Add DataHub MCP recall context - #341
Conversation
…text # Conflicts: # docs/api/openapi.json # server/cmd/mnemo-server/main.go # server/internal/handler/handler.go # server/internal/handler/memory.go
King-Dylan
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7edcc9e8c5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if resp.StatusCode < 200 || resp.StatusCode >= 300 { | ||
| return nil, fmt.Errorf("MCP %s returned HTTP %d: %s", method, resp.StatusCode, compactSnippet(string(responseBody), 300)) |
There was a problem hiding this comment.
When a stateful MCP server restarts or expires a session, requests carrying the old Mcp-Session-Id return HTTP 404; the Streamable HTTP spec requires the client to start a new initialize without that session. This path only returns an error and leaves initialized/sessionID unchanged, so every later recall keeps sending the stale session ID and DataHub context remains broken until the mem9 process is restarted.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
From the DataHub side - I don't think this 404 path is actually reachable against DataHub's MCP server. It runs stateless (stateless_http=True), so there's no server-side session to expire or 404 on; a stale Mcp-Session-Id with a valid token just proceeds. So the re-initialize-on-404 handling likely isn't needed.
The failure that does happen is auth-related, and it differs by deployment: self-hosted has no HTTP auth layer and surfaces a bad/expired token as a JSON-RPC error inside a 200 body (which post() already catches), while DataHub Cloud's managed endpoint rejects at the HTTP layer with 401 (or 400 when OAuth2 isn't enabled). Might be worth special-casing 401/400 so a Cloud token problem gives a clear signal — but the stateless 404 case here shouldn't bite.
…text # Conflicts: # docs/api/openapi.json # server/internal/handler/memory.go
lakshay-nasa
left a comment
lakshay-nasa
left a comment
There was a problem hiding this comment.
Thanks for putting this together. Read through the DataHub integration and dropped a few notes inline — most are minor, one's worth a closer look (how curated names/descriptions get resolved on the DataHub side).
| } | ||
|
|
||
| func dataHubEntityTitle(entity map[string]any) string { | ||
| return firstString(entity, "name", "displayName", "qualifiedName", "properties.name", "properties.displayName", "properties.qualifiedName", "urn") |
There was a problem hiding this comment.
Reading through the normalization here - one thing worth knowing from the DataHub side: the human-set display values live under editableProperties, and DataHub expects clients to prefer those over the ingested properties.* ones (resolution order is editableProperties → properties → top-level). This title waterfall checks properties.name but not editableProperties.name, and just above in dataHubEntityItem the description checks editableProperties.description last rather than first.
The effect is that for assets a team has curated in DataHub, you'd surface the raw ingested name/description instead of the display values they actually set - which tends to be exactly the well-managed datasets people search for. Putting editableProperties.* first for both should cover it; it falls through harmlessly on types that don't have it (dashboards/charts only carry editableProperties.description, not a name).
Since the tests mock the MCP responses, this reads green today - a fixture with editableProperties populated would surface it.
| ID: firstString(entity, "urn"), | ||
| Title: title, | ||
| Snippet: compactSnippet(snippet, 800), | ||
| URL: firstString(entity, "url", "properties.externalUrl"), |
There was a problem hiding this comment.
Small thing: url only gets populated on DataHub Cloud – self-hosted/OSS doesn't inject it, so this will be empty there. Probably worth treating as optional in whatever renders external_context. (properties.externalUrl as a fallback is a nice touch — it's sometimes set, though often not either.)
| "max_results": lineageLimit, | ||
| "offset": 0, | ||
| }) | ||
| if err != nil { |
There was a problem hiding this comment.
Minor: a get_lineage failure gets swallowed here with no log. The best-effort behavior is right — better to return partial context than fail recall — but a debug-level log would make a systematically broken lineage path noticeable instead of silently empty.
| if resp.StatusCode < 200 || resp.StatusCode >= 300 { | ||
| return nil, fmt.Errorf("MCP %s returned HTTP %d: %s", method, resp.StatusCode, compactSnippet(string(responseBody), 300)) |
There was a problem hiding this comment.
From the DataHub side - I don't think this 404 path is actually reachable against DataHub's MCP server. It runs stateless (stateless_http=True), so there's no server-side session to expire or 404 on; a stale Mcp-Session-Id with a valid token just proceeds. So the re-initialize-on-404 handling likely isn't needed.
The failure that does happen is auth-related, and it differs by deployment: self-hosted has no HTTP auth layer and surfaces a bad/expired token as a JSON-RPC error inside a 200 body (which post() already catches), while DataHub Cloud's managed endpoint rejects at the HTTP layer with 401 (or 400 when OAuth2 isn't enabled). Might be worth special-casing 401/400 so a Cloud token problem gives a clear signal — but the stateless 404 case here shouldn't bite.
2 participants
Summary
external_contextto memory list/search responses for data-asset questions.search,get_entities, and one-hop upstream/downstreamget_lineagecalls.Executive Revenueand prefer dataset URNs when selecting the lineage root so recall consistently returns the right lineage target.Purpose
This PR is the mem9-side DataHub integration path only. mem9 remains the source of truth for user and agent memory, while DataHub remains the source of truth for data assets, metadata, ownership, quality, and lineage. At recall time, agents can combine mem9 long-term memory with governed DataHub metadata context without copying DataHub into mem9 or writing mem9 data back to DataHub yet.
Refs #340
Demo Boundary
The standalone demo has been extracted into a separate repository:
This PR intentionally keeps only the integration and recall behavior inside
mem9.Next Step
Follow-up work should add an optional mem9 -> DataHub publish path for curated data-asset observations. That path should classify and reconcile mem9 memories that clearly refer to DataHub entities, map them to DataHub URNs, and publish selected metadata-grade signals such as dashboard or data issue observations, freshness or quality annotations, owner hints, glossary notes, or custom DataHub aspects. Raw personal memory and unrelated agent memory should stay in mem9.
Validation
python3 -m json.tool docs/api/openapi.json >/dev/nullcd server && go test -count=1 -run 'TestDataHubMCP|TestShouldQueryDataHubContext|TestFormatDataHubSearchQuery|TestNormalizeDataHubTextContent|TestNormalizeDataHubLineageToolResult|TestExtractSSEData|TestCollectDataHubURNs|TestTrimExternalContextItems|TestCompactSnippet|TestListMemories_IncludesDataHubExternalContextForDataQuestions|TestListMemories_ExternalContextFalseOverrideWins|TestListMemories_ExternalContextRepeatedFalseOverrideWins|TestLoad_DataHubMCP' ./internal/service ./internal/handler ./internal/configcd server && go test -race -count=1 ./internal/service/...git diff --checkExecutive Revenue,mart.revenue, and both upstream and downstream lineage summaries.