@lleverage-ai/agent-sdk 0.0.3 → 0.0.4

package/dist/agent.js

@@ -42,6 +42,46 @@ class InterruptSignal extends Error {
 function isInterruptSignal(error) {
     return error instanceof InterruptSignal;
 }
+/**
+ * Outermost tool wrapper that intercepts flow-control signals.
+ *
+ * When a tool throws `InterruptSignal`, this wrapper catches it before the AI
+ * SDK can, stores it in the shared `signalState`, and returns a placeholder
+ * string. Combined with a custom `stopWhen` condition, this cleanly stops
+ * generation and allows `generate()` to inspect `signalState` in the normal
+ * return path (not the catch block).
+ *
+ * @internal
+ */
+function wrapToolsWithSignalCatching(tools, signalState) {
+    const wrapped = {};
+    for (const [name, toolDef] of Object.entries(tools)) {
+        if (!toolDef.execute) {
+            wrapped[name] = toolDef;
+            continue;
+        }
+        const originalExecute = toolDef.execute;
+        wrapped[name] = {
+            ...toolDef,
+            execute: async (input, options) => {
+                try {
+                    return await originalExecute.call(toolDef, input, options);
+                }
+                catch (error) {
+                    if (isInterruptSignal(error)) {
+                        if (signalState.interrupt) {
+                            throw error; // Already have a signal — let AI SDK handle this one
+                        }
+                        signalState.interrupt = error;
+                        return "[Interrupt requested]";
+                    }
+                    throw error;
+                }
+            },
+        };
+    }
+    return wrapped;
+}
 /**
  * File edit tool names that get auto-approved in acceptEdits mode.
  * @internal
@@ -457,17 +497,21 @@ function wrapToolsWithHooks(tools, hookRegistration, agent, sessionId) {
                     return output;
                 }
                 catch (error) {
-                    // Invoke PostToolUseFailure hooks
-                    if (hookRegistration?.PostToolUseFailure?.length) {
-                        const failureInput = {
-                            hook_event_name: "PostToolUseFailure",
-                            session_id: sessionId,
-                            cwd: process.cwd(),
-                            tool_name: name,
-                            tool_input: input,
-                            error: error instanceof Error ? error : new Error(String(error)),
-                        };
-                        await invokeMatchingHooks(hookRegistration.PostToolUseFailure, name, failureInput, toolUseId, agent);
+                    // Skip PostToolUseFailure for flow-control signals — these are not
+                    // actual failures but intentional control flow (interrupt).
+                    if (!isInterruptSignal(error)) {
+                        // Invoke PostToolUseFailure hooks
+                        if (hookRegistration?.PostToolUseFailure?.length) {
+                            const failureInput = {
+                                hook_event_name: "PostToolUseFailure",
+                                session_id: sessionId,
+                                cwd: process.cwd(),
+                                tool_name: name,
+                                tool_input: input,
+                                error: error instanceof Error ? error : new Error(String(error)),
+                            };
+                            await invokeMatchingHooks(hookRegistration.PostToolUseFailure, name, failureInput, toolUseId, agent);
+                        }
                     }
                     throw error;
                 }
@@ -541,7 +585,8 @@ export function createAgent(options) {
     // Process middleware to get hooks (middleware hooks come before explicit hooks)
     const middleware = options.middleware ?? [];
     const middlewareHooks = applyMiddleware(middleware);
-    const mergedHooks = mergeHooks(middlewareHooks, options.hooks);
+    const pluginHooks = (options.plugins ?? []).filter((p) => p.hooks).map((p) => p.hooks);
+    const mergedHooks = mergeHooks(middlewareHooks, ...pluginHooks, options.hooks);
     // Create options with merged hooks for all hook lookups
     const effectiveHooks = mergedHooks;
     // Permission mode (mutable for setPermissionMode)
@@ -572,6 +617,18 @@ export function createAgent(options) {
     }
     // Initialize task manager for background task tracking
     const taskManager = new TaskManager();
+    // Background task completion options
+    const waitForBackgroundTasks = options.waitForBackgroundTasks ?? true;
+    const formatTaskCompletion = options.formatTaskCompletion ??
+        ((task) => {
+            const command = task.metadata?.command ?? "unknown command";
+            return `[Background task completed: ${task.id}]\nCommand: ${command}\nOutput:\n${task.result ?? "(no output)"}`;
+        });
+    const formatTaskFailure = options.formatTaskFailure ??
+        ((task) => {
+            const command = task.metadata?.command ?? "unknown command";
+            return `[Background task failed: ${task.id}]\nCommand: ${command}\nError: ${task.error ?? "Unknown error"}`;
+        });
     // Determine plugin loading mode
     // Track whether it was explicitly set to distinguish from default
     const explicitPluginLoading = options.pluginLoading !== undefined;
@@ -828,6 +885,7 @@ export function createAgent(options) {
         // only advertises tools the agent will actually expose
         const filteredTools = filterToolsByAllowed((() => {
             const allTools = { ...coreTools };
+            Object.assign(allTools, runtimeTools);
             Object.assign(allTools, mcpManager.getToolSet());
             if (toolRegistry) {
                 Object.assign(allTools, toolRegistry.getLoadedTools());
@@ -878,10 +936,14 @@ export function createAgent(options) {
         // Build prompt using prompt builder
         return promptBuilder.build(context);
     };
-    // Helper to get current active tools (core + MCP + dynamically loaded from registry)
+    // Runtime tools added/removed dynamically by plugins at runtime
+    const runtimeTools = {};
+    // Helper to get current active tools (core + runtime + MCP + dynamically loaded from registry)
     const getActiveToolSet = (threadId) => {
         // Start with core tools
         const allTools = { ...coreTools };
+        // Add runtime tools (added by plugins at runtime)
+        Object.assign(allTools, runtimeTools);
         // Add MCP tools from plugin registrations
         const mcpTools = mcpManager.getToolSet();
         Object.assign(allTools, mcpTools);
@@ -910,6 +972,8 @@ export function createAgent(options) {
     const getActiveToolSetWithStreaming = (streamingContext, threadId, step) => {
         // Start with core tools
         const allTools = { ...coreTools };
+        // Add runtime tools (added by plugins at runtime)
+        Object.assign(allTools, runtimeTools);
         // Process plugins - invoke function-based tools with streaming context
         for (const plugin of options.plugins ?? []) {
             if (plugin.tools) {
@@ -1222,6 +1286,34 @@ export function createAgent(options) {
             usage: step.usage,
         }));
     }
+    /**
+     * Get the next actionable task prompt from the background task queue.
+     *
+     * Waits for a task to reach terminal state, skips already-consumed and killed
+     * tasks, formats the result as a prompt, removes the task, and returns it.
+     * Returns null when no more tasks need processing.
+     */
+    async function getNextTaskPrompt() {
+        while (taskManager.hasActiveTasks() || taskManager.hasTerminalTasks()) {
+            const completedTask = await taskManager.waitForNextCompletion();
+            // Dedup: skip if already consumed via task_output tool
+            if (!taskManager.getTask(completedTask.id)) {
+                continue;
+            }
+            // Skip killed tasks (user already knows)
+            if (completedTask.status === "killed") {
+                taskManager.removeTask(completedTask.id);
+                continue;
+            }
+            // Format as follow-up prompt
+            const prompt = completedTask.status === "completed"
+                ? formatTaskCompletion(completedTask)
+                : formatTaskFailure(completedTask);
+            taskManager.removeTask(completedTask.id);
+            return prompt;
+        }
+        return null;
+    }
     const agent = {
         id,
         options,
@@ -1251,9 +1343,17 @@ export function createAgent(options) {
                     lastBuiltMessages = messages;
                     const maxSteps = options.maxSteps ?? 10;
                     const startStep = checkpoint?.step ?? 0;
+                    // Shared signal state: flow-control signals (interrupt) thrown by tools
+                    // are caught by the outermost wrapper and stored here. A custom
+                    // stopWhen condition stops generation after the current step completes.
+                    const signalState = {};
                     // Build initial params - use active tools (core + dynamically loaded + task)
-                    // Apply hooks AFTER adding task tool so task tool is also wrapped
-                    const activeTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    // Apply hooks AFTER adding task tool so task tool is also wrapped.
+                    // Then wrap with signal catching as the outermost layer so that
+                    // InterruptSignal is intercepted before the AI SDK can catch it and
+                    // convert it to a tool-error result.
+                    const hookedTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    const activeTools = wrapToolsWithSignalCatching(hookedTools, signalState);
                     // Build prompt context and generate system prompt
                     const promptContext = buildPromptContext(messages, effectiveGenOptions.threadId);
                     const systemPrompt = getSystemPrompt(promptContext);
@@ -1268,6 +1368,9 @@ export function createAgent(options) {
                         providerOptions: effectiveGenOptions.providerOptions,
                         headers: effectiveGenOptions.headers,
                     };
+                    // Stop condition: stop when an interrupt signal was caught, OR when
+                    // the step count reaches maxSteps (whichever comes first).
+                    const signalStopCondition = () => signalState.interrupt != null;
                     // Execute generation
                     const response = await generateText({
                         model: retryState.currentModel,
@@ -1278,13 +1381,53 @@ export function createAgent(options) {
                         temperature: initialParams.temperature,
                         stopSequences: initialParams.stopSequences,
                         abortSignal: initialParams.abortSignal,
-                        stopWhen: stepCountIs(maxSteps),
+                        stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
                         // Passthrough AI SDK options
                         output: effectiveGenOptions.output,
                         // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
                         providerOptions: initialParams.providerOptions,
                         headers: initialParams.headers,
                     });
+                    // Check for intercepted interrupt signal (cooperative path)
+                    if (signalState.interrupt) {
+                        const interrupt = signalState.interrupt.interrupt;
+                        // Save the interrupt to checkpoint
+                        if (effectiveGenOptions.threadId && options.checkpointer) {
+                            const checkpoint = await options.checkpointer.load(effectiveGenOptions.threadId);
+                            if (checkpoint) {
+                                const updatedCheckpoint = updateCheckpoint(checkpoint, {
+                                    pendingInterrupt: interrupt,
+                                });
+                                await options.checkpointer.save(updatedCheckpoint);
+                            }
+                        }
+                        // Emit InterruptRequested hook
+                        const interruptRequestedHooks = effectiveHooks?.InterruptRequested ?? [];
+                        if (interruptRequestedHooks.length > 0) {
+                            const hookInput = {
+                                hook_event_name: "InterruptRequested",
+                                session_id: effectiveGenOptions.threadId ?? "default",
+                                cwd: process.cwd(),
+                                interrupt_id: interrupt.id,
+                                interrupt_type: interrupt.type,
+                                tool_call_id: interrupt.toolCallId,
+                                tool_name: interrupt.toolName,
+                                request: interrupt.request,
+                            };
+                            await invokeHooksWithTimeout(interruptRequestedHooks, hookInput, null, agent);
+                        }
+                        // Return interrupted result with partial results from the response
+                        const interruptedResult = {
+                            status: "interrupted",
+                            interrupt,
+                            partial: {
+                                text: response.text,
+                                steps: mapSteps(response.steps),
+                                usage: response.usage,
+                            },
+                        };
+                        return interruptedResult;
+                    }
                     // Only access output if an output schema was provided
                     // (accessing response.output throws AI_NoOutputGeneratedError otherwise)
                     let output;
@@ -1342,7 +1485,39 @@ export function createAgent(options) {
                             finalResult = updatedResult;
                         }
                     }
-                    return finalResult;
+                    // --- Background task completion loop ---
+                    if (!waitForBackgroundTasks) {
+                        return finalResult;
+                    }
+                    // When checkpointing is active, the checkpoint already contains the
+                    // full conversation history (saved above). Passing explicit messages
+                    // would cause buildMessages() to load checkpoint messages AND append
+                    // the same messages again, causing duplication.
+                    const hasCheckpointing = !!(effectiveGenOptions.threadId && options.checkpointer);
+                    let lastResult = finalResult;
+                    let runningMessages = hasCheckpointing
+                        ? []
+                        : [...messages, { role: "assistant", content: finalResult.text }];
+                    let followUpPrompt = await getNextTaskPrompt();
+                    while (followUpPrompt !== null) {
+                        lastResult = await agent.generate({
+                            ...genOptions,
+                            prompt: followUpPrompt,
+                            messages: hasCheckpointing ? undefined : runningMessages,
+                        });
+                        if (lastResult.status === "interrupted") {
+                            return lastResult;
+                        }
+                        if (!hasCheckpointing) {
+                            runningMessages = [
+                                ...runningMessages,
+                                { role: "user", content: followUpPrompt },
+                                { role: "assistant", content: lastResult.text },
+                            ];
+                        }
+                        followUpPrompt = await getNextTaskPrompt();
+                    }
+                    return lastResult;
                 }
                 catch (error) {
                     // Check if this is an InterruptSignal (new interrupt system)
@@ -1521,9 +1696,13 @@ export function createAgent(options) {
                     const { messages, checkpoint } = await buildMessages(effectiveGenOptions);
                     const maxSteps = options.maxSteps ?? 10;
                     const startStep = checkpoint?.step ?? 0;
+                    // Signal state for cooperative signal catching in streaming mode
+                    const signalState = {};
                     // Build initial params - use active tools (core + dynamically loaded + task)
-                    // Apply hooks AFTER adding task tool so task tool is also wrapped
-                    const activeTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    // Apply hooks AFTER adding task tool so task tool is also wrapped.
+                    // Then wrap with signal catching as the outermost layer.
+                    const hookedTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    const activeTools = wrapToolsWithSignalCatching(hookedTools, signalState);
                     // Build prompt context and generate system prompt
                     const promptContext = buildPromptContext(messages, effectiveGenOptions.threadId);
                     const systemPrompt = getSystemPrompt(promptContext);
@@ -1538,6 +1717,9 @@ export function createAgent(options) {
                         providerOptions: effectiveGenOptions.providerOptions,
                         headers: effectiveGenOptions.headers,
                     };
+                    // Stop condition: stop when an interrupt signal was caught, OR when
+                    // the step count reaches maxSteps.
+                    const signalStopCondition = () => signalState.interrupt != null;
                     // Execute stream
                     const response = streamText({
                         model: retryState.currentModel,
@@ -1548,7 +1730,7 @@ export function createAgent(options) {
                         temperature: initialParams.temperature,
                         stopSequences: initialParams.stopSequences,
                         abortSignal: initialParams.abortSignal,
-                        stopWhen: stepCountIs(maxSteps),
+                        stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
                         // Passthrough AI SDK options
                         output: genOptions.output,
                         // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
@@ -1632,7 +1814,36 @@ export function createAgent(options) {
                         await invokeHooksWithTimeout(postGenerateHooks, postGenerateInput, null, agent);
                         // Note: updatedResult is not applied for streaming since the stream has already been sent
                     }
-                    // Success - break out of retry loop
+                    // --- Background task completion loop ---
+                    if (!waitForBackgroundTasks || signalState.interrupt) {
+                        return;
+                    }
+                    const hasCheckpointing = !!(effectiveGenOptions.threadId && options.checkpointer);
+                    let currentMessages = hasCheckpointing
+                        ? []
+                        : [...messages, { role: "assistant", content: text }];
+                    let followUpPrompt = await getNextTaskPrompt();
+                    while (followUpPrompt !== null) {
+                        const followUpGen = agent.stream({
+                            ...genOptions,
+                            prompt: followUpPrompt,
+                            messages: hasCheckpointing ? undefined : currentMessages,
+                        });
+                        let followUpText = "";
+                        for await (const part of followUpGen) {
+                            yield part;
+                            if (part.type === "text-delta")
+                                followUpText += part.text;
+                        }
+                        if (!hasCheckpointing) {
+                            currentMessages = [
+                                ...currentMessages,
+                                { role: "user", content: followUpPrompt },
+                                { role: "assistant", content: followUpText },
+                            ];
+                        }
+                        followUpPrompt = await getNextTaskPrompt();
+                    }
                     return;
                 }
                 catch (error) {
@@ -1682,9 +1893,13 @@ export function createAgent(options) {
                     const { messages, checkpoint } = await buildMessages(effectiveGenOptions);
                     const maxSteps = options.maxSteps ?? 10;
                     const startStep = checkpoint?.step ?? 0;
+                    // Signal state for cooperative signal catching in streaming mode
+                    const signalState = {};
                     // Build initial params - use active tools (core + dynamically loaded + task)
-                    // Apply hooks AFTER adding task tool so task tool is also wrapped
-                    const activeTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    // Apply hooks AFTER adding task tool so task tool is also wrapped.
+                    // Then wrap with signal catching as the outermost layer.
+                    const hookedTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    const activeTools = wrapToolsWithSignalCatching(hookedTools, signalState);
                     // Build prompt context and generate system prompt
                     const promptContext = buildPromptContext(messages, effectiveGenOptions.threadId);
                     const systemPrompt = getSystemPrompt(promptContext);
@@ -1699,11 +1914,18 @@ export function createAgent(options) {
                         providerOptions: effectiveGenOptions.providerOptions,
                         headers: effectiveGenOptions.headers,
                     };
+                    // Capture currentModel for use in the callback closure
+                    const modelToUse = retryState.currentModel;
+                    // Stop condition: stop when an interrupt signal was caught, OR when
+                    // the step count reaches maxSteps.
+                    const signalStopCondition = () => signalState.interrupt != null;
                     // Track step count for incremental checkpointing
                     let currentStepCount = 0;
-                    // Execute stream
+                    // Execute streamText OUTSIDE createUIMessageStream so errors propagate
+                    // to the retry loop (if streamText throws synchronously on creation,
+                    // e.g. rate limit, the catch block handles retry/fallback).
                     const result = streamText({
-                        model: retryState.currentModel,
+                        model: modelToUse,
                         system: initialParams.system,
                         messages: initialParams.messages,
                         tools: initialParams.tools,
@@ -1711,7 +1933,7 @@ export function createAgent(options) {
                         temperature: initialParams.temperature,
                         stopSequences: initialParams.stopSequences,
                         abortSignal: initialParams.abortSignal,
-                        stopWhen: stepCountIs(maxSteps),
+                        stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
                         // Passthrough AI SDK options
                         output: effectiveGenOptions.output,
                         // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
@@ -1774,10 +1996,93 @@ export function createAgent(options) {
                             }
                         },
                     });
-                    // Return AI SDK compatible response for use with useChat
-                    // Note: Not passing originalMessages since ModelMessage[] != UIMessage[]
-                    // The AI SDK will reconstruct messages from the stream
-                    return result.toUIMessageStreamResponse();
+                    // Use createUIMessageStream to control stream lifecycle for background task follow-ups
+                    const stream = createUIMessageStream({
+                        execute: async ({ writer }) => {
+                            // Merge initial generation into the stream
+                            writer.merge(result.toUIMessageStream());
+                            // Wait for initial generation to complete to get final text
+                            const text = await result.text;
+                            // --- Background task completion loop ---
+                            if (waitForBackgroundTasks) {
+                                // Track accumulated steps for checkpoint saves
+                                const initialSteps = await result.steps;
+                                let accumulatedStepCount = initialSteps.length;
+                                let currentMessages = [
+                                    ...messages,
+                                    { role: "assistant", content: text },
+                                ];
+                                let followUpPrompt = await getNextTaskPrompt();
+                                while (followUpPrompt !== null) {
+                                    // Stream follow-up generation into the same writer
+                                    const followUpResult = streamText({
+                                        model: modelToUse,
+                                        system: initialParams.system,
+                                        messages: [
+                                            ...currentMessages,
+                                            { role: "user", content: followUpPrompt },
+                                        ],
+                                        tools: initialParams.tools,
+                                        maxOutputTokens: initialParams.maxTokens,
+                                        temperature: initialParams.temperature,
+                                        stopSequences: initialParams.stopSequences,
+                                        abortSignal: initialParams.abortSignal,
+                                        stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
+                                        // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
+                                        providerOptions: initialParams.providerOptions,
+                                        headers: initialParams.headers,
+                                    });
+                                    writer.merge(followUpResult.toUIMessageStream());
+                                    const followUpText = await followUpResult.text;
+                                    currentMessages = [
+                                        ...currentMessages,
+                                        { role: "user", content: followUpPrompt },
+                                        { role: "assistant", content: followUpText },
+                                    ];
+                                    // --- Post-completion bookkeeping for follow-ups ---
+                                    const followUpSteps = await followUpResult.steps;
+                                    accumulatedStepCount += followUpSteps.length;
+                                    // Checkpoint save
+                                    if (effectiveGenOptions.threadId && options.checkpointer) {
+                                        await saveCheckpoint(effectiveGenOptions.threadId, currentMessages, startStep + accumulatedStepCount);
+                                    }
+                                    // Context manager update
+                                    const followUpUsage = await followUpResult.usage;
+                                    if (options.contextManager?.updateUsage && followUpUsage) {
+                                        options.contextManager.updateUsage({
+                                            inputTokens: followUpUsage.inputTokens,
+                                            outputTokens: followUpUsage.outputTokens,
+                                            totalTokens: followUpUsage.totalTokens,
+                                        });
+                                    }
+                                    // PostGenerate hooks
+                                    const followUpPostGenerateHooks = effectiveHooks?.PostGenerate ?? [];
+                                    if (followUpPostGenerateHooks.length > 0) {
+                                        const followUpFinishReason = await followUpResult.finishReason;
+                                        const followUpHookResult = {
+                                            status: "complete",
+                                            text: followUpText,
+                                            usage: followUpUsage,
+                                            finishReason: followUpFinishReason,
+                                            output: undefined,
+                                            steps: mapSteps(followUpSteps),
+                                        };
+                                        const followUpPostGenerateInput = {
+                                            hook_event_name: "PostGenerate",
+                                            session_id: effectiveGenOptions.threadId ?? "default",
+                                            cwd: process.cwd(),
+                                            options: effectiveGenOptions,
+                                            result: followUpHookResult,
+                                        };
+                                        await invokeHooksWithTimeout(followUpPostGenerateHooks, followUpPostGenerateInput, null, agent);
+                                    }
+                                    followUpPrompt = await getNextTaskPrompt();
+                                }
+                            }
+                        },
+                    });
+                    // Convert the stream to a Response
+                    return createUIMessageStreamResponse({ stream });
                 }
                 catch (error) {
                     // Normalize error to AgentError
@@ -1816,9 +2121,13 @@ export function createAgent(options) {
                     const { messages, checkpoint } = await buildMessages(effectiveGenOptions);
                     const maxSteps = options.maxSteps ?? 10;
                     const startStep = checkpoint?.step ?? 0;
+                    // Signal state for cooperative signal catching in streaming mode
+                    const signalState = {};
                     // Build initial params - use active tools (core + dynamically loaded + task)
-                    // Apply hooks AFTER adding task tool so task tool is also wrapped
-                    const activeTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    // Apply hooks AFTER adding task tool so task tool is also wrapped.
+                    // Then wrap with signal catching as the outermost layer.
+                    const hookedTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSet(effectiveGenOptions.threadId)), effectiveGenOptions.threadId);
+                    const activeTools = wrapToolsWithSignalCatching(hookedTools, signalState);
                     // Build prompt context and generate system prompt
                     const promptContext = buildPromptContext(messages, effectiveGenOptions.threadId);
                     const systemPrompt = getSystemPrompt(promptContext);
@@ -1835,6 +2144,9 @@ export function createAgent(options) {
                     };
                     // Track step count for incremental checkpointing
                     let currentStepCount = 0;
+                    // Stop condition: stop when an interrupt signal was caught, OR when
+                    // the step count reaches maxSteps.
+                    const signalStopCondition = () => signalState.interrupt != null;
                     // Execute stream
                     const result = streamText({
                         model: retryState.currentModel,
@@ -1845,7 +2157,7 @@ export function createAgent(options) {
                         temperature: initialParams.temperature,
                         stopSequences: initialParams.stopSequences,
                         abortSignal: initialParams.abortSignal,
-                        stopWhen: stepCountIs(maxSteps),
+                        stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
                         // Passthrough AI SDK options
                         output: effectiveGenOptions.output,
                         // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
@@ -1968,9 +2280,13 @@ export function createAgent(options) {
                             }
                             // Create streaming context for tools
                             const streamingContext = { writer };
+                            // Signal state for cooperative signal catching in streaming mode
+                            const signalState = {};
                             // Build tools with streaming context and task tool
-                            // Apply hooks AFTER adding task tool so task tool is also wrapped
-                            const streamingTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSetWithStreaming(streamingContext, effectiveGenOptions.threadId), streamingContext), effectiveGenOptions.threadId);
+                            // Apply hooks AFTER adding task tool so task tool is also wrapped.
+                            // Then wrap with signal catching as the outermost layer.
+                            const hookedStreamingTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSetWithStreaming(streamingContext, effectiveGenOptions.threadId), streamingContext), effectiveGenOptions.threadId);
+                            const streamingTools = wrapToolsWithSignalCatching(hookedStreamingTools, signalState);
                             // Build prompt context and generate system prompt
                             const promptContext = buildPromptContext(messages, effectiveGenOptions.threadId);
                             const systemPrompt = getSystemPrompt(promptContext);
@@ -1988,6 +2304,9 @@ export function createAgent(options) {
                             };
                             // Track step count for incremental checkpointing
                             let currentStepCount = 0;
+                            // Stop condition: stop when a flow-control signal was caught, OR when
+                            // the step count reaches maxSteps.
+                            const signalStopCondition = () => signalState.interrupt != null;
                             // Execute stream
                             const result = streamText({
                                 model: modelToUse,
@@ -1998,7 +2317,7 @@ export function createAgent(options) {
                                 temperature: initialParams.temperature,
                                 stopSequences: initialParams.stopSequences,
                                 abortSignal: initialParams.abortSignal,
-                                stopWhen: stepCountIs(maxSteps),
+                                stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
                                 // Passthrough AI SDK options
                                 output: effectiveGenOptions.output,
                                 // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
@@ -2066,6 +2385,84 @@ export function createAgent(options) {
                             });
                             // Merge the streamText output into the UI message stream
                             writer.merge(result.toUIMessageStream());
+                            // Wait for initial generation to complete to get final text
+                            const text = await result.text;
+                            // --- Background task completion loop ---
+                            if (waitForBackgroundTasks) {
+                                // Track accumulated steps for checkpoint saves
+                                const initialSteps = await result.steps;
+                                let accumulatedStepCount = initialSteps.length;
+                                let currentMessages = [
+                                    ...messages,
+                                    { role: "assistant", content: text },
+                                ];
+                                let followUpPrompt = await getNextTaskPrompt();
+                                while (followUpPrompt !== null) {
+                                    // Stream follow-up generation into the same writer
+                                    const followUpResult = streamText({
+                                        model: modelToUse,
+                                        system: initialParams.system,
+                                        messages: [
+                                            ...currentMessages,
+                                            { role: "user", content: followUpPrompt },
+                                        ],
+                                        tools: initialParams.tools,
+                                        maxOutputTokens: initialParams.maxTokens,
+                                        temperature: initialParams.temperature,
+                                        stopSequences: initialParams.stopSequences,
+                                        abortSignal: initialParams.abortSignal,
+                                        stopWhen: [signalStopCondition, stepCountIs(maxSteps)],
+                                        // biome-ignore lint/suspicious/noExplicitAny: Type cast needed for AI SDK compatibility
+                                        providerOptions: initialParams.providerOptions,
+                                        headers: initialParams.headers,
+                                    });
+                                    writer.merge(followUpResult.toUIMessageStream());
+                                    const followUpText = await followUpResult.text;
+                                    currentMessages = [
+                                        ...currentMessages,
+                                        { role: "user", content: followUpPrompt },
+                                        { role: "assistant", content: followUpText },
+                                    ];
+                                    // --- Post-completion bookkeeping for follow-ups ---
+                                    const followUpSteps = await followUpResult.steps;
+                                    accumulatedStepCount += followUpSteps.length;
+                                    // Checkpoint save
+                                    if (effectiveGenOptions.threadId && options.checkpointer) {
+                                        await saveCheckpoint(effectiveGenOptions.threadId, currentMessages, startStep + accumulatedStepCount);
+                                    }
+                                    // Context manager update
+                                    const followUpUsage = await followUpResult.usage;
+                                    if (options.contextManager?.updateUsage && followUpUsage) {
+                                        options.contextManager.updateUsage({
+                                            inputTokens: followUpUsage.inputTokens,
+                                            outputTokens: followUpUsage.outputTokens,
+                                            totalTokens: followUpUsage.totalTokens,
+                                        });
+                                    }
+                                    // PostGenerate hooks
+                                    const followUpPostGenerateHooks = effectiveHooks?.PostGenerate ?? [];
+                                    if (followUpPostGenerateHooks.length > 0) {
+                                        const followUpFinishReason = await followUpResult.finishReason;
+                                        const followUpHookResult = {
+                                            status: "complete",
+                                            text: followUpText,
+                                            usage: followUpUsage,
+                                            finishReason: followUpFinishReason,
+                                            output: undefined,
+                                            steps: mapSteps(followUpSteps),
+                                        };
+                                        const followUpPostGenerateInput = {
+                                            hook_event_name: "PostGenerate",
+                                            session_id: effectiveGenOptions.threadId ?? "default",
+                                            cwd: process.cwd(),
+                                            options: effectiveGenOptions,
+                                            result: followUpHookResult,
+                                        };
+                                        await invokeHooksWithTimeout(followUpPostGenerateHooks, followUpPostGenerateInput, null, agent);
+                                    }
+                                    followUpPrompt = await getNextTaskPrompt();
+                                }
+                            }
                         },
                     });
                     // Convert the stream to a Response
@@ -2095,6 +2492,14 @@ export function createAgent(options) {
         getActiveTools() {
             return getActiveToolSet();
         },
+        addRuntimeTools(tools) {
+            Object.assign(runtimeTools, tools);
+        },
+        removeRuntimeTools(toolNames) {
+            for (const name of toolNames) {
+                delete runtimeTools[name];
+            }
+        },
         loadTools(toolNames) {
             if (!toolRegistry) {
                 // No registry in eager mode - all tools already loaded