illuma-agents 1.0.37 → 1.0.39

package/src/graphs/Graph.ts

@@ -37,6 +37,7 @@ import {
   Providers,
   StepTypes,
   MessageTypes,
+  Constants,
 } from '@/common';
 import {
   formatAnthropicArtifactContent,
@@ -141,6 +142,12 @@ export abstract class Graph<
   /** Set of invoked tool call IDs from non-message run steps completed mid-run, if any */
   invokedToolIds?: Set<string>;
   handlerRegistry: HandlerRegistry | undefined;
+  /**
+   * Tool session contexts for automatic state persistence across tool invocations.
+   * Keyed by tool name (e.g., Constants.EXECUTE_CODE).
+   * Currently supports code execution session tracking (session_id, files).
+   */
+  sessions: t.ToolSessionMap = new Map();
 }
 
 export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
@@ -487,11 +494,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       finalInstructions != null &&
       finalInstructions &&
       provider === Providers.ANTHROPIC &&
-      ((
-        (clientOptions as t.AnthropicClientOptions).clientOptions
-          ?.defaultHeaders as Record<string, string> | undefined
-      )?.['anthropic-beta']?.includes('prompt-caching') ??
-        false)
+      (clientOptions as t.AnthropicClientOptions).promptCache === true
     ) {
       finalInstructions = {
         content: [
@@ -528,6 +531,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       errorHandler: (data, metadata) =>
         StandardGraph.handleToolCallErrorStatic(this, data, metadata),
       toolRegistry: agentContext?.toolRegistry,
+      sessions: this.sessions,
     });
   }
 
@@ -662,6 +666,106 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
     }
   }
 
+  /**
+   * Execute model invocation with structured output.
+   * Uses withStructuredOutput to force the model to return JSON conforming to the schema.
+   * Disables streaming and returns a validated JSON response.
+   */
+  private async attemptStructuredInvoke(
+    {
+      currentModel,
+      finalMessages,
+      schema,
+      structuredOutputConfig,
+    }: {
+      currentModel: t.ChatModelInstance;
+      finalMessages: BaseMessage[];
+      schema: Record<string, unknown>;
+      structuredOutputConfig: t.StructuredOutputConfig;
+    },
+    config?: RunnableConfig
+  ): Promise<{
+    structuredResponse: Record<string, unknown>;
+    rawMessage?: AIMessageChunk;
+  }> {
+    const model = this.overrideModel ?? currentModel;
+    if (!model) {
+      throw new Error('No model found');
+    }
+
+    const {
+      name = 'StructuredResponse',
+      includeRaw = false,
+      handleErrors = true,
+      maxRetries = 2,
+    } = structuredOutputConfig;
+
+    // Use withStructuredOutput to bind the schema
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const structuredModel = (model as any).withStructuredOutput(schema, {
+      name,
+      includeRaw,
+      strict: structuredOutputConfig.strict !== false,
+    });
+
+    let lastError: Error | undefined;
+    let attempts = 0;
+
+    while (attempts <= maxRetries) {
+      try {
+        const result = await structuredModel.invoke(finalMessages, config);
+
+        // Handle includeRaw response format
+        if (includeRaw && result.raw && result.parsed) {
+          return {
+            structuredResponse: result.parsed as Record<string, unknown>,
+            rawMessage: result.raw as AIMessageChunk,
+          };
+        }
+
+        // Direct response
+        return {
+          structuredResponse: result as Record<string, unknown>,
+        };
+      } catch (error) {
+        lastError = error as Error;
+        attempts++;
+
+        // If error handling is disabled, throw immediately
+        if (handleErrors === false) {
+          throw error;
+        }
+
+        // If we've exhausted retries, throw
+        if (attempts > maxRetries) {
+          throw new Error(
+            `Structured output failed after ${maxRetries + 1} attempts: ${lastError.message}`
+          );
+        }
+
+        // Add error message to conversation for retry
+        const errorMessage =
+          typeof handleErrors === 'string'
+            ? handleErrors
+            : `The response did not match the expected schema. Error: ${lastError.message}. Please try again with a valid response.`;
+
+        console.warn(
+          `[Graph] Structured output attempt ${attempts} failed: ${lastError.message}. Retrying...`
+        );
+
+        // Add the error as a human message for context
+        finalMessages = [
+          ...finalMessages,
+          new HumanMessage({
+            content: `[VALIDATION ERROR]\n${errorMessage}`,
+          }),
+        ];
+      }
+    }
+
+    throw lastError ?? new Error('Structured output failed');
+  }
+
   cleanupSignalListener(currentModel?: t.ChatModel): void {
     if (!this.signal) {
       return;
@@ -860,14 +964,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         const anthropicOptions = agentContext.clientOptions as
           | t.AnthropicClientOptions
           | undefined;
-        const defaultHeaders = anthropicOptions?.clientOptions
-          ?.defaultHeaders as Record<string, string> | undefined;
-        const anthropicBeta = defaultHeaders?.['anthropic-beta'];
-
-        if (
-          typeof anthropicBeta === 'string' &&
-          anthropicBeta.includes('prompt-caching')
-        ) {
+        if (anthropicOptions?.promptCache === true) {
           finalMessages = addCacheControl<BaseMessage>(finalMessages);
         }
       } else if (agentContext.provider === Providers.BEDROCK) {
@@ -951,6 +1048,56 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         config
       );
 
+      // Check if structured output mode is enabled
+      if (
+        agentContext.isStructuredOutputMode &&
+        agentContext.structuredOutput
+      ) {
+        const schema = agentContext.getStructuredOutputSchema();
+        if (!schema) {
+          throw new Error('Structured output schema is not configured');
+        }
+
+        try {
+          // Use structured output invocation (non-streaming)
+          const { structuredResponse, rawMessage } =
+            await this.attemptStructuredInvoke(
+              {
+                currentModel: model as t.ChatModelInstance,
+                finalMessages,
+                schema,
+                structuredOutputConfig: agentContext.structuredOutput,
+              },
+              config
+            );
+
+          // Emit structured output event
+          await safeDispatchCustomEvent(
+            GraphEvents.ON_STRUCTURED_OUTPUT,
+            {
+              structuredResponse,
+              schema,
+              raw: rawMessage,
+            },
+            config
+          );
+
+          agentContext.currentUsage = rawMessage
+            ? this.getUsageMetadata(rawMessage)
+            : undefined;
+          this.cleanupSignalListener();
+
+          // Return both the structured response and the raw message
+          return {
+            messages: rawMessage ? [rawMessage] : [],
+            structuredResponse,
+          };
+        } catch (structuredError) {
+          console.error('[Graph] Structured output failed:', structuredError);
+          throw structuredError;
+        }
+      }
+
       try {
         result = await this.attemptInvoke(
           {
@@ -1366,6 +1513,60 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
       throw new Error(`No run step found for stepId ${stepId}`);
     }
 
+    /**
+     * Extract and store code execution session context from artifacts.
+     * Each file is stamped with its source session_id to support multi-session file tracking.
+     * When the same filename appears in a later execution, the newer version replaces the old.
+     */
+    const toolName = output.name;
+    if (
+      toolName === Constants.EXECUTE_CODE ||
+      toolName === Constants.PROGRAMMATIC_TOOL_CALLING
+    ) {
+      const artifact = output.artifact as t.CodeExecutionArtifact | undefined;
+      const newFiles = artifact?.files ?? [];
+      const hasNewFiles = newFiles.length > 0;
+
+      if (
+        hasNewFiles &&
+        artifact?.session_id != null &&
+        artifact.session_id !== ''
+      ) {
+        /**
+         * Stamp each new file with its source session_id.
+         * This enables files from different executions (parallel or sequential)
+         * to be tracked and passed to subsequent calls.
+         */
+        const filesWithSession: t.FileRefs = newFiles.map((file) => ({
+          ...file,
+          session_id: artifact.session_id,
+        }));
+
+        const existingSession = this.sessions.get(Constants.EXECUTE_CODE) as
+          | t.CodeSessionContext
+          | undefined;
+        const existingFiles = existingSession?.files ?? [];
+
+        /**
+         * Merge files, preferring latest versions by name.
+         * If a file with the same name exists, replace it with the new version.
+         * This handles cases where files are edited/recreated in subsequent executions.
+         */
+        const newFileNames = new Set(filesWithSession.map((f) => f.name));
+        const filteredExisting = existingFiles.filter(
+          (f) => !newFileNames.has(f.name)
+        );
+
+        this.sessions.set(Constants.EXECUTE_CODE, {
+          /** Keep latest session_id for reference/fallback */
+          session_id: artifact.session_id,
+          /** Accumulated files with latest versions preferred */
+          files: [...filteredExisting, ...filesWithSession],
+          lastUpdated: Date.now(),
+        });
+      }
+    }
+
     const dispatchedOutput =
       typeof output.content === 'string'
         ? output.content