Merge pull request #427 from Kilo-Org/catrielmuller/351-sub-agent-progress

fix(vscode): display sub-agent task tool list in webview

Author: catrielmuller

packages/kilo-vscode/src/KiloProvider.ts

@@ -268,6 +268,9 @@ export class KiloProvider implements vscode.WebviewViewProvider {
         case "loadMessages":
           await this.handleLoadMessages(message.sessionID)
           break
+        case "syncSession":
+          await this.handleSyncSession(message.sessionID)
+          break
         case "loadSessions":
           await this.handleLoadSessions()
           break
@@ -584,6 +587,44 @@ export class KiloProvider implements vscode.WebviewViewProvider {
     }
   }
 
+  /**
+   * Handle syncing a child session (e.g. spawned by the task tool).
+   * Tracks the session for SSE events and fetches its messages.
+   */
+  private async handleSyncSession(sessionID: string): Promise<void> {
+    if (!this.httpClient) return
+    if (this.trackedSessionIds.has(sessionID)) return
+
+    this.trackedSessionIds.add(sessionID)
+
+    try {
+      const workspaceDir = this.getWorkspaceDirectory(sessionID)
+      const messagesData = await this.httpClient.getMessages(sessionID, workspaceDir)
+
+      const messages = messagesData.map((m) => ({
+        id: m.info.id,
+        sessionID: m.info.sessionID,
+        role: m.info.role,
+        parts: m.parts,
+        createdAt: new Date(m.info.time.created).toISOString(),
+        cost: m.info.cost,
+        tokens: m.info.tokens,
+      }))
+
+      for (const message of messages) {
+        this.connectionService.recordMessageSessionId(message.id, message.sessionID)
+      }
+
+      this.postMessage({
+        type: "messagesLoaded",
+        sessionID,
+        messages,
+      })
+    } catch (err) {
+      console.error("[Kilo New] KiloProvider: Failed to sync child session:", err)
+    }
+  }
+
   /**
    * Handle loading all sessions.
    */

packages/kilo-vscode/webview-ui/src/App.tsx

@@ -50,19 +50,14 @@ export const DataBridge: Component<{ children: any }> = (props) => {
 
   const data = createMemo(() => {
     const id = session.currentSessionID()
+    const msgs = session.allMessages()
+    const parts = session.allParts()
     return {
-      session: session.sessions().map((s) => ({ ...s, id: s.id, role: "user" as const })),
+      session: session.sessions().map((s) => ({ ...s, id: s.id, role: "user" as const })) as unknown as any[],
       session_status: {} as Record<string, any>,
       session_diff: {} as Record<string, any[]>,
-      message: id ? { [id]: session.messages() as unknown as SDKMessage[] } : {},
-      part: id
-        ? Object.fromEntries(
-            session
-              .messages()
-              .map((msg) => [msg.id, session.getParts(msg.id) as unknown as SDKPart[]])
-              .filter(([, parts]) => (parts as SDKPart[]).length > 0),
-          )
-        : {},
+      message: msgs as unknown as Record<string, SDKMessage[]>,
+      part: parts as unknown as Record<string, SDKPart[]>,
       permission: id ? { [id]: session.permissions() as unknown as any[] } : {},
     }
   })
@@ -71,8 +66,12 @@ export const DataBridge: Component<{ children: any }> = (props) => {
     session.respondToPermission(input.permissionID, input.response)
   }
 
+  const sync = (sessionID: string) => {
+    session.syncSession(sessionID)
+  }
+
   return (
-    <DataProvider data={data()} directory="" onPermissionRespond={respond}>
+    <DataProvider data={data()} directory="" onPermissionRespond={respond} onSyncSession={sync}>
       {props.children}
     </DataProvider>
   )

packages/kilo-vscode/webview-ui/src/context/session.tsx

@@ -102,6 +102,12 @@ interface SessionContextValue {
   // Messages for current session
   messages: Accessor<Message[]>
 
+  // All messages keyed by sessionID (includes child sessions)
+  allMessages: () => Record<string, Message[]>
+
+  // All parts keyed by messageID (includes child sessions)
+  allParts: () => Record<string, Part[]>
+
   // Parts for a specific message
   getParts: (messageID: string) => Part[]
 
@@ -143,6 +149,7 @@ interface SessionContextValue {
   selectSession: (id: string) => void
   deleteSession: (id: string) => void
   renameSession: (id: string, title: string) => void
+  syncSession: (sessionID: string) => void
 }
 
 const SessionContext = createContext<SessionContextValue>()
@@ -790,6 +797,14 @@ export const SessionProvider: ParentComponent = (props) => {
     return store.parts[messageID] || []
   }
 
+  const allMessages = () => store.messages
+
+  const allParts = () => store.parts
+
+  function syncSession(sessionID: string) {
+    vscode.postMessage({ type: "syncSession", sessionID })
+  }
+
   const todos = () => {
     const id = currentSessionID()
     return id ? store.todos[id] || [] : []
@@ -865,6 +880,8 @@ export const SessionProvider: ParentComponent = (props) => {
     selectAgent,
     getSessionAgent: (sessionID: string) => store.agentSelections[sessionID] ?? defaultAgent(),
     getSessionModel: (sessionID: string) => store.modelSelections[sessionID] ?? provider.defaultSelection(),
+    allMessages,
+    allParts,
     sendMessage,
     abort,
     compact,
@@ -877,6 +894,7 @@ export const SessionProvider: ParentComponent = (props) => {
     selectSession,
     deleteSession,
     renameSession,
+    syncSession,
   }
 
   return <SessionContext.Provider value={value}>{props.children}</SessionContext.Provider>

packages/kilo-vscode/webview-ui/src/types/messages.ts

@@ -729,6 +729,11 @@ export interface ResetAllSettingsRequest {
   type: "resetAllSettings"
 }
 
+export interface SyncSessionRequest {
+  type: "syncSession"
+  sessionID: string
+}
+
 // Agent Manager worktree messages
 export interface CreateWorktreeSessionRequest {
   type: "agentManager.createWorktreeSession"
@@ -773,6 +778,7 @@ export type WebviewMessage =
   | UpdateConfigMessage
   | RequestNotificationSettingsMessage
   | ResetAllSettingsRequest
+  | SyncSessionRequest
   | CreateWorktreeSessionRequest
 
 // ============================================