@townco/agent 0.1.73 → 0.1.75

package/dist/acp-server/adapter.js

@@ -258,6 +258,9 @@ export class AgentAcpAdapter {
         const subagentsMetadata = this.getSubagentsMetadata();
         response._meta = {
             ...response._meta,
+            // Include model and system prompt for Town Hall comparison feature
+            model: this.agent.definition.model,
+            systemPrompt: this.agent.definition.systemPrompt ?? null,
             ...(this.agentDescription
                 ? { agentDescription: this.agentDescription }
                 : {}),
@@ -276,11 +279,14 @@ export class AgentAcpAdapter {
     }
     async newSession(params) {
         const sessionId = Math.random().toString(36).substring(2);
+        // Extract configOverrides from _meta if provided (Town Hall comparison feature)
+        const configOverrides = params._meta?.configOverrides;
         this.sessions.set(sessionId, {
             pendingPrompt: null,
             messages: [],
             context: [],
             requestParams: params,
+            configOverrides,
         });
         // Note: Initial message is sent by the HTTP transport when SSE connection is established
         // This ensures the message is delivered after the client is ready to receive it
@@ -475,6 +481,7 @@ export class AgentAcpAdapter {
     async prompt(params) {
         const promptSpan = telemetry.startSpan("adapter.prompt", {
             "session.id": params.sessionId,
+            "agent.session_id": params.sessionId,
         });
         const spanContext = promptSpan
             ? trace.setSpan(context.active(), promptSpan)
@@ -535,6 +542,11 @@ export class AgentAcpAdapter {
             .filter((p) => p.type === "text")
             .map((p) => p.text)
             .join("\n");
+        // Store user message in span attribute for debugger UI
+        const promptSpan = trace.getActiveSpan();
+        if (promptSpan) {
+            promptSpan.setAttribute("user.message", userMessageText);
+        }
         logger.info("User message received", {
             sessionId: params.sessionId,
             messagePreview: userMessageText.slice(0, 100),
@@ -654,11 +666,26 @@ export class AgentAcpAdapter {
             if (session.requestParams._meta) {
                 invokeParams.sessionMeta = session.requestParams._meta;
             }
+            // Pass config overrides if defined (Town Hall comparison feature)
+            if (session.configOverrides) {
+                invokeParams.configOverrides = session.configOverrides;
+            }
             const generator = this.agent.invoke(invokeParams);
+            // Track the invocation span for parenting hook spans
+            let invocationSpan = null;
             // Manually iterate to capture the return value
             let iterResult = await generator.next();
             while (!iterResult.done) {
                 const msg = iterResult.value;
+                // Capture the invocation span so we can use it for parenting hook spans
+                if ("sessionUpdate" in msg &&
+                    msg.sessionUpdate === "__invocation_span" &&
+                    "invocationSpan" in msg) {
+                    invocationSpan = msg.invocationSpan;
+                    // Skip this internal message - don't yield it
+                    iterResult = await generator.next();
+                    continue;
+                }
                 // Capture tool overhead info if provided by harness
                 if ("sessionUpdate" in msg &&
                     msg.sessionUpdate === "tool_overhead_info") {
@@ -739,6 +766,12 @@ export class AgentAcpAdapter {
                         "icon" in toolCallMsg._meta
                         ? String(toolCallMsg._meta.icon)
                         : undefined;
+                    const verbiage = toolCallMsg._meta &&
+                        typeof toolCallMsg._meta === "object" &&
+                        "verbiage" in toolCallMsg._meta &&
+                        typeof toolCallMsg._meta.verbiage === "object"
+                        ? toolCallMsg._meta.verbiage
+                        : undefined;
                     const batchId = toolCallMsg._meta &&
                         typeof toolCallMsg._meta === "object" &&
                         "batchId" in toolCallMsg._meta
@@ -771,6 +804,7 @@ export class AgentAcpAdapter {
                         title: toolCallMsg.title || "Tool",
                         ...(prettyName ? { prettyName } : {}),
                         ...(icon ? { icon } : {}),
+                        ...(verbiage ? { verbiage } : {}),
                         ...(subline ? { subline } : {}),
                         kind: toolCallMsg.kind || "other",
                         status: toolCallMsg.status || "pending",
@@ -1003,7 +1037,7 @@ export class AgentAcpAdapter {
                                 toolResultsTokens: midTurnSnapshot.context_size.toolResultsTokens,
                             });
                             // Execute hooks mid-turn to check if compaction is needed
-                            const midTurnContextEntries = await this.executeHooksIfConfigured(session, params.sessionId, "mid_turn");
+                            const midTurnContextEntries = await this.executeHooksIfConfigured(session, params.sessionId, "mid_turn", invocationSpan);
                             // Append new context entries returned by hooks (e.g., compaction)
                             if (midTurnContextEntries.length > 0) {
                                 logger.info(`Appending ${midTurnContextEntries.length} new context entries from mid_turn hooks`, {
@@ -1154,20 +1188,25 @@ export class AgentAcpAdapter {
      * Execute hooks if configured for this agent
      * Returns new context entries that should be appended to session.context
      */
-    async executeHooksIfConfigured(session, sessionId, executionPoint) {
+    async executeHooksIfConfigured(session, sessionId, executionPoint, parentSpan) {
         // Check if hooks are configured and session persistence is enabled
         const hooks = this.agent.definition.hooks;
         if (this.noSession || !hooks || hooks.length === 0) {
             return [];
         }
+        // If a parent span is provided (for mid-turn hooks), set it as active context
+        const baseContext = parentSpan
+            ? trace.setSpan(context.active(), parentSpan)
+            : context.active();
         const hookSpan = telemetry.startSpan("adapter.executeHooks", {
             "hooks.executionPoint": executionPoint,
             "hooks.count": hooks.length,
             "session.id": sessionId,
-        });
+            "agent.session_id": sessionId,
+        }, baseContext);
         const spanContext = hookSpan
-            ? trace.setSpan(context.active(), hookSpan)
-            : context.active();
+            ? trace.setSpan(baseContext, hookSpan)
+            : baseContext;
         return context.with(spanContext, async () => {
             try {
                 return await this._executeHooksImpl(session, sessionId, executionPoint, hooks);

package/dist/acp-server/http.js

@@ -45,6 +45,9 @@ function compressIfNeeded(rawMsg) {
 }
 // Use PGlite in-memory database for LISTEN/NOTIFY
 const pg = new PGlite();
+// Store for oversized responses that can't go through PostgreSQL NOTIFY
+// Key: request ID, Value: response object
+const oversizedResponses = new Map();
 // Helper to create safe channel names from untrusted IDs
 function safeChannelName(prefix, id) {
     const hash = createHash("sha256").update(id).digest("hex").slice(0, 16);
@@ -133,26 +136,21 @@ export function makeHttpTransport(agent, agentDir, agentName) {
                     const escapedPayload = payload.replace(/'/g, "''");
                     // Check if even compressed payload is too large
                     if (compressedSize > 7500) {
-                        logger.error("Response payload too large even after compression", {
+                        logger.info("Response payload too large for NOTIFY, using direct storage", {
                             requestId: rawMsg.id,
                             originalSize,
                             compressedSize,
                         });
-                        // Send error response
-                        const errorResponse = {
+                        // Store the response for direct retrieval by the HTTP handler
+                        oversizedResponses.set(rawMsg.id, rawMsg);
+                        // Send a small marker through NOTIFY to signal the response is ready
+                        const markerResponse = {
                             jsonrpc: "2.0",
                             id: rawMsg.id,
-                            error: {
-                                code: -32603,
-                                message: "Response payload too large even after compression",
-                                data: {
-                                    originalSize,
-                                    compressedSize,
-                                },
-                            },
+                            _oversized: true,
                         };
-                        const errorPayload = JSON.stringify(errorResponse).replace(/'/g, "''");
-                        await pg.query(`NOTIFY ${channel}, '${errorPayload}'`);
+                        const markerPayload = JSON.stringify(markerResponse).replace(/'/g, "''");
+                        await pg.query(`NOTIFY ${channel}, '${markerPayload}'`);
                         continue;
                     }
                     try {
@@ -555,6 +553,24 @@ export function makeHttpTransport(agent, agentDir, agentName) {
                         };
                     }
                 }
+                // Check if this is an oversized response marker
+                if (rawResponse &&
+                    typeof rawResponse === "object" &&
+                    "_oversized" in rawResponse &&
+                    rawResponse._oversized === true &&
+                    "id" in rawResponse) {
+                    // Retrieve the actual response from the oversized storage
+                    const actualResponse = oversizedResponses.get(String(rawResponse.id));
+                    if (actualResponse) {
+                        oversizedResponses.delete(String(rawResponse.id));
+                        logger.debug("Retrieved oversized response from storage", { id });
+                        responseResolver(actualResponse);
+                        return;
+                    }
+                    logger.error("Oversized response not found in storage", {
+                        id: rawResponse.id,
+                    });
+                }
                 responseResolver(rawResponse);
             });
             // Write NDJSON line into the ACP inbound stream

package/dist/runner/agent-runner.d.ts

@@ -1,4 +1,5 @@
 import type { PromptRequest, PromptResponse, SessionNotification } from "@agentclientprotocol/sdk";
+import type { Span } from "@opentelemetry/api";
 import { z } from "zod";
 import type { ContentBlock } from "../acp-server/session-storage.js";
 export declare const zAgentRunnerParams: z.ZodObject<{
@@ -8,7 +9,7 @@ export declare const zAgentRunnerParams: z.ZodObject<{
     suggestedPrompts: z.ZodOptional<z.ZodArray<z.ZodString>>;
     systemPrompt: z.ZodNullable<z.ZodString>;
     model: z.ZodString;
-    tools: z.ZodOptional<z.ZodArray<z.ZodUnion<readonly [z.ZodUnion<readonly [z.ZodLiteral<"todo_write">, z.ZodLiteral<"get_weather">, z.ZodLiteral<"web_search">, z.ZodLiteral<"filesystem">, z.ZodLiteral<"generate_image">, z.ZodLiteral<"browser">]>, z.ZodObject<{
+    tools: z.ZodOptional<z.ZodArray<z.ZodUnion<readonly [z.ZodUnion<readonly [z.ZodLiteral<"todo_write">, z.ZodLiteral<"get_weather">, z.ZodLiteral<"web_search">, z.ZodLiteral<"town_web_search">, z.ZodLiteral<"filesystem">, z.ZodLiteral<"generate_image">, z.ZodLiteral<"browser">]>, z.ZodObject<{
         type: z.ZodLiteral<"custom">;
         modulePath: z.ZodString;
     }, z.core.$strip>, z.ZodObject<{
@@ -66,10 +67,16 @@ export interface SessionMessage {
     content: ContentBlock[];
     timestamp: string;
 }
+export interface ConfigOverrides {
+    model?: string;
+    systemPrompt?: string;
+    tools?: string[];
+}
 export type InvokeRequest = Omit<PromptRequest, "_meta"> & {
     messageId: string;
     sessionMeta?: Record<string, unknown>;
     contextMessages?: SessionMessage[];
+    configOverrides?: ConfigOverrides;
 };
 export interface TokenUsage {
     inputTokens?: number;
@@ -125,6 +132,9 @@ export type ExtendedSessionUpdate = (SessionNotification["update"] & {
     sessionUpdate: "tool_overhead_info";
     toolOverheadTokens: number;
     mcpOverheadTokens: number;
+} | {
+    sessionUpdate: "__invocation_span";
+    invocationSpan: Span;
 } | AgentMessageChunkWithTokens | HookNotificationUpdate;
 /** Describes an object that can run an agent definition */
 export interface AgentRunner {

package/dist/runner/hooks/predefined/tool-response-compactor.js

@@ -1,8 +1,13 @@
-import { ChatAnthropic } from "@langchain/anthropic";
-import { HumanMessage, SystemMessage } from "@langchain/core/messages";
+import Anthropic from "@anthropic-ai/sdk";
 import { createLogger } from "../../../logger.js";
+import { telemetry } from "../../../telemetry/index.js";
 import { countToolResultTokens } from "../../../utils/token-counter.js";
 const logger = createLogger("tool-response-compactor");
+// Create Anthropic client directly (not using LangChain)
+// This ensures compaction LLM calls don't get captured by LangGraph's streaming
+const anthropic = new Anthropic({
+    apiKey: process.env.ANTHROPIC_API_KEY,
+});
 // Haiku 4.5 for compaction (fast and cost-effective)
 const COMPACTION_MODEL = "claude-haiku-4-5-20251001";
 const COMPACTION_MODEL_CONTEXT = 200000; // Haiku context size for calculating truncation limits
@@ -243,10 +248,6 @@ export const toolResponseCompactor = async (ctx) => {
  * Recursive LLM compaction with adaptive retries
  */
 async function compactWithLLM(rawOutput, toolName, toolInput, conversationContext, targetTokens) {
-    const model = new ChatAnthropic({
-        model: COMPACTION_MODEL,
-        temperature: 0,
-    });
     // Step 1: Understand what we're looking for (only need to do this once)
     const analysisPrompt = `You are helping to manage context size in an agent conversation.
 
@@ -260,12 +261,38 @@ ${conversationContext}
 Based on the tool input and conversation context, what key information is the user looking for from this tool response?
 
 Provide a concise list (3-5 bullet points) of the most important elements to extract.`;
-    const analysisResponse = await model.invoke([
-        new SystemMessage("You are a helpful assistant analyzing information needs."),
-        new HumanMessage(analysisPrompt),
-    ]);
-    const keyRequirements = typeof analysisResponse.content === "string"
-        ? analysisResponse.content
+    // Create OTEL span for analysis call
+    const analysisSpan = telemetry.startSpan("compaction.analysis", {
+        "gen_ai.operation.name": "chat",
+        "gen_ai.provider.name": "anthropic",
+        "gen_ai.request.model": COMPACTION_MODEL,
+        "gen_ai.system_instructions": "You are a helpful assistant analyzing information needs.",
+    });
+    let analysisResponse;
+    try {
+        analysisResponse = await telemetry.withActiveSpanAsync(analysisSpan, () => anthropic.messages.create({
+            model: COMPACTION_MODEL,
+            max_tokens: 1024,
+            temperature: 0,
+            system: "You are a helpful assistant analyzing information needs.",
+            messages: [
+                {
+                    role: "user",
+                    content: analysisPrompt,
+                },
+            ],
+        }));
+        // Record token usage
+        telemetry.recordTokenUsage(analysisResponse.usage.input_tokens, analysisResponse.usage.output_tokens, analysisSpan);
+        telemetry.endSpan(analysisSpan);
+    }
+    catch (error) {
+        telemetry.endSpan(analysisSpan, error);
+        throw error;
+    }
+    const firstContent = analysisResponse.content[0];
+    const keyRequirements = firstContent && firstContent.type === "text"
+        ? firstContent.text
         : "Extract relevant information";
     logger.info("Identified key requirements for compaction", {
         requirements: keyRequirements.substring(0, 200),
@@ -323,14 +350,41 @@ Your task: Further compact this data by:
 
 Return ONLY valid JSON (no explanation text).`;
         }
-        const compactionResponse = await model.invoke([
-            new SystemMessage("You are a helpful assistant compacting data."),
-            new HumanMessage(compactionPrompt),
-        ]);
+        // Create OTEL span for compaction call
+        const compactionSpan = telemetry.startSpan("compaction.compact", {
+            "gen_ai.operation.name": "chat",
+            "gen_ai.provider.name": "anthropic",
+            "gen_ai.request.model": COMPACTION_MODEL,
+            "gen_ai.system_instructions": "You are a helpful assistant compacting data.",
+            "compaction.attempt": attempt + 1,
+            "compaction.target_tokens": targetTokens,
+            "compaction.current_tokens": currentTokens,
+        });
+        let compactionResponse;
+        try {
+            compactionResponse = await telemetry.withActiveSpanAsync(compactionSpan, () => anthropic.messages.create({
+                model: COMPACTION_MODEL,
+                max_tokens: 4096,
+                temperature: 0,
+                system: "You are a helpful assistant compacting data.",
+                messages: [
+                    {
+                        role: "user",
+                        content: compactionPrompt,
+                    },
+                ],
+            }));
+            // Record token usage
+            telemetry.recordTokenUsage(compactionResponse.usage.input_tokens, compactionResponse.usage.output_tokens, compactionSpan);
+            telemetry.endSpan(compactionSpan);
+        }
+        catch (error) {
+            telemetry.endSpan(compactionSpan, error);
+            throw error;
+        }
         // Extract and parse JSON
-        const responseText = typeof compactionResponse.content === "string"
-            ? compactionResponse.content
-            : JSON.stringify(compactionResponse.content);
+        const firstContent = compactionResponse.content[0];
+        const responseText = firstContent && firstContent.type === "text" ? firstContent.text : "";
         const jsonMatch = responseText.match(/```(?:json)?\n([\s\S]*?)\n```/) || [
             null,
             responseText,

package/dist/runner/langchain/index.js

@@ -1,6 +1,6 @@
 import { MultiServerMCPClient } from "@langchain/mcp-adapters";
 import { context, propagation, trace } from "@opentelemetry/api";
-import { loadAuthCredentials } from "@townco/core/auth";
+import { getShedAuth } from "@townco/core/auth";
 import { AIMessageChunk, createAgent, ToolMessage, tool, } from "langchain";
 import { z } from "zod";
 import { SUBAGENT_MODE_KEY } from "../../acp-server/adapter";
@@ -15,7 +15,7 @@ import { makeGenerateImageTool } from "./tools/generate_image";
 import { SUBAGENT_TOOL_NAME } from "./tools/subagent";
 import { hashQuery, queryToToolCallId, subagentEvents, } from "./tools/subagent-connections";
 import { TODO_WRITE_TOOL_NAME, todoWrite } from "./tools/todo";
-import { makeWebSearchTools } from "./tools/web_search";
+import { makeTownWebSearchTools, makeWebSearchTools } from "./tools/web_search";
 const _logger = createLogger("agent-runner");
 const getWeather = tool(({ city }) => `It's always sunny in ${city}!`, {
     name: "get_weather",
@@ -30,6 +30,7 @@ export const TOOL_REGISTRY = {
     todo_write: todoWrite,
     get_weather: getWeather,
     web_search: () => makeWebSearchTools(),
+    town_web_search: () => makeTownWebSearchTools(),
     filesystem: () => makeFilesystemTools(process.cwd()),
     generate_image: () => makeGenerateImageTool(),
     browser: () => makeBrowserTools(),
@@ -96,10 +97,6 @@ export class LangchainAgent {
             // Clear the buffer after flushing
             pendingToolCallNotifications.length = 0;
         }
-        // Set session_id as a base attribute so all spans in this invocation include it
-        telemetry.setBaseAttributes({
-            "agent.session_id": req.sessionId,
-        });
         const subagentUpdateQueue = [];
         let subagentUpdateResolver = null;
         const subagentMessagesQueue = [];
@@ -176,11 +173,17 @@ export class LangchainAgent {
                 };
             }
         }
+        // Add agent.session_id as a base attribute so it propagates to all child spans
+        // We'll clear this in a finally block to prevent cross-contamination
+        telemetry.setBaseAttributes({
+            "agent.session_id": req.sessionId,
+        });
         // 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.message_id": req.messageId,
+            "agent.session_id": req.sessionId,
         }, parentContext);
         // Create a context with the invocation span as active
         // This will be used when creating child spans (tool calls)
@@ -192,7 +195,29 @@ export class LangchainAgent {
             sessionId: req.sessionId,
             messageId: req.messageId,
         });
+        // Yield the invocation span to the adapter so it can use it for parenting hook spans
+        if (invocationSpan) {
+            yield {
+                sessionUpdate: "__invocation_span",
+                invocationSpan,
+            };
+        }
+        // Declare otelCallbacks outside try block so it's accessible in catch
+        let otelCallbacks = null;
         try {
+            // Determine effective model early so we can detect provider for callbacks
+            // Use override model if provided (Town Hall comparison feature)
+            const effectiveModel = req.configOverrides?.model ?? this.definition.model;
+            const provider = detectProvider(effectiveModel);
+            // Create OTEL callbacks for instrumentation early so we can use them during tool wrapping
+            // Track iteration index across LLM calls in this invocation
+            const iterationIndexRef = { current: 0 };
+            otelCallbacks = makeOtelCallbacks({
+                provider,
+                model: effectiveModel,
+                parentContext: invocationContext,
+                iterationIndexRef,
+            });
             // Track todo_write tool call IDs to suppress their tool_call notifications
             const todoWriteToolCallIds = new Set();
             // --------------------------------------------------------------------------
@@ -391,19 +416,34 @@ export class LangchainAgent {
                 : wrappedTools;
             // Wrap tools with tracing so each tool executes within its own span context.
             // This ensures subagent spans are children of the Task tool span.
-            const finalTools = filteredTools.map((t) => wrapToolWithTracing(t));
+            // Pass the context getter so tools can nest under the current iteration span.
+            let finalTools = filteredTools.map((t) => wrapToolWithTracing(t, otelCallbacks?.getCurrentIterationContext ??
+                (() => invocationContext)));
+            // Apply tool overrides if provided (Town Hall comparison feature)
+            if (req.configOverrides?.tools && req.configOverrides.tools.length > 0) {
+                const allowedToolNames = new Set(req.configOverrides.tools);
+                finalTools = finalTools.filter((t) => allowedToolNames.has(t.name));
+                _logger.debug("Applied tool override filter", {
+                    requested: req.configOverrides.tools,
+                    filtered: finalTools.map((t) => t.name),
+                });
+            }
             // 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 model = createModelFromString(effectiveModel);
             const agentConfig = {
                 model,
                 tools: finalTools,
             };
-            if (this.definition.systemPrompt) {
-                agentConfig.systemPrompt = this.definition.systemPrompt;
+            // Use override system prompt if provided (Town Hall comparison feature)
+            const effectiveSystemPrompt = req.configOverrides?.systemPrompt !== undefined
+                ? req.configOverrides.systemPrompt
+                : this.definition.systemPrompt;
+            if (effectiveSystemPrompt) {
+                agentConfig.systemPrompt = effectiveSystemPrompt;
             }
             // Inject system prompt with optional TodoWrite instructions
             const hasTodoWrite = builtInNames.includes("todo_write");
@@ -411,8 +451,6 @@ export class LangchainAgent {
                 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;
             // Helper to convert content blocks to LangChain format
@@ -503,12 +541,6 @@ export class LangchainAgent {
                     },
                 ];
             }
-            // Create OTEL callbacks for instrumentation
-            const otelCallbacks = makeOtelCallbacks({
-                provider,
-                model: this.definition.model,
-                parentContext: invocationContext,
-            });
             // 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 }, {
@@ -639,6 +671,7 @@ export class LangchainAgent {
                             const matchingTool = finalTools.find((t) => t.name === toolCall.name);
                             let prettyName = matchingTool?.prettyName;
                             const icon = matchingTool?.icon;
+                            const verbiage = matchingTool?.verbiage;
                             // For the Task tool, use the displayName (or agentName as fallback) as the prettyName
                             if (toolCall.name === SUBAGENT_TOOL_NAME &&
                                 toolCall.args &&
@@ -680,6 +713,7 @@ export class LangchainAgent {
                                         messageId: req.messageId,
                                         ...(prettyName ? { prettyName } : {}),
                                         ...(icon ? { icon } : {}),
+                                        ...(verbiage ? { verbiage } : {}),
                                         ...(batchId ? { batchId } : {}),
                                     },
                                 });
@@ -698,6 +732,7 @@ export class LangchainAgent {
                                         messageId: req.messageId,
                                         ...(prettyName ? { prettyName } : {}),
                                         ...(icon ? { icon } : {}),
+                                        ...(verbiage ? { verbiage } : {}),
                                         ...(batchId ? { batchId } : {}),
                                     },
                                 });
@@ -828,7 +863,9 @@ export class LangchainAgent {
                                         toolUseBlock.name &&
                                         !preliminaryToolCallIds.has(toolUseBlock.id)) {
                                         preliminaryToolCallIds.add(toolUseBlock.id);
-                                        pendingToolCallNotifications.push({
+                                        // Yield the preliminary notification immediately (don't buffer)
+                                        // This allows the UI to show the "selecting" state before params arrive
+                                        yield {
                                             sessionUpdate: "tool_call",
                                             toolCallId: toolUseBlock.id,
                                             title: toolUseBlock.name,
@@ -836,7 +873,7 @@ export class LangchainAgent {
                                             status: "pending",
                                             rawInput: {}, // Args not available yet
                                             _meta: { messageId: req.messageId },
-                                        });
+                                        };
                                     }
                                 }
                                 else if (part.type === "input_json_delta") {
@@ -846,8 +883,8 @@ export class LangchainAgent {
                                     throw new Error(`Unhandled AIMessageChunk content block type: ${part.type}\n${JSON.stringify(part)}`);
                                 }
                             }
-                            // Don't flush here - these are preliminary tool_use blocks
-                            // We'll flush when we get the full tool calls in "updates" mode
+                            // Preliminary tool_use blocks are yielded immediately above
+                            // Full tool calls with params will come later in "updates" mode
                         }
                         else {
                             throw new Error(`Unhandled AIMessageChunk content type: ${typeof aiMessage.content}`);
@@ -921,6 +958,8 @@ export class LangchainAgent {
             if (subagentUpdateResolver) {
                 subagentUpdateResolver = null;
             }
+            // Clean up any remaining iteration span
+            otelCallbacks?.cleanup();
             // Log successful completion
             telemetry.log("info", "Agent invocation completed", {
                 sessionId: req.sessionId,
@@ -937,6 +976,8 @@ export class LangchainAgent {
             // Clean up subagent event listeners on error
             subagentEvents.off("connection", onSubagentConnection);
             subagentEvents.off("messages", onSubagentMessages);
+            // Clean up any remaining iteration span
+            otelCallbacks?.cleanup();
             // Log error and end span with error status
             telemetry.log("error", "Agent invocation failed", {
                 error: error instanceof Error ? error.message : String(error),
@@ -945,6 +986,10 @@ export class LangchainAgent {
             telemetry.endSpan(invocationSpan, error instanceof Error ? error : new Error(String(error)));
             throw error;
         }
+        finally {
+            // Clear agent.session_id from base attributes to prevent cross-contamination
+            telemetry.clearBaseAttribute("agent.session_id");
+        }
     }
 }
 const modelRequestSchema = z.object({
@@ -956,17 +1001,17 @@ const makeMcpToolsClient = (mcpConfigs) => {
     const mcpServers = mcpConfigs?.map((config) => {
         if (typeof config === "string") {
             // String configs use the centralized MCP proxy with auth
-            const credentials = loadAuthCredentials();
-            if (!credentials) {
-                throw new Error("Not logged in. Run 'town login' first to use cloud MCP servers.");
+            const shedAuth = getShedAuth();
+            if (!shedAuth) {
+                throw new Error("Not logged in. Run 'town login' or set SHED_API_KEY to use cloud MCP servers.");
             }
-            const proxyUrl = process.env.MCP_PROXY_URL ?? `${credentials.shed_url}/mcp_proxy`;
+            const proxyUrl = process.env.MCP_PROXY_URL ?? `${shedAuth.shedUrl}/mcp_proxy`;
             return [
                 config,
                 {
                     url: `${proxyUrl}?server=${config}`,
                     headers: {
-                        Authorization: `Bearer ${credentials.access_token}`,
+                        Authorization: `Bearer ${shedAuth.accessToken}`,
                     },
                 },
             ];
@@ -1059,18 +1104,22 @@ export { makeSubagentsTool } from "./tools/subagent.js";
  * This ensures the tool executes within its own span context,
  * so any child operations (like subagent spawning) become children
  * of the tool span rather than the parent invocation span.
+ * @param originalTool The tool to wrap
+ * @param getIterationContext Function that returns the current iteration context
  */
-function wrapToolWithTracing(originalTool) {
+function wrapToolWithTracing(originalTool, getIterationContext) {
     const wrappedFunc = async (input) => {
         const toolInputJson = JSON.stringify(input);
+        // Get the current iteration context so the tool span is created as a child
+        const iterationContext = getIterationContext();
         const toolSpan = telemetry.startSpan("agent.tool_call", {
             "tool.name": originalTool.name,
             "tool.input": toolInputJson,
-        });
+        }, iterationContext);
         // Create a context with the tool span as active
         const spanContext = toolSpan
-            ? trace.setSpan(context.active(), toolSpan)
-            : context.active();
+            ? trace.setSpan(iterationContext, toolSpan)
+            : iterationContext;
         try {
             // Execute within the tool span's context
             const result = await context.with(spanContext, () => originalTool.invoke(input));