@townco/agent 0.1.77 → 0.1.79

package/dist/acp-server/adapter.js

@@ -1,7 +1,7 @@
 import * as acp from "@agentclientprotocol/sdk";
 import { context, trace } from "@opentelemetry/api";
 import { createLogger } from "../logger.js";
-import { HookExecutor, loadHookCallback } from "../runner/hooks";
+import { getModelContextWindow, HookExecutor, loadHookCallback, } from "../runner/hooks";
 import { telemetry } from "../telemetry/index.js";
 import { calculateContextSize, } from "../utils/context-size-calculator.js";
 import { countToolResultTokens } from "../utils/token-counter.js";
@@ -278,16 +278,18 @@ export class AgentAcpAdapter {
         return response;
     }
     async newSession(params) {
+        // Generate a unique session ID for this session
         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, {
+        const sessionData = {
             pendingPrompt: null,
             messages: [],
             context: [],
             requestParams: params,
             configOverrides,
-        });
+        };
+        this.sessions.set(sessionId, sessionData);
         // 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
         return {
@@ -599,7 +601,9 @@ export class AgentAcpAdapter {
             // Calculate context size - no LLM call yet, so only estimated values
             const context_size = calculateContextSize(contextMessages, this.agent.definition.systemPrompt ?? undefined, undefined, // No LLM-reported tokens yet
             this.currentToolOverheadTokens, // Include tool overhead
-            this.currentMcpOverheadTokens);
+            this.currentMcpOverheadTokens, // Include MCP overhead
+            getModelContextWindow(this.agent.definition.model), // Model context window for UI
+            true);
             const contextSnapshot = createContextSnapshot(session.messages.length - 1, // Exclude the newly added user message (it will be passed separately via prompt)
             new Date().toISOString(), previousContext, context_size);
             session.context.push(contextSnapshot);
@@ -826,6 +830,11 @@ export class AgentAcpAdapter {
                             toolCallBlock.status =
                                 updateMsg.status;
                         }
+                        // Update rawInput if provided (for preliminary -> full tool call flow)
+                        if (updateMsg.rawInput &&
+                            Object.keys(updateMsg.rawInput).length > 0) {
+                            toolCallBlock.rawInput = updateMsg.rawInput;
+                        }
                         if (updateMsg.rawOutput) {
                             toolCallBlock.rawOutput = updateMsg.rawOutput;
                         }
@@ -903,11 +912,42 @@ export class AgentAcpAdapter {
                             const hooks = this.agent.definition.hooks ?? [];
                             if (hooks.some((h) => h.type === "tool_response")) {
                                 const latestContext = session.context[session.context.length - 1];
-                                const currentContextTokens = latestContext?.context_size.llmReportedInputTokens ??
-                                    latestContext?.context_size.totalEstimated ??
-                                    0;
+                                // Use max of estimated and LLM-reported tokens (same logic as UI)
+                                const currentContextTokens = Math.max(latestContext?.context_size.totalEstimated ?? 0, latestContext?.context_size.llmReportedInputTokens ?? 0);
+                                // Send the context size to the UI so it shows the same value we're using for hook evaluation
+                                if (latestContext?.context_size) {
+                                    const contextSizeForUI = {
+                                        ...latestContext.context_size,
+                                        totalEstimated: currentContextTokens,
+                                    };
+                                    this.connection.sessionUpdate({
+                                        sessionId: params.sessionId,
+                                        update: {
+                                            sessionUpdate: "agent_message_chunk",
+                                            content: {
+                                                type: "text",
+                                                text: "",
+                                            },
+                                            _meta: {
+                                                context_size: contextSizeForUI,
+                                            },
+                                        },
+                                    });
+                                }
                                 const outputTokens = countToolResultTokens(rawOutput);
-                                const hookExecutor = new HookExecutor(hooks, this.agent.definition.model, (callbackRef) => loadHookCallback(callbackRef, this.agentDir));
+                                // Create notification callback to stream hook events in real-time
+                                const sendHookNotification = (notification) => {
+                                    this.connection.sessionUpdate({
+                                        sessionId: params.sessionId,
+                                        update: {
+                                            sessionUpdate: "hook_notification",
+                                            id: `hook_${Date.now()}_${Math.random().toString(36).substring(2, 9)}`,
+                                            notification,
+                                            messageId,
+                                        },
+                                    });
+                                };
+                                const hookExecutor = new HookExecutor(hooks, this.agent.definition.model, (callbackRef) => loadHookCallback(callbackRef, this.agentDir), sendHookNotification);
                                 const hookResult = await hookExecutor.executeToolResponseHooks({
                                     messages: session.messages,
                                     context: session.context,
@@ -919,6 +959,8 @@ export class AgentAcpAdapter {
                                     rawOutput,
                                     outputTokens,
                                 });
+                                // Note: Notifications are now sent in real-time via the callback
+                                // The hookResult.notifications array is kept for backwards compatibility
                                 // Apply modifications if hook returned them
                                 if (hookResult.modifiedOutput) {
                                     rawOutput = hookResult.modifiedOutput;
@@ -1019,7 +1061,8 @@ export class AgentAcpAdapter {
                             // Calculate context size - tool result is now in the message, but hasn't been sent to LLM yet
                             const context_size = calculateContextSize(contextMessages, this.agent.definition.systemPrompt ?? undefined, undefined, // Tool result hasn't been sent to LLM yet, so no new LLM-reported tokens
                             this.currentToolOverheadTokens, // Include tool overhead
-                            this.currentMcpOverheadTokens);
+                            this.currentMcpOverheadTokens, // Include MCP overhead
+                            getModelContextWindow(this.agent.definition.model));
                             // Create snapshot with a pointer to the partial message (not a full copy!)
                             const midTurnSnapshot = {
                                 timestamp: new Date().toISOString(),
@@ -1172,12 +1215,26 @@ export class AgentAcpAdapter {
                 }
             }
             // Calculate context size with LLM-reported tokens from this turn
+            // Exclude tool results - they're only sent during the turn they were received,
+            // not in subsequent turns (only messages are sent)
             const context_size = calculateContextSize(contextMessages, this.agent.definition.systemPrompt ?? undefined, turnTokenUsage.inputTokens, // Final LLM-reported tokens from this turn
             this.currentToolOverheadTokens, // Include tool overhead
-            this.currentMcpOverheadTokens);
+            this.currentMcpOverheadTokens, // Include MCP overhead
+            getModelContextWindow(this.agent.definition.model), // Model context window for UI
+            true);
             const contextSnapshot = createContextSnapshot(session.messages.length, new Date().toISOString(), previousContext, context_size);
             session.context.push(contextSnapshot);
             await this.saveSessionToDisk(params.sessionId, session);
+            // Send final context_size to UI (with tool results excluded)
+            // This ensures the UI shows the correct context size at turn-end
+            this.connection.sessionUpdate({
+                sessionId: params.sessionId,
+                update: {
+                    sessionUpdate: "agent_message_chunk",
+                    content: { type: "text", text: "" },
+                    _meta: { context_size },
+                },
+            });
         }
         session.pendingPrompt = null;
         return {
@@ -1222,7 +1279,18 @@ export class AgentAcpAdapter {
             contextEntries: session.context.length,
             totalMessages: session.messages.length,
         });
-        const hookExecutor = new HookExecutor(hooks, this.agent.definition.model, (callbackRef) => loadHookCallback(callbackRef, this.agentDir));
+        // Create notification callback to stream hook events in real-time
+        const sendHookNotification = (notification) => {
+            this.connection.sessionUpdate({
+                sessionId,
+                update: {
+                    sessionUpdate: "hook_notification",
+                    id: `hook_${Date.now()}_${Math.random().toString(36).substring(2, 9)}`,
+                    notification,
+                },
+            });
+        };
+        const hookExecutor = new HookExecutor(hooks, this.agent.definition.model, (callbackRef) => loadHookCallback(callbackRef, this.agentDir), sendHookNotification);
         // Create read-only session view for hooks
         const readonlySession = {
             messages: session.messages,
@@ -1233,23 +1301,45 @@ export class AgentAcpAdapter {
         const latestContext = session.context.length > 0
             ? session.context[session.context.length - 1]
             : undefined;
-        // Prefer LLM-reported tokens (most accurate), fall back to our estimate
-        const actualInputTokens = latestContext?.context_size.llmReportedInputTokens ??
-            latestContext?.context_size.totalEstimated ??
-            0;
+        // Use max of estimated and LLM-reported tokens (same logic as UI)
+        // This is conservative - we want to trigger compaction earlier rather than later
+        const actualInputTokens = Math.max(latestContext?.context_size.totalEstimated ?? 0, latestContext?.context_size.llmReportedInputTokens ?? 0);
         logger.debug("Using tokens for hook execution", {
             llmReported: latestContext?.context_size.llmReportedInputTokens,
             estimated: latestContext?.context_size.totalEstimated,
             used: actualInputTokens,
         });
-        const hookResult = await hookExecutor.executeHooks(readonlySession, actualInputTokens);
-        // Send hook notifications to client
-        for (const notification of hookResult.notifications) {
+        // Send the context size to the UI so it shows the same value we're using for hook evaluation
+        // This ensures the UI percentage matches what the hook sees
+        if (latestContext?.context_size) {
+            // Create an updated context_size with the actualInputTokens we computed
+            // so UI can calculate the same percentage
+            const contextSizeForUI = {
+                ...latestContext.context_size,
+                // Override with the max value we computed (what hooks actually use)
+                totalEstimated: actualInputTokens,
+            };
             this.connection.sessionUpdate({
                 sessionId,
-                update: notification,
+                update: {
+                    sessionUpdate: "agent_message_chunk",
+                    content: {
+                        type: "text",
+                        text: "",
+                    },
+                    _meta: {
+                        context_size: contextSizeForUI,
+                    },
+                },
+            });
+            logger.debug("Sent context_size update to UI before hook execution", {
+                actualInputTokens,
+                modelContextWindow: contextSizeForUI.modelContextWindow,
             });
         }
+        const hookResult = await hookExecutor.executeHooks(readonlySession, actualInputTokens);
+        // Note: Notifications are now sent in real-time via the callback
+        // The hookResult.notifications array is kept for backwards compatibility
         // Return new context entries (will be appended by caller)
         return hookResult.newContextEntries;
     }