From 7696fca8266dcd10e360fc6b31a94e6dc1f55766 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 28 Mar 2026 15:31:13 -0500 Subject: [PATCH] fix: restore MCP telemetry compatibility and outcomes Signed-off-by: phernandez --- .../api/v2/routers/search_router.py | 5 +- src/basic_memory/api/v2/utils.py | 11 +- src/basic_memory/mcp/tools/build_context.py | 2 +- src/basic_memory/mcp/tools/edit_note.py | 2 +- src/basic_memory/mcp/tools/read_note.py | 30 +++-- src/basic_memory/mcp/tools/search.py | 6 +- src/basic_memory/mcp/tools/utils.py | 118 +++++++++++++++--- src/basic_memory/mcp/tools/write_note.py | 2 +- src/basic_memory/services/entity_service.py | 16 ++- src/basic_memory/services/search_service.py | 9 +- src/basic_memory/telemetry.py | 12 +- tests/api/v2/test_utils_telemetry.py | 5 +- tests/mcp/test_client_telemetry.py | 106 +++++++++++++++- tests/mcp/test_tool_read_note.py | 57 +++++++++ tests/mcp/test_tool_telemetry.py | 10 +- tests/test_telemetry.py | 38 +++++- 16 files changed, 370 insertions(+), 59 deletions(-) diff --git a/src/basic_memory/api/v2/routers/search_router.py b/src/basic_memory/api/v2/routers/search_router.py index eb098729..37ebfd58 100644 --- a/src/basic_memory/api/v2/routers/search_router.py +++ b/src/basic_memory/api/v2/routers/search_router.py @@ -57,7 +57,10 @@ async def search( page_size=page_size, retrieval_mode=query.retrieval_mode.value, has_query=bool( - (query.text and query.text.strip()) or query.title or query.permalink or query.permalink_match + (query.text and query.text.strip()) + or query.title + or query.permalink + or query.permalink_match ), has_filters=bool(query.note_types or query.entity_types or query.metadata_filters), ): diff --git a/src/basic_memory/api/v2/utils.py b/src/basic_memory/api/v2/utils.py index c72a0c1b..e11d0ee5 100644 --- a/src/basic_memory/api/v2/utils.py +++ b/src/basic_memory/api/v2/utils.py @@ -39,7 +39,9 @@ async def to_graph_context( entity_ids_needed: set[int] = set() for context_item in context_result.results: for item in ( - [context_item.primary_result] + context_item.observations + context_item.related_results + [context_item.primary_result] + + context_item.observations + + context_item.related_results ): if item.type == SearchItemType.ENTITY: # Entity's own ID for its external_id @@ -100,9 +102,7 @@ async def to_graph_context( created_at=item.created_at, ) case SearchItemType.RELATION: - from_title = ( - entity_title_lookup.get(item.from_id) if item.from_id else None - ) # pyright: ignore + from_title = entity_title_lookup.get(item.from_id) if item.from_id else None # pyright: ignore to_title = entity_title_lookup.get(item.to_id) if item.to_id else None from_ext_id = ( entity_external_id_lookup.get(item.from_id) if item.from_id else None @@ -154,7 +154,8 @@ async def to_graph_context( generated_at=context_result.metadata.generated_at, primary_count=context_result.metadata.primary_count, related_count=context_result.metadata.related_count, - total_results=context_result.metadata.primary_count + context_result.metadata.related_count, + total_results=context_result.metadata.primary_count + + context_result.metadata.related_count, total_relations=context_result.metadata.total_relations, total_observations=context_result.metadata.total_observations, ) diff --git a/src/basic_memory/mcp/tools/build_context.py b/src/basic_memory/mcp/tools/build_context.py index 28a7265a..4131b365 100644 --- a/src/basic_memory/mcp/tools/build_context.py +++ b/src/basic_memory/mcp/tools/build_context.py @@ -206,7 +206,7 @@ async def build_context( "mcp.tool.build_context", entrypoint="mcp", tool_name="build_context", - project_name=project, + requested_project=project, workspace_id=workspace, depth=depth or 1, timeframe=timeframe, diff --git a/src/basic_memory/mcp/tools/edit_note.py b/src/basic_memory/mcp/tools/edit_note.py index 49323a2d..c0bb3c46 100644 --- a/src/basic_memory/mcp/tools/edit_note.py +++ b/src/basic_memory/mcp/tools/edit_note.py @@ -275,7 +275,7 @@ async def edit_note( "mcp.tool.edit_note", entrypoint="mcp", tool_name="edit_note", - project_name=project, + requested_project=project, workspace_id=workspace, edit_operation=operation, output_format=output_format, diff --git a/src/basic_memory/mcp/tools/read_note.py b/src/basic_memory/mcp/tools/read_note.py index 3d336ba3..09dc8916 100644 --- a/src/basic_memory/mcp/tools/read_note.py +++ b/src/basic_memory/mcp/tools/read_note.py @@ -16,8 +16,8 @@ from basic_memory.mcp.project_context import ( resolve_project_and_path, ) from basic_memory.mcp.server import mcp +from basic_memory.mcp.tools.search import search_notes from basic_memory.schemas.memory import memory_url_path -from basic_memory.schemas.search import SearchQuery from basic_memory.utils import validate_project_path @@ -144,7 +144,7 @@ async def read_note( "mcp.tool.read_note", entrypoint="mcp", tool_name="read_note", - project_name=project, + requested_project=project, workspace_id=workspace, output_format=output_format, page=page, @@ -199,12 +199,11 @@ async def read_note( ) # Import here to avoid circular import - from basic_memory.mcp.clients import KnowledgeClient, ResourceClient, SearchClient + from basic_memory.mcp.clients import KnowledgeClient, ResourceClient # Use typed clients for API calls knowledge_client = KnowledgeClient(client, active_project.external_id) resource_client = ResourceClient(client, active_project.external_id) - search_client = SearchClient(client, active_project.external_id) async def _read_json_payload(entity_id: str) -> dict: with telemetry.scope( @@ -214,7 +213,9 @@ async def read_note( phase="shape_response", ): entity = await knowledge_client.get_entity(entity_id) - response = await resource_client.read(entity_id, page=page, page_size=page_size) + response = await resource_client.read( + entity_id, page=page, page_size=page_size + ) content_text = response.text body_content, parsed_frontmatter = _parse_opening_frontmatter(content_text) return { @@ -241,15 +242,22 @@ async def read_note( return results if isinstance(results, list) else [] async def _search_candidates(identifier_text: str, *, title_only: bool) -> dict: - query = SearchQuery(title=identifier_text) if title_only else SearchQuery( - text=identifier_text - ) - response = await search_client.search( - query.model_dump(mode="json", exclude_none=True), + # Trigger: direct entity resolution failed for the caller's identifier. + # Why: search_notes applies the same memory:// normalization and tool-level + # query handling as the rest of MCP routing, which raw client calls skip. + # Outcome: unresolved memory URLs still fall back through normalized search. + search_type = "title" if title_only else "text" + response = await search_notes( + project=active_project.name, + workspace=workspace, + query=identifier_text, + search_type=search_type, page=page, page_size=page_size, + output_format="json", + context=context, ) - return response.model_dump(mode="json") + return response if isinstance(response, dict) else {} def _result_title(item: dict) -> str: return str(item.get("title") or "") diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py index 9f33f9a9..94c99aa9 100644 --- a/src/basic_memory/mcp/tools/search.py +++ b/src/basic_memory/mcp/tools/search.py @@ -528,7 +528,7 @@ async def search_notes( "mcp.tool.search_notes", entrypoint="mcp", tool_name="search_notes", - project_name=project, + requested_project=project, workspace_id=workspace, search_type=search_type or "default", output_format=output_format, @@ -537,7 +537,9 @@ async def search_notes( has_query=bool(query and query.strip()), note_type_filter_count=len(note_types), entity_type_filter_count=len(entity_types), - has_filters=bool(metadata_filters or tags or status or note_types or entity_types or after_date), + has_filters=bool( + metadata_filters or tags or status or note_types or entity_types or after_date + ), has_tags_filter=bool(tags), has_status_filter=bool(status), ): diff --git a/src/basic_memory/mcp/tools/utils.py b/src/basic_memory/mcp/tools/utils.py index 6f847c10..ded5a7bd 100644 --- a/src/basic_memory/mcp/tools/utils.py +++ b/src/basic_memory/mcp/tools/utils.py @@ -5,6 +5,7 @@ to the Basic Memory API, with improved error handling and logging. """ import typing +from contextlib import contextmanager from typing import Optional from httpx import Response, URL, AsyncClient, HTTPStatusError @@ -27,6 +28,58 @@ from basic_memory import telemetry from basic_memory.config import ConfigManager +def _classify_http_outcome(status_code: int) -> str: + """Map HTTP status codes to a low-cardinality outcome label.""" + if 200 <= status_code < 300: + return "success" + if 300 <= status_code < 400: # pragma: no cover + return "redirect" + if 400 <= status_code < 500: + return "client_error" + if 500 <= status_code < 600: + return "server_error" + return "unknown" # pragma: no cover + + +class _RequestSpan: + """Small adapter for attaching outcome metadata to a live request span.""" + + def __init__(self, active_span: typing.Any | None): + self._active_span = active_span + + def record_response(self, response: Response) -> None: + self._set_attributes( + { + "status_code": response.status_code, + "is_success": response.is_success, + "outcome": _classify_http_outcome(response.status_code), + } + ) + + def record_transport_error(self, exc: Exception) -> None: + self._set_attributes( + { + "is_success": False, + "outcome": "transport_error", + "error_type": type(exc).__name__, + } + ) + + def _set_attributes(self, attrs: dict[str, typing.Any]) -> None: + if self._active_span is None: + return + + set_attributes = getattr(self._active_span, "set_attributes", None) + if callable(set_attributes): + set_attributes(attrs) + return + + set_attribute = getattr(self._active_span, "set_attribute", None) + if callable(set_attribute): + for key, value in attrs.items(): + set_attribute(key, value) + + def get_error_message( status_code: int, url: URL | str, method: str, msg: Optional[str] = None ) -> str: @@ -136,6 +189,7 @@ def _resolve_error_message( return get_error_message(status_code, url, method) +@contextmanager def _request_scope( method: str, *, @@ -146,16 +200,18 @@ def _request_scope( has_body: bool = False, ): """Create the shared MCP transport span used by all HTTP helpers.""" - return telemetry.scope( - "mcp.http.request", - method=method, - client_name=client_name, - operation=operation, - path_template=path_template, - phase="request", - has_query=bool(params), - has_body=has_body, - ) + attrs = { + "method": method, + "client_name": client_name, + "operation": operation, + "path_template": path_template, + "phase": "request", + "has_query": bool(params), + "has_body": has_body, + } + with telemetry.contextualize(**attrs): + with telemetry.started_span("mcp.http.request", **attrs) as active_span: + yield _RequestSpan(active_span) async def call_get( @@ -194,6 +250,7 @@ async def call_get( """ logger.debug(f"Calling GET '{url}' params: '{params}'") error_message = None + request_span: _RequestSpan | None = None try: with _request_scope( @@ -202,7 +259,7 @@ async def call_get( operation=operation, path_template=path_template, params=params, - ): + ) as request_span: response = await client.get( url, params=params, @@ -213,6 +270,7 @@ async def call_get( timeout=timeout, extensions=extensions, ) + request_span.record_response(response) if response.is_success: return response @@ -239,6 +297,10 @@ async def call_get( except HTTPStatusError as e: raise ToolError(error_message) from e + except Exception as e: + if request_span is not None: + request_span.record_transport_error(e) + raise async def call_put( @@ -285,6 +347,7 @@ async def call_put( """ logger.debug(f"Calling PUT '{url}'") error_message = None + request_span: _RequestSpan | None = None try: with _request_scope( @@ -294,7 +357,7 @@ async def call_put( path_template=path_template, params=params, has_body=any(value is not None for value in (content, data, files, json)), - ): + ) as request_span: response = await client.put( url, content=content, @@ -309,6 +372,7 @@ async def call_put( timeout=timeout, extensions=extensions, ) + request_span.record_response(response) if response.is_success: return response @@ -336,6 +400,10 @@ async def call_put( except HTTPStatusError as e: raise ToolError(error_message) from e + except Exception as e: + if request_span is not None: + request_span.record_transport_error(e) + raise async def call_patch( @@ -381,6 +449,7 @@ async def call_patch( ToolError: If the request fails with an appropriate error message """ logger.debug(f"Calling PATCH '{url}'") + request_span: _RequestSpan | None = None try: with _request_scope( @@ -390,7 +459,7 @@ async def call_patch( path_template=path_template, params=params, has_body=any(value is not None for value in (content, data, files, json)), - ): + ) as request_span: response = await client.patch( url, content=content, @@ -405,6 +474,7 @@ async def call_patch( timeout=timeout, extensions=extensions, ) + request_span.record_response(response) if response.is_success: return response @@ -437,6 +507,10 @@ async def call_patch( error_message = _resolve_error_message(status_code, url, "PATCH", response_data) raise ToolError(error_message) from e + except Exception as e: + if request_span is not None: + request_span.record_transport_error(e) + raise async def call_post( @@ -483,6 +557,7 @@ async def call_post( """ logger.debug(f"Calling POST '{url}'") error_message = None + request_span: _RequestSpan | None = None try: with _request_scope( @@ -492,7 +567,7 @@ async def call_post( path_template=path_template, params=params, has_body=any(value is not None for value in (content, data, files, json)), - ): + ) as request_span: response = await client.post( url=url, content=content, @@ -507,7 +582,8 @@ async def call_post( timeout=timeout, extensions=extensions, ) - logger.debug(f"response: {response.json()}") + request_span.record_response(response) + logger.debug(f"response: {_extract_response_data(response)}") if response.is_success: return response @@ -534,6 +610,10 @@ async def call_post( except HTTPStatusError as e: raise ToolError(error_message) from e + except Exception as e: + if request_span is not None: + request_span.record_transport_error(e) + raise async def resolve_entity_id(client: AsyncClient, project_external_id: str, identifier: str) -> str: @@ -604,6 +684,7 @@ async def call_delete( """ logger.debug(f"Calling DELETE '{url}'") error_message = None + request_span: _RequestSpan | None = None try: with _request_scope( @@ -612,7 +693,7 @@ async def call_delete( operation=operation, path_template=path_template, params=params, - ): + ) as request_span: response = await client.delete( url=url, params=params, @@ -623,6 +704,7 @@ async def call_delete( timeout=timeout, extensions=extensions, ) + request_span.record_response(response) if response.is_success: return response @@ -649,3 +731,7 @@ async def call_delete( except HTTPStatusError as e: raise ToolError(error_message) from e + except Exception as e: + if request_span is not None: + request_span.record_transport_error(e) + raise diff --git a/src/basic_memory/mcp/tools/write_note.py b/src/basic_memory/mcp/tools/write_note.py index f832565d..2ec997e9 100644 --- a/src/basic_memory/mcp/tools/write_note.py +++ b/src/basic_memory/mcp/tools/write_note.py @@ -153,7 +153,7 @@ async def write_note( "mcp.tool.write_note", entrypoint="mcp", tool_name="write_note", - project_name=project, + requested_project=project, workspace_id=workspace, note_type=note_type, overwrite=effective_overwrite, diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 61cbf133..7df735ee 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -433,7 +433,9 @@ class EntityService(BaseService[EntityModel]): action="update", phase="upsert_entity", ): - entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) + entity = await self.upsert_entity_from_markdown( + file_path, entity_markdown, is_new=False + ) with telemetry.scope( "entity_service.update.update_checksum", @@ -465,7 +467,9 @@ class EntityService(BaseService[EntityModel]): action="fast_write", phase="resolve_entity", ): - existing = await self.repository.get_by_external_id(external_id) if external_id else None + existing = ( + await self.repository.get_by_external_id(external_id) if external_id else None + ) # Trigger: external_id already exists # Why: avoid duplicate entities when title-derived paths change @@ -698,7 +702,9 @@ class EntityService(BaseService[EntityModel]): action="reindex", phase="upsert_entity", ): - updated = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) + updated = await self.upsert_entity_from_markdown( + file_path, entity_markdown, is_new=False + ) with telemetry.scope( "entity_service.reindex.update_checksum", domain="entity_service", @@ -1019,7 +1025,9 @@ class EntityService(BaseService[EntityModel]): action="edit", phase="upsert_entity", ): - entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) + entity = await self.upsert_entity_from_markdown( + file_path, entity_markdown, is_new=False + ) with telemetry.scope( "entity_service.edit.update_checksum", diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 6057482f..b1a77051 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -174,10 +174,7 @@ class SearchService: retrieval_mode = query.retrieval_mode or SearchRetrievalMode.FTS strict_search_text = query.text has_query = bool( - strict_search_text - or query.title - or query.permalink - or query.permalink_match + strict_search_text or query.title or query.permalink or query.permalink_match ) has_filters = bool( metadata_filters @@ -658,7 +655,9 @@ class SearchService: ) ) if len(obs_content_stems) > MAX_CONTENT_STEMS_SIZE: # pragma: no cover - obs_content_stems = obs_content_stems[:MAX_CONTENT_STEMS_SIZE] # pragma: no cover + obs_content_stems = obs_content_stems[ + :MAX_CONTENT_STEMS_SIZE + ] # pragma: no cover rows_to_index.append( SearchIndexRow( id=obs.id, diff --git a/src/basic_memory/telemetry.py b/src/basic_memory/telemetry.py index adc48b35..35f6cfb5 100644 --- a/src/basic_memory/telemetry.py +++ b/src/basic_memory/telemetry.py @@ -159,13 +159,20 @@ operation = scope @contextmanager def span(name: str, **attrs: Any) -> Iterator[None]: """Create a manual Logfire span when telemetry is enabled.""" + with started_span(name, **attrs): + yield + + +@contextmanager +def started_span(name: str, **attrs: Any) -> Iterator[Any | None]: + """Create a manual Logfire span and expose the active span handle when available.""" logfire = _load_logfire() if logfire is None or not _STATE.configured: # pragma: no cover yield # pragma: no cover return # pragma: no cover - with logfire.span(name, **_filter_attributes(attrs)): - yield + with logfire.span(name, **_filter_attributes(attrs)) as active_span: + yield active_span __all__ = [ @@ -177,5 +184,6 @@ __all__ = [ "reset_telemetry_state", "scope", "span", + "started_span", "telemetry_enabled", ] diff --git a/tests/api/v2/test_utils_telemetry.py b/tests/api/v2/test_utils_telemetry.py index 6262c638..5f02f8ea 100644 --- a/tests/api/v2/test_utils_telemetry.py +++ b/tests/api/v2/test_utils_telemetry.py @@ -32,7 +32,10 @@ async def test_to_search_results_emits_hydration_spans(monkeypatch) -> None: class FakeEntityService: async def get_entities_by_id(self, ids): - return [SimpleNamespace(permalink="notes/root"), SimpleNamespace(permalink="notes/child")] + return [ + SimpleNamespace(permalink="notes/root"), + SimpleNamespace(permalink="notes/child"), + ] now = datetime.now(timezone.utc) results = [ diff --git a/tests/mcp/test_client_telemetry.py b/tests/mcp/test_client_telemetry.py index 8aa88211..609261ef 100644 --- a/tests/mcp/test_client_telemetry.py +++ b/tests/mcp/test_client_telemetry.py @@ -7,9 +7,11 @@ from contextlib import contextmanager import httpx import pytest +from mcp.server.fastmcp.exceptions import ToolError knowledge_client_module = importlib.import_module("basic_memory.mcp.clients.knowledge") search_client_module = importlib.import_module("basic_memory.mcp.clients.search") +utils_module = importlib.import_module("basic_memory.mcp.tools.utils") def _capture_spans(): @@ -20,13 +22,27 @@ def _capture_spans(): spans.append((name, attrs)) yield - return spans, fake_span + @contextmanager + def fake_started_span(name: str, **attrs): + spans.append((name, attrs)) + + class FakeStartedSpan: + def set_attribute(self, key: str, value) -> None: + attrs[key] = value + + def set_attributes(self, new_attrs: dict) -> None: + attrs.update(new_attrs) + + yield FakeStartedSpan() + + return spans, fake_span, fake_started_span @pytest.mark.asyncio async def test_knowledge_client_resolve_entity_emits_client_and_http_spans(monkeypatch) -> None: - spans, fake_span = _capture_spans() + spans, fake_span, fake_started_span = _capture_spans() monkeypatch.setattr(knowledge_client_module.telemetry, "span", fake_span) + monkeypatch.setattr(utils_module.telemetry, "started_span", fake_started_span) async def handler(request: httpx.Request) -> httpx.Response: assert request.method == "POST" @@ -50,13 +66,17 @@ async def test_knowledge_client_resolve_entity_emits_client_and_http_spans(monke "phase": "request", "has_query": False, "has_body": True, + "status_code": 200, + "is_success": True, + "outcome": "success", } @pytest.mark.asyncio async def test_search_client_emits_client_and_http_spans(monkeypatch) -> None: - spans, fake_span = _capture_spans() + spans, fake_span, fake_started_span = _capture_spans() monkeypatch.setattr(search_client_module.telemetry, "span", fake_span) + monkeypatch.setattr(utils_module.telemetry, "started_span", fake_started_span) async def handler(request: httpx.Request) -> httpx.Response: assert request.method == "POST" @@ -88,4 +108,84 @@ async def test_search_client_emits_client_and_http_spans(monkeypatch) -> None: "phase": "request", "has_query": True, "has_body": True, + "status_code": 200, + "is_success": True, + "outcome": "success", } + + +@pytest.mark.asyncio +async def test_call_get_emits_http_outcome_for_client_errors(monkeypatch) -> None: + spans, _, fake_started_span = _capture_spans() + monkeypatch.setattr(utils_module.telemetry, "started_span", fake_started_span) + + async def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "GET" + return httpx.Response(404, json={"detail": "missing"}) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="https://example.test") as client: + with pytest.raises(ToolError, match="missing"): + await utils_module.call_get( + client, + "/missing", + client_name="knowledge", + operation="get_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}", + ) + + assert spans == [ + ( + "mcp.http.request", + { + "method": "GET", + "client_name": "knowledge", + "operation": "get_entity", + "path_template": "/v2/projects/{project_id}/knowledge/entities/{entity_id}", + "phase": "request", + "has_query": False, + "has_body": False, + "status_code": 404, + "is_success": False, + "outcome": "client_error", + }, + ) + ] + + +@pytest.mark.asyncio +async def test_call_get_emits_transport_error_outcome(monkeypatch) -> None: + spans, _, fake_started_span = _capture_spans() + monkeypatch.setattr(utils_module.telemetry, "started_span", fake_started_span) + + async def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("boom", request=request) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="https://example.test") as client: + with pytest.raises(httpx.ConnectError, match="boom"): + await utils_module.call_get( + client, + "/boom", + client_name="knowledge", + operation="get_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}", + ) + + assert spans == [ + ( + "mcp.http.request", + { + "method": "GET", + "client_name": "knowledge", + "operation": "get_entity", + "path_template": "/v2/projects/{project_id}/knowledge/entities/{entity_id}", + "phase": "request", + "has_query": False, + "has_body": False, + "is_success": False, + "outcome": "transport_error", + "error_type": "ConnectError", + }, + ) + ] diff --git a/tests/mcp/test_tool_read_note.py b/tests/mcp/test_tool_read_note.py index d485f1ec..5fa4dbfb 100644 --- a/tests/mcp/test_tool_read_note.py +++ b/tests/mcp/test_tool_read_note.py @@ -264,6 +264,63 @@ async def test_read_note_memory_url_with_project_prefix(app, test_project): assert "Testing memory:// URL handling with project prefix" in content +@pytest.mark.asyncio +async def test_read_note_memory_url_fallback_uses_search_tool_normalization( + monkeypatch, app, test_project +): + """Fallback search should go back through search_notes for memory:// normalization.""" + await write_note( + project=test_project.name, + title="Memory URL Fallback Note", + directory="test", + content="Fallback note content", + ) + + import importlib + + read_note_module = importlib.import_module("basic_memory.mcp.tools.read_note") + clients_mod = importlib.import_module("basic_memory.mcp.clients") + OriginalKnowledgeClient = clients_mod.KnowledgeClient + + fallback_memory_url = f"memory://{test_project.name}/test/memory-url-fallback-note" + search_calls: list[tuple[str, str, str | None]] = [] + + class SelectiveKnowledgeClient(OriginalKnowledgeClient): + async def resolve_entity(self, identifier: str, *, strict: bool = False) -> int: + if strict and identifier.endswith("test/memory-url-fallback-note"): + raise RuntimeError("force direct lookup failure") + return await super().resolve_entity(identifier, strict=strict) + + async def fake_search_notes_fn(*, query, search_type, project, **kwargs): + search_calls.append((search_type, query, project)) + return { + "results": [ + { + "title": "Memory URL Fallback Note", + "permalink": "test/memory-url-fallback-note", + "content": "", + "type": "entity", + "score": 1.0, + "file_path": "test/Memory URL Fallback Note.md", + } + ], + "current_page": 1, + "page_size": 10, + } + + monkeypatch.setattr(clients_mod, "KnowledgeClient", SelectiveKnowledgeClient) + monkeypatch.setattr(read_note_module, "search_notes", fake_search_notes_fn) + + result = await read_note(fallback_memory_url) + + assert search_calls == [ + ("title", fallback_memory_url, test_project.name), + ("text", fallback_memory_url, test_project.name), + ] + assert "I couldn't find an exact match" in result + assert "Memory URL Fallback Note" in result + + class TestReadNoteSecurityValidation: """Test read_note security validation features.""" diff --git a/tests/mcp/test_tool_telemetry.py b/tests/mcp/test_tool_telemetry.py index 26e39aaa..5c0fedf4 100644 --- a/tests/mcp/test_tool_telemetry.py +++ b/tests/mcp/test_tool_telemetry.py @@ -56,7 +56,7 @@ async def test_write_note_emits_root_operation_and_project_context( { "entrypoint": "mcp", "tool_name": "write_note", - "project_name": test_project.name, + "requested_project": test_project.name, "workspace_id": None, "note_type": "note", "overwrite": False, @@ -107,7 +107,7 @@ async def test_read_note_emits_root_operation_and_project_context( { "entrypoint": "mcp", "tool_name": "read_note", - "project_name": test_project.name, + "requested_project": test_project.name, "workspace_id": None, "output_format": "json", "page": 1, @@ -164,7 +164,7 @@ async def test_search_notes_emits_root_operation_and_project_context( { "entrypoint": "mcp", "tool_name": "search_notes", - "project_name": test_project.name, + "requested_project": test_project.name, "workspace_id": None, "search_type": "text", "output_format": "json", @@ -223,7 +223,7 @@ async def test_edit_note_emits_root_operation_and_project_context( { "entrypoint": "mcp", "tool_name": "edit_note", - "project_name": test_project.name, + "requested_project": test_project.name, "workspace_id": None, "edit_operation": "append", "output_format": "json", @@ -282,7 +282,7 @@ async def test_build_context_emits_root_operation_and_project_context( { "entrypoint": "mcp", "tool_name": "build_context", - "project_name": test_project.name, + "requested_project": test_project.name, "workspace_id": None, "depth": 2, "timeframe": "7d", diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index 3bcf0f5a..cd133df4 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -33,7 +33,15 @@ class FakeLogfire: @contextmanager def span(self, name: str, **attrs): self.span_calls.append((name, attrs)) - yield + + class FakeStartedSpan: + def set_attribute(self, key: str, value) -> None: + attrs[key] = value + + def set_attributes(self, new_attrs: dict) -> None: + attrs.update(new_attrs) + + yield FakeStartedSpan() def test_configure_telemetry_disabled_is_noop() -> None: @@ -140,6 +148,34 @@ def test_span_uses_logfire_when_enabled(monkeypatch) -> None: assert fake_logfire.span_calls == [("mcp.tool.write_note", {"project_name": "main"})] +def test_started_span_exposes_mutable_logfire_handle(monkeypatch) -> None: + fake_logfire = FakeLogfire() + telemetry.reset_telemetry_state() + monkeypatch.setattr(telemetry, "_load_logfire", lambda: fake_logfire) + telemetry.configure_telemetry( + "basic-memory-mcp", + environment="dev", + enable_logfire=True, + ) + + with telemetry.started_span("mcp.http.request", method="GET") as span: + assert span is not None + span.set_attribute("status_code", 200) + span.set_attributes({"is_success": True, "outcome": "success"}) + + assert fake_logfire.span_calls == [ + ( + "mcp.http.request", + { + "method": "GET", + "status_code": 200, + "is_success": True, + "outcome": "success", + }, + ) + ] + + def test_operation_creates_span_and_log_context(monkeypatch) -> None: fake_logfire = FakeLogfire() records: list[dict] = []