11import type { BackgroundAgent } from "@/types" ;
2- import type { TaskStartedEvent , TaskProgressEvent , TaskNotificationEvent } from "@/types" ;
2+ import type { TaskProgressEvent , TaskNotificationEvent } from "@/types" ;
33
44type Listener = ( sessionId : string ) => void ;
55
6+ interface AsyncAgentInfo {
7+ toolUseId : string ;
8+ agentId : string ;
9+ description : string ;
10+ outputFile : string ;
11+ }
12+
613/**
714 * Shared store for event-driven background agent tracking.
815 *
9- * Both useClaude (active session) and BackgroundSessionStore (backgrounded
10- * sessions) push SDK task lifecycle events here. The useBackgroundAgents
11- * hook subscribes via useSyncExternalStore.
16+ * Only tracks BACKGROUND (async) agents — foreground agents use the
17+ * existing parentToolMap/subagentSteps system in useClaude.
18+ *
19+ * Registration: from tool_result with isAsync: true (definitive async signal).
20+ * Updates: from task_progress events (live metrics) and task-notification XML
21+ * in user messages (completion).
1222 */
1323class BackgroundAgentStore {
1424 private agents = new Map < string , Map < string , BackgroundAgent > > ( ) ;
@@ -43,33 +53,37 @@ class BackgroundAgentStore {
4353 this . notify ( sessionId ) ;
4454 }
4555
46- handleTaskStarted ( sessionId : string , event : TaskStartedEvent ) : void {
47- if ( ! event . tool_use_id ) return ;
56+ /**
57+ * Register a background agent from tool_result with isAsync: true.
58+ * This is the only entry point — task_started fires for ALL agents
59+ * (foreground + background), so we don't use it.
60+ */
61+ registerAsyncAgent ( sessionId : string , info : AsyncAgentInfo ) : void {
4862 let map = this . agents . get ( sessionId ) ;
4963 if ( ! map ) {
5064 map = new Map ( ) ;
5165 this . agents . set ( sessionId , map ) ;
5266 }
53- // Don't overwrite if already exists (e.g. duplicate event)
54- if ( map . has ( event . tool_use_id ) ) return ;
67+ if ( map . has ( info . toolUseId ) ) return ;
5568
56- map . set ( event . tool_use_id , {
57- agentId : event . task_id ,
58- description : event . description ,
69+ map . set ( info . toolUseId , {
70+ agentId : info . agentId ,
71+ description : info . description ,
5972 prompt : "" ,
60- outputFile : "" ,
73+ outputFile : info . outputFile ,
6174 launchedAt : Date . now ( ) ,
6275 status : "running" ,
6376 activity : [ ] ,
64- toolUseId : event . tool_use_id ,
65- taskId : event . task_id ,
77+ toolUseId : info . toolUseId ,
78+ taskId : info . agentId ,
6679 } ) ;
6780 this . notify ( sessionId ) ;
6881 }
6982
7083 handleTaskProgress ( sessionId : string , event : TaskProgressEvent ) : void {
7184 if ( ! event . tool_use_id ) return ;
7285 const agent = this . agents . get ( sessionId ) ?. get ( event . tool_use_id ) ;
86+ // Only update agents we've registered (i.e. background agents)
7387 if ( ! agent ) return ;
7488
7589 agent . usage = {
@@ -109,6 +123,37 @@ class BackgroundAgentStore {
109123 this . notify ( sessionId ) ;
110124 }
111125
126+ /**
127+ * Parse task completion from user text messages containing <task-notification> XML.
128+ * The SDK delivers task completion as a user text message, NOT as a system event.
129+ */
130+ handleUserMessage ( sessionId : string , content : string ) : void {
131+ if ( ! content . includes ( "<task-notification>" ) ) return ;
132+
133+ const toolUseId = extractXmlTag ( content , "tool-use-id" ) ;
134+ if ( ! toolUseId ) return ;
135+
136+ const agent = this . agents . get ( sessionId ) ?. get ( toolUseId ) ;
137+ if ( ! agent ) return ;
138+
139+ const status = extractXmlTag ( content , "status" ) ;
140+ agent . status = status === "completed" ? "completed" : "error" ;
141+ agent . result = extractXmlTag ( content , "summary" ) || undefined ;
142+
143+ const tokens = extractXmlTag ( content , "total_tokens" ) ;
144+ const tools = extractXmlTag ( content , "tool_uses" ) ;
145+ const duration = extractXmlTag ( content , "duration_ms" ) ;
146+ if ( tokens ) {
147+ agent . usage = {
148+ totalTokens : parseInt ( tokens , 10 ) || 0 ,
149+ toolUses : parseInt ( tools ?? "0" , 10 ) || 0 ,
150+ durationMs : parseInt ( duration ?? "0" , 10 ) || 0 ,
151+ } ;
152+ }
153+
154+ this . notify ( sessionId ) ;
155+ }
156+
112157 dismissAgent ( sessionId : string , agentId : string ) : void {
113158 const map = this . agents . get ( sessionId ) ;
114159 if ( ! map ) return ;
@@ -122,4 +167,11 @@ class BackgroundAgentStore {
122167 }
123168}
124169
170+ /** Extract text content of an XML-like tag from a string. */
171+ function extractXmlTag ( text : string , tag : string ) : string | null {
172+ const re = new RegExp ( `<${ tag } ${ tag } ) ;
173+ const match = re . exec ( text ) ;
174+ return match ? match [ 1 ] . trim ( ) : null ;
175+ }
176+
125177export const bgAgentStore = new BackgroundAgentStore ( ) ;
0 commit comments