@theia/ai-chat 1.76.0-next.18 → 1.76.0-next.23

package/src/browser/agent-delegation-tool.ts

@@ -24,8 +24,11 @@ import {
     ChatChangeEvent,
     ChatRequest,
     ChatResponseContent,
+    ChatContextManager,
     ChatService,
     ChatServiceFactory,
+    ChatSession,
+    ChatSessionStatus,
     ChatToolContext,
     MutableChatModel,
     MutableChatRequestModel,
@@ -57,9 +60,12 @@ export class AgentDelegationTool implements ToolProvider {
             name: AgentDelegationTool.ID,
             description:
                 'Delegate a task or question to a specific AI agent. IMPORTANT: When you delegate a task or question to a specific AI agent using this tool, ' +
-                'remember that each sub-agent operates solely within its specialized capabilities and tools and does not have access to previous conversation context ' +
-                ' or external systems. Therefore, it is crucial to provide all necessary context and detailed information directly within your request to ensure accurate ' +
-                'and effective task completion. ' +
+                'remember that each sub-agent operates solely within its specialized capabilities and tools and, unless you resume a previous delegation via sessionId, ' +
+                'does not have access to previous conversation context or external systems. Therefore, it is crucial to provide all necessary context and detailed ' +
+                'information directly within your request to ensure accurate and effective task completion. ' +
+                'Each result ends with a "[delegation sessionId: <id>]" line. To send a follow-up request to the same agent with its previous conversation retained ' +
+                '(e.g. to request fixes after reviewing its result), pass that id as sessionId. Prefer resuming for follow-ups on the same task; ' +
+                'start a fresh session (omit sessionId) for unrelated tasks, as a resumed session grows with every round. ' +
                 'You may optionally pass a taskContextId to make a specific task context (e.g. a plan) available to the delegated agent via its system prompt.',
             parameters: {
                 type: 'object',
@@ -78,6 +84,12 @@ export class AgentDelegationTool implements ToolProvider {
                         type: 'string',
                         description: 'Optional task context ID to make available to the delegated agent. The agent will see the task context in its system prompt.',
                     },
+                    sessionId: {
+                        type: 'string',
+                        description: 'Optional ID of a previous delegation session to continue, as reported in the "[delegation sessionId: <id>]" line of an earlier result. ' +
+                            'The delegated agent keeps its previous conversation context, so you do not need to repeat context you already sent it. ' +
+                            'Only a session that you delegated yourself can be resumed, and the agentId must match the original delegation.',
+                    },
                 },
                 required: ['agentId', 'prompt'],
             },
@@ -102,7 +114,7 @@ export class AgentDelegationTool implements ToolProvider {
 
         try {
             const args = JSON.parse(arg_string);
-            const { agentId, prompt, taskContextId } = args;
+            const { agentId, prompt, taskContextId, sessionId } = args;
 
             if (!agentId || !prompt) {
                 const errorMsg = 'Both agentId and prompt parameters are required.';
@@ -121,131 +133,183 @@ export class AgentDelegationTool implements ToolProvider {
                 return errorMsg;
             }
 
-            let newSession;
+            let session: ChatSession;
             let childModelDisposable: Disposable | undefined;
-            try {
+            if (sessionId) {
                 const chatService = this.getChatService();
+                // Restore from storage if necessary: delegated sessions are auto-saved, so after a reload the
+                // delegating session's history still carries the id while the session itself is no longer in memory.
+                const existingSession = await chatService.getOrRestoreSession(sessionId);
+                if (!existingSession) {
+                    const errorMsg = `Delegation session '${sessionId}' not found. It may have been deleted. ` +
+                        'Send the request again without sessionId to start a fresh session — the agent will not have any prior context.';
+                    // warn, not error: a stale session id is expected in a normal agentic loop (e.g. the user deleted the session)
+                    this.logger.warn(errorMsg);
+                    return errorMsg;
+                }
+                if (existingSession.parentSessionId !== ctx.request.session.id) {
+                    // Without this, any session id would be accepted, including an unrelated top-level user chat,
+                    // and the delegated request would be injected into it.
+                    const errorMsg = `Delegation session '${sessionId}' was not delegated by this session and cannot be resumed here. ` +
+                        'Only pass a sessionId reported by one of your own earlier delegations, or omit it to start a fresh session.';
+                    this.logger.error(errorMsg);
+                    return errorMsg;
+                }
+                if (ChatSessionStatus.isInProgress(existingSession.model.status)) {
+                    const errorMsg = `Delegation session '${sessionId}' is still processing a previous request. ` +
+                        'Wait for its result before sending a follow-up, or omit sessionId to start a fresh session.';
+                    // warn, not error: expected when the caller issues parallel requests against the same session
+                    this.logger.warn(errorMsg);
+                    return errorMsg;
+                }
+                if (existingSession.pinnedAgent && existingSession.pinnedAgent.id !== agentId) {
+                    const errorMsg = `Delegation session '${sessionId}' belongs to agent '${existingSession.pinnedAgent.id}', not '${agentId}'. ` +
+                        'Pass the matching agentId, or omit sessionId to start a fresh session.';
+                    this.logger.error(errorMsg);
+                    return errorMsg;
+                }
 
-                // Store the current active session to restore it after delegation
-                const currentActiveSession = chatService.getActiveSession();
+                // Deliberately keep rootSessionId/parentSessionId as set at creation time; a resumed
+                // session stays attached to its original delegation hierarchy.
+                if (taskContextId) {
+                    this.setTaskContext(existingSession.model.context, taskContextId);
+                }
 
-                newSession = chatService.createSession(
-                    undefined,
-                    { focus: false },
-                    agent
-                );
+                session = existingSession;
+                // Re-establish bubbling for this request; the previous disposable was disposed when the
+                // original delegation completed, and the parent request/response are new objects now.
+                childModelDisposable = this.setupChildSessionBubbling(existingSession.model as MutableChatModel, ctx.request.session, ctx.response);
+            } else {
+                try {
+                    const chatService = this.getChatService();
 
-                // Set root session ID to enable task context sharing across delegation chains
-                // Root is either the current root (for nested delegation) or current session (for first-level delegation)
-                const rootId = ctx.rootSessionId || ctx.request.session.id;
-                newSession.rootSessionId = rootId;
-                newSession.model.rootSessionId = rootId;
+                    // Store the current active session to restore it after delegation
+                    const currentActiveSession = chatService.getActiveSession();
 
-                // Track the immediate parent (the delegating session) so the UI can render the full
-                // delegation hierarchy, not just a flat list under the root.
-                const parentId = ctx.request.session.id;
-                newSession.parentSessionId = parentId;
-                newSession.model.parentSessionId = parentId;
+                    const newSession = chatService.createSession(
+                        undefined,
+                        { focus: false },
+                        agent
+                    );
 
-                if (taskContextId) {
-                    newSession.model.context.addVariables({
-                        variable: TASK_CONTEXT_VARIABLE,
-                        arg: taskContextId
-                    });
-                }
+                    // Set root session ID to enable task context sharing across delegation chains
+                    // Root is either the current root (for nested delegation) or current session (for first-level delegation)
+                    const rootId = ctx.rootSessionId || ctx.request.session.id;
+                    newSession.rootSessionId = rootId;
+                    newSession.model.rootSessionId = rootId;
+
+                    // Track the immediate parent (the delegating session) so the UI can render the full
+                    // delegation hierarchy, not just a flat list under the root.
+                    const parentId = ctx.request.session.id;
+                    newSession.parentSessionId = parentId;
+                    newSession.model.parentSessionId = parentId;
+
+                    if (taskContextId) {
+                        newSession.model.context.addVariables({
+                            variable: TASK_CONTEXT_VARIABLE,
+                            arg: taskContextId
+                        });
+                    }
 
-                // Immediately restore the original active session to avoid confusing the user
-                if (currentActiveSession) {
-                    chatService.setActiveSession(currentActiveSession.id, { focus: false });
-                }
+                    // Immediately restore the original active session to avoid confusing the user
+                    if (currentActiveSession) {
+                        chatService.setActiveSession(currentActiveSession.id, { focus: false });
+                    }
 
-                // Setup bubbling of child session events to parent session
-                childModelDisposable = this.setupChildSessionBubbling(newSession.model as MutableChatModel, ctx.request.session, ctx.response);
-            } catch (sessionError) {
-                const errorMsg = `Failed to create chat session for agent '${agentId}': ${sessionError instanceof Error ? sessionError.message : sessionError}`;
-                this.logger.error(errorMsg, sessionError);
-                return errorMsg;
+                    // Setup bubbling of child session events to parent session
+                    childModelDisposable = this.setupChildSessionBubbling(newSession.model as MutableChatModel, ctx.request.session, ctx.response);
+                    session = newSession;
+                } catch (sessionError) {
+                    const errorMsg = `Failed to create chat session for agent '${agentId}': ${sessionError instanceof Error ? sessionError.message : sessionError}`;
+                    this.logger.error(errorMsg, sessionError);
+                    return errorMsg;
+                }
             }
 
-            // Send the request
-            const chatRequest: ChatRequest = {
-                text: `@${agentId} ${prompt}`,
-            };
-
-            let response: ChatRequestInvocation | undefined;
+            // Send the request. The `finally` releases the event bubbling on every outcome: a resumed
+            // session is long-lived, so leaking a listener per failed round would forward its events
+            // into stale parent responses.
             try {
-                if (ctx.cancellationToken?.isCancellationRequested) {
-                    return 'Operation cancelled by user';
-                }
+                const chatRequest: ChatRequest = {
+                    text: `@${agentId} ${prompt}`,
+                };
 
-                const chatService = this.getChatService();
-                response = await chatService.sendRequest(
-                    newSession.id,
-                    chatRequest
-                );
-
-                if (ctx.cancellationToken) {
-                    ctx.cancellationToken.onCancellationRequested(
-                        async () => {
-                            if (response) {
-                                ((await response?.requestCompleted) as MutableChatRequestModel).cancel();
-                            }
-                        }
-                    );
-                }
-            } catch (sendError) {
-                const errorMsg = `Failed to send request to agent '${agentId}': ${sendError instanceof Error ? sendError.message : sendError}`;
-                this.logger.error(errorMsg, sendError);
-                return errorMsg;
-            }
+                let response: ChatRequestInvocation | undefined;
+                try {
+                    if (ctx.cancellationToken?.isCancellationRequested) {
+                        return 'Operation cancelled by user';
+                    }
 
-            if (response) {
-                // Store the invocation in the registry so the renderer can access it
-                if (ctx.toolCallId) {
-                    this.pendingDelegations.set(ctx.toolCallId, {
-                        prompt,
-                        invocation: response
-                    });
-                    // Clean up when the delegated session is deleted
                     const chatService = this.getChatService();
-                    const toolCallId = ctx.toolCallId;
-                    const sessionEventListener = chatService.onSessionEvent(event => {
-                        if (event.type === 'deleted' && event.sessionId === newSession.id) {
-                            this.pendingDelegations.delete(toolCallId);
-                            sessionEventListener.dispose();
-                        }
-                    });
+                    response = await chatService.sendRequest(
+                        session.id,
+                        chatRequest
+                    );
+
+                    if (ctx.cancellationToken) {
+                        ctx.cancellationToken.onCancellationRequested(
+                            async () => {
+                                if (response) {
+                                    ((await response?.requestCompleted) as MutableChatRequestModel).cancel();
+                                }
+                            }
+                        );
+                    }
+                } catch (sendError) {
+                    const errorMsg = `Failed to send request to agent '${agentId}': ${sendError instanceof Error ? sendError.message : sendError}`;
+                    this.logger.error(errorMsg, sendError);
+                    return errorMsg;
                 }
 
-                try {
-                    // Wait for completion to return the final result as tool output
-                    const result = await response.responseCompleted;
-                    const filteredContent = result.response.content.filter(c => !ThinkingChatResponseContent.is(c) && !ToolCallChatResponseContent.is(c));
-                    const stringResult = filteredContent
-                        .map(c => ChatResponseContent.hasAsString(c) ? c.asString() : undefined)
-                        .filter((text): text is string => text !== undefined && text !== '')
-                        .join('\n\n');
-
-                    // Clean up the parent-child link after completion (event bubbling is no longer needed)
-                    childModelDisposable?.dispose();
-
-                    // Return the raw text to the top-level Agent, as a tool result
-                    return stringResult;
-                } catch (completionError) {
-                    if (
-                        completionError instanceof Error &&
-                        completionError.message.includes('cancelled')
-                    ) {
-                        return 'Operation cancelled by user';
+                if (response) {
+                    // Store the invocation in the registry so the renderer can access it
+                    if (ctx.toolCallId) {
+                        this.pendingDelegations.set(ctx.toolCallId, {
+                            prompt,
+                            invocation: response
+                        });
+                        // Clean up when the delegated session is deleted
+                        const chatService = this.getChatService();
+                        const toolCallId = ctx.toolCallId;
+                        const sessionEventListener = chatService.onSessionEvent(event => {
+                            if (event.type === 'deleted' && event.sessionId === session.id) {
+                                this.pendingDelegations.delete(toolCallId);
+                                sessionEventListener.dispose();
+                            }
+                        });
                     }
-                    const errorMsg = `Failed to complete response from agent '${agentId}': ${completionError instanceof Error ? completionError.message : completionError}`;
-                    this.logger.error(errorMsg, completionError);
+
+                    try {
+                        // Wait for completion to return the final result as tool output
+                        const result = await response.responseCompleted;
+                        const filteredContent = result.response.content.filter(c => !ThinkingChatResponseContent.is(c) && !ToolCallChatResponseContent.is(c));
+                        const stringResult = filteredContent
+                            .map(c => ChatResponseContent.hasAsString(c) ? c.asString() : undefined)
+                            .filter((text): text is string => text !== undefined && text !== '')
+                            .join('\n\n');
+
+                        // Return the raw text to the top-level Agent, as a tool result. The session id
+                        // allows the caller to send follow-up requests into the same session.
+                        const sessionIdSuffix = `[delegation sessionId: ${session.id}]`;
+                        return stringResult ? `${stringResult}\n\n${sessionIdSuffix}` : sessionIdSuffix;
+                    } catch (completionError) {
+                        if (
+                            completionError instanceof Error &&
+                            completionError.message.includes('cancelled')
+                        ) {
+                            return 'Operation cancelled by user';
+                        }
+                        const errorMsg = `Failed to complete response from agent '${agentId}': ${completionError instanceof Error ? completionError.message : completionError}`;
+                        this.logger.error(errorMsg, completionError);
+                        return errorMsg;
+                    }
+                } else {
+                    const errorMsg = `Delegation to agent '${agentId}' has failed: no response returned.`;
+                    this.logger.error(errorMsg);
                     return errorMsg;
                 }
-            } else {
-                const errorMsg = `Delegation to agent '${agentId}' has failed: no response returned.`;
-                this.logger.error(errorMsg);
-                return errorMsg;
+            } finally {
+                childModelDisposable?.dispose();
             }
         } catch (error) {
             this.logger.error('Failed to delegate to agent', error);
@@ -255,6 +319,24 @@ export class AgentDelegationTool implements ToolProvider {
         }
     }
 
+    /**
+     * Makes the given task context the session's task context, replacing a previously set one.
+     * `addVariables` deduplicates on variable id *and* argument, so adding a different task context
+     * id would otherwise leave the session with both in its system prompt.
+     */
+    private setTaskContext(context: ChatContextManager, taskContextId: string): void {
+        const staleIndices = context.getVariables().reduce<number[]>((indices, variable, index) => {
+            if (variable.variable.id === TASK_CONTEXT_VARIABLE.id && variable.arg !== taskContextId) {
+                indices.push(index);
+            }
+            return indices;
+        }, []);
+        if (staleIndices.length > 0) {
+            context.deleteVariables(...staleIndices);
+        }
+        context.addVariables({ variable: TASK_CONTEXT_VARIABLE, arg: taskContextId });
+    }
+
     /**
      * Sets up all event bubbling from a delegated child session to the parent session:
      * - Interaction forwarding: child interactionNeeded events are forwarded to the parent response