illuma-agents 1.0.2 → 1.0.4

package/dist/esm/graphs/Graph.mjs

@@ -1,29 +1,29 @@
 import { nanoid } from 'nanoid';
 import { concat } from '@langchain/core/utils/stream';
 import { ChatVertexAI } from '@langchain/google-vertexai';
-import { StateGraph, START, END } from '@langchain/langgraph';
-import { dispatchCustomEvent } from '@langchain/core/callbacks/dispatch';
+import { Annotation, messagesStateReducer, StateGraph, START, END } from '@langchain/langgraph';
+import { RunnableLambda } from '@langchain/core/runnables';
 import { SystemMessage, AIMessageChunk, ToolMessage } from '@langchain/core/messages';
 import { GraphNodeKeys, ContentTypes, Providers, GraphEvents, StepTypes } from '../common/enum.mjs';
-import { getChatModelClass, manualToolStreamProviders } from '../llm/providers.mjs';
-import { ToolNode, toolsCondition } from '../tools/ToolNode.mjs';
-import { convertMessagesToContent, formatAnthropicArtifactContent, formatArtifactPayload, modifyDeltaProperties } from '../messages/core.mjs';
+import { convertMessagesToContent, modifyDeltaProperties, formatAnthropicArtifactContent, formatArtifactPayload } from '../messages/core.mjs';
 import { createPruneMessages } from '../messages/prune.mjs';
+import { addCacheControl, addBedrockCacheControl } from '../messages/cache.mjs';
+import { formatContentStrings } from '../messages/content.mjs';
 import { resetIfNotEmpty, joinKeys } from '../utils/graph.mjs';
 import { isOpenAILike, isGoogleLike } from '../utils/llm.mjs';
 import { sleep } from '../utils/run.mjs';
 import 'js-tiktoken/lite';
+import { getChatModelClass, manualToolStreamProviders } from '../llm/providers.mjs';
+import { ToolNode, toolsCondition } from '../tools/ToolNode.mjs';
 import { ChatOpenAI, AzureChatOpenAI } from '../llm/openai/index.mjs';
+import { safeDispatchCustomEvent } from '../utils/events.mjs';
+import { AgentContext } from '../agents/AgentContext.mjs';
 import { createFakeStreamingLLM } from '../llm/fake.mjs';
 
 /* eslint-disable no-console */
 // src/graphs/Graph.ts
 const { AGENT, TOOLS } = GraphNodeKeys;
 class Graph {
-    lastToken;
-    tokenTypeSwitch;
-    reasoningKey = 'reasoning_content';
-    currentTokenType = ContentTypes.TEXT;
     messageStepHasToolCalls = new Map();
     messageIdsByStepKey = new Map();
     prelimMessageIdsByStepKey = new Map();
@@ -32,74 +32,37 @@ class Graph {
     stepKeyIds = new Map();
     contentIndexMap = new Map();
     toolCallStepIds = new Map();
-    currentUsage;
-    indexTokenCountMap = {};
-    maxContextTokens;
-    pruneMessages;
-    /** The amount of time that should pass before another consecutive API call */
-    streamBuffer;
-    tokenCounter;
     signal;
     /** Set of invoked tool call IDs from non-message run steps completed mid-run, if any */
     invokedToolIds;
     handlerRegistry;
 }
 class StandardGraph extends Graph {
-    graphState;
-    clientOptions;
-    boundModel;
-    /** The last recorded timestamp that a stream API call was invoked */
-    lastStreamCall;
-    systemMessage;
+    overrideModel;
+    /** Optional compile options passed into workflow.compile() */
+    compileOptions;
     messages = [];
     runId;
-    tools;
-    toolMap;
     startIndex = 0;
-    provider;
-    toolEnd;
     signal;
-    constructor({ runId, tools, signal, toolMap, provider, streamBuffer, instructions, reasoningKey, clientOptions, toolEnd = false, additional_instructions = '', }) {
+    /** Map of agent contexts by agent ID */
+    agentContexts = new Map();
+    /** Default agent ID to use */
+    defaultAgentId;
+    constructor({ 
+    // parent-level graph inputs
+    runId, signal, agents, tokenCounter, indexTokenCountMap, }) {
         super();
         this.runId = runId;
-        this.tools = tools;
         this.signal = signal;
-        this.toolEnd = toolEnd;
-        this.toolMap = toolMap;
-        this.provider = provider;
-        this.streamBuffer = streamBuffer;
-        this.clientOptions = clientOptions;
-        this.graphState = this.createGraphState();
-        this.boundModel = this.initializeModel();
-        if (reasoningKey) {
-            this.reasoningKey = reasoningKey;
+        if (agents.length === 0) {
+            throw new Error('At least one agent configuration is required');
         }
-        let finalInstructions = instructions;
-        if (additional_instructions) {
-            finalInstructions =
-                finalInstructions != null && finalInstructions
-                    ? `${finalInstructions}\n\n${additional_instructions}`
-                    : additional_instructions;
-        }
-        if (finalInstructions != null &&
-            finalInstructions &&
-            provider === Providers.ANTHROPIC &&
-            (clientOptions.clientOptions
-                ?.defaultHeaders?.['anthropic-beta']?.includes('prompt-caching') ??
-                false)) {
-            finalInstructions = {
-                content: [
-                    {
-                        type: 'text',
-                        text: instructions,
-                        cache_control: { type: 'ephemeral' },
-                    },
-                ],
-            };
-        }
-        if (finalInstructions != null && finalInstructions !== '') {
-            this.systemMessage = new SystemMessage(finalInstructions);
+        for (const agentConfig of agents) {
+            const agentContext = AgentContext.fromConfig(agentConfig, tokenCounter, indexTokenCountMap);
+            this.agentContexts.set(agentConfig.agentId, agentContext);
         }
+        this.defaultAgentId = agents[0].agentId;
     }
     /* Init */
     resetValues(keepContent) {
@@ -112,16 +75,12 @@ class StandardGraph extends Graph {
         this.stepKeyIds = resetIfNotEmpty(this.stepKeyIds, new Map());
         this.toolCallStepIds = resetIfNotEmpty(this.toolCallStepIds, new Map());
         this.messageIdsByStepKey = resetIfNotEmpty(this.messageIdsByStepKey, new Map());
-        this.messageStepHasToolCalls = resetIfNotEmpty(this.prelimMessageIdsByStepKey, new Map());
+        this.messageStepHasToolCalls = resetIfNotEmpty(this.messageStepHasToolCalls, new Map());
         this.prelimMessageIdsByStepKey = resetIfNotEmpty(this.prelimMessageIdsByStepKey, new Map());
-        this.currentTokenType = resetIfNotEmpty(this.currentTokenType, ContentTypes.TEXT);
-        this.lastToken = resetIfNotEmpty(this.lastToken, undefined);
-        this.tokenTypeSwitch = resetIfNotEmpty(this.tokenTypeSwitch, undefined);
-        this.indexTokenCountMap = resetIfNotEmpty(this.indexTokenCountMap, {});
-        this.currentUsage = resetIfNotEmpty(this.currentUsage, undefined);
-        this.tokenCounter = resetIfNotEmpty(this.tokenCounter, undefined);
-        this.maxContextTokens = resetIfNotEmpty(this.maxContextTokens, undefined);
         this.invokedToolIds = resetIfNotEmpty(this.invokedToolIds, undefined);
+        for (const context of this.agentContexts.values()) {
+            context.reset();
+        }
     }
     /* Run Step Processing */
     getRunStep(stepId) {
@@ -131,6 +90,27 @@ class StandardGraph extends Graph {
         }
         return undefined;
     }
+    getAgentContext(metadata) {
+        if (!metadata) {
+            throw new Error('No metadata provided to retrieve agent context');
+        }
+        const currentNode = metadata.langgraph_node;
+        if (!currentNode) {
+            throw new Error('No langgraph_node in metadata to retrieve agent context');
+        }
+        let agentId;
+        if (currentNode.startsWith(AGENT)) {
+            agentId = currentNode.substring(AGENT.length);
+        }
+        else if (currentNode.startsWith(TOOLS)) {
+            agentId = currentNode.substring(TOOLS.length);
+        }
+        const agentContext = this.agentContexts.get(agentId ?? '');
+        if (!agentContext) {
+            throw new Error(`No agent context found for agent ID ${agentId}`);
+        }
+        return agentContext;
+    }
     getStepKey(metadata) {
         if (!metadata)
             return '';
@@ -176,11 +156,12 @@ class StandardGraph extends Graph {
             metadata.langgraph_step,
             metadata.checkpoint_ns,
         ];
-        if (this.currentTokenType === ContentTypes.THINK ||
-            this.currentTokenType === 'think_and_text') {
+        const agentContext = this.getAgentContext(metadata);
+        if (agentContext.currentTokenType === ContentTypes.THINK ||
+            agentContext.currentTokenType === 'think_and_text') {
             keyList.push('reasoning');
         }
-        else if (this.tokenTypeSwitch === 'content') {
+        else if (agentContext.tokenTypeSwitch === 'content') {
             keyList.push('post-reasoning');
         }
         if (this.invokedToolIds != null && this.invokedToolIds.size > 0) {
@@ -199,112 +180,158 @@ class StandardGraph extends Graph {
         return convertMessagesToContent(this.messages.slice(this.startIndex));
     }
     /* Graph */
-    createGraphState() {
-        return {
-            messages: {
-                value: (x, y) => {
-                    if (!x.length) {
-                        if (this.systemMessage) {
-                            x.push(this.systemMessage);
-                        }
-                        this.startIndex = x.length + y.length;
-                    }
-                    const current = x.concat(y);
-                    this.messages = current;
-                    return current;
-                },
-                default: () => [],
-            },
-        };
+    createSystemRunnable({ provider, clientOptions, instructions, additional_instructions, }) {
+        let finalInstructions = instructions;
+        if (additional_instructions != null && additional_instructions !== '') {
+            finalInstructions =
+                finalInstructions != null && finalInstructions
+                    ? `${finalInstructions}\n\n${additional_instructions}`
+                    : additional_instructions;
+        }
+        if (finalInstructions != null &&
+            finalInstructions &&
+            provider === Providers.ANTHROPIC &&
+            (clientOptions.clientOptions
+                ?.defaultHeaders?.['anthropic-beta']?.includes('prompt-caching') ??
+                false)) {
+            finalInstructions = {
+                content: [
+                    {
+                        type: 'text',
+                        text: instructions,
+                        cache_control: { type: 'ephemeral' },
+                    },
+                ],
+            };
+        }
+        if (finalInstructions != null && finalInstructions !== '') {
+            const systemMessage = new SystemMessage(finalInstructions);
+            return RunnableLambda.from((messages) => {
+                return [systemMessage, ...messages];
+            }).withConfig({ runName: 'prompt' });
+        }
     }
-    initializeTools() {
-        // return new ToolNode<t.BaseGraphState>(this.tools);
+    initializeTools({ currentTools, currentToolMap, }) {
         return new ToolNode({
-            tools: this.tools || [],
-            toolMap: this.toolMap,
+            tools: currentTools ?? [],
+            toolMap: currentToolMap,
             toolCallStepIds: this.toolCallStepIds,
             errorHandler: (data, metadata) => StandardGraph.handleToolCallErrorStatic(this, data, metadata),
         });
     }
-    initializeModel() {
-        const ChatModelClass = getChatModelClass(this.provider);
-        const model = new ChatModelClass(this.clientOptions);
-        if (isOpenAILike(this.provider) &&
+    initializeModel({ provider, tools, clientOptions, }) {
+        const ChatModelClass = getChatModelClass(provider);
+        const model = new ChatModelClass(clientOptions ?? {});
+        if (isOpenAILike(provider) &&
             (model instanceof ChatOpenAI || model instanceof AzureChatOpenAI)) {
-            model.temperature = this.clientOptions
+            model.temperature = clientOptions
                 .temperature;
-            model.topP = this.clientOptions.topP;
-            model.frequencyPenalty = this.clientOptions
+            model.topP = clientOptions.topP;
+            model.frequencyPenalty = clientOptions
                 .frequencyPenalty;
-            model.presencePenalty = this.clientOptions
+            model.presencePenalty = clientOptions
                 .presencePenalty;
-            model.n = this.clientOptions.n;
+            model.n = clientOptions.n;
         }
-        else if (this.provider === Providers.VERTEXAI &&
+        else if (provider === Providers.VERTEXAI &&
             model instanceof ChatVertexAI) {
-            model.temperature = this.clientOptions
+            model.temperature = clientOptions
                 .temperature;
-            model.topP = this.clientOptions
-                .topP;
-            model.topK = this.clientOptions
-                .topK;
-            model.topLogprobs = this.clientOptions
+            model.topP = clientOptions.topP;
+            model.topK = clientOptions.topK;
+            model.topLogprobs = clientOptions
                 .topLogprobs;
-            model.frequencyPenalty = this.clientOptions
+            model.frequencyPenalty = clientOptions
                 .frequencyPenalty;
-            model.presencePenalty = this.clientOptions
+            model.presencePenalty = clientOptions
                 .presencePenalty;
-            model.maxOutputTokens = this.clientOptions
+            model.maxOutputTokens = clientOptions
                 .maxOutputTokens;
         }
-        if (!this.tools || this.tools.length === 0) {
+        if (!tools || tools.length === 0) {
             return model;
         }
-        return model.bindTools(this.tools);
+        return model.bindTools(tools);
     }
     overrideTestModel(responses, sleep, toolCalls) {
-        this.boundModel = createFakeStreamingLLM({
+        this.overrideModel = createFakeStreamingLLM({
             responses,
             sleep,
             toolCalls,
         });
     }
-    getNewModel({ provider, clientOptions, omitOptions, }) {
+    getNewModel({ provider, clientOptions, }) {
         const ChatModelClass = getChatModelClass(provider);
-        const options = omitOptions && clientOptions == null
-            ? Object.assign(Object.fromEntries(Object.entries(this.clientOptions).filter(([key]) => !omitOptions.has(key))), clientOptions)
-            : (clientOptions ?? this.clientOptions);
-        return new ChatModelClass(options);
+        return new ChatModelClass(clientOptions ?? {});
     }
-    storeUsageMetadata(finalMessage) {
+    getUsageMetadata(finalMessage) {
         if (finalMessage &&
             'usage_metadata' in finalMessage &&
             finalMessage.usage_metadata != null) {
-            this.currentUsage = finalMessage.usage_metadata;
+            return finalMessage.usage_metadata;
+        }
+    }
+    /** Execute model invocation with streaming support */
+    async attemptInvoke({ currentModel, finalMessages, provider, tools, }, config) {
+        const model = this.overrideModel ?? currentModel;
+        if (!model) {
+            throw new Error('No model found');
+        }
+        if ((tools?.length ?? 0) > 0 && manualToolStreamProviders.has(provider)) {
+            if (!model.stream) {
+                throw new Error('Model does not support stream');
+            }
+            const stream = await model.stream(finalMessages, config);
+            let finalChunk;
+            for await (const chunk of stream) {
+                await safeDispatchCustomEvent(GraphEvents.CHAT_MODEL_STREAM, { chunk, emitted: true }, config);
+                finalChunk = finalChunk ? concat(finalChunk, chunk) : chunk;
+            }
+            finalChunk = modifyDeltaProperties(provider, finalChunk);
+            return { messages: [finalChunk] };
+        }
+        else {
+            const finalMessage = await model.invoke(finalMessages, config);
+            if ((finalMessage.tool_calls?.length ?? 0) > 0) {
+                finalMessage.tool_calls = finalMessage.tool_calls?.filter((tool_call) => !!tool_call.name);
+            }
+            return { messages: [finalMessage] };
         }
     }
-    cleanupSignalListener() {
+    cleanupSignalListener(currentModel) {
         if (!this.signal) {
             return;
         }
-        if (!this.boundModel) {
+        const model = this.overrideModel ?? currentModel;
+        if (!model) {
             return;
         }
-        const client = this.boundModel?.exposedClient;
+        const client = model?.exposedClient;
         if (!client?.abortHandler) {
             return;
         }
         this.signal.removeEventListener('abort', client.abortHandler);
         client.abortHandler = undefined;
     }
-    createCallModel() {
+    createCallModel(agentId = 'default', currentModel) {
         return async (state, config) => {
-            const { provider = '' } = config?.configurable ?? {};
-            if (this.boundModel == null) {
+            /**
+             * Get agent context - it must exist by this point
+             */
+            const agentContext = this.agentContexts.get(agentId);
+            if (!agentContext) {
+                throw new Error(`Agent context not found for agentId: ${agentId}`);
+            }
+            const model = this.overrideModel ?? currentModel;
+            if (!model) {
                 throw new Error('No Graph model found');
             }
-            if (!config || !provider) {
-                throw new Error(`No ${config ? 'provider' : 'config'} provided`);
+            if (!config) {
+                throw new Error('No config provided');
+            }
+            // Ensure token calculations are complete before proceeding
+            if (agentContext.tokenCalculationPromise) {
+                await agentContext.tokenCalculationPromise;
             }
             if (!config.signal) {
                 config.signal = this.signal;
@@ -312,120 +339,208 @@ class StandardGraph extends Graph {
             this.config = config;
             const { messages } = state;
             let messagesToUse = messages;
-            if (!this.pruneMessages &&
-                this.tokenCounter &&
-                this.maxContextTokens != null &&
-                this.indexTokenCountMap[0] != null) {
-                const isAnthropicWithThinking = (this.provider === Providers.ANTHROPIC &&
-                    this.clientOptions.thinking !=
+            if (!agentContext.pruneMessages &&
+                agentContext.tokenCounter &&
+                agentContext.maxContextTokens != null &&
+                agentContext.indexTokenCountMap[0] != null) {
+                const isAnthropicWithThinking = (agentContext.provider === Providers.ANTHROPIC &&
+                    agentContext.clientOptions.thinking !=
                         null) ||
-                    (this.provider === Providers.BEDROCK &&
-                        this.clientOptions
+                    (agentContext.provider === Providers.BEDROCK &&
+                        agentContext.clientOptions
                             .additionalModelRequestFields?.['thinking'] != null) ||
-                    (this.provider === Providers.OPENAI &&
-                        this.clientOptions.modelKwargs
+                    (agentContext.provider === Providers.OPENAI &&
+                        agentContext.clientOptions.modelKwargs
                             ?.thinking?.type === 'enabled');
-                this.pruneMessages = createPruneMessages({
-                    provider: this.provider,
-                    indexTokenCountMap: this.indexTokenCountMap,
-                    maxTokens: this.maxContextTokens,
-                    tokenCounter: this.tokenCounter,
+                agentContext.pruneMessages = createPruneMessages({
                     startIndex: this.startIndex,
+                    provider: agentContext.provider,
+                    tokenCounter: agentContext.tokenCounter,
+                    maxTokens: agentContext.maxContextTokens,
                     thinkingEnabled: isAnthropicWithThinking,
+                    indexTokenCountMap: agentContext.indexTokenCountMap,
                 });
             }
-            if (this.pruneMessages) {
-                const { context, indexTokenCountMap } = this.pruneMessages({
+            if (agentContext.pruneMessages) {
+                const { context, indexTokenCountMap } = agentContext.pruneMessages({
                     messages,
-                    usageMetadata: this.currentUsage,
+                    usageMetadata: agentContext.currentUsage,
                     // startOnMessageType: 'human',
                 });
-                this.indexTokenCountMap = indexTokenCountMap;
+                agentContext.indexTokenCountMap = indexTokenCountMap;
                 messagesToUse = context;
             }
-            const finalMessages = messagesToUse;
+            let finalMessages = messagesToUse;
+            if (agentContext.useLegacyContent) {
+                finalMessages = formatContentStrings(finalMessages);
+            }
             const lastMessageX = finalMessages.length >= 2
                 ? finalMessages[finalMessages.length - 2]
                 : null;
             const lastMessageY = finalMessages.length >= 1
                 ? finalMessages[finalMessages.length - 1]
                 : null;
-            if (provider === Providers.BEDROCK &&
+            if (agentContext.provider === Providers.BEDROCK &&
                 lastMessageX instanceof AIMessageChunk &&
                 lastMessageY instanceof ToolMessage &&
                 typeof lastMessageX.content === 'string') {
                 finalMessages[finalMessages.length - 2].content = '';
             }
             const isLatestToolMessage = lastMessageY instanceof ToolMessage;
-            if (isLatestToolMessage && provider === Providers.ANTHROPIC) {
+            if (isLatestToolMessage &&
+                agentContext.provider === Providers.ANTHROPIC) {
                 formatAnthropicArtifactContent(finalMessages);
             }
             else if (isLatestToolMessage &&
-                (isOpenAILike(provider) || isGoogleLike(provider))) {
+                (isOpenAILike(agentContext.provider) ||
+                    isGoogleLike(agentContext.provider))) {
                 formatArtifactPayload(finalMessages);
             }
-            if (this.lastStreamCall != null && this.streamBuffer != null) {
-                const timeSinceLastCall = Date.now() - this.lastStreamCall;
-                if (timeSinceLastCall < this.streamBuffer) {
-                    const timeToWait = Math.ceil((this.streamBuffer - timeSinceLastCall) / 1000) * 1000;
+            if (agentContext.provider === Providers.ANTHROPIC) {
+                const anthropicOptions = agentContext.clientOptions;
+                const defaultHeaders = anthropicOptions?.clientOptions
+                    ?.defaultHeaders;
+                const anthropicBeta = defaultHeaders?.['anthropic-beta'];
+                if (typeof anthropicBeta === 'string' &&
+                    anthropicBeta.includes('prompt-caching')) {
+                    finalMessages = addCacheControl(finalMessages);
+                }
+            }
+            else if (agentContext.provider === Providers.BEDROCK) {
+                const bedrockOptions = agentContext.clientOptions;
+                if (bedrockOptions?.promptCache === true) {
+                    finalMessages = addBedrockCacheControl(finalMessages);
+                }
+            }
+            if (agentContext.lastStreamCall != null &&
+                agentContext.streamBuffer != null) {
+                const timeSinceLastCall = Date.now() - agentContext.lastStreamCall;
+                if (timeSinceLastCall < agentContext.streamBuffer) {
+                    const timeToWait = Math.ceil((agentContext.streamBuffer - timeSinceLastCall) / 1000) *
+                        1000;
                     await sleep(timeToWait);
                 }
             }
-            this.lastStreamCall = Date.now();
+            agentContext.lastStreamCall = Date.now();
             let result;
-            if ((this.tools?.length ?? 0) > 0 &&
-                manualToolStreamProviders.has(provider)) {
-                const stream = await this.boundModel.stream(finalMessages, config);
-                let finalChunk;
-                for await (const chunk of stream) {
-                    dispatchCustomEvent(GraphEvents.CHAT_MODEL_STREAM, { chunk }, config);
-                    if (!finalChunk) {
-                        finalChunk = chunk;
+            const fallbacks = agentContext.clientOptions?.fallbacks ??
+                [];
+            if (finalMessages.length === 0) {
+                throw new Error(JSON.stringify({
+                    type: 'empty_messages',
+                    info: 'Message pruning removed all messages as none fit in the context window. Please increase the context window size or make your message shorter.',
+                }));
+            }
+            try {
+                result = await this.attemptInvoke({
+                    currentModel: model,
+                    finalMessages,
+                    provider: agentContext.provider,
+                    tools: agentContext.tools,
+                }, config);
+            }
+            catch (primaryError) {
+                let lastError = primaryError;
+                for (const fb of fallbacks) {
+                    try {
+                        let model = this.getNewModel({
+                            provider: fb.provider,
+                            clientOptions: fb.clientOptions,
+                        });
+                        const bindableTools = agentContext.tools;
+                        model = (!bindableTools || bindableTools.length === 0
+                            ? model
+                            : model.bindTools(bindableTools));
+                        result = await this.attemptInvoke({
+                            currentModel: model,
+                            finalMessages,
+                            provider: fb.provider,
+                            tools: agentContext.tools,
+                        }, config);
+                        lastError = undefined;
+                        break;
                     }
-                    else {
-                        finalChunk = concat(finalChunk, chunk);
+                    catch (e) {
+                        lastError = e;
+                        continue;
                     }
                 }
-                finalChunk = modifyDeltaProperties(this.provider, finalChunk);
-                result = { messages: [finalChunk] };
-            }
-            else {
-                const finalMessage = (await this.boundModel.invoke(finalMessages, config));
-                if ((finalMessage.tool_calls?.length ?? 0) > 0) {
-                    finalMessage.tool_calls = finalMessage.tool_calls?.filter((tool_call) => {
-                        if (!tool_call.name) {
-                            return false;
-                        }
-                        return true;
-                    });
+                if (lastError !== undefined) {
+                    throw lastError;
                 }
-                result = { messages: [finalMessage] };
             }
-            this.storeUsageMetadata(result.messages?.[0]);
+            if (!result) {
+                throw new Error('No result after model invocation');
+            }
+            agentContext.currentUsage = this.getUsageMetadata(result.messages?.[0]);
             this.cleanupSignalListener();
             return result;
         };
     }
-    createWorkflow() {
+    createAgentNode(agentId) {
+        const agentContext = this.agentContexts.get(agentId);
+        if (!agentContext) {
+            throw new Error(`Agent context not found for agentId: ${agentId}`);
+        }
+        let currentModel = this.initializeModel({
+            tools: agentContext.tools,
+            provider: agentContext.provider,
+            clientOptions: agentContext.clientOptions,
+        });
+        if (agentContext.systemRunnable) {
+            currentModel = agentContext.systemRunnable.pipe(currentModel);
+        }
+        const agentNode = `${AGENT}${agentId}`;
+        const toolNode = `${TOOLS}${agentId}`;
         const routeMessage = (state, config) => {
             this.config = config;
-            return toolsCondition(state, this.invokedToolIds);
+            return toolsCondition(state, toolNode, this.invokedToolIds);
         };
-        const workflow = new StateGraph({
-            channels: this.graphState,
-        })
-            .addNode(AGENT, this.createCallModel())
-            .addNode(TOOLS, this.initializeTools())
-            .addEdge(START, AGENT)
-            .addConditionalEdges(AGENT, routeMessage)
-            .addEdge(TOOLS, this.toolEnd ? END : AGENT);
-        return workflow.compile();
+        const StateAnnotation = Annotation.Root({
+            messages: Annotation({
+                reducer: messagesStateReducer,
+                default: () => [],
+            }),
+        });
+        const workflow = new StateGraph(StateAnnotation)
+            .addNode(agentNode, this.createCallModel(agentId, currentModel))
+            .addNode(toolNode, this.initializeTools({
+            currentTools: agentContext.tools,
+            currentToolMap: agentContext.toolMap,
+        }))
+            .addEdge(START, agentNode)
+            .addConditionalEdges(agentNode, routeMessage)
+            .addEdge(toolNode, agentContext.toolEnd ? END : agentNode);
+        // Cast to unknown to avoid tight coupling to external types; options are opt-in
+        return workflow.compile(this.compileOptions);
+    }
+    createWorkflow() {
+        /** Use the default (first) agent for now */
+        const agentNode = this.createAgentNode(this.defaultAgentId);
+        const StateAnnotation = Annotation.Root({
+            messages: Annotation({
+                reducer: (a, b) => {
+                    if (!a.length) {
+                        this.startIndex = a.length + b.length;
+                    }
+                    const result = messagesStateReducer(a, b);
+                    this.messages = result;
+                    return result;
+                },
+                default: () => [],
+            }),
+        });
+        const workflow = new StateGraph(StateAnnotation)
+            .addNode(this.defaultAgentId, agentNode, { ends: [END] })
+            .addEdge(START, this.defaultAgentId)
+            .compile();
+        return workflow;
     }
     /* Dispatchers */
     /**
      * Dispatches a run step to the client, returns the step ID
      */
-    dispatchRunStep(stepKey, stepDetails) {
+    async dispatchRunStep(stepKey, stepDetails) {
         if (!this.config) {
             throw new Error('No config provided');
         }
@@ -453,17 +568,21 @@ class StandardGraph extends Graph {
         }
         this.contentData.push(runStep);
         this.contentIndexMap.set(stepId, runStep.index);
-        dispatchCustomEvent(GraphEvents.ON_RUN_STEP, runStep, this.config);
+        await safeDispatchCustomEvent(GraphEvents.ON_RUN_STEP, runStep, this.config);
         return stepId;
     }
-    handleToolCallCompleted(data, metadata, omitOutput) {
+    async handleToolCallCompleted(data, metadata, omitOutput) {
         if (!this.config) {
             throw new Error('No config provided');
         }
         if (!data.output) {
             return;
         }
-        const { input, output } = data;
+        const { input, output: _output } = data;
+        if (_output?.lg_name === 'Command') {
+            return;
+        }
+        const output = _output;
         const { tool_call_id } = output;
         const stepId = this.toolCallStepIds.get(tool_call_id) ?? '';
         if (!stepId) {
@@ -484,7 +603,9 @@ class StandardGraph extends Graph {
             output: omitOutput === true ? '' : dispatchedOutput,
             progress: 1,
         };
-        this.handlerRegistry?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)?.handle(GraphEvents.ON_RUN_STEP_COMPLETED, {
+        await this.handlerRegistry
+            ?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)
+            ?.handle(GraphEvents.ON_RUN_STEP_COMPLETED, {
             result: {
                 id: stepId,
                 index: runStep.index,
@@ -497,7 +618,7 @@ class StandardGraph extends Graph {
      * Static version of handleToolCallError to avoid creating strong references
      * that prevent garbage collection
      */
-    static handleToolCallErrorStatic(graph, data, metadata) {
+    static async handleToolCallErrorStatic(graph, data, metadata) {
         if (!graph.config) {
             throw new Error('No config provided');
         }
@@ -521,7 +642,7 @@ class StandardGraph extends Graph {
             output: `Error processing tool${error?.message != null ? `: ${error.message}` : ''}`,
             progress: 1,
         };
-        graph.handlerRegistry
+        await graph.handlerRegistry
             ?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)
             ?.handle(GraphEvents.ON_RUN_STEP_COMPLETED, {
             result: {
@@ -536,10 +657,10 @@ class StandardGraph extends Graph {
      * Instance method that delegates to the static method
      * Kept for backward compatibility
      */
-    handleToolCallError(data, metadata) {
-        StandardGraph.handleToolCallErrorStatic(this, data, metadata);
+    async handleToolCallError(data, metadata) {
+        await StandardGraph.handleToolCallErrorStatic(this, data, metadata);
     }
-    dispatchRunStepDelta(id, delta) {
+    async dispatchRunStepDelta(id, delta) {
         if (!this.config) {
             throw new Error('No config provided');
         }
@@ -550,9 +671,9 @@ class StandardGraph extends Graph {
             id,
             delta,
         };
-        dispatchCustomEvent(GraphEvents.ON_RUN_STEP_DELTA, runStepDelta, this.config);
+        await safeDispatchCustomEvent(GraphEvents.ON_RUN_STEP_DELTA, runStepDelta, this.config);
     }
-    dispatchMessageDelta(id, delta) {
+    async dispatchMessageDelta(id, delta) {
         if (!this.config) {
             throw new Error('No config provided');
         }
@@ -560,9 +681,9 @@ class StandardGraph extends Graph {
             id,
             delta,
         };
-        dispatchCustomEvent(GraphEvents.ON_MESSAGE_DELTA, messageDelta, this.config);
+        await safeDispatchCustomEvent(GraphEvents.ON_MESSAGE_DELTA, messageDelta, this.config);
     }
-    dispatchReasoningDelta = (stepId, delta) => {
+    dispatchReasoningDelta = async (stepId, delta) => {
         if (!this.config) {
             throw new Error('No config provided');
         }
@@ -570,7 +691,7 @@ class StandardGraph extends Graph {
             id: stepId,
             delta,
         };
-        dispatchCustomEvent(GraphEvents.ON_REASONING_DELTA, reasoningDelta, this.config);
+        await safeDispatchCustomEvent(GraphEvents.ON_REASONING_DELTA, reasoningDelta, this.config);
     };
 }