feat: structured lineage and column reach for dbt_lineage and impact_analysis

packages/opencode/src/altimate/mcp-server/server.ts

@@ -43,6 +43,81 @@ export async function serverInstructions(): Promise<string> {
   }
 }
 
+// altimate_change start — lineage tools answer with structured content as well as text
+// Lineage tools also return their metadata as structured content, so a client that draws lineage
+// (the Claude Code plugin's lineage pane) reads edges and downstream lists without parsing text.
+// Claude Code shows structured content to the model in place of the text, so `report` carries the
+// (already truncated) text (dbt_lineage: its header, its mappings travel as data), only the lineage
+// fields are passed, and the whole is capped.
+const STRUCTURED_FIELDS: Record<string, readonly string[]> = {
+  dbt_lineage: ["model_name", "model_unique_id", "confidence", "confidence_factors", "upstream", "columns", "edges", "error"],
+  impact_analysis: [
+    "model",
+    "column",
+    "change_type",
+    "severity",
+    "direct_count",
+    "transitive_count",
+    "test_count",
+    "downstream",
+    "exposures",
+    "column_reach",
+    "column_reach_error",
+    "exposures_error",
+    "error",
+  ],
+}
+export const STRUCTURED_TOOLS: ReadonlySet<string> = new Set(Object.keys(STRUCTURED_FIELDS))
+const STRUCTURED_OUTPUT_SCHEMA = {
+  type: "object" as const,
+  properties: {
+    report: { type: "string", description: "The tool's text report" },
+    truncated: { type: "boolean", description: "True when the lineage fields were left out for size; the report is kept, shortened if it is too large on its own" },
+  },
+  required: ["report"],
+  additionalProperties: true,
+}
+// The structured result reaches the model in Claude Code, so it stays within the same order of
+// size as a text result: past this, only the report is sent. Claude Code caps an MCP result at
+// 25k tokens by default and JSON of ids and edges runs ~3-4 characters a token, so this stays under that.
+const STRUCTURED_MAX_BYTES = 60_000
+
+export function structuredResult(tool: string, metadata: Record<string, unknown>, report: string): Record<string, unknown> {
+  // A tool that builds its own structured result (dbt_lineage from the dbt integration) hands it
+  // over whole; otherwise the lineage fields of its metadata are passed.
+  const own = metadata.structuredContent
+  if (own && typeof own === "object" && !Array.isArray(own)) {
+    // The text repeats every edge that `columnLineage` already carries: only its header goes with them.
+    const head = report.split(/\n(?:Column edges \(|Trace of )/)[0]!.trim()
+    const full = { ...(own as Record<string, unknown>), report: head }
+    return Buffer.byteLength(JSON.stringify(full)) <= STRUCTURED_MAX_BYTES ? full : capReport(report, true)
+  }
+  const fields = Object.fromEntries((STRUCTURED_FIELDS[tool] ?? []).filter((key) => key in metadata).map((key) => [key, metadata[key]]))
+  // dbt_lineage's text prints the same mappings and edges as `columns` and `edges`: keep its header.
+  const head = tool === "dbt_lineage" && "columns" in fields ? report.split(/\n(?:Column Mappings|Lineage Edges):/)[0]!.trim() : report
+  const full = { ...fields, report: head }
+  return Buffer.byteLength(JSON.stringify(full)) <= STRUCTURED_MAX_BYTES ? full : capReport(report, true)
+}
+/** An error report held to the same cap as a result, so a huge exception message cannot exceed it. */
+export function structuredError(message: string): Record<string, unknown> {
+  return capReport(message, false)
+}
+// `dropped`: the result's other fields were left out, so it is marked truncated even when the
+// report itself fits.
+function capReport(message: string, dropped: boolean): Record<string, unknown> {
+  // Measured as sent: JSON escaping can grow the text (a newline is two bytes, a control
+  // character six), so the cap applies to the serialized result, shortening until it fits.
+  const size = (value: Record<string, unknown>) => Buffer.byteLength(JSON.stringify(value))
+  const whole = dropped ? { report: message, truncated: true } : { report: message }
+  if (size(whole) <= STRUCTURED_MAX_BYTES) return whole
+  let report = message
+  while (report.length > 0 && size({ report, truncated: true }) > STRUCTURED_MAX_BYTES) {
+    report = report.slice(0, Math.floor(report.length * 0.9))
+  }
+  return { report, truncated: true }
+}
+// altimate_change end
+
 /** Spilled results: a resource_link next to the text, which keeps the preview, row count and path. */
 function resultLinks(metadata: Record<string, unknown>) {
   // A batched sql_execute call reports one saved file per spilled query.
@@ -97,6 +172,7 @@ export async function createServer(options: ServeOptions = {}) {
       name,
       description: descriptions[name] ?? tool.description,
       inputSchema: tool.jsonSchema as McpTool["inputSchema"],
+      ...(STRUCTURED_TOOLS.has(name) && { outputSchema: STRUCTURED_OUTPUT_SCHEMA }),
     })), ...(workspace?.tools ?? [])],
   }))
   // Ask the user once per session when the model weighed two readings and the client can ask; any
@@ -146,19 +222,19 @@ export async function createServer(options: ServeOptions = {}) {
       const clarified = request.params.name === "answer_commit" ? await clarifyReading(request.params.arguments, result, extra.signal) : undefined
       if (clarified) return { content: [{ type: "text" as const, text: clarified }], isError: true }
       const isError = resultFailed(result)
-      // A tool that builds a structured result (dbt_lineage) hands it over beside its text.
-      const structured = result.metadata?.structuredContent
       return {
         content: [{ type: "text" as const, text: result.output }, ...resultLinks(result.metadata)],
-        ...(structured && typeof structured === "object" && !isError
-          ? { structuredContent: structured as Record<string, unknown> }
-          : {}),
         isError,
+        ...(STRUCTURED_TOOLS.has(request.params.name) && {
+          structuredContent: structuredResult(request.params.name, result.metadata, result.output),
+        }),
       }
     } catch (error) {
+      const message = error instanceof Error ? error.message : String(error)
       return {
-        content: [{ type: "text" as const, text: error instanceof Error ? error.message : String(error) }],
+        content: [{ type: "text" as const, text: message }],
         isError: true,
+        ...(STRUCTURED_TOOLS.has(request.params.name) && { structuredContent: structuredError(message) }),
       }
     }
   }