@memberjunction/ai-agents 5.38.0 → 5.40.0

package/dist/base-agent.js

@@ -11,7 +11,8 @@
  * @since 2.49.0
  */
 import { FileStorageEngineBase } from '@memberjunction/core-entities';
-import { Metadata, RunView, LogStatus, LogStatusEx, LogError, LogErrorEx, IsVerboseLoggingEnabled } from '@memberjunction/core';
+import { Metadata, RunView, LogStatus, LogStatusEx, LogError, LogErrorEx, IsVerboseLoggingEnabled, DatabaseProviderBase } from '@memberjunction/core';
+import { AgentRunWatchdog } from './agent-run-watchdog.js';
 import { AIPromptRunner } from '@memberjunction/ai-prompts';
 import { BaseAgentType } from './agent-types/base-agent-type.js';
 import { CopyScalarsAndArrays, JSONValidator, SafeExpressionEvaluator, UUIDsEqual } from '@memberjunction/global';
@@ -26,6 +27,7 @@ import { AgentRunner } from './AgentRunner.js';
 import { PayloadManager } from './PayloadManager.js';
 import { ScratchpadManager } from './ScratchpadManager.js';
 import { ArtifactToolManager } from './ArtifactToolManager.js';
+import { PipelineExecutor, PipelineToolRegistry, ActionInvocable, ArtifactToolInvocable, BuildPipelineToolDocs, formatFinalOutput, summarizePipelineStages, } from './pipeline/index.js';
 import { AgentDataPreloader } from './AgentDataPreloader.js';
 import { ClientToolRequestManager } from './ClientToolRequestManager.js';
 import { ConversationMessageResolver } from './utils/ConversationMessageResolver.js';
@@ -94,6 +96,11 @@ export class BaseAgent {
         /**
          * Queue map to chain database saves sequentially per step entity.
          * Prevents UPDATE queries running before INSERT queries on quick steps.
+         *
+         * Keyed by the step ENTITY INSTANCE, not its `ID`: a new step's `ID` is empty at create time and
+         * only gets populated during its INSERT `Save()`, so keying by `ID` would file the create and the
+         * finalize under different buckets and defeat the chain — letting the UPDATE race ahead of the
+         * INSERT on millisecond-fast steps (e.g. pipelines), which left them stuck at `Running`.
          */
         this._stepSavePromises = new Map();
         /**
@@ -285,6 +292,22 @@ export class BaseAgent {
      * @private
      */
     static { this.MAX_CONSECUTIVE_FAILED_STEPS = 10; }
+    /**
+     * Maximum consecutive *unproductive* retry steps before forcing termination.
+     *
+     * An unproductive retry is a 'Retry' next-step that carries an errorMessage — i.e. one
+     * produced by {@link BaseAgentType.createRetryStep} because the model's output could not
+     * be parsed or failed structural validation (e.g. the LLM returned conversational prose
+     * instead of the required JSON envelope). These do NOT count as 'Failed' steps, so they
+     * bypass {@link MAX_CONSECUTIVE_FAILED_STEPS} entirely and — without this guard — loop
+     * until the far-higher absolute iteration cap (effectively forever, burning time and tokens).
+     *
+     * Legitimate yield/await retries (pipeline / client-tools / sub-agent re-entry) are created
+     * via createNextStep('Retry', …) WITHOUT an errorMessage, so they do not increment this
+     * counter. Any productive (non-unproductive-retry) step resets it.
+     * @private
+     */
+    static { this.MAX_CONSECUTIVE_UNPRODUCTIVE_RETRIES = 10; }
     /**
      * Returns the active metadata provider for this agent run. Subclasses MUST
      * use this getter (rather than `new Metadata()` or `Metadata.Provider`) so
@@ -1209,6 +1232,7 @@ export class BaseAgent {
         let currentNextStep = null;
         let stepCount = 0;
         let consecutiveFailedSteps = 0;
+        let consecutiveUnproductiveRetries = 0;
         while (continueExecution) {
             // Check for cancellation before each step
             if (params.cancellationToken?.aborted) {
@@ -1248,6 +1272,37 @@ export class BaseAgent {
             else if (nextStep.step !== 'Failed') {
                 consecutiveFailedSteps = 0;
             }
+            // Track consecutive *unproductive* retries to prevent infinite loops that the
+            // consecutive-failed-steps net above cannot catch. A model that repeatedly returns
+            // output we can't parse/validate (e.g. conversational prose instead of the required
+            // JSON envelope) yields a stream of 'Retry' steps — never 'Failed' — so the failed-step
+            // counter resets every turn and never trips. Such retries are produced via
+            // createRetryStep(), which always sets an errorMessage; legitimate yield/await retries
+            // (pipeline / client-tools / sub-agent re-entry) carry no errorMessage and are exempt.
+            const isUnproductiveRetry = nextStep.step === 'Retry' && !nextStep.terminate && !!nextStep.errorMessage;
+            if (isUnproductiveRetry) {
+                consecutiveUnproductiveRetries++;
+                if (consecutiveUnproductiveRetries >= BaseAgent.MAX_CONSECUTIVE_UNPRODUCTIVE_RETRIES) {
+                    this.logError(`⛔ Agent '${params.agent.Name}' reached maximum consecutive unproductive retries ` +
+                        `(${BaseAgent.MAX_CONSECUTIVE_UNPRODUCTIVE_RETRIES}). The model is repeatedly returning output ` +
+                        `that cannot be parsed or validated. Forcing termination to prevent infinite loop.`, {
+                        agent: params.agent,
+                        category: 'ExecutionSafetyNet',
+                        metadata: {
+                            consecutiveUnproductiveRetries,
+                            lastError: nextStep.errorMessage
+                        }
+                    });
+                    nextStep.step = 'Failed';
+                    nextStep.terminate = true;
+                    nextStep.errorMessage = `Agent terminated after ${consecutiveUnproductiveRetries} consecutive unproductive retries ` +
+                        `(model repeatedly returned output that could not be parsed or validated). ` +
+                        `Last error: ${nextStep.errorMessage || 'Unknown'}`;
+                }
+            }
+            else {
+                consecutiveUnproductiveRetries = 0;
+            }
             // Check if we should continue or terminate
             if (nextStep.terminate) {
                 continueExecution = false;
@@ -1813,6 +1868,21 @@ export class BaseAgent {
             else if (this._artifactToolManager.HasArtifacts()) {
                 this.logStatus(`[ArtifactTools] Artifacts present but tools disabled by agent config (includeArtifactToolsDocs=false)`, true, params);
             }
+            // Inject pipeline tool docs when pipelines are enabled and at least one source exists.
+            // A pipeline's first step must be a source (Action or artifact tool); with none
+            // available pipelines are impossible, so BuildPipelineToolDocs returns '' and the
+            // template's `{{ _PIPELINE_TOOLS }}` block stays empty.
+            const pipelineDocsEnabled = agentTypePromptParams?.includePipelineDocs !== false;
+            if (pipelineDocsEnabled) {
+                const sourceNames = [
+                    ...this.getEffectiveActionsForValidation(params.agent.ID).map((a) => a.Name),
+                    ...this._artifactToolManager.GetAvailableToolNames(),
+                ];
+                const pipelineDocs = BuildPipelineToolDocs(sourceNames);
+                if (pipelineDocs) {
+                    promptParams.data['_PIPELINE_TOOLS'] = pipelineDocs;
+                }
+            }
             // Pass file artifacts as candidate native file inputs.
             // The AIPromptRunner will check these against the resolved driver's
             // FileCapabilities and attach qualifying files as native content blocks.
@@ -3305,9 +3375,10 @@ The context is now within limits. Please retry your request with the recovered c
         const body = toolResults.map((r, i) => {
             const heading = `### ${i + 1}. ${r.artifactId}.${r.tool}(${JSON.stringify(r.input)})`;
             if (r.result.success) {
-                const data = typeof r.result.data === 'string'
+                const raw = typeof r.result.data === 'string'
                     ? r.result.data
                     : JSON.stringify(r.result.data, null, 2);
+                const data = this.capStandaloneToolResultText(raw);
                 return `${heading}\n\`\`\`json\n${data}\n\`\`\``;
             }
             return `${heading}\n**Error:** ${r.result.errorMessage}`;
@@ -3331,6 +3402,154 @@ The context is now within limits. Please retry your request with the recovered c
         };
         params.conversationMessages.push(message);
     }
+    /**
+     * Character budget (~4 chars/token) for a SINGLE standalone artifact-tool result injected into
+     * the conversation. A `get_full` on a large artifact can otherwise dump the whole thing into
+     * context and overflow the model's window — the exact failure pipelines exist to avoid. Override
+     * in a subclass to tune. Pipelines are unaffected: their intermediate results never flow through
+     * here, and the executor already caps a pipeline's final output.
+     *
+     * @protected
+     */
+    get maxStandaloneToolResultChars() {
+        return 100_000; // ~25k tokens
+    }
+    /**
+     * Bound a standalone tool result to {@link maxStandaloneToolResultChars}: return a head slice
+     * plus a redirect that teaches the agent to page (`get_rows`) or reduce (`pipeline`) instead of
+     * reading a whole large artifact. Mirrors how read/search tools cap output at the tool boundary.
+     *
+     * @protected
+     */
+    capStandaloneToolResultText(text) {
+        const budget = this.maxStandaloneToolResultChars;
+        if (text.length <= budget) {
+            return text;
+        }
+        const omitted = text.length - budget;
+        return (text.slice(0, budget) +
+            `\n\n…[truncated ${omitted.toLocaleString()} chars. This artifact is too large to read whole ` +
+            `(~${Math.round(text.length / 4000)}k tokens) — reading it in full overflows the context window. ` +
+            `Instead: page it with get_rows(start, count), or run a pipeline that filters/aggregates it ` +
+            `server-side (where / select / groupBy → only the small final result returns to you).]`);
+    }
+    /**
+     * Builds a per-run {@link PipelineToolRegistry} that unifies the three pipeline-able
+     * substrates behind one namespace: built-in transforms, the agent's effective Actions, and
+     * the run's artifact tools. Transforms register first so their reserved names win; a source
+     * whose name collides with a transform is skipped for pipeline use (still callable normally)
+     * and logged, rather than aborting the whole pipeline.
+     *
+     * @protected
+     */
+    buildPipelineRegistry(params) {
+        const registry = new PipelineToolRegistry();
+        const register = (invocable) => {
+            try {
+                registry.Register(invocable);
+            }
+            catch (e) {
+                this.logStatus(`[Pipeline] Skipped tool "${invocable.toolName}": ${e.message}`, true, params);
+            }
+        };
+        // Operators (where/select/map/…) are pure code-defined verbs, not registry tools — only
+        // capabilities (Actions + artifact tools) live here as pipeline sources/stages.
+        // Actions — each wrapped to run via the existing single-action execution path.
+        this.getEffectiveActionsForValidation(params.agent.ID).forEach((actionEntity) => register(new ActionInvocable(actionEntity.Name, (p) => this.ExecuteSingleAction(params, { name: actionEntity.Name, params: p }, actionEntity, params.contextUser))));
+        // Artifact tools — one invocable per distinct tool name; `artifactId` is supplied as a
+        // call-time param so the same `{ tool, params }` step shape works across all substrates.
+        this._artifactToolManager.GetAvailableToolNames().forEach((toolName) => register(new ArtifactToolInvocable(toolName, async (tool, p) => {
+            const stored = await this._artifactToolManager.ExecuteSingleToolCall({
+                artifactId: String(p.artifactId ?? ''),
+                tool,
+                input: p,
+            });
+            return stored.result;
+        })));
+        return registry;
+    }
+    /**
+     * Runs a tool pipeline as a single `Tool` step in the run tree (sibling of the prompt step that
+     * requested it, matching artifact-tool steps). ALL pipeline observability lives in this step's
+     * `OutputData` — the per-stage breakdown, totals, bytes saved, and the tool chain — so there are
+     * no dedicated pipeline entities and no extra SQL I/O; the run tree alone carries everything a
+     * debug UI needs.
+     *
+     * @protected
+     */
+    async executePipelineAsStep(pipeline, params) {
+        const stepEntity = await this.createStepEntity({
+            stepType: 'Tool',
+            stepName: `Pipeline: ${pipeline.steps.length} step(s)`,
+            contextUser: params.contextUser,
+            inputData: { steps: pipeline.steps },
+        });
+        const registry = this.buildPipelineRegistry(params);
+        // The executor converts stage-level errors into a failed RESULT (it doesn't throw for those),
+        // but an unexpected throw — e.g. a tool returning a non-serializable value (BigInt/circular)
+        // that trips JSON.stringify in the executor's byte-accounting — must NEVER leave this step
+        // stuck on 'Running'. Catch it and materialize a failed result so finalize always runs and the
+        // failure surfaces as a 'Failed' step (answering "do pipeline errors show as errors?": yes).
+        let result;
+        try {
+            result = await new PipelineExecutor(registry).Execute(pipeline.steps);
+        }
+        catch (e) {
+            result = {
+                success: false,
+                finalOutput: null,
+                steps: [],
+                error: `Pipeline crashed: ${e?.message ?? String(e)}`,
+                contextBytesSaved: 0,
+            };
+        }
+        // A pipeline is ONE run-step — not a parent + a child step per stage. It runs server-side in a
+        // single fast pass, so the full per-stage breakdown + totals live in this step's OutputData for
+        // a debug UI to visualize — no separate entities, no extra DB writes.
+        await this.finalizeStepEntity(stepEntity, result.success, result.success ? undefined : result.error, {
+            success: result.success,
+            toolChain: summarizePipelineStages(result.steps),
+            steps: result.steps,
+            contextBytesSaved: result.contextBytesSaved,
+            totalBytesStreamed: result.steps.reduce((sum, s) => sum + s.outputSize, 0),
+            totalDurationMs: result.steps.reduce((sum, s) => sum + s.durationMs, 0),
+            failedStepIndex: result.failedStepIndex,
+        });
+        return result;
+    }
+    /**
+     * Pushes the pipeline's final output (or its failure message) into the conversation for the
+     * LLM's next turn, mirroring the artifact-tool "inject once, then expire" pattern. Only the
+     * final output is surfaced — intermediate step outputs never enter the context window.
+     *
+     * @protected
+     */
+    injectPipelineResultMessage(params, result) {
+        const diagnostic = result.success && result.diagnostic
+            ? `\n⚠ Empty result — ${result.diagnostic}`
+            : '';
+        // Identify which pipeline this result belongs to (stage chain, e.g. `get_rows → where →
+        // select`). Without it, multiple pipeline results across turns are indistinguishable once
+        // compacted — mirrors how artifact-tool results name their tool/artifact.
+        const label = summarizePipelineStages(result.steps);
+        const content = result.success
+            ? `Pipeline result [${label}] (final stage value — intermediate stages stayed out of context, ~${result.contextBytesSaved} bytes saved):\n\`\`\`\n${formatFinalOutput(result.finalOutput)}\n\`\`\`${diagnostic}`
+            : `Pipeline failed [${label}].\n${result.error}`;
+        const message = {
+            role: 'user',
+            content,
+            metadata: {
+                turnAdded: this._promptTurnCount,
+                messageType: 'tool-result',
+                expirationTurns: 3,
+                expirationMode: 'Compact',
+                compactMode: 'First N Chars',
+                compactLength: 500,
+                compactPromptId: '',
+            },
+        };
+        params.conversationMessages.push(message);
+    }
     /**
      * Creates a chat message containing sub-agent execution results.
      *
@@ -3523,7 +3742,8 @@ The context is now within limits. Please retry your request with the recovered c
             { docsFlag: 'includeForEachDocs', responseTypeKey: 'forEach' },
             { docsFlag: 'includeWhileDocs', responseTypeKey: 'while' },
             { docsFlag: 'includeScratchpadDocs', responseTypeKey: 'scratchpad' },
-            { docsFlag: 'includeArtifactToolsDocs', responseTypeKey: 'artifactToolCalls' }
+            { docsFlag: 'includeArtifactToolsDocs', responseTypeKey: 'artifactToolCalls' },
+            { docsFlag: 'includePipelineDocs', responseTypeKey: 'pipeline' }
         ];
         for (const { docsFlag, responseTypeKey } of alignmentMappings) {
             // Check if the user explicitly set this response type property
@@ -4468,6 +4688,13 @@ The context is now within limits. Please retry your request with the recovered c
             const errorMessage = JSON.stringify(CopyScalarsAndArrays(this._agentRun.LatestResult));
             throw new Error(`Failed to create agent run record: Details: ${errorMessage}`);
         }
+        // Hand the now-persisted run (it has a stable ID) to the watchdog so a process restart,
+        // crash, or failed terminal-state write can't leave it stuck 'Running' forever. Only the
+        // server-side DB provider can heartbeat via SQL; client/non-DB providers simply opt out.
+        const runProvider = params.provider || this._activeProvider;
+        if (runProvider instanceof DatabaseProviderBase && params.contextUser) {
+            AgentRunWatchdog.Instance.Track(this._agentRun.ID, runProvider, params.contextUser);
+        }
         // Invoke callback if provided
         if (modifiedParams.onAgentRunCreated) {
             try {
@@ -4672,15 +4899,16 @@ The context is now within limits. Please retry your request with the recovered c
      * @protected
      */
     queueStepSave(stepEntity) {
-        const id = stepEntity.ID;
-        const previousSave = this._stepSavePromises.get(id) ?? Promise.resolve();
+        // Chain on the entity INSTANCE (stable), NOT stepEntity.ID — the ID is empty until the INSERT
+        // Save() assigns it, so an ID-keyed chain breaks for fast create→finalize sequences.
+        const previousSave = this._stepSavePromises.get(stepEntity) ?? Promise.resolve();
         const currentSave = previousSave.then(() => stepEntity.Save()).then((ok) => {
             if (!ok) {
-                LogError(`Failed to save agent run step record ${id}: ${stepEntity.LatestResult?.CompleteMessage ?? 'unknown error'}`);
+                LogError(`Failed to save agent run step record ${stepEntity.ID || '(unsaved)'}: ${stepEntity.LatestResult?.CompleteMessage ?? 'unknown error'}`);
             }
             return ok;
         });
-        this._stepSavePromises.set(id, currentSave);
+        this._stepSavePromises.set(stepEntity, currentSave);
         this._pendingSaves.push(currentSave);
     }
     /**
@@ -5183,6 +5411,14 @@ The context is now within limits. Please retry your request with the recovered c
             else if (this._artifactToolManager.HasArtifacts()) {
                 this.logStatus(`[ArtifactTools] LLM did not use artifact tools this turn (artifacts available but not accessed)`, true, params);
             }
+            // Execute a tool pipeline if provided (zero turn cost — processed inline). Each step's
+            // output is threaded into the next server-side; only the final step's output returns to
+            // the LLM, so intermediate payloads never enter the context window.
+            if (initialNextStep.pipeline?.steps?.length) {
+                this.logStatus(`[Pipeline] LLM requested a ${initialNextStep.pipeline.steps.length}-stage pipeline: ${initialNextStep.pipeline.steps.map(s => s.tool ?? Object.keys(s)[0]).join(' | ')}`, true, params);
+                const pipelineResult = await this.executePipelineAsStep(initialNextStep.pipeline, params);
+                this.injectPipelineResultMessage(params, pipelineResult);
+            }
             // now that we have processed the payload, we can process the next step which does validation and changes the next step if
             // validation fails
             const updatedNextStep = await this.processNextStep(initialNextStep, params, config.agentType, promptResult, finalPayload, stepEntity);
@@ -5440,9 +5676,13 @@ The context is now within limits. Please retry your request with the recovered c
             }
         });
         // Add assistant message indicating we're executing a sub-agent
+        // Recorded as a `user`-role environment annotation (not an `assistant` turn) for the same
+        // reason as the action record above: the model's real output is the JSON envelope, and
+        // storing framework prose as an `assistant` turn trains strong in-context models to imitate
+        // the prose and drift off the required JSON format. See the note at the action-record push.
         params.conversationMessages.push({
-            role: 'assistant',
-            content: `I'm delegating this task to the "${subAgentRequest.name}" agent.\n\nReason: ${subAgentRequest.message}`
+            role: 'user',
+            content: `[You delegated this task to the "${subAgentRequest.name}" agent. Reason: ${subAgentRequest.message}]`
         });
         // Prepare input data for the step
         const inputData = {
@@ -5842,9 +6082,11 @@ The context is now within limits. Please retry your request with the recovered c
                 hierarchicalStep: this.buildHierarchicalStep(stepCount + 1, this._parentStepCounts)
             }
         });
+        // `user`-role environment annotation (not an `assistant` turn) — see the note on the
+        // single-delegation push above for why framework prose must not be stored as assistant turns.
         params.conversationMessages.push({
-            role: 'assistant',
-            content: `I'm delegating this task to the parallel sub-agent "${request.name}".\n\nReason: ${request.message}`
+            role: 'user',
+            content: `[You delegated this task to the parallel sub-agent "${request.name}". Reason: ${request.message}]`
         });
         return { request: request, subAgentEntity, relationship };
     }
@@ -6138,9 +6380,13 @@ The context is now within limits. Please retry your request with the recovered c
             }
         });
         // Add assistant message indicating we're executing a related sub-agent
+        // Recorded as a `user`-role environment annotation (not an `assistant` turn) for the same
+        // reason as the action record above: the model's real output is the JSON envelope, and
+        // storing framework prose as an `assistant` turn trains strong in-context models to imitate
+        // the prose and drift off the required JSON format. See the note at the action-record push.
         params.conversationMessages.push({
-            role: 'assistant',
-            content: `I'm delegating this task to the "${subAgentRequest.name}" agent.\n\nReason: ${subAgentRequest.message}`
+            role: 'user',
+            content: `[You delegated this task to the "${subAgentRequest.name}" agent. Reason: ${subAgentRequest.message}]`
         });
         // Prepare input data for the step
         const inputData = {
@@ -6579,12 +6825,23 @@ The context is now within limits. Please retry your request with the recovered c
                 },
                 displayMode: 'live' // Only show in live mode
             });
-            // Build detailed action execution message with parameters using markdown formatting
-            // This creates a permanent, lightweight record of what was requested
+            // Build a detailed record of the action(s) invoked, with parameters, in markdown.
+            // This is a permanent, lightweight memory of what was requested.
+            //
+            // IMPORTANT — this record is injected as a `user`-role environment annotation, NOT an
+            // `assistant` turn. The model's actual output is the JSON envelope, but we don't store
+            // that raw JSON; we store this human-readable summary instead. If it were recorded as an
+            // `assistant` turn, then after a few action-heavy turns the model's entire visible
+            // assistant history would be prose like "I'm executing the X action with parameters: …",
+            // and strong in-context learners (e.g. Gemini Flash) imitate that demonstrated pattern
+            // over the system-prompt instruction — drifting into prose and breaking JSON parsing,
+            // which (pre-guardrail) looped forever. Phrasing it in second person under the `user`
+            // role keeps the memory while removing the false assistant-prose exemplar. The
+            // human-facing narration is emitted separately via onProgress above.
             let actionMessage;
             if (actions.length === 1) {
                 const aa = actions[0];
-                actionMessage = `I'm executing the **${aa.name}** action`;
+                actionMessage = `[You invoked the **${aa.name}** action`;
                 // Add parameters if they exist
                 if (aa.params && Object.keys(aa.params).length > 0) {
                     const paramsList = Object.entries(aa.params)
@@ -6593,14 +6850,14 @@ The context is now within limits. Please retry your request with the recovered c
                         return `• **${key}**: ${displayValue}`;
                     })
                         .join('\n');
-                    actionMessage += ` with parameters:\n${paramsList}`;
+                    actionMessage += ` with parameters:\n${paramsList}\n]`;
                 }
                 else {
-                    actionMessage += '.';
+                    actionMessage += '.]';
                 }
             }
             else {
-                actionMessage = `I'm executing **${actions.length} actions** in parallel:\n\n` + actions.map((aa, index) => {
+                actionMessage = `[You invoked **${actions.length} actions** in parallel:\n\n` + actions.map((aa, index) => {
                     let actionText = `${index + 1}. **${aa.name}**`;
                     // Add parameters if they exist
                     if (aa.params && Object.keys(aa.params).length > 0) {
@@ -6613,12 +6870,13 @@ The context is now within limits. Please retry your request with the recovered c
                         actionText += `\n${paramsList}`;
                     }
                     return actionText;
-                }).join('\n\n');
+                }).join('\n\n') + '\n]';
             }
             if (addConversationMessage) {
-                // Add assistant message (no metadata - this is a permanent record)
+                // Record as a `user`-role environment annotation (no metadata - permanent record).
+                // See the note above on why this is NOT an `assistant` turn.
                 params.conversationMessages.push({
-                    role: 'assistant',
+                    role: 'user',
                     content: actionMessage
                 });
             }
@@ -7901,6 +8159,8 @@ The context is now within limits. Please retry your request with the recovered c
             this._agentRun.TotalTokensUsed = tokenStats.totalTokens;
             this._agentRun.TotalPromptTokensUsed = tokenStats.promptTokens;
             this._agentRun.TotalCompletionTokensUsed = tokenStats.completionTokens;
+            this._agentRun.TotalCacheReadTokensUsed = tokenStats.cacheReadTokens;
+            this._agentRun.TotalCacheWriteTokensUsed = tokenStats.cacheWriteTokens;
             this._agentRun.TotalCost = tokenStats.totalCost;
             await this._agentRun.Save();
         }
@@ -7927,6 +8187,8 @@ The context is now within limits. Please retry your request with the recovered c
             this._agentRun.TotalTokensUsed = tokenStats.totalTokens;
             this._agentRun.TotalPromptTokensUsed = tokenStats.promptTokens;
             this._agentRun.TotalCompletionTokensUsed = tokenStats.completionTokens;
+            this._agentRun.TotalCacheReadTokensUsed = tokenStats.cacheReadTokens;
+            this._agentRun.TotalCacheWriteTokensUsed = tokenStats.cacheWriteTokens;
             this._agentRun.TotalCost = tokenStats.totalCost;
             await this._agentRun.Save();
         }
@@ -8008,6 +8270,8 @@ The context is now within limits. Please retry your request with the recovered c
             this._agentRun.TotalTokensUsed = tokenStats.totalTokens;
             this._agentRun.TotalPromptTokensUsed = tokenStats.promptTokens;
             this._agentRun.TotalCompletionTokensUsed = tokenStats.completionTokens;
+            this._agentRun.TotalCacheReadTokensUsed = tokenStats.cacheReadTokens;
+            this._agentRun.TotalCacheWriteTokensUsed = tokenStats.cacheWriteTokens;
             this._agentRun.TotalCost = tokenStats.totalCost;
             const ok = await this._agentRun.Save();
             if (!ok) {
@@ -8046,15 +8310,19 @@ The context is now within limits. Please retry your request with the recovered c
         let totalTokens = 0;
         let promptTokens = 0;
         let completionTokens = 0;
+        let cacheReadTokens = 0;
+        let cacheWriteTokens = 0;
         let totalCost = 0;
         // Iterate through the agent run's steps to sum up tokens
         if (this._agentRun?.Steps) {
             for (const step of this._agentRun.Steps) {
                 if (step.StepType === 'Prompt' && step.PromptRun) {
-                    // Add tokens from prompt runs
+                    // Add tokens from prompt runs (rollup fields include any nested child prompt runs)
                     totalTokens += step.PromptRun.TokensUsedRollup || 0;
                     promptTokens += step.PromptRun.TokensPromptRollup || 0;
                     completionTokens += step.PromptRun.TokensCompletionRollup || 0;
+                    cacheReadTokens += step.PromptRun.TokensCacheReadRollup || 0;
+                    cacheWriteTokens += step.PromptRun.TokensCacheWriteRollup || 0;
                     totalCost += step.PromptRun.TotalCost || 0;
                 }
                 else if (step.StepType === 'Sub-Agent' && step.SubAgentRun) {
@@ -8062,11 +8330,13 @@ The context is now within limits. Please retry your request with the recovered c
                     totalTokens += step.SubAgentRun.TotalTokensUsed || 0;
                     promptTokens += step.SubAgentRun.TotalPromptTokensUsed || 0;
                     completionTokens += step.SubAgentRun.TotalCompletionTokensUsed || 0;
+                    cacheReadTokens += step.SubAgentRun.TotalCacheReadTokensUsed || 0;
+                    cacheWriteTokens += step.SubAgentRun.TotalCacheWriteTokensUsed || 0;
                     totalCost += step.SubAgentRun.TotalCost || 0;
                 }
             }
         }
-        return { totalTokens, promptTokens, completionTokens, totalCost };
+        return { totalTokens, promptTokens, completionTokens, cacheReadTokens, cacheWriteTokens, totalCost };
     }
     /**
      * Gets the count of how many times a specific action has been executed in this agent run.