@illuma-ai/agents 1.0.90 → 1.0.94

package/src/agents/AgentContext.ts

@@ -49,6 +49,7 @@ export class AgentContext {
       structured_output: structuredOutputSnake,
       discoveredTools,
       summarizeCallback,
+      persistedSummary,
     } = agentConfig;
 
     // Normalize structured output: support both camelCase and snake_case inputs
@@ -56,7 +57,10 @@ export class AgentContext {
     let structuredOutput: t.StructuredOutputConfig | undefined;
     if (structuredOutputCamel) {
       structuredOutput = structuredOutputCamel;
-    } else if (structuredOutputSnake?.enabled && structuredOutputSnake.schema) {
+    } else if (
+      structuredOutputSnake?.enabled === true &&
+      structuredOutputSnake.schema != null
+    ) {
       // Convert snake_case input to StructuredOutputConfig
       structuredOutput = {
         schema: structuredOutputSnake.schema,
@@ -89,12 +93,10 @@ export class AgentContext {
       dynamicContext,
       structuredOutput,
       discoveredTools,
+      summarizeCallback,
+      persistedSummary,
     });
 
-    if (summarizeCallback) {
-      agentContext.summarizeCallback = summarizeCallback;
-    }
-
     if (tokenCounter) {
       // Initialize system runnable BEFORE async tool token calculation
       // This ensures system message tokens are in instructionTokens before
@@ -102,6 +104,7 @@ export class AgentContext {
       agentContext.initializeSystemRunnable();
 
       const tokenMap = indexTokenCountMap || {};
+      agentContext.baseIndexTokenCountMap = { ...tokenMap };
       agentContext.indexTokenCountMap = tokenMap;
       agentContext.tokenCalculationPromise = agentContext
         .calculateInstructionTokens(tokenCounter)
@@ -113,6 +116,7 @@ export class AgentContext {
           console.error('Error calculating instruction tokens:', err);
         });
     } else if (indexTokenCountMap) {
+      agentContext.baseIndexTokenCountMap = { ...indexTokenCountMap };
       agentContext.indexTokenCountMap = indexTokenCountMap;
     }
 
@@ -120,17 +124,19 @@ export class AgentContext {
   }
 
   /** Agent identifier */
-  agentId: string;
+  agentId!: string;
   /** Human-readable name for this agent (used in handoff context). Falls back to agentId if not provided. */
   name?: string;
   /** Description of what this agent does (used to enrich handoff tool descriptions). */
   description?: string;
   /** Provider for this specific agent */
-  provider: Providers;
+  provider!: Providers;
   /** Client options for this agent */
   clientOptions?: t.ClientOptions;
   /** Token count map indexed by message position */
   indexTokenCountMap: Record<string, number | undefined> = {};
+  /** Canonical pre-run token map used to restore token accounting on reset */
+  baseIndexTokenCountMap: Record<string, number> = {};
   /** Maximum context tokens for this agent */
   maxContextTokens?: number;
   /** Current usage metadata for this agent */
@@ -147,6 +153,8 @@ export class AgentContext {
   lastStreamCall?: number;
   /** Tools available to this agent */
   tools?: t.GraphTools;
+  /** Graph-managed tools (e.g., handoff tools created by MultiAgentGraph) that bypass event-driven dispatch */
+  graphTools?: t.GraphTools;
   /** Tool map for this agent */
   toolMap?: t.ToolMap;
   /**
@@ -177,6 +185,8 @@ export class AgentContext {
   lastToken?: string;
   /** Token type switch state */
   tokenTypeSwitch?: 'reasoning' | 'content';
+  /** Tracks how many reasoning→text transitions have occurred (ensures unique post-reasoning step keys) */
+  reasoningTransitionCount = 0;
   /** Current token type being processed */
   currentTokenType: ContentTypes.TEXT | ContentTypes.THINK | 'think_and_text' =
     ContentTypes.TEXT;
@@ -209,13 +219,13 @@ export class AgentContext {
     artifacts: number;
     memory: number;
   } = {
-      branding: 0,
-      toolRouting: 0,
-      agentInstructions: 0,
-      mcpInstructions: 0,
-      artifacts: 0,
-      memory: 0,
-    };
+    branding: 0,
+    toolRouting: 0,
+    agentInstructions: 0,
+    mcpInstructions: 0,
+    artifacts: 0,
+    memory: 0,
+  };
   /**
    * Handoff context when this agent receives control via handoff.
    * Contains source and parallel execution info for system message context.
@@ -232,11 +242,10 @@ export class AgentContext {
    * instead of streaming text.
    */
   structuredOutput?: t.StructuredOutputConfig;
-  /**
-   * Optional callback for summarizing messages that were pruned from context.
-   * When provided, discarded messages are summarized instead of silently dropped.
-   */
-  summarizeCallback?: (messagesToRefine: BaseMessage[]) => Promise<string | undefined>;
+  /** Optional callback to summarize discarded messages during context pruning */
+  summarizeCallback?: (messages: BaseMessage[]) => Promise<string | undefined>;
+  /** Pre-existing summary loaded from persistent storage, injected into context on new turns */
+  persistedSummary?: string;
 
   constructor({
     agentId,
@@ -260,6 +269,8 @@ export class AgentContext {
     useLegacyContent,
     structuredOutput,
     discoveredTools,
+    summarizeCallback,
+    persistedSummary,
   }: {
     agentId: string;
     name?: string;
@@ -282,6 +293,10 @@ export class AgentContext {
     useLegacyContent?: boolean;
     structuredOutput?: t.StructuredOutputConfig;
     discoveredTools?: string[];
+    summarizeCallback?: (
+      messages: BaseMessage[]
+    ) => Promise<string | undefined>;
+    persistedSummary?: string;
   }) {
     this.agentId = agentId;
     this.name = name;
@@ -299,6 +314,8 @@ export class AgentContext {
     this.additionalInstructions = additionalInstructions;
     this.dynamicContext = dynamicContext;
     this.structuredOutput = structuredOutput;
+    this.summarizeCallback = summarizeCallback;
+    this.persistedSummary = persistedSummary;
     if (reasoningKey) {
       this.reasoningKey = reasoningKey;
     }
@@ -324,8 +341,11 @@ export class AgentContext {
    * and return a validated JSON response.
    */
   get isStructuredOutputMode(): boolean {
+    // Runtime safety: schema can be null/undefined via API despite type saying required
     return (
-      this.structuredOutput != null && this.structuredOutput.schema != null
+      this.structuredOutput != null &&
+      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+      this.structuredOutput.schema != null
     );
   }
 
@@ -341,17 +361,25 @@ export class AgentContext {
     const schema = { ...this.structuredOutput.schema };
 
     // Ensure type is set
-    if (!schema.type && schema.properties) {
+    if (schema.type == null && schema.properties != null) {
       schema.type = 'object';
     }
 
     // Add title from config name
-    if (this.structuredOutput.name && !schema.title) {
+    if (
+      this.structuredOutput.name != null &&
+      this.structuredOutput.name !== '' &&
+      schema.title == null
+    ) {
       schema.title = this.structuredOutput.name;
     }
 
     // Add description from config
-    if (this.structuredOutput.description && !schema.description) {
+    if (
+      this.structuredOutput.description != null &&
+      this.structuredOutput.description !== '' &&
+      schema.description == null
+    ) {
       schema.description = this.structuredOutput.description;
     }
 
@@ -636,35 +664,6 @@ export class AgentContext {
       }
     }
 
-    // Handle Bedrock prompt caching (Converse API)
-    // Adds cachePoint block after text content for system message caching
-    // NOTE: Both Claude and Nova models support cachePoint in system and messages
-    // (Nova does NOT support cachePoint in tools - that check is in bedrock/index.ts)
-    if (this.provider === Providers.BEDROCK) {
-      const bedrockOptions = this.clientOptions as
-        | t.BedrockAnthropicInput
-        | undefined;
-      const modelId = bedrockOptions?.model?.toLowerCase() ?? '';
-      const supportsCaching =
-        modelId.includes('claude') ||
-        modelId.includes('anthropic') ||
-        modelId.includes('nova');
-
-      if (bedrockOptions?.promptCache === true && supportsCaching) {
-        finalInstructions = {
-          content: [
-            {
-              type: 'text',
-              text: instructionsString,
-            },
-            {
-              cachePoint: { type: 'default' },
-            },
-          ],
-        };
-      }
-    }
-
     const systemMessage = new SystemMessage(finalInstructions);
 
     // Update token counts (subtract old, add new)
@@ -690,14 +689,32 @@ export class AgentContext {
     this.cachedSystemRunnable = undefined;
     this.systemRunnableStale = true;
     this.lastToken = undefined;
-    this.indexTokenCountMap = {};
+    this.indexTokenCountMap = { ...this.baseIndexTokenCountMap };
     this.currentUsage = undefined;
     this.pruneMessages = undefined;
     this.lastStreamCall = undefined;
     this.tokenTypeSwitch = undefined;
+    this.reasoningTransitionCount = 0;
     this.currentTokenType = ContentTypes.TEXT;
     this.discoveredToolNames.clear();
     this.handoffContext = undefined;
+
+    if (this.tokenCounter) {
+      this.initializeSystemRunnable();
+      const baseTokenMap = { ...this.baseIndexTokenCountMap };
+      this.indexTokenCountMap = baseTokenMap;
+      this.tokenCalculationPromise = this.calculateInstructionTokens(
+        this.tokenCounter
+      )
+        .then(() => {
+          this.updateTokenMapWithInstructions(baseTokenMap);
+        })
+        .catch((err) => {
+          console.error('Error calculating instruction tokens:', err);
+        });
+    } else {
+      this.tokenCalculationPromise = undefined;
+    }
   }
 
   /**
@@ -729,8 +746,15 @@ export class AgentContext {
     tokenCounter: t.TokenCounter
   ): Promise<void> {
     let toolTokens = 0;
-    this.toolsDetail = []; // Reset per-tool breakdown
+    // Track names to avoid double-counting when a tool appears in both
+    // this.tools (bound StructuredTool instances) and this.toolDefinitions
+    // (MCP / event-driven schemas).
+    const countedToolNames = new Set<string>();
+
+    // Reset per-tool breakdown
+    this.toolsDetail = [];
 
+    // Count tokens for bound tools (StructuredTool instances with .schema)
     if (this.tools && this.tools.length > 0) {
       for (const tool of this.tools) {
         const genericTool = tool as Record<string, unknown>;
@@ -738,23 +762,50 @@ export class AgentContext {
           genericTool.schema != null &&
           typeof genericTool.schema === 'object'
         ) {
+          const toolName = (genericTool.name as string | undefined) ?? '';
           const jsonSchema = toJsonSchema(
             genericTool.schema,
-            (genericTool.name as string | undefined) ?? '',
+            toolName,
             (genericTool.description as string | undefined) ?? ''
           );
-          const toolName = (genericTool.name as string) || 'unknown';
           const tokens = tokenCounter(
             new SystemMessage(JSON.stringify(jsonSchema))
           );
+          if (toolName) {
+            countedToolNames.add(toolName);
+          }
 
           // Track per-tool breakdown
-          this.toolsDetail.push({ name: toolName, tokens });
+          this.toolsDetail.push({ name: toolName || 'unknown', tokens });
           toolTokens += tokens;
         }
       }
     }
 
+    // Count tokens for tool definitions (MCP / event-driven tools).
+    // These are sent to the provider API as tool schemas alongside bound tools.
+    // Both can be populated simultaneously (graph tools + MCP tools).
+    if (this.toolDefinitions && this.toolDefinitions.length > 0) {
+      for (const def of this.toolDefinitions) {
+        if (countedToolNames.has(def.name)) {
+          continue; // Already counted via this.tools
+        }
+        const schema = {
+          name: def.name,
+          description: def.description ?? '',
+          parameters: def.parameters ?? {},
+        };
+        const defTokens = tokenCounter(
+          new SystemMessage(JSON.stringify(schema))
+        );
+        this.toolsDetail.push({
+          name: def.name || 'unknown',
+          tokens: defTokens,
+        });
+        toolTokens += defTokens;
+      }
+    }
+
     // Store total tool tokens for breakdown reporting
     this.toolTokensTotal = toolTokens;
 
@@ -811,7 +862,7 @@ export class AgentContext {
       artifacts: number;
       memory: number;
     };
-    } {
+  } {
     // Calculate sum of prompt components
     const promptsTotal =
       this.promptBreakdown.branding +
@@ -923,17 +974,22 @@ export class AgentContext {
     }
 
     /** Traditional mode: filter actual tool instances */
-    if (!this.tools || !this.toolRegistry) {
-      return this.tools;
+    const filtered =
+      !this.tools || !this.toolRegistry
+        ? this.tools
+        : this.filterToolsForBinding(this.tools);
+
+    if (this.graphTools && this.graphTools.length > 0) {
+      return [...(filtered ?? []), ...this.graphTools];
     }
 
-    return this.filterToolsForBinding(this.tools);
+    return filtered;
   }
 
-  /** Creates schema-only tools from toolDefinitions for event-driven mode */
+  /** Creates schema-only tools from toolDefinitions for event-driven mode, merged with native tools */
   private getEventDrivenToolsForBinding(): t.GraphTools {
     if (!this.toolDefinitions) {
-      return [];
+      return this.graphTools ?? [];
     }
 
     const defsToInclude = this.toolDefinitions.filter((def) => {
@@ -950,7 +1006,19 @@ export class AgentContext {
       return true;
     });
 
-    return createSchemaOnlyTools(defsToInclude) as t.GraphTools;
+    const schemaTools = createSchemaOnlyTools(defsToInclude) as t.GraphTools;
+
+    const allTools = [...schemaTools];
+
+    if (this.graphTools && this.graphTools.length > 0) {
+      allTools.push(...this.graphTools);
+    }
+
+    if (this.tools && this.tools.length > 0) {
+      allTools.push(...this.tools);
+    }
+
+    return allTools;
   }
 
   /** Filters tool instances for binding based on registry config */