@townco/ui 0.1.109 → 0.1.111

package/dist/sdk/transports/http.js

@@ -25,6 +25,7 @@ export class HttpTransport {
     options;
     isReceivingMessages = false;
     isInReplayMode = false; // True during session replay, ignores non-replay streaming
+    pendingReplayUpdates = []; // Queue session updates during replay for late-subscribing callbacks
     agentInfo;
     constructor(options) {
         // Validate options at the boundary using Zod
@@ -565,6 +566,61 @@ export class HttpTransport {
     }
     onSessionUpdate(callback) {
         this.sessionUpdateCallbacks.add(callback);
+        // Replay any queued session updates from replay that arrived before callbacks were registered
+        // IMPORTANT: Use notifySessionUpdate to send to ALL registered callbacks, not just this one.
+        // This ensures that when multiple hooks register (e.g., message handler, useToolCalls),
+        // they ALL receive the queued updates, not just the first one to register.
+        if (this.pendingReplayUpdates.length > 0) {
+            logger.info("Replaying queued session updates to ALL callbacks", {
+                count: this.pendingReplayUpdates.length,
+                callbackCount: this.sessionUpdateCallbacks.size,
+                updateTypes: this.pendingReplayUpdates.map((u) => u.type),
+                updateDetails: this.pendingReplayUpdates.map((u) => {
+                    if (u.type === "tool_call") {
+                        return { type: "tool_call", toolCallId: u.toolCall?.id };
+                    }
+                    if (u.type === "sources") {
+                        return {
+                            type: "sources",
+                            count: u.sources?.length,
+                            ids: u.sources?.map((s) => s.id),
+                        };
+                    }
+                    if (u.type === "generic" && u.message) {
+                        const firstContent = u.message.content?.[0];
+                        return {
+                            type: "message",
+                            role: u.message.role,
+                            contentLength: firstContent?.type === "text" ? firstContent.text?.length : 0,
+                        };
+                    }
+                    return { type: u.type };
+                }),
+            });
+            // Copy and clear the queue immediately to prevent re-triggering when other callbacks register
+            const updatesToReplay = [...this.pendingReplayUpdates];
+            this.pendingReplayUpdates = [];
+            // Use setTimeout to ensure all callbacks have a chance to register before replaying
+            setTimeout(() => {
+                logger.info("Starting replay of queued updates to all callbacks", {
+                    updateCount: updatesToReplay.length,
+                    callbackCount: this.sessionUpdateCallbacks.size,
+                });
+                for (const update of updatesToReplay) {
+                    try {
+                        logger.info("Replaying update to all callbacks", {
+                            type: update.type,
+                        });
+                        // Notify ALL registered callbacks, not just the one that triggered this replay
+                        this.notifySessionUpdate(update);
+                    }
+                    catch (error) {
+                        logger.error("Error replaying session update", { error });
+                    }
+                }
+                logger.info("Finished replay of queued updates");
+            }, 0);
+        }
         return () => {
             this.sessionUpdateCallbacks.delete(callback);
         };
@@ -810,6 +866,28 @@ export class HttpTransport {
         try {
             const message = JSON.parse(data);
             logger.debug("Received SSE message", { message });
+            // Check if this is a sources message (custom extension to ACP)
+            const isSourcesMessage = message &&
+                typeof message === "object" &&
+                message.method === "session/update" &&
+                message.params?.update?.sessionUpdate === "sources";
+            if (isSourcesMessage) {
+                // Use console.warn directly for gui-console capture
+                console.warn("🟢 RECEIVED SOURCES SSE MESSAGE", {
+                    sourcesCount: message.params?.update?.sources?.length,
+                    isInReplayMode: this.isInReplayMode,
+                    callbackCount: this.sessionUpdateCallbacks.size,
+                });
+                // Handle sources directly without ACP schema validation
+                try {
+                    this.handleSessionNotification(message.params);
+                    console.warn("🟢 AFTER handleSessionNotification for sources");
+                }
+                catch (error) {
+                    console.error("🔴 ERROR in handleSessionNotification for sources", error);
+                }
+                return;
+            }
             // Validate the message is an ACP agent outgoing message
             const parseResult = acp.agentOutgoingMessageSchema.safeParse(message);
             if (!parseResult.success) {
@@ -843,18 +921,29 @@ export class HttpTransport {
      * Handle a session notification from the agent
      */
     handleSessionNotification(params) {
+        // Extract update type early to check if this is a sources notification
+        const paramsAny = params;
+        const updateAny = paramsAny.update;
+        const isSourcesNotification = updateAny?.sessionUpdate === "sources";
         // Skip processing if stream has been cancelled/completed
-        if (this.streamComplete) {
+        // BUT always allow sources through - they arrive after stream completion during replay
+        if (this.streamComplete && !isSourcesNotification) {
             logger.debug("Skipping session notification - stream complete/cancelled");
             return;
         }
+        if (this.streamComplete && isSourcesNotification) {
+            console.warn("🟢 Processing sources notification after stream complete");
+        }
         logger.debug("handleSessionNotification called", { params });
         // Extract content from the update
         const paramsExtended = params;
         const update = paramsExtended.update;
         const sessionId = this.currentSessionId || params.sessionId;
-        logger.debug("Update session type", {
+        logger.warn("📥 SSE UPDATE RECEIVED", {
             sessionUpdate: update?.sessionUpdate,
+            hasUpdate: !!update,
+            isInReplayMode: this.isInReplayMode,
+            callbackCount: this.sessionUpdateCallbacks.size,
         });
         // Handle sandbox file changes
         // Type assertion needed because TypeScript doesn't recognize this as a valid session update type
@@ -1064,9 +1153,18 @@ export class HttpTransport {
                 isReplay,
                 isInReplayMode: this.isInReplayMode,
             });
-            // During replay, notify directly since there's no active receive() consumer
+            // During replay, handle specially
             if (isReplay || this.isInReplayMode) {
-                this.notifySessionUpdate(sessionUpdate);
+                // If no callbacks registered yet, queue for late-subscribing hooks
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    logger.debug("Queueing tool_call for late-subscribing callbacks", {
+                        toolCallId: toolCall.id,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                }
+                else {
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
                 // Queue tool call as a chunk for ordered processing during live streaming
@@ -1273,8 +1371,16 @@ export class HttpTransport {
             // Check if this is replay (tool_call_update doesn't have isReplay in _meta,
             // but we can check if we're in replay mode)
             if (this.isInReplayMode) {
-                // During replay, notify directly since there's no active receive() consumer
-                this.notifySessionUpdate(sessionUpdate);
+                // If no callbacks registered yet, queue for late-subscribing hooks
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    logger.debug("Queueing tool_call_update for late-subscribing callbacks", {
+                        toolCallId: toolCallUpdate.id,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                }
+                else {
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
                 // Queue tool call update as a chunk for ordered processing
@@ -1406,7 +1512,16 @@ export class HttpTransport {
             };
             // During replay, notify directly; otherwise queue for ordered processing
             if (this.isInReplayMode) {
-                this.notifySessionUpdate(sessionUpdate);
+                // If no callbacks registered yet, queue for late-subscribing hooks
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    logger.debug("Queueing tool_output for late-subscribing callbacks", {
+                        toolCallId: toolOutput.id,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                }
+                else {
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
                 // Queue tool output as a chunk for ordered processing
@@ -1433,21 +1548,56 @@ export class HttpTransport {
             update.sessionUpdate === "sources") {
             // Sources notification - citation sources from tool calls
             const sourcesUpdate = update;
-            logger.debug("Received sources notification", {
+            console.warn("🔵 SOURCES in handleSessionNotification", {
                 sourcesCount: sourcesUpdate.sources.length,
+                isInReplayMode: this.isInReplayMode,
+                callbackCount: this.sessionUpdateCallbacks.size,
+                firstSourceId: sourcesUpdate.sources[0]?.id,
+                firstSourceToolCallId: sourcesUpdate.sources[0]?.toolCallId,
             });
-            // Create a sources chunk for the message queue
-            const sourcesChunk = {
+            // Create a sources session update
+            const sessionUpdate = {
                 type: "sources",
+                sessionId,
+                status: "active",
                 sources: sourcesUpdate.sources,
             };
-            // Queue for ordered processing
-            const resolver = this.chunkResolvers.shift();
-            if (resolver) {
-                resolver(sourcesChunk);
+            // During replay, handle sources specially
+            if (this.isInReplayMode) {
+                // If no callbacks are registered yet (React hooks haven't subscribed),
+                // queue the sources to be replayed when they do subscribe
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    console.warn("🔵 QUEUEING sources for late-subscribing callbacks", {
+                        sourcesCount: sourcesUpdate.sources.length,
+                        queueLengthBefore: this.pendingReplayUpdates.length,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                    console.warn("🔵 Queue length after:", {
+                        queueLengthAfter: this.pendingReplayUpdates.length,
+                    });
+                }
+                else {
+                    console.warn("🔵 NOTIFYING sources immediately (callbacks registered)", {
+                        sourcesCount: sourcesUpdate.sources.length,
+                        callbackCount: this.sessionUpdateCallbacks.size,
+                    });
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
-                this.messageQueue.push(sourcesChunk);
+                // Create a sources chunk for the message queue
+                const sourcesChunk = {
+                    type: "sources",
+                    sources: sourcesUpdate.sources,
+                };
+                // Queue for ordered processing
+                const resolver = this.chunkResolvers.shift();
+                if (resolver) {
+                    resolver(sourcesChunk);
+                }
+                else {
+                    this.messageQueue.push(sourcesChunk);
+                }
             }
         }
         else if (update?.sessionUpdate === "agent_message_chunk") {
@@ -1539,8 +1689,18 @@ export class HttpTransport {
                             timestamp: new Date().toISOString(),
                         },
                     };
-                    // Notify as a complete message (for session replay or initial message)
-                    this.notifySessionUpdate(messageSessionUpdate);
+                    // During replay, queue messages along with tool_calls and sources
+                    // so everything is processed in the correct order
+                    if (isReplay && this.sessionUpdateCallbacks.size === 0) {
+                        logger.debug("Queueing assistant message for late-subscribing callbacks", {
+                            textLength: contentObj.text.length,
+                        });
+                        this.pendingReplayUpdates.push(messageSessionUpdate);
+                    }
+                    else {
+                        // Notify as a complete message (for session replay or initial message)
+                        this.notifySessionUpdate(messageSessionUpdate);
+                    }
                 }
             }
             // Send session update for:
@@ -1554,6 +1714,11 @@ export class HttpTransport {
         else if (update?.sessionUpdate === "user_message_chunk") {
             // Handle user message chunks (could be from replay or new messages)
             logger.debug("Received user_message_chunk", { update });
+            // Check if this is a replay
+            const isReplay = update._meta &&
+                typeof update._meta === "object" &&
+                "isReplay" in update._meta &&
+                update._meta.isReplay === true;
             const content = update.content;
             if (content && typeof content === "object") {
                 const contentObj = content;
@@ -1575,13 +1740,28 @@ export class HttpTransport {
                             timestamp: new Date().toISOString(),
                         },
                     };
-                    logger.debug("Notifying session update for user message");
-                    this.notifySessionUpdate(sessionUpdate);
+                    // During replay, queue messages along with tool_calls and sources
+                    // so everything is processed in the correct order
+                    if (isReplay && this.sessionUpdateCallbacks.size === 0) {
+                        logger.debug("Queueing user message for late-subscribing callbacks", {
+                            textLength: contentObj.text.length,
+                        });
+                        this.pendingReplayUpdates.push(sessionUpdate);
+                    }
+                    else {
+                        logger.debug("Notifying session update for user message");
+                        this.notifySessionUpdate(sessionUpdate);
+                    }
                 }
             }
         }
         else {
             // Handle other session updates
+            logger.warn("⚠️ UNHANDLED SESSION UPDATE - falling through to generic", {
+                sessionUpdate: update?.sessionUpdate,
+                updateKeys: update ? Object.keys(update) : [],
+                hasSources: update && "sources" in update,
+            });
             const sessionUpdate = {
                 type: "generic",
                 sessionId,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@townco/ui",
-  "version": "0.1.109",
+  "version": "0.1.111",
   "type": "module",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -49,7 +49,7 @@
     "@radix-ui/react-slot": "^1.2.4",
     "@radix-ui/react-tabs": "^1.1.13",
     "@radix-ui/react-tooltip": "^1.2.8",
-    "@townco/core": "0.0.87",
+    "@townco/core": "0.0.89",
     "@types/mdast": "^4.0.4",
     "@uiw/react-json-view": "^2.0.0-alpha.39",
     "class-variance-authority": "^0.7.1",
@@ -67,7 +67,7 @@
     "zustand": "^5.0.8"
   },
   "devDependencies": {
-    "@townco/tsconfig": "0.1.106",
+    "@townco/tsconfig": "0.1.108",
     "@types/node": "^24.10.0",
     "@types/react": "^19.2.2",
     "@types/unist": "^3.0.3",