From 54bc20cb7f10762fdc55b8d5fe0a3829a9f41fdd Mon Sep 17 00:00:00 2001 From: bendtherules Date: Mon, 27 Apr 2026 18:31:25 +0530 Subject: [PATCH] feat(langfuse): add tool-level spans with input/output to traces Wrap each MCP tool handler in withSpan()/withSpanContext() to create individual named spans per tool call: - ask262_search_spec_sections - ask262_get_section_content - ask262_evaluate_in_engine262 Each span includes: - langfuse.observation.input: JSON of tool arguments (query, sectionIds, code) - tool name and relevant metadata as span attributes This makes the Langfuse trace hierarchy show which MCP tool was called and what input it received, nested under the root HTTP/stdio request span. --- src/mcp-server-http.ts | 84 ++++++++++++++++++++++++++++++----------- src/mcp-server-stdio.ts | 42 +++++++++++++++++---- 2 files changed, 98 insertions(+), 28 deletions(-) diff --git a/src/mcp-server-http.ts b/src/mcp-server-http.ts index 419edab..6b525e2 100644 --- a/src/mcp-server-http.ts +++ b/src/mcp-server-http.ts @@ -36,7 +36,7 @@ import { import { DEFAULT_PORT, STORAGE_DIR as STORAGE_DIR_REL } from "./constants.js"; import { createEmbeddings } from "./lib/embeddings-factory.js"; import { LogOperation, logger } from "./lib/logger.js"; -import { withSpan } from "./lib/tracing.js"; +import { getSessionMetadata, setupTracing, withSpan } from "./lib/tracing.js"; // Resolve storage path relative to this script's directory const __filename = fileURLToPath(import.meta.url); @@ -90,12 +90,22 @@ async function createMcpServer() { }, }, async ({ query }) => { - const result = await searchSpecTool({ query }); - return { - content: [{ type: "text", text: JSON.stringify(result, null, 2) }], - structuredContent: result, - isError: false, - }; + return await withSpan( + "ask262_search_spec_sections", + { + "langfuse.observation.input": JSON.stringify({ query }), + tool: searchSpecToolName, + query, + }, + async () => { + const result = await searchSpecTool({ query }); + return { + content: [{ type: "text", text: JSON.stringify(result, null, 2) }], + structuredContent: result, + isError: false, + }; + }, + ); }, ); @@ -113,12 +123,25 @@ async function createMcpServer() { }, }, async ({ sectionIds, recursive }) => { - const result = await getSectionContentTool({ sectionIds, recursive }); - return { - content: [{ type: "text", text: JSON.stringify(result, null, 2) }], - structuredContent: result, - isError: false, - }; + return await withSpan( + "ask262_get_section_content", + { + "langfuse.observation.input": JSON.stringify({ + sectionIds, + recursive, + }), + tool: sectionContentToolName, + section_count: sectionIds.length, + }, + async () => { + const result = await getSectionContentTool({ sectionIds, recursive }); + return { + content: [{ type: "text", text: JSON.stringify(result, null, 2) }], + structuredContent: result, + isError: false, + }; + }, + ); }, ); @@ -136,14 +159,28 @@ async function createMcpServer() { }, }, async ({ code }) => { - const result = await evaluateTool({ code }); - const isError = result.error !== undefined; - const text = isError ? result.error : JSON.stringify(result, null, 2); - return { - content: [{ type: "text", text }], - structuredContent: result, - isError, - }; + return await withSpan( + "ask262_evaluate_in_engine262", + { + "langfuse.observation.input": JSON.stringify({ + code: code.slice(0, 200), + }), + tool: evaluateToolName, + code_length: code.length, + }, + async () => { + const result = await evaluateTool({ code }); + const isError = result.error !== undefined; + const text = isError + ? result.error + : JSON.stringify(result, null, 2); + return { + content: [{ type: "text", text }], + structuredContent: result, + isError, + }; + }, + ); }, ); @@ -156,6 +193,9 @@ async function createMcpServer() { } export async function main() { + // Initialize tracing (registers Langfuse processor when enabled) + setupTracing(); + // Initialize HTTP server logger const log = await logger.forComponent("http-server"); @@ -217,6 +257,7 @@ export async function main() { } // Handle request within trace context (passing trace ID from header if available) + const sessionMetadata = getSessionMetadata("http"); return await withSpan( LogOperation.HANDLING_MCP_HTTP_REQUEST, { method: c.req.method, client_ip: clientIp }, @@ -257,6 +298,7 @@ export async function main() { } }, traceId, + sessionMetadata, ); }); diff --git a/src/mcp-server-stdio.ts b/src/mcp-server-stdio.ts index a9a1599..94ce499 100644 --- a/src/mcp-server-stdio.ts +++ b/src/mcp-server-stdio.ts @@ -37,7 +37,12 @@ import { import { STORAGE_DIR as STORAGE_DIR_REL } from "./constants.js"; import { createEmbeddings } from "./lib/embeddings-factory.js"; import { LogOperation, logger } from "./lib/logger.js"; -import { createProcessScopedTrace, withSpanContext } from "./lib/tracing.js"; +import { + createProcessScopedTrace, + getSessionMetadata, + setupTracing, + withSpanContext, +} from "./lib/tracing.js"; // Resolve storage path relative to this script's directory const __filename = fileURLToPath(import.meta.url); @@ -90,6 +95,9 @@ export interface SearchSpecMCPOutput extends McpToolOutputBase { const embeddings = createEmbeddings(); export async function main() { + // Initialize tracing (registers Langfuse processor when enabled) + setupTracing(); + // Initialize stdio server logger const log = await logger.forComponent("stdio-server"); @@ -139,8 +147,12 @@ export async function main() { async ({ query }: SearchSpecMCPInput): Promise => { return await withSpanContext( sessionTraceId, - "vector_search", - { tool: searchSpecToolName, query }, + "ask262_search_spec_sections", + { + "langfuse.observation.input": JSON.stringify({ query }), + tool: searchSpecToolName, + query, + }, async () => { const result = await searchSpecTool({ query }); return { @@ -149,6 +161,7 @@ export async function main() { isError: false, }; }, + getSessionMetadata("stdio"), ); }, ); @@ -172,8 +185,15 @@ export async function main() { }: GetSectionContentMCPInput): Promise => { return await withSpanContext( sessionTraceId, - "section_fetch", - { tool: sectionContentToolName, section_count: sectionIds.length }, + "ask262_get_section_content", + { + "langfuse.observation.input": JSON.stringify({ + sectionIds, + recursive, + }), + tool: sectionContentToolName, + section_count: sectionIds.length, + }, async () => { const result = await getSectionContentTool({ sectionIds, recursive }); return { @@ -182,6 +202,7 @@ export async function main() { isError: false, }; }, + getSessionMetadata("stdio"), ); }, ); @@ -202,8 +223,14 @@ export async function main() { async ({ code }: EvaluateToolMCPInput): Promise => { return await withSpanContext( sessionTraceId, - "code_execution", - { tool: evaluateToolName, code_length: code.length }, + "ask262_evaluate_in_engine262", + { + "langfuse.observation.input": JSON.stringify({ + code: code.slice(0, 200), + }), + tool: evaluateToolName, + code_length: code.length, + }, async () => { const result = await evaluateTool({ code }); const isError = result.error !== undefined; @@ -214,6 +241,7 @@ export async function main() { isError, }; }, + getSessionMetadata("stdio"), ); }, );