@townco/agent 0.1.49 → 0.1.51

package/dist/runner/langchain/index.js

@@ -1,10 +1,13 @@
 import { MultiServerMCPClient } from "@langchain/mcp-adapters";
+import { context, propagation, trace } from "@opentelemetry/api";
 import { createLogger } from "@townco/core";
 import { AIMessageChunk, createAgent, ToolMessage, tool, } from "langchain";
 import { z } from "zod";
 import { SUBAGENT_MODE_KEY } from "../../acp-server/adapter";
+import { telemetry } from "../../telemetry/index.js";
 import { loadCustomToolModule, } from "../tool-loader.js";
-import { createModelFromString } from "./model-factory.js";
+import { createModelFromString, detectProvider } from "./model-factory.js";
+import { makeOtelCallbacks } from "./otel-callbacks.js";
 import { makeFilesystemTools } from "./tools/filesystem";
 import { TASK_TOOL_NAME } from "./tools/subagent";
 import { TODO_WRITE_TOOL_NAME, todoWrite } from "./tools/todo";
@@ -51,369 +54,568 @@ async function loadCustomTools(modulePaths) {
 }
 export class LangchainAgent {
     definition;
+    toolSpans = new Map();
     constructor(params) {
         this.definition = params;
     }
     async *invoke(req) {
-        // Track todo_write tool call IDs to suppress their tool_call notifications
-        const todoWriteToolCallIds = new Set();
-        // --------------------------------------------------------------------------
-        // Resolve tools: built-ins (string) + custom ({ type: "custom", modulePath })
-        // + filesystem ({ type: "filesystem", working_directory? })
-        // --------------------------------------------------------------------------
-        const enabledTools = [];
-        const toolDefs = this.definition.tools ?? [];
-        const builtInNames = [];
-        const customToolPaths = [];
-        for (const t of toolDefs) {
-            if (typeof t === "string") {
-                builtInNames.push(t);
-            }
-            else if (t && typeof t === "object" && "type" in t) {
-                const type = t.type;
-                if (type === "custom" &&
-                    "modulePath" in t &&
-                    typeof t.modulePath === "string") {
-                    customToolPaths.push(t.modulePath);
-                }
-                else if (type === "filesystem") {
-                    const wd = t.working_directory ??
-                        process.cwd();
-                    enabledTools.push(...makeFilesystemTools(wd));
+        // Derive the parent OTEL context for this invocation.
+        // If this is a subagent and the parent process propagated an OTEL trace
+        // context via sessionMeta.otelTraceContext, use that as the parent;
+        // otherwise, fall back to the current active context.
+        let parentContext = context.active();
+        const meta = req.sessionMeta;
+        if (meta?.otelTraceContext) {
+            parentContext = propagation.extract(parentContext, meta.otelTraceContext);
+        }
+        // Track turn-level token usage from API responses
+        const turnTokenUsage = {
+            inputTokens: 0,
+            outputTokens: 0,
+            totalTokens: 0,
+        };
+        const countedMessageIds = new Set();
+        // Start telemetry span for entire invocation
+        const invocationSpan = telemetry.startSpan("agent.invoke", {
+            "agent.model": this.definition.model,
+            "agent.subagent": meta?.[SUBAGENT_MODE_KEY] === true,
+            "agent.session_id": req.sessionId,
+            "agent.message_id": req.messageId,
+        }, parentContext);
+        // Create a context with the invocation span as active
+        // This will be used when creating child spans (tool calls)
+        const invocationContext = invocationSpan
+            ? trace.setSpan(parentContext, invocationSpan)
+            : parentContext;
+        telemetry.log("info", "Agent invocation started", {
+            model: this.definition.model,
+            sessionId: req.sessionId,
+            messageId: req.messageId,
+        });
+        try {
+            // Track todo_write tool call IDs to suppress their tool_call notifications
+            const todoWriteToolCallIds = new Set();
+            // --------------------------------------------------------------------------
+            // Resolve tools: built-ins (string) + custom ({ type: "custom", modulePath })
+            // + filesystem ({ type: "filesystem", working_directory? })
+            // --------------------------------------------------------------------------
+            const enabledTools = [];
+            const toolDefs = this.definition.tools ?? [];
+            const builtInNames = [];
+            const customToolPaths = [];
+            for (const t of toolDefs) {
+                if (typeof t === "string") {
+                    builtInNames.push(t);
                 }
-                else if (type === "direct") {
-                    // Handle direct tool objects (imported in code)
-                    // biome-ignore lint/suspicious/noExplicitAny: mlai unsure how to best type this
-                    const addedTool = tool(t.fn, {
-                        name: t.name,
-                        description: t.description,
-                        schema: t.schema,
-                    });
-                    addedTool.prettyName = t.prettyName;
-                    addedTool.icon = t.icon;
-                    enabledTools.push(addedTool);
+                else if (t && typeof t === "object" && "type" in t) {
+                    const type = t.type;
+                    if (type === "custom" &&
+                        "modulePath" in t &&
+                        typeof t.modulePath === "string") {
+                        customToolPaths.push(t.modulePath);
+                    }
+                    else if (type === "filesystem") {
+                        const wd = t.working_directory ??
+                            process.cwd();
+                        enabledTools.push(...makeFilesystemTools(wd));
+                    }
+                    else if (type === "direct") {
+                        // Handle direct tool objects (imported in code)
+                        // biome-ignore lint/suspicious/noExplicitAny: mlai unsure how to best type this
+                        const addedTool = tool(t.fn, {
+                            name: t.name,
+                            description: t.description,
+                            schema: t.schema,
+                        });
+                        addedTool.prettyName = t.prettyName;
+                        addedTool.icon = t.icon;
+                        enabledTools.push(addedTool);
+                    }
                 }
             }
-        }
-        // Built-in tools from registry
-        for (const name of builtInNames) {
-            const entry = TOOL_REGISTRY[name];
-            if (!entry) {
-                throw new Error(`Unknown built-in tool "${name}"`);
-            }
-            if (typeof entry === "function") {
-                const result = entry();
-                if (Array.isArray(result)) {
-                    enabledTools.push(...result);
+            // Built-in tools from registry
+            for (const name of builtInNames) {
+                const entry = TOOL_REGISTRY[name];
+                if (!entry) {
+                    throw new Error(`Unknown built-in tool "${name}"`);
+                }
+                if (typeof entry === "function") {
+                    const result = entry();
+                    if (Array.isArray(result)) {
+                        enabledTools.push(...result);
+                    }
+                    else {
+                        enabledTools.push(result);
+                    }
                 }
                 else {
-                    enabledTools.push(result);
+                    enabledTools.push(entry);
                 }
             }
+            // Custom tools loaded from modulePaths
+            if (customToolPaths.length > 0) {
+                const customTools = await loadCustomTools(customToolPaths);
+                enabledTools.push(...customTools);
+            }
+            // MCP tools
+            if ((this.definition.mcps?.length ?? 0) > 0) {
+                enabledTools.push(...(await makeMcpToolsClient(this.definition.mcps).getTools()));
+            }
+            // Wrap tools with response compaction if hook is configured
+            const hooks = this.definition.hooks ?? [];
+            const hasToolResponseHook = hooks.some((h) => h.type === "tool_response");
+            const noSession = req.sessionMeta?.[SUBAGENT_MODE_KEY] === true; // Subagents don't have session storage
+            // Track cumulative tool output tokens in this turn for proper context calculation
+            let cumulativeToolOutputTokens = 0;
+            let wrappedTools = enabledTools;
+            if (hasToolResponseHook && !noSession) {
+                const { countToolResultTokens } = await import("../../utils/token-counter.js");
+                const { toolResponseCompactor } = await import("../hooks/predefined/tool-response-compactor.js");
+                wrappedTools = enabledTools.map((originalTool) => {
+                    const wrappedFunc = async (input) => {
+                        // Execute the original tool
+                        const result = await originalTool.invoke(input);
+                        // Check if result should be compacted
+                        const resultStr = typeof result === "string" ? result : JSON.stringify(result);
+                        const rawOutput = { content: resultStr };
+                        const outputTokens = countToolResultTokens(rawOutput);
+                        // Skip compaction for small results (under 10k tokens)
+                        if (outputTokens < 10000) {
+                            // Still track this in cumulative total
+                            cumulativeToolOutputTokens += outputTokens;
+                            return result;
+                        }
+                        _logger.info("Tool wrapper: compacting large tool result", {
+                            toolName: originalTool.name,
+                            originalTokens: outputTokens,
+                            cumulativeToolOutputTokens,
+                        });
+                        // Calculate current context including all tool outputs so far in this turn
+                        // This ensures we account for multiple large tool calls in the same turn
+                        const baseContextTokens = turnTokenUsage.inputTokens || 10000;
+                        const currentTokens = baseContextTokens + cumulativeToolOutputTokens;
+                        const maxTokens = 200000; // Claude's limit
+                        // Build proper hook context with all required fields
+                        const hookContext = {
+                            session: {
+                                messages: req.contextMessages || [],
+                                context: [],
+                                requestParams: {
+                                    hookSettings: hooks.find((h) => h.type === "tool_response")
+                                        ?.setting,
+                                },
+                            },
+                            currentTokens,
+                            maxTokens,
+                            percentage: (currentTokens / maxTokens) * 100,
+                            model: this.definition.model,
+                            toolResponse: {
+                                toolCallId: "pending",
+                                toolName: originalTool.name,
+                                toolInput: input,
+                                rawOutput,
+                                outputTokens,
+                            },
+                        };
+                        // Call the tool response compactor directly
+                        const hookResult = await toolResponseCompactor(hookContext);
+                        // Extract modified output from metadata
+                        if (hookResult?.metadata?.modifiedOutput) {
+                            const modifiedOutput = hookResult.metadata
+                                .modifiedOutput;
+                            const compactedTokens = countToolResultTokens(modifiedOutput);
+                            // Update cumulative total with the compacted size (not original!)
+                            cumulativeToolOutputTokens += compactedTokens;
+                            _logger.info("Tool wrapper: compaction complete", {
+                                toolName: originalTool.name,
+                                originalTokens: outputTokens,
+                                compactedTokens,
+                                reduction: `${((1 - compactedTokens / outputTokens) * 100).toFixed(1)}%`,
+                                totalCumulativeTokens: cumulativeToolOutputTokens,
+                            });
+                            return typeof result === "string"
+                                ? modifiedOutput.content
+                                : JSON.stringify(modifiedOutput);
+                        }
+                        // No compaction happened, count original size
+                        cumulativeToolOutputTokens += outputTokens;
+                        return result;
+                    };
+                    // Create new tool with wrapped function
+                    const wrappedTool = tool(wrappedFunc, {
+                        name: originalTool.name,
+                        description: originalTool.description,
+                        schema: originalTool.schema,
+                    });
+                    // Preserve metadata
+                    wrappedTool.prettyName = originalTool.prettyName;
+                    wrappedTool.icon = originalTool.icon;
+                    return wrappedTool;
+                });
+            }
+            // Filter tools if running in subagent mode
+            const isSubagent = req.sessionMeta?.[SUBAGENT_MODE_KEY] === true;
+            const finalTools = isSubagent
+                ? wrappedTools.filter((t) => t.name !== TODO_WRITE_TOOL_NAME && t.name !== TASK_TOOL_NAME)
+                : wrappedTools;
+            // Create the model instance using the factory
+            // This detects the provider from the model string:
+            // - "gemini-2.0-flash" → Google Generative AI
+            // - "vertex-gemini-2.0-flash" → Vertex AI (strips prefix)
+            // - "claude-sonnet-4-5-20250929" → Anthropic
+            const model = createModelFromString(this.definition.model);
+            const agentConfig = {
+                model,
+                tools: finalTools,
+            };
+            if (this.definition.systemPrompt) {
+                agentConfig.systemPrompt = this.definition.systemPrompt;
+            }
+            // Inject system prompt with optional TodoWrite instructions
+            const hasTodoWrite = builtInNames.includes("todo_write");
+            if (hasTodoWrite) {
+                agentConfig.systemPrompt = `${agentConfig.systemPrompt ?? ""}\n\n${TODO_WRITE_INSTRUCTIONS}`;
+            }
+            const agent = createAgent(agentConfig);
+            // Add logging callbacks for model requests
+            const provider = detectProvider(this.definition.model);
+            // Build messages from context history if available, otherwise use just the prompt
+            let messages;
+            if (req.contextMessages && req.contextMessages.length > 0) {
+                // Use context messages (already resolved from context entries)
+                // Convert to LangChain format
+                messages = req.contextMessages.map((msg) => ({
+                    type: msg.role === "user" ? "human" : "ai",
+                    // Extract text from content blocks
+                    content: msg.content
+                        .filter((block) => block.type === "text")
+                        .map((block) => block.text)
+                        .join(""),
+                }));
+                // Add the current prompt as the final human message
+                const currentPromptText = req.prompt
+                    .filter((promptMsg) => promptMsg.type === "text")
+                    .map((promptMsg) => promptMsg.text)
+                    .join("\n");
+                messages.push({
+                    type: "human",
+                    content: currentPromptText,
+                });
+            }
             else {
-                enabledTools.push(entry);
+                // Fallback: No context history, use just the prompt
+                messages = req.prompt
+                    .filter((promptMsg) => promptMsg.type === "text")
+                    .map((promptMsg) => ({
+                    type: "human",
+                    content: promptMsg.text,
+                }));
             }
-        }
-        // Custom tools loaded from modulePaths
-        if (customToolPaths.length > 0) {
-            const customTools = await loadCustomTools(customToolPaths);
-            enabledTools.push(...customTools);
-        }
-        // MCP tools
-        if ((this.definition.mcps?.length ?? 0) > 0) {
-            enabledTools.push(...(await makeMcpToolsClient(this.definition.mcps).getTools()));
-        }
-        // Filter tools if running in subagent mode
-        const isSubagent = req.sessionMeta?.[SUBAGENT_MODE_KEY] === true;
-        const finalTools = isSubagent
-            ? enabledTools.filter((t) => t.name !== TODO_WRITE_TOOL_NAME && t.name !== TASK_TOOL_NAME)
-            : enabledTools;
-        // Create the model instance using the factory
-        // This detects the provider from the model string:
-        // - "gemini-2.0-flash" → Google Generative AI
-        // - "vertex-gemini-2.0-flash" → Vertex AI (strips prefix)
-        // - "claude-sonnet-4-5-20250929" → Anthropic
-        const model = createModelFromString(this.definition.model);
-        const agentConfig = {
-            model,
-            tools: finalTools,
-        };
-        if (this.definition.systemPrompt) {
-            agentConfig.systemPrompt = this.definition.systemPrompt;
-        }
-        // Inject system prompt with optional TodoWrite instructions
-        const hasTodoWrite = builtInNames.includes("todo_write");
-        if (hasTodoWrite) {
-            agentConfig.systemPrompt = `${agentConfig.systemPrompt ?? ""}\n\n${TODO_WRITE_INSTRUCTIONS}`;
-        }
-        const agent = createAgent(agentConfig);
-        // Build messages from context history if available, otherwise use just the prompt
-        let messages;
-        if (req.contextMessages && req.contextMessages.length > 0) {
-            // Use context messages (already resolved from context entries)
-            // Convert to LangChain format
-            messages = req.contextMessages.map((msg) => ({
-                type: msg.role === "user" ? "human" : "ai",
-                // Extract text from content blocks
-                content: msg.content
-                    .filter((block) => block.type === "text")
-                    .map((block) => block.text)
-                    .join(""),
-            }));
-            // Add the current prompt as the final human message
-            const currentPromptText = req.prompt
-                .filter((promptMsg) => promptMsg.type === "text")
-                .map((promptMsg) => promptMsg.text)
-                .join("\n");
-            messages.push({
-                type: "human",
-                content: currentPromptText,
+            // Create OTEL callbacks for instrumentation
+            const otelCallbacks = makeOtelCallbacks({
+                provider,
+                model: this.definition.model,
+                parentContext: invocationContext,
             });
-        }
-        else {
-            // Fallback: No context history, use just the prompt
-            messages = req.prompt
-                .filter((promptMsg) => promptMsg.type === "text")
-                .map((promptMsg) => ({
-                type: "human",
-                content: promptMsg.text,
+            // Create the stream within the invocation context so AsyncLocalStorage
+            // propagates the context to all tool executions and callbacks
+            const stream = context.with(invocationContext, () => agent.stream({ messages }, {
+                streamMode: ["updates", "messages"],
+                recursionLimit: 200,
+                callbacks: [otelCallbacks],
             }));
-        }
-        const stream = agent.stream({ messages }, {
-            streamMode: ["updates", "messages"],
-            recursionLimit: 200,
-        });
-        for await (const [streamMode, chunk] of await stream) {
-            if (streamMode === "updates") {
-                const updatesChunk = modelRequestSchema.safeParse(chunk);
-                if (!updatesChunk.success) {
-                    // Other kinds of updates are either handled in the 'messages'
-                    // streamMode (tool calls), or we don't care about them so far (not
-                    // known yet).
-                    continue;
-                }
-                const updatesMessages = updatesChunk.data.model_request.messages;
-                if (!updatesMessages.every((m) => m instanceof AIMessageChunk)) {
-                    throw new Error(`Unhandled updates message chunk types: ${JSON.stringify(updatesMessages)}`);
-                }
-                for (const msg of updatesMessages) {
-                    // Extract token usage metadata if available
-                    const tokenUsage = msg.usage_metadata
-                        ? {
-                            inputTokens: msg.usage_metadata.input_tokens,
-                            outputTokens: msg.usage_metadata.output_tokens,
-                            totalTokens: msg.usage_metadata.total_tokens,
+            for await (const streamItem of await stream) {
+                const [streamMode, chunk] = streamItem;
+                if (streamMode === "updates") {
+                    const updatesChunk = modelRequestSchema.safeParse(chunk);
+                    if (!updatesChunk.success) {
+                        // Other kinds of updates are either handled in the 'messages'
+                        // streamMode (tool calls), or we don't care about them so far (not
+                        // known yet).
+                        continue;
+                    }
+                    const updatesMessages = updatesChunk.data.model_request.messages;
+                    if (!updatesMessages.every((m) => m instanceof AIMessageChunk)) {
+                        throw new Error(`Unhandled updates message chunk types: ${JSON.stringify(updatesMessages)}`);
+                    }
+                    for (const msg of updatesMessages) {
+                        // Extract token usage metadata if available
+                        const tokenUsage = msg.usage_metadata
+                            ? {
+                                inputTokens: msg.usage_metadata.input_tokens,
+                                outputTokens: msg.usage_metadata.output_tokens,
+                                totalTokens: msg.usage_metadata.total_tokens,
+                            }
+                            : undefined;
+                        // Record token usage in telemetry
+                        if (tokenUsage) {
+                            telemetry.recordTokenUsage(tokenUsage.inputTokens ?? 0, tokenUsage.outputTokens ?? 0, invocationSpan);
                         }
-                        : undefined;
-                    for (const toolCall of msg.tool_calls ?? []) {
-                        if (toolCall.id == null) {
-                            throw new Error(`Tool call is missing id: ${JSON.stringify(toolCall)}`);
+                        // Accumulate token usage (deduplicate by message ID)
+                        if (tokenUsage && msg.id && !countedMessageIds.has(msg.id)) {
+                            turnTokenUsage.inputTokens += tokenUsage.inputTokens ?? 0;
+                            turnTokenUsage.outputTokens += tokenUsage.outputTokens ?? 0;
+                            turnTokenUsage.totalTokens += tokenUsage.totalTokens ?? 0;
+                            countedMessageIds.add(msg.id);
+                        }
+                        for (const toolCall of msg.tool_calls ?? []) {
+                            if (toolCall.id == null) {
+                                throw new Error(`Tool call is missing id: ${JSON.stringify(toolCall)}`);
+                            }
+                            // Create tool span within the invocation context
+                            // This makes the tool span a child of the invocation span
+                            const toolSpan = context.with(invocationContext, () => telemetry.startSpan("agent.tool_call", {
+                                "tool.name": toolCall.name,
+                                "tool.id": toolCall.id,
+                            }));
+                            this.toolSpans.set(toolCall.id, toolSpan);
+                            telemetry.log("info", `Tool call started: ${toolCall.name}`, {
+                                toolCallId: toolCall.id,
+                                toolName: toolCall.name,
+                                toolArgs: JSON.stringify(toolCall.args),
+                            });
+                            // TODO: re-add this suppression of the todo_write tool call when we
+                            // are rendering the agent-plan update in the UIs
+                            // If this is a todo_write tool call, yield an agent-plan update
+                            //if (toolCall.name === "todo_write" && toolCall.args?.todos) {
+                            //  const entries = toolCall.args.todos.flatMap((todo: unknown) => {
+                            //    const validation = todoItemSchema.safeParse(todo);
+                            //    if (!validation.success) {
+                            //      // Invalid todo - filter it out
+                            //      return [];
+                            //    }
+                            //    return [
+                            //      {
+                            //        content: validation.data.content,
+                            //        status: validation.data.status,
+                            //        priority: "medium" as const,
+                            //      },
+                            //    ];
+                            //  });
+                            //  yield {
+                            //    sessionUpdate: "plan",
+                            //    entries: entries,
+                            //  };
+                            //  // Track this tool call ID to suppress tool_call notifications
+                            //  todoWriteToolCallIds.add(toolCall.id);
+                            //  continue;
+                            //}
+                            const matchingTool = finalTools.find((t) => t.name === toolCall.name);
+                            const prettyName = matchingTool?.prettyName;
+                            const icon = matchingTool?.icon;
+                            yield {
+                                sessionUpdate: "tool_call",
+                                toolCallId: toolCall.id,
+                                title: toolCall.name,
+                                kind: "other",
+                                status: "pending",
+                                rawInput: toolCall.args,
+                                ...(tokenUsage ? { tokenUsage } : {}),
+                                _meta: {
+                                    messageId: req.messageId,
+                                    ...(prettyName ? { prettyName } : {}),
+                                    ...(icon ? { icon } : {}),
+                                },
+                            };
+                            yield {
+                                sessionUpdate: "tool_call_update",
+                                toolCallId: toolCall.id,
+                                status: "in_progress",
+                                ...(tokenUsage ? { tokenUsage } : {}),
+                                _meta: { messageId: req.messageId },
+                            };
                         }
-                        // TODO: re-add this suppression of the todo_write tool call when we
-                        // are rendering the agent-plan update in the UIs
-                        // If this is a todo_write tool call, yield an agent-plan update
-                        //if (toolCall.name === "todo_write" && toolCall.args?.todos) {
-                        //  const entries = toolCall.args.todos.flatMap((todo: unknown) => {
-                        //    const validation = todoItemSchema.safeParse(todo);
-                        //    if (!validation.success) {
-                        //      // Invalid todo - filter it out
-                        //      return [];
-                        //    }
-                        //    return [
-                        //      {
-                        //        content: validation.data.content,
-                        //        status: validation.data.status,
-                        //        priority: "medium" as const,
-                        //      },
-                        //    ];
-                        //  });
-                        //  yield {
-                        //    sessionUpdate: "plan",
-                        //    entries: entries,
-                        //  };
-                        //  // Track this tool call ID to suppress tool_call notifications
-                        //  todoWriteToolCallIds.add(toolCall.id);
-                        //  continue;
-                        //}
-                        const matchingTool = finalTools.find((t) => t.name === toolCall.name);
-                        const prettyName = matchingTool?.prettyName;
-                        const icon = matchingTool?.icon;
-                        yield {
-                            sessionUpdate: "tool_call",
-                            toolCallId: toolCall.id,
-                            title: toolCall.name,
-                            kind: "other",
-                            status: "pending",
-                            rawInput: toolCall.args,
-                            ...(tokenUsage ? { tokenUsage } : {}),
-                            _meta: {
-                                messageId: req.messageId,
-                                ...(prettyName ? { prettyName } : {}),
-                                ...(icon ? { icon } : {}),
-                            },
-                        };
-                        yield {
-                            sessionUpdate: "tool_call_update",
-                            toolCallId: toolCall.id,
-                            status: "in_progress",
-                            ...(tokenUsage ? { tokenUsage } : {}),
-                            _meta: { messageId: req.messageId },
-                        };
                     }
                 }
-            }
-            else if (streamMode === "messages") {
-                const aiMessage = chunk[0];
-                if (aiMessage instanceof AIMessageChunk) {
-                    // Extract token usage metadata if available
-                    const messageTokenUsage = aiMessage.usage_metadata
-                        ? {
-                            inputTokens: aiMessage.usage_metadata.input_tokens,
-                            outputTokens: aiMessage.usage_metadata.output_tokens,
-                            totalTokens: aiMessage.usage_metadata.total_tokens,
+                else if (streamMode === "messages") {
+                    const aiMessage = chunk[0];
+                    if (aiMessage instanceof AIMessageChunk) {
+                        // Extract token usage metadata if available
+                        const messageTokenUsage = aiMessage.usage_metadata
+                            ? {
+                                inputTokens: aiMessage.usage_metadata.input_tokens,
+                                outputTokens: aiMessage.usage_metadata.output_tokens,
+                                totalTokens: aiMessage.usage_metadata.total_tokens,
+                            }
+                            : undefined;
+                        // Accumulate token usage (deduplicate by message ID)
+                        if (messageTokenUsage &&
+                            aiMessage.id &&
+                            !countedMessageIds.has(aiMessage.id)) {
+                            turnTokenUsage.inputTokens += messageTokenUsage.inputTokens ?? 0;
+                            turnTokenUsage.outputTokens +=
+                                messageTokenUsage.outputTokens ?? 0;
+                            turnTokenUsage.totalTokens += messageTokenUsage.totalTokens ?? 0;
+                            countedMessageIds.add(aiMessage.id);
                         }
-                        : undefined;
-                    if (messageTokenUsage) {
-                        const contentType = typeof aiMessage.content;
-                        const contentIsArray = Array.isArray(aiMessage.content);
-                        const contentLength = contentIsArray
-                            ? aiMessage.content.length
-                            : typeof aiMessage.content === "string"
+                        if (messageTokenUsage) {
+                            const contentType = typeof aiMessage.content;
+                            const contentIsArray = Array.isArray(aiMessage.content);
+                            const contentLength = contentIsArray
                                 ? aiMessage.content.length
-                                : -1;
-                        _logger.debug("messageTokenUsage", {
-                            messageTokenUsage,
-                            contentType,
-                            isArray: contentIsArray,
-                            length: contentLength,
-                        });
-                    }
-                    // If we have tokenUsage but no content, send a token-only chunk
-                    if (messageTokenUsage &&
-                        (typeof aiMessage.content === "string"
-                            ? aiMessage.content === ""
-                            : Array.isArray(aiMessage.content) &&
-                                aiMessage.content.length === 0)) {
-                        _logger.debug("sending token-only chunk", {
-                            messageTokenUsage,
-                        });
-                        const msgToYield = {
-                            sessionUpdate: "agent_message_chunk",
-                            content: {
-                                type: "text",
-                                text: "", // Empty text, just carrying tokenUsage
-                            },
-                            _meta: {
-                                tokenUsage: messageTokenUsage,
-                            },
-                        };
-                        yield msgToYield;
-                        continue; // Skip the rest of the processing for this chunk
-                    }
-                    if (typeof aiMessage.content === "string") {
-                        const msgToYield = messageTokenUsage
-                            ? {
+                                : typeof aiMessage.content === "string"
+                                    ? aiMessage.content.length
+                                    : -1;
+                            _logger.debug("messageTokenUsage", {
+                                messageTokenUsage,
+                                contentType,
+                                isArray: contentIsArray,
+                                length: contentLength,
+                            });
+                        }
+                        // If we have tokenUsage but no content, send a token-only chunk
+                        if (messageTokenUsage &&
+                            (typeof aiMessage.content === "string"
+                                ? aiMessage.content === ""
+                                : Array.isArray(aiMessage.content) &&
+                                    aiMessage.content.length === 0)) {
+                            _logger.debug("sending token-only chunk", {
+                                messageTokenUsage,
+                            });
+                            const msgToYield = {
                                 sessionUpdate: "agent_message_chunk",
                                 content: {
                                     type: "text",
-                                    text: aiMessage.content,
+                                    text: "", // Empty text, just carrying tokenUsage
                                 },
                                 _meta: {
                                     tokenUsage: messageTokenUsage,
                                 },
-                            }
-                            : {
-                                sessionUpdate: "agent_message_chunk",
-                                content: {
-                                    type: "text",
-                                    text: aiMessage.content,
-                                },
                             };
-                        yield msgToYield;
+                            yield msgToYield;
+                            continue; // Skip the rest of the processing for this chunk
+                        }
+                        if (typeof aiMessage.content === "string") {
+                            const msgToYield = messageTokenUsage
+                                ? {
+                                    sessionUpdate: "agent_message_chunk",
+                                    content: {
+                                        type: "text",
+                                        text: aiMessage.content,
+                                    },
+                                    _meta: {
+                                        tokenUsage: messageTokenUsage,
+                                    },
+                                }
+                                : {
+                                    sessionUpdate: "agent_message_chunk",
+                                    content: {
+                                        type: "text",
+                                        text: aiMessage.content,
+                                    },
+                                };
+                            yield msgToYield;
+                        }
+                        else if (Array.isArray(aiMessage.content)) {
+                            for (const part of aiMessage.content) {
+                                if (part.type === "text" && typeof part.text === "string") {
+                                    const msgToYield = messageTokenUsage
+                                        ? {
+                                            sessionUpdate: "agent_message_chunk",
+                                            content: {
+                                                type: "text",
+                                                text: part.text,
+                                            },
+                                            _meta: {
+                                                tokenUsage: messageTokenUsage,
+                                            },
+                                        }
+                                        : {
+                                            sessionUpdate: "agent_message_chunk",
+                                            content: {
+                                                type: "text",
+                                                text: part.text,
+                                            },
+                                        };
+                                    yield msgToYield;
+                                }
+                                else if (part.type === "tool_use") {
+                                    // We don't care about tool use chunks -- do nothing
+                                }
+                                else if (part.type === "input_json_delta") {
+                                    // We don't care about tool use input delta chunks -- do nothing
+                                }
+                                else {
+                                    throw new Error(`Unhandled AIMessageChunk content block type: ${part.type}\n${JSON.stringify(part)}`);
+                                }
+                            }
+                        }
+                        else {
+                            throw new Error(`Unhandled AIMessageChunk content type: ${typeof aiMessage.content}`);
+                        }
                     }
-                    else if (Array.isArray(aiMessage.content)) {
-                        for (const part of aiMessage.content) {
-                            if (part.type === "text" && typeof part.text === "string") {
-                                const msgToYield = messageTokenUsage
-                                    ? {
-                                        sessionUpdate: "agent_message_chunk",
-                                        content: {
-                                            type: "text",
-                                            text: part.text,
-                                        },
-                                        _meta: {
-                                            tokenUsage: messageTokenUsage,
-                                        },
-                                    }
-                                    : {
-                                        sessionUpdate: "agent_message_chunk",
+                    else if (aiMessage instanceof ToolMessage) {
+                        if (typeof aiMessage.content === "string") {
+                            if (todoWriteToolCallIds.has(aiMessage.tool_call_id)) {
+                                // Skip tool_call_update for todo_write tools
+                                continue;
+                            }
+                            // End telemetry span for this tool call
+                            const toolSpan = this.toolSpans.get(aiMessage.tool_call_id);
+                            if (toolSpan) {
+                                telemetry.log("info", "Tool call completed", {
+                                    toolCallId: aiMessage.tool_call_id,
+                                });
+                                telemetry.endSpan(toolSpan);
+                                this.toolSpans.delete(aiMessage.tool_call_id);
+                            }
+                            // Send status update (metadata only, no content)
+                            yield {
+                                sessionUpdate: "tool_call_update",
+                                toolCallId: aiMessage.tool_call_id,
+                                status: "completed",
+                                _meta: { messageId: req.messageId },
+                            };
+                            // Send tool output separately (via direct SSE, bypassing PostgreSQL NOTIFY)
+                            yield {
+                                sessionUpdate: "tool_output",
+                                toolCallId: aiMessage.tool_call_id,
+                                content: [
+                                    {
+                                        type: "content",
                                         content: {
                                             type: "text",
-                                            text: part.text,
+                                            text: aiMessage.content,
                                         },
-                                    };
-                                yield msgToYield;
-                            }
-                            else if (part.type === "tool_use") {
-                                // We don't care about tool use chunks -- do nothing
-                            }
-                            else if (part.type === "input_json_delta") {
-                                // We don't care about tool use input delta chunks -- do nothing
-                            }
-                            else {
-                                throw new Error(`Unhandled AIMessageChunk content block type: ${part.type}\n${JSON.stringify(part)}`);
-                            }
+                                    },
+                                ],
+                                rawOutput: { content: aiMessage.content },
+                                _meta: { messageId: req.messageId },
+                            };
                         }
-                    }
-                    else {
-                        throw new Error(`Unhandled AIMessageChunk content type: ${typeof aiMessage.content}`);
-                    }
-                }
-                else if (aiMessage instanceof ToolMessage) {
-                    if (typeof aiMessage.content === "string") {
-                        if (todoWriteToolCallIds.has(aiMessage.tool_call_id)) {
-                            // Skip tool_call_update for todo_write tools
-                            continue;
+                        else {
+                            throw new Error(`Unhandled ToolMessage content type: ${typeof aiMessage.content}`);
                         }
-                        // Send status update (metadata only, no content)
-                        yield {
-                            sessionUpdate: "tool_call_update",
-                            toolCallId: aiMessage.tool_call_id,
-                            status: "completed",
-                            _meta: { messageId: req.messageId },
-                        };
-                        // Send tool output separately (via direct SSE, bypassing PostgreSQL NOTIFY)
-                        yield {
-                            sessionUpdate: "tool_output",
-                            toolCallId: aiMessage.tool_call_id,
-                            content: [
-                                {
-                                    type: "content",
-                                    content: {
-                                        type: "text",
-                                        text: aiMessage.content,
-                                    },
-                                },
-                            ],
-                            rawOutput: { content: aiMessage.content },
-                            _meta: { messageId: req.messageId },
-                        };
                     }
                     else {
-                        throw new Error(`Unhandled ToolMessage content type: ${typeof aiMessage.content}`);
+                        throw new Error(`Unhandled message chunk type: ${JSON.stringify(aiMessage)}`);
                     }
                 }
                 else {
-                    throw new Error(`Unhandled message chunk type: ${JSON.stringify(aiMessage)}`);
+                    throw new Error(`Unhandled stream mode: ${streamMode}`);
                 }
             }
-            else {
-                throw new Error(`Unhandled stream mode: ${streamMode}`);
-            }
+            // Log successful completion
+            telemetry.log("info", "Agent invocation completed", {
+                sessionId: req.sessionId,
+            });
+            telemetry.endSpan(invocationSpan);
+            return {
+                stopReason: "end_turn",
+                _meta: {
+                    tokenUsage: turnTokenUsage,
+                },
+            };
+        }
+        catch (error) {
+            // Log error and end span with error status
+            telemetry.log("error", "Agent invocation failed", {
+                error: error instanceof Error ? error.message : String(error),
+                sessionId: req.sessionId,
+            });
+            telemetry.endSpan(invocationSpan, error instanceof Error ? error : new Error(String(error)));
+            throw error;
         }
-        return {
-            stopReason: "end_turn",
-        };
     }
 }
 const modelRequestSchema = z.object({