@lleverage-ai/agent-sdk 0.0.3 → 0.0.5

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,62 @@ 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 checkpoint with messages AND the pending interrupt.
+                        // The normal completion path saves messages at the end of generate(),
+                        // but when interrupted we return early. Without saving here, resume()
+                        // cannot find the checkpoint (or finds one without messages/interrupt).
+                        if (effectiveGenOptions.threadId && options.checkpointer) {
+                            const checkpointThreadId = forkedSessionId ?? effectiveGenOptions.threadId;
+                            const finalMessages = [
+                                ...messages,
+                                { role: "assistant", content: response.text },
+                            ];
+                            const savedCheckpoint = await saveCheckpoint(checkpointThreadId, finalMessages, startStep + response.steps.length);
+                            if (savedCheckpoint) {
+                                const withInterrupt = updateCheckpoint(savedCheckpoint, {
+                                    pendingInterrupt: interrupt,
+                                });
+                                await options.checkpointer.save(withInterrupt);
+                                threadCheckpoints.set(checkpointThreadId, withInterrupt);
+                            }
+                        }
+                        // 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,20 +1494,54 @@ 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)
                     if (isInterruptSignal(error)) {
                         const interrupt = error.interrupt;
-                        // Save the interrupt to checkpoint
+                        // Save checkpoint with messages AND the pending interrupt (catch-block path).
+                        // lastBuiltMessages holds the messages built before generateText was called.
                         if (effectiveGenOptions.threadId && options.checkpointer) {
-                            const checkpoint = await options.checkpointer.load(effectiveGenOptions.threadId);
-                            if (checkpoint) {
-                                const updatedCheckpoint = updateCheckpoint(checkpoint, {
+                            const savedCheckpoint = await saveCheckpoint(effectiveGenOptions.threadId, lastBuiltMessages ?? [], 0);
+                            if (savedCheckpoint) {
+                                const withInterrupt = updateCheckpoint(savedCheckpoint, {
                                     pendingInterrupt: interrupt,
                                 });
-                                await options.checkpointer.save(updatedCheckpoint);
+                                await options.checkpointer.save(withInterrupt);
+                                threadCheckpoints.set(effectiveGenOptions.threadId, withInterrupt);
                             }
                         }
                         // Emit InterruptRequested hook
@@ -1521,9 +1707,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 +1728,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 +1741,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 +1825,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 +1904,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 +1925,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 +1944,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 +2007,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 +2132,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 +2155,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 +2168,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 +2291,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 +2315,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 +2328,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 +2396,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 +2503,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
@@ -2131,11 +2547,11 @@ export function createAgent(options) {
             if (interrupt.id !== interruptId) {
                 throw new Error(`Cannot resume: interrupt ID mismatch. Expected ${interrupt.id}, got ${interruptId}`);
             }
-            // Store the response for the tool wrapper to use
-            const toolCallId = interrupt.toolCallId;
-            if (toolCallId) {
-                pendingResponses.set(toolCallId, response);
-            }
+            // Store the response keyed by interrupt ID (format: "int_<toolCallId>").
+            // The interrupt() function in the tool wrapper looks up responses using
+            // this exact key format, so we must use interrupt.id — NOT the raw
+            // toolCallId which would never match.
+            pendingResponses.set(interrupt.id, response);
             // Emit InterruptResolved hook
             const interruptResolvedHooks = effectiveHooks?.InterruptResolved ?? [];
             if (interruptResolvedHooks.length > 0) {
@@ -2194,20 +2610,19 @@ export function createAgent(options) {
                         });
                     }
                     catch (error) {
-                        toolResultOutput = {
-                            error: true,
-                            message: error instanceof Error ? error.message : String(error),
-                        };
+                        toolResultOutput = `Tool execution failed: ${error instanceof Error ? error.message : String(error)}`;
                     }
                 }
                 else {
                     // Denied: Create a synthetic denial result
-                    toolResultOutput = {
-                        denied: true,
-                        message: `Tool "${interrupt.toolName}" was denied by user${approvalResponse.reason ? `: ${approvalResponse.reason}` : ""}`,
-                    };
+                    toolResultOutput = `Tool "${interrupt.toolName}" was denied by user${approvalResponse.reason ? `: ${approvalResponse.reason}` : ""}`;
                 }
-                // Build the tool result message
+                // Build the tool result message with proper ToolResultOutput format.
+                // The AI SDK validates messages against modelMessageSchema which
+                // requires output to be { type: 'text', value } or { type: 'json', value }.
+                const approvalOutput = typeof toolResultOutput === "string"
+                    ? { type: "text", value: toolResultOutput }
+                    : { type: "json", value: toolResultOutput };
                 const toolResultMessage = {
                     role: "tool",
                     content: [
@@ -2215,7 +2630,7 @@ export function createAgent(options) {
                             type: "tool-result",
                             toolCallId: interrupt.toolCallId,
                             toolName: interrupt.toolName,
-                            output: toolResultOutput,
+                            output: approvalOutput,
                         },
                     ],
                 };
@@ -2232,7 +2647,7 @@ export function createAgent(options) {
                 });
                 await options.checkpointer.save(updatedCheckpoint);
                 // Clean up the response from our maps
-                pendingResponses.delete(interrupt.toolCallId);
+                pendingResponses.delete(interrupt.id);
                 approvalDecisions.delete(interrupt.toolCallId);
                 // Continue generation from the updated checkpoint
                 return agent.generate({
@@ -2241,19 +2656,128 @@ export function createAgent(options) {
                     prompt: undefined,
                 });
             }
-            // For custom interrupts, the tool's interrupt() call will receive the response
-            // Store it and re-run generation - the tool will get the response
-            // Clear the interrupt from checkpoint before continuing
-            const updatedCheckpoint = updateCheckpoint(checkpoint, {
-                pendingInterrupt: undefined,
-            });
-            await options.checkpointer.save(updatedCheckpoint);
-            // Continue generation - the wrapped tool will pick up the response from pendingResponses
-            return agent.generate({
-                threadId,
-                ...genOptions,
-                prompt: undefined,
-            });
+            // For custom interrupts (e.g. ask_user), manually execute the tool so
+            // that its interrupt() call receives the stored response deterministically.
+            // Re-running generate() would rely on the model re-calling the same tool,
+            // but tool call IDs change each generation so pendingResponses would never
+            // be matched.
+            {
+                const customToolCallId = interrupt.toolCallId;
+                const customToolName = interrupt.toolName;
+                if (!customToolCallId || !customToolName) {
+                    throw new Error("Cannot resume custom interrupt: missing toolCallId or toolName on interrupt");
+                }
+                // Build the assistant message with the original tool call
+                const customAssistantMessage = {
+                    role: "assistant",
+                    content: [
+                        {
+                            type: "tool-call",
+                            toolCallId: customToolCallId,
+                            toolName: customToolName,
+                            input: interrupt.request,
+                        },
+                    ],
+                };
+                // Collect all tools (same pattern as the approval path above)
+                const customTools = { ...coreTools };
+                for (const plugin of options.plugins ?? []) {
+                    if (plugin.tools && typeof plugin.tools !== "function") {
+                        Object.assign(customTools, plugin.tools);
+                    }
+                }
+                const customMcpTools = mcpManager.getToolSet();
+                Object.assign(customTools, customMcpTools);
+                const customTool = customTools[customToolName];
+                if (!customTool || !customTool.execute) {
+                    throw new Error(`Cannot resume: tool "${customToolName}" not found or has no execute function`);
+                }
+                let customToolResult;
+                try {
+                    // Execute the tool, providing an interrupt function that returns
+                    // the stored user response. This mirrors what happens inside the
+                    // permission-mode tool wrapper when pendingResponses has a match.
+                    customToolResult = await customTool.execute(interrupt.request, {
+                        toolCallId: customToolCallId,
+                        messages: checkpoint.messages,
+                        abortSignal: genOptions?.signal,
+                        interrupt: async (request) => {
+                            // First call: return the stored user response (mirrors the
+                            // permission-mode wrapper when pendingResponses has a match).
+                            if (pendingResponses.has(interrupt.id)) {
+                                const stored = pendingResponses.get(interrupt.id);
+                                pendingResponses.delete(interrupt.id);
+                                return stored;
+                            }
+                            // Subsequent calls: no stored response — throw InterruptSignal
+                            // so the tool can pause again (e.g. multi-step wizards).
+                            const newInterruptData = createInterrupt({
+                                id: `int_${customToolCallId}`,
+                                threadId,
+                                type: "custom",
+                                toolCallId: customToolCallId,
+                                toolName: customToolName,
+                                request,
+                                step: checkpoint.step,
+                            });
+                            throw new InterruptSignal(newInterruptData);
+                        },
+                    });
+                }
+                catch (executeError) {
+                    if (isInterruptSignal(executeError)) {
+                        // Tool threw another interrupt — persist it and return interrupted
+                        const newInterrupt = executeError.interrupt;
+                        const reInterruptCheckpoint = updateCheckpoint(checkpoint, {
+                            pendingInterrupt: newInterrupt,
+                        });
+                        await options.checkpointer.save(reInterruptCheckpoint);
+                        threadCheckpoints.set(threadId, reInterruptCheckpoint);
+                        return {
+                            status: "interrupted",
+                            interrupt: newInterrupt,
+                            partial: undefined,
+                        };
+                    }
+                    customToolResult = `Tool execution failed: ${executeError instanceof Error ? executeError.message : String(executeError)}`;
+                }
+                // Build tool result message with proper ToolResultOutput format.
+                const customOutput = typeof customToolResult === "string"
+                    ? { type: "text", value: customToolResult }
+                    : { type: "json", value: customToolResult };
+                const customToolResultMessage = {
+                    role: "tool",
+                    content: [
+                        {
+                            type: "tool-result",
+                            toolCallId: customToolCallId,
+                            toolName: customToolName,
+                            output: customOutput,
+                        },
+                    ],
+                };
+                // Update checkpoint with tool call + result, clear interrupt
+                const customUpdatedMessages = [
+                    ...checkpoint.messages,
+                    customAssistantMessage,
+                    customToolResultMessage,
+                ];
+                const customUpdatedCheckpoint = updateCheckpoint(checkpoint, {
+                    messages: customUpdatedMessages,
+                    pendingInterrupt: undefined,
+                    step: checkpoint.step + 1,
+                });
+                await options.checkpointer.save(customUpdatedCheckpoint);
+                threadCheckpoints.set(threadId, customUpdatedCheckpoint);
+                // Clean up
+                pendingResponses.delete(interrupt.id);
+                // Continue generation from the updated checkpoint
+                return agent.generate({
+                    threadId,
+                    ...genOptions,
+                    prompt: undefined,
+                });
+            }
         },
         async dispose() {
             // Kill all running background tasks