@townco/agent 0.1.112 → 0.1.114

package/dist/runner/langchain/tools/subagent-connections.js

@@ -1,22 +1,14 @@
 import { createHash } from "node:crypto";
-import { EventEmitter } from "node:events";
 import { createLogger } from "@townco/core";
+import { getInvocationContext } from "../../session-context.js";
 const logger = createLogger("subagent-connections");
-/**
- * Event emitter for subagent connection events.
- * The runner listens to these events and emits tool_call_update.
- */
-export const subagentEvents = new EventEmitter();
 /**
  * Maps query hash to toolCallId.
  * Set by the runner when it sees a subagent tool_call.
+ * This is still global because the registration happens in the parent runner
+ * before subagent execution begins.
  */
 export const queryToToolCallId = new Map();
-/**
- * Maps query hash to resolved toolCallId (preserved after connection event).
- * Used to correlate messages when the tool completes.
- */
-const queryToResolvedToolCallId = new Map();
 /**
  * Generate a hash from the query string for correlation.
  */
@@ -29,61 +21,45 @@ export function hashQuery(query) {
     return hash;
 }
 /**
- * Called by the subagent tool when connection is established.
- * Emits an event that the runner can listen to.
+ * Called by the subagent tool during execution to emit incremental messages.
+ * Emits an event with the messages for live streaming to the UI.
+ * Uses the invocation context's EventEmitter to ensure messages go to the correct parent.
+ * @param completed - If true, signals that the subagent stream has finished
  */
-export function emitSubagentConnection(queryHash, connectionInfo) {
-    logger.info("emitSubagentConnection called", {
-        queryHash,
-        port: connectionInfo.port,
-        sessionId: connectionInfo.sessionId,
-        registeredHashes: Array.from(queryToToolCallId.keys()),
-    });
+export function emitSubagentMessages(queryHash, messages, completed = false) {
     const toolCallId = queryToToolCallId.get(queryHash);
-    if (toolCallId) {
-        logger.info("Found toolCallId for queryHash, emitting connection event", {
-            queryHash,
-            toolCallId,
-            port: connectionInfo.port,
-            sessionId: connectionInfo.sessionId,
-        });
-        subagentEvents.emit("connection", {
-            toolCallId,
-            ...connectionInfo,
-        });
-        // Preserve the toolCallId for message emission, but remove from pending lookup
-        queryToResolvedToolCallId.set(queryHash, toolCallId);
-        queryToToolCallId.delete(queryHash);
-    }
-    else {
-        logger.warn("No toolCallId found for queryHash", {
+    const invocationCtx = getInvocationContext();
+    if (!invocationCtx) {
+        logger.warn("✗ No invocation context found - cannot emit subagent messages", {
             queryHash,
-            registeredHashes: Array.from(queryToToolCallId.keys()),
+            hasToolCallId: !!toolCallId,
         });
+        return;
     }
-}
-/**
- * Called by the subagent tool when it completes with accumulated messages.
- * Emits an event with the messages for session storage.
- */
-export function emitSubagentMessages(queryHash, messages) {
-    const toolCallId = queryToResolvedToolCallId.get(queryHash);
     if (toolCallId) {
-        logger.info("Emitting subagent messages for storage", {
+        const firstMessage = messages[0];
+        logger.info("✓ Emitting subagent messages for live streaming", {
             queryHash,
             toolCallId,
             messageCount: messages.length,
+            hasContent: firstMessage ? firstMessage.content.length > 0 : false,
+            hasToolCalls: firstMessage ? firstMessage.toolCalls.length > 0 : false,
+            completed,
+            invocationId: invocationCtx.invocationId,
         });
-        subagentEvents.emit("messages", {
+        // Emit to the parent's invocation-scoped EventEmitter
+        invocationCtx.subagentEventEmitter.emit("messages", {
             toolCallId,
             messages,
+            completed,
         });
-        // Clean up the resolved mapping
-        queryToResolvedToolCallId.delete(queryHash);
     }
     else {
-        logger.warn("No resolved toolCallId for messages emission", {
+        logger.warn("✗ No toolCallId found for messages emission (RACE CONDITION)", {
             queryHash,
+            registeredHashes: Array.from(queryToToolCallId.keys()),
+            mapSize: queryToToolCallId.size,
+            invocationId: invocationCtx.invocationId,
         });
     }
 }