@illuma-ai/agents 1.0.89 → 1.0.93

package/dist/esm/tools/ToolNode.mjs

@@ -5,7 +5,7 @@ import 'nanoid';
 import '../messages/core.mjs';
 import { processToolOutput } from '../utils/toonFormat.mjs';
 import { RunnableCallable } from '../utils/run.mjs';
-import 'js-tiktoken';
+import 'ai-tokenizer';
 import 'zod-to-json-schema';
 import { safeDispatchCustomEvent } from '../utils/events.mjs';
 
@@ -32,7 +32,7 @@ function extractStringContent(content) {
             if (typeof block === 'string') {
                 textParts.push(block);
             }
-            else if (block && typeof block === 'object') {
+            else if (block != null && typeof block === 'object') {
                 // Handle {type: 'text', text: '...'} format
                 const obj = block;
                 if (obj.type === 'text' && typeof obj.text === 'string') {
@@ -60,6 +60,8 @@ class ToolNode extends RunnableCallable {
     toolCallStepIds;
     errorHandler;
     toolUsageCount;
+    /** Maps toolCallId → turn captured in runTool, used by handleRunToolCompletions */
+    toolCallTurns = new Map();
     /** Tool registry for filtering (lazy computation of programmatic maps) */
     toolRegistry;
     /** Cached programmatic tools (computed once on first PTC call) */
@@ -68,13 +70,15 @@ class ToolNode extends RunnableCallable {
     sessions;
     /** When true, dispatches ON_TOOL_EXECUTE events instead of invoking tools directly */
     eventDrivenMode = false;
-    /** Tool definitions for event-driven mode */
-    toolDefinitions;
     /** Agent ID for event-driven mode */
     agentId;
+    /** Tool names that bypass event dispatch and execute directly (e.g., graph-managed handoff tools) */
+    directToolNames;
     /** HITL tool approval configuration */
     toolApprovalConfig;
-    constructor({ tools, toolMap, name, tags, errorHandler, toolCallStepIds, handleToolErrors, loadRuntimeTools, toolRegistry, sessions, eventDrivenMode, toolDefinitions, agentId, toolApprovalConfig, }) {
+    /** Buffer for recovering truncated tool call arguments from streaming data */
+    streamingToolCallBuffer;
+    constructor({ tools, toolMap, name, tags, errorHandler, toolCallStepIds, handleToolErrors, loadRuntimeTools, toolRegistry, sessions, eventDrivenMode, agentId, directToolNames, streamingToolCallBuffer, toolApprovalConfig, }) {
         super({ name, tags, func: (input, config) => this.run(input, config) });
         this.toolMap = toolMap ?? new Map(tools.map((tool) => [tool.name, tool]));
         this.toolCallStepIds = toolCallStepIds;
@@ -85,8 +89,9 @@ class ToolNode extends RunnableCallable {
         this.toolRegistry = toolRegistry;
         this.sessions = sessions;
         this.eventDrivenMode = eventDrivenMode ?? false;
-        this.toolDefinitions = toolDefinitions;
         this.agentId = agentId;
+        this.directToolNames = directToolNames;
+        this.streamingToolCallBuffer = streamingToolCallBuffer;
         this.toolApprovalConfig = toolApprovalConfig;
     }
     /**
@@ -128,8 +133,13 @@ class ToolNode extends RunnableCallable {
             return false;
         }
         const { defaultPolicy, overrides, executionContext } = this.toolApprovalConfig;
-        // Scheduled executions bypass all approval checks
-        if (executionContext === 'scheduled') {
+        // Scheduled and browser executions bypass all approval checks
+        if (executionContext === 'scheduled' || executionContext === 'browser') {
+            return false;
+        }
+        // Graph-managed tools (handoff/transfer tools) bypass HITL approval —
+        // these are internal routing mechanisms, not user-facing tool executions
+        if (this.directToolNames?.has(toolName)) {
             return false;
         }
         // Determine the effective policy for this tool
@@ -187,8 +197,19 @@ class ToolNode extends RunnableCallable {
             }
             const turn = this.toolUsageCount.get(call.name) ?? 0;
             this.toolUsageCount.set(call.name, turn + 1);
-            const args = call.args;
+            if (call.id != null && call.id !== '') {
+                this.toolCallTurns.set(call.id, turn);
+            }
+            let args = call.args;
             const stepId = this.toolCallStepIds?.get(call.id);
+            // Recover truncated tool call arguments from the streaming buffer.
+            // When max_tokens truncates a tool call mid-JSON, parsePartialJson may lose
+            // content that was already streamed to the UI. The buffer has the raw accumulated
+            // arg string, so we can extract field values that were dropped.
+            if (call.id && this.streamingToolCallBuffer?.has(call.id)) {
+                args = this.recoverTruncatedArgs(call.name, call.id, args);
+                this.streamingToolCallBuffer.clear(call.id);
+            }
             // Build invoke params - LangChain extracts non-schema fields to config.toolCall
             let invokeParams = {
                 ...call,
@@ -217,11 +238,14 @@ class ToolNode extends RunnableCallable {
              * Each file uses its own session_id (supporting multi-session file tracking).
              * Both session_id and _injected_files are injected directly to invokeParams
              * (not inside args) so they bypass Zod schema validation and reach config.toolCall.
+             *
+             * session_id is always injected when available (even without tracked files)
+             * so the CodeExecutor can fall back to the /files endpoint for session continuity.
              */
             if (call.name === Constants.EXECUTE_CODE ||
                 call.name === Constants.PROGRAMMATIC_TOOL_CALLING) {
                 const codeSession = this.sessions?.get(Constants.EXECUTE_CODE);
-                if (codeSession?.session_id) {
+                if (codeSession != null && codeSession.session_id !== '') {
                     /**
                      * Always inject session_id so retries reuse the same workspace.
                      * Also inject file refs when files exist from previous executions.
@@ -259,14 +283,17 @@ class ToolNode extends RunnableCallable {
                     // Human denied the tool call - return a denial message
                     return new ToolMessage({
                         status: 'error',
-                        content: `Tool call "${call.name}" was denied by the user.${approvalResponse.feedback ? ` Reason: ${approvalResponse.feedback}` : ''} Please acknowledge the denial and proceed without executing this tool.`,
+                        content: `Tool call "${call.name}" was denied by the user.${approvalResponse.feedback != null && approvalResponse.feedback !== '' ? ` Reason: ${approvalResponse.feedback}` : ''} Please acknowledge the denial and proceed without executing this tool.`,
                         name: call.name,
                         tool_call_id: call.id ?? '',
                     });
                 }
                 // Human approved - optionally use modified args
                 if (approvalResponse.modifiedArgs) {
-                    invokeParams = { ...invokeParams, args: approvalResponse.modifiedArgs };
+                    invokeParams = {
+                        ...invokeParams,
+                        args: approvalResponse.modifiedArgs,
+                    };
                 }
             }
             const output = await tool.invoke(invokeParams, config);
@@ -370,22 +397,264 @@ class ToolNode extends RunnableCallable {
         }
     }
     /**
-     * Execute all tool calls via ON_TOOL_EXECUTE event dispatch.
-     * Used in event-driven mode where the host handles actual tool execution.
+     * Recover truncated tool call arguments using the raw streaming buffer.
+     *
+     * When parsePartialJson drops content (e.g., a large "content" field truncated
+     * at max_tokens), this method extracts the field value from the raw accumulated
+     * arg string and merges it back into the parsed args object.
+     *
+     * This is generic — it checks ALL string fields in the raw buffer, not just
+     * content_tool fields. Any tool with a truncated string value benefits.
      */
-    async executeViaEvent(toolCalls, config, 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    input) {
+    /**
+     * Recover truncated tool call arguments using the raw streaming buffer.
+     *
+     * Strategy:
+     * 1. If args are completely empty → try parsePartialJson on the raw buffer
+     * 2. Otherwise → field-level recovery: extract missing fields from raw buffer
+     *
+     * @param toolName - Tool name (for logging)
+     * @param toolCallId - The tool call ID
+     * @param args - The parsed args (potentially incomplete from parsePartialJson)
+     * @returns The args with recovered fields merged in
+     */
+    /* eslint-disable @typescript-eslint/no-explicit-any */
+    recoverTruncatedArgs(toolName, toolCallId, args) {
+        /* eslint-enable @typescript-eslint/no-explicit-any */
+        if (!this.streamingToolCallBuffer)
+            return args;
+        const rawArgs = this.streamingToolCallBuffer.getRawArgs(toolCallId);
+        if (!rawArgs)
+            return args;
+        const rawLen = rawArgs.length;
+        // If args is completely empty (JSON parsing aborted before any key),
+        // attempt to parse the raw accumulated string
+        if (args == null ||
+            (typeof args === 'object' && Object.keys(args).length === 0)) {
+            try {
+                // eslint-disable-next-line @typescript-eslint/no-require-imports
+                const { parsePartialJson } = require('@langchain/core/utils/json');
+                const recovered = parsePartialJson(rawArgs);
+                if (recovered != null &&
+                    typeof recovered === 'object' &&
+                    Object.keys(recovered).length > 0) {
+                    // eslint-disable-next-line no-console
+                    console.warn(`[TruncationRecovery] Tool=${toolName}, callId=${toolCallId}, rawBufferLen=${rawLen}, ` +
+                        `parsedFields=[], recoveredFields=[${Object.keys(recovered).join(',')}] (full parse)`);
+                    return recovered;
+                }
+            }
+            catch {
+                // parsePartialJson failed — fall through to field-level extraction
+            }
+        }
+        // Field-level recovery: extract missing fields from the raw buffer
+        const parsedArgs = typeof args === 'object' ? { ...args } : {};
+        const recoveredFields = [];
+        // Extract field names from the raw JSON string
+        const fieldPattern = /"([^"]+)"\s*:/g;
+        let match;
+        const rawFieldNames = [];
+        while ((match = fieldPattern.exec(rawArgs)) !== null) {
+            rawFieldNames.push(match[1]);
+        }
+        for (const fieldName of rawFieldNames) {
+            if (parsedArgs[fieldName] == null ||
+                parsedArgs[fieldName] === '' ||
+                parsedArgs[fieldName] === undefined) {
+                const rawValue = this.streamingToolCallBuffer.extractFieldValue(toolCallId, fieldName);
+                if (rawValue != null && rawValue !== '') {
+                    parsedArgs[fieldName] = rawValue;
+                    recoveredFields.push(fieldName);
+                }
+            }
+        }
+        if (recoveredFields.length > 0) {
+            // eslint-disable-next-line no-console
+            console.warn(`[TruncationRecovery] Tool=${toolName}, callId=${toolCallId}, rawBufferLen=${rawLen}, ` +
+                `parsedFields=[${Object.keys(args ?? {}).join(',')}], recoveredFields=[${recoveredFields.join(',')}]`);
+        }
+        return parsedArgs;
+    }
+    /**
+     * Builds code session context for injection into event-driven tool calls.
+     * Mirrors the session injection logic in runTool() for direct execution.
+     */
+    getCodeSessionContext() {
+        if (!this.sessions) {
+            return undefined;
+        }
+        const codeSession = this.sessions.get(Constants.EXECUTE_CODE);
+        if (!codeSession) {
+            return undefined;
+        }
+        const context = {
+            session_id: codeSession.session_id,
+        };
+        if (codeSession.files && codeSession.files.length > 0) {
+            context.files = codeSession.files.map((file) => ({
+                session_id: file.session_id ?? codeSession.session_id,
+                id: file.id,
+                name: file.name,
+            }));
+        }
+        return context;
+    }
+    /**
+     * Extracts code execution session context from tool results and stores in Graph.sessions.
+     * Mirrors the session storage logic in handleRunToolCompletions for direct execution.
+     */
+    storeCodeSessionFromResults(results, requests) {
+        if (!this.sessions) {
+            return;
+        }
+        for (let i = 0; i < results.length; i++) {
+            const result = results[i];
+            if (result.status !== 'success' || result.artifact == null) {
+                continue;
+            }
+            const request = requests.find((r) => r.id === result.toolCallId);
+            if (request?.name !== Constants.EXECUTE_CODE &&
+                request?.name !== Constants.PROGRAMMATIC_TOOL_CALLING) {
+                continue;
+            }
+            const artifact = result.artifact;
+            if (artifact?.session_id == null || artifact.session_id === '') {
+                continue;
+            }
+            const newFiles = artifact.files ?? [];
+            const existingSession = this.sessions.get(Constants.EXECUTE_CODE);
+            const existingFiles = existingSession?.files ?? [];
+            if (newFiles.length > 0) {
+                const filesWithSession = newFiles.map((file) => ({
+                    ...file,
+                    session_id: artifact.session_id,
+                }));
+                const newFileNames = new Set(filesWithSession.map((f) => f.name));
+                const filteredExisting = existingFiles.filter((f) => !newFileNames.has(f.name));
+                this.sessions.set(Constants.EXECUTE_CODE, {
+                    session_id: artifact.session_id,
+                    files: [...filteredExisting, ...filesWithSession],
+                    lastUpdated: Date.now(),
+                });
+            }
+            else {
+                this.sessions.set(Constants.EXECUTE_CODE, {
+                    session_id: artifact.session_id,
+                    files: existingFiles,
+                    lastUpdated: Date.now(),
+                });
+            }
+        }
+    }
+    /**
+     * Post-processes standard runTool outputs: dispatches ON_RUN_STEP_COMPLETED
+     * and stores code session context. Mirrors the completion handling in
+     * dispatchToolEvents for the event-driven path.
+     *
+     * By handling completions here in graph context (rather than in the
+     * stream consumer via ToolEndHandler), the race between the stream
+     * consumer and graph execution is eliminated.
+     */
+    handleRunToolCompletions(calls, outputs, config) {
+        for (let i = 0; i < calls.length; i++) {
+            const call = calls[i];
+            const output = outputs[i];
+            const turn = this.toolCallTurns.get(call.id) ?? 0;
+            if (isCommand(output)) {
+                continue;
+            }
+            const toolMessage = output;
+            const toolCallId = call.id ?? '';
+            // Skip error ToolMessages when errorHandler already dispatched ON_RUN_STEP_COMPLETED
+            // via handleToolCallErrorStatic. Without this check, errors would be double-dispatched.
+            if (toolMessage.status === 'error' && this.errorHandler != null) {
+                continue;
+            }
+            // Store code session context from tool results
+            if (this.sessions &&
+                (call.name === Constants.EXECUTE_CODE ||
+                    call.name === Constants.PROGRAMMATIC_TOOL_CALLING)) {
+                const artifact = toolMessage.artifact;
+                if (artifact?.session_id != null && artifact.session_id !== '') {
+                    const newFiles = artifact.files ?? [];
+                    const existingSession = this.sessions.get(Constants.EXECUTE_CODE);
+                    const existingFiles = existingSession?.files ?? [];
+                    if (newFiles.length > 0) {
+                        const filesWithSession = newFiles.map((file) => ({
+                            ...file,
+                            session_id: artifact.session_id,
+                        }));
+                        const newFileNames = new Set(filesWithSession.map((f) => f.name));
+                        const filteredExisting = existingFiles.filter((f) => !newFileNames.has(f.name));
+                        this.sessions.set(Constants.EXECUTE_CODE, {
+                            session_id: artifact.session_id,
+                            files: [...filteredExisting, ...filesWithSession],
+                            lastUpdated: Date.now(),
+                        });
+                    }
+                    else {
+                        this.sessions.set(Constants.EXECUTE_CODE, {
+                            session_id: artifact.session_id,
+                            files: existingFiles,
+                            lastUpdated: Date.now(),
+                        });
+                    }
+                }
+            }
+            // Dispatch ON_RUN_STEP_COMPLETED via custom event (same path as dispatchToolEvents)
+            const stepId = this.toolCallStepIds?.get(toolCallId) ?? '';
+            if (!stepId) {
+                continue;
+            }
+            const contentString = typeof toolMessage.content === 'string'
+                ? toolMessage.content
+                : JSON.stringify(toolMessage.content);
+            const tool_call = {
+                args: typeof call.args === 'string'
+                    ? call.args
+                    : JSON.stringify(call.args ?? {}),
+                name: call.name,
+                id: toolCallId,
+                output: contentString,
+                progress: 1,
+            };
+            safeDispatchCustomEvent(GraphEvents.ON_RUN_STEP_COMPLETED, {
+                result: {
+                    id: stepId,
+                    index: turn,
+                    type: 'tool_call',
+                    tool_call,
+                },
+            }, config);
+        }
+    }
+    /**
+     * Dispatches tool calls to the host via ON_TOOL_EXECUTE event and returns raw ToolMessages.
+     * Core logic for event-driven execution, separated from output shaping.
+     */
+    async dispatchToolEvents(toolCalls, config) {
         const requests = toolCalls.map((call) => {
             const turn = this.toolUsageCount.get(call.name) ?? 0;
             this.toolUsageCount.set(call.name, turn + 1);
-            return {
+            // Recover truncated args from streaming buffer (same as runTool path)
+            let args = call.args;
+            if (call.id && this.streamingToolCallBuffer?.has(call.id)) {
+                args = this.recoverTruncatedArgs(call.name, call.id, args);
+                this.streamingToolCallBuffer.clear(call.id);
+            }
+            const request = {
                 id: call.id,
                 name: call.name,
-                args: call.args,
+                args,
                 stepId: this.toolCallStepIds?.get(call.id),
                 turn,
             };
+            if (call.name === Constants.EXECUTE_CODE ||
+                call.name === Constants.PROGRAMMATIC_TOOL_CALLING) {
+                request.codeSessionContext = this.getCodeSessionContext();
+            }
+            return request;
         });
         const results = await new Promise((resolve, reject) => {
             const request = {
@@ -399,10 +668,17 @@ class ToolNode extends RunnableCallable {
             };
             safeDispatchCustomEvent(GraphEvents.ON_TOOL_EXECUTE, request, config);
         });
-        const outputs = results.map((result) => {
+        this.storeCodeSessionFromResults(results, requests);
+        return results.map((result) => {
             const request = requests.find((r) => r.id === result.toolCallId);
             const toolName = request?.name ?? 'unknown';
             const stepId = this.toolCallStepIds?.get(result.toolCallId) ?? '';
+            if (!stepId) {
+                // eslint-disable-next-line no-console
+                console.warn(`[ToolNode] toolCallStepIds missing entry for toolCallId=${result.toolCallId} (tool=${toolName}). ` +
+                    'This indicates a race between the stream consumer and graph execution. ' +
+                    `Map size: ${this.toolCallStepIds?.size ?? 0}`);
+            }
             let toolMessage;
             let contentString;
             if (result.status === 'error') {
@@ -447,16 +723,28 @@ class ToolNode extends RunnableCallable {
             safeDispatchCustomEvent(GraphEvents.ON_RUN_STEP_COMPLETED, runStepCompletedData, config);
             return toolMessage;
         });
+    }
+    /**
+     * Execute all tool calls via ON_TOOL_EXECUTE event dispatch.
+     * Used in event-driven mode where the host handles actual tool execution.
+     */
+    async executeViaEvent(toolCalls, config, 
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    input) {
+        const outputs = await this.dispatchToolEvents(toolCalls, config);
         return (Array.isArray(input) ? outputs : { messages: outputs });
     }
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     async run(input, config) {
+        this.toolCallTurns.clear();
         let outputs;
         if (this.isSendInput(input)) {
-            if (this.eventDrivenMode) {
+            const isDirectTool = this.directToolNames?.has(input.lg_tool_call.name);
+            if (this.eventDrivenMode && isDirectTool !== true) {
                 return this.executeViaEvent([input.lg_tool_call], config, input);
             }
             outputs = [await this.runTool(input.lg_tool_call, config)];
+            this.handleRunToolCompletions([input.lg_tool_call], outputs, config);
         }
         else {
             let messages;
@@ -500,9 +788,26 @@ class ToolNode extends RunnableCallable {
                     !(call.id?.startsWith('srvtoolu_') ?? false));
             }) ?? [];
             if (this.eventDrivenMode && filteredCalls.length > 0) {
-                return this.executeViaEvent(filteredCalls, config, input);
+                if (!this.directToolNames || this.directToolNames.size === 0) {
+                    return this.executeViaEvent(filteredCalls, config, input);
+                }
+                const directCalls = filteredCalls.filter((c) => this.directToolNames.has(c.name));
+                const eventCalls = filteredCalls.filter((c) => !this.directToolNames.has(c.name));
+                const directOutputs = directCalls.length > 0
+                    ? await Promise.all(directCalls.map((call) => this.runTool(call, config)))
+                    : [];
+                if (directCalls.length > 0 && directOutputs.length > 0) {
+                    this.handleRunToolCompletions(directCalls, directOutputs, config);
+                }
+                const eventOutputs = eventCalls.length > 0
+                    ? await this.dispatchToolEvents(eventCalls, config)
+                    : [];
+                outputs = [...directOutputs, ...eventOutputs];
+            }
+            else {
+                outputs = await Promise.all(filteredCalls.map((call) => this.runTool(call, config)));
+                this.handleRunToolCompletions(filteredCalls, outputs, config);
             }
-            outputs = await Promise.all(filteredCalls.map((call) => this.runTool(call, config)));
         }
         if (!outputs.some(isCommand)) {
             return (Array.isArray(input) ? outputs : { messages: outputs });
@@ -622,17 +927,15 @@ function areToolCallsInvoked(message, invokedToolIds) {
     return (message.tool_calls?.every((toolCall) => toolCall.id != null && invokedToolIds.has(toolCall.id)) ?? false);
 }
 function toolsCondition(state, toolNode, invokedToolIds) {
-    const message = Array.isArray(state)
-        ? state[state.length - 1]
-        : state.messages[state.messages.length - 1];
-    if ('tool_calls' in message &&
+    const messages = Array.isArray(state) ? state : state.messages;
+    const message = messages[messages.length - 1];
+    if (message &&
+        'tool_calls' in message &&
         (message.tool_calls?.length ?? 0) > 0 &&
         !areToolCallsInvoked(message, invokedToolIds)) {
         return toolNode;
     }
-    else {
-        return END;
-    }
+    return END;
 }
 
 export { ToolNode, toolsCondition };