@rudderjs/ai 1.2.0 → 1.4.0

package/dist/agent.js

@@ -1,8 +1,9 @@
 import { z } from 'zod';
 import { AiRegistry } from './registry.js';
-import { isPauseForClientToolsChunk, toolDefinition, toolToSchema } from './tool.js';
+import { isPauseForClientToolsChunk, pauseForClientTools, toolDefinition, toolToSchema } from './tool.js';
 import { attachmentsToContentParts, getMessageText } from './attachment.js';
 import { QueuedPromptBuilder } from './queue-job.js';
+import { resolveAutoPersistSpec, runWithPersistence, runWithPersistenceStreaming, } from './conversation-persistence.js';
 import { runOnConfig, runOnChunk, runOnBeforeToolCall, runOnAfterToolCall, runSequential, runOnUsage, runOnAbort, runOnError, } from './middleware.js';
 // ─── AI Observer (lazy accessor) ─────────────────────────
 function _getAiObservers() {
@@ -55,6 +56,60 @@ export class Agent {
     temperature() { return undefined; }
     /** Max tokens for response */
     maxTokens() { return undefined; }
+    /**
+     * Declarative prompt-cache configuration.
+     *
+     * Override on a subclass to mark stable parts of the prompt as cacheable
+     * — provider adapters translate to native primitives (Anthropic
+     * `cache_control`, OpenAI `prompt_cache_key`, Google `cachedContent`)
+     * so cache hits can save 50–90% on input tokens for long system prompts,
+     * tool definitions, or stable conversation context.
+     *
+     * Returning `undefined` (the default) means no caching. Per-call override
+     * via `agent.prompt(input, { cache: false })` disables caching for that
+     * call; passing a {@link CacheableConfig} for `cache` replaces the agent
+     * default for that call.
+     *
+     * @example
+     * class SupportAgent extends Agent {
+     *   instructions() { return LONG_SYSTEM_PROMPT }
+     *   tools()        { return [tool1, tool2, tool3] }
+     *   cacheable() {
+     *     return { instructions: true, tools: true }
+     *   }
+     * }
+     */
+    cacheable() { return undefined; }
+    /**
+     * Opt into auto-persisted conversation behavior. Override on a subclass
+     * to declare *which* user owns the thread and (optionally) which
+     * specific thread, and the framework will load history before each
+     * `prompt()`/`stream()` call and append the new turn after it — without
+     * any caller having to remember `forUser()` / `continue()`.
+     *
+     * Returning `false` (the default) disables auto-persist; the agent runs
+     * stateless. Returning a {@link ConversationalSpec} opts in:
+     *
+     * @example
+     * class ChatAgent extends Agent {
+     *   conversational() {
+     *     return { user: Auth.user()?.id }   // null user → falsy → opt-out
+     *   }
+     * }
+     *
+     * await new ChatAgent().prompt('Hi')          // auto-loads + auto-saves
+     *
+     * **Precedence (high → low):**
+     * 1. Explicit `agent.forUser(id).prompt()` / `agent.continue(id).prompt()`
+     * 2. Per-call `prompt(input, { conversation: false | {...} })`
+     * 3. This method's return value
+     *
+     * Async returns are supported — useful when the user identity is fetched
+     * from an async DI binding.
+     */
+    conversational() {
+        return false;
+    }
     /**
      * Default for `AgentPromptOptions.parallelTools`. When `true` (default),
      * multiple tool calls within a single step run their `execute()` functions
@@ -64,11 +119,15 @@ export class Agent {
     parallelTools() { return true; }
     /** Run the agent with a prompt (non-streaming) */
     async prompt(input, options) {
+        const spec = await resolveAutoPersistSpec(() => this.conversational(), options?.conversation);
+        if (spec) {
+            return runWithPersistence(spec, this.constructor.name, resolveConversationStore, input, options, (effOptions) => runAgentLoop(this, input, effOptions));
+        }
         return runAgentLoop(this, input, options);
     }
     /** Run the agent with a prompt (streaming) */
     stream(input, options) {
-        return runAgentLoopStreaming(this, input, options);
+        return runStreamWithMaybeAutoPersist(this, input, options);
     }
     /** Queue the prompt for background execution */
     queue(input, options) {
@@ -83,17 +142,201 @@ export class Agent {
         return new ConversableAgent(this).continue(conversationId);
     }
     asTool(options) {
+        if (options.suspendable && !options.streaming) {
+            throw new Error('[RudderJS AI] asTool: `suspendable` requires `streaming: true` (or a projector). Silent suspend would leave the parent UI with no progress signal between sub-agent invocations.');
+        }
         const schema = options.inputSchema ?? z.object({ prompt: z.string() });
         const promptOf = options.prompt ?? ((input) => input.prompt);
         const modelOutput = options.modelOutput ?? ((response) => response.text);
+        if (!options.streaming) {
+            // 1.2.0 zero-config path — single prompt() call, single AgentResponse out.
+            return toolDefinition({
+                name: options.name,
+                description: options.description,
+                inputSchema: schema,
+            })
+                .server((input) => this.prompt(promptOf(input)))
+                .modelOutput(modelOutput);
+        }
+        const project = options.streaming === true ? defaultSubAgentProjector : options.streaming;
+        const innerAgent = this; // eslint-disable-line @typescript-eslint/no-this-alias
+        const agentName = options.name;
+        const suspendable = options.suspendable;
+        const generatorExecute = async function* (input) {
+            const userPrompt = promptOf(input);
+            yield { kind: 'agent_start', agentName };
+            const streamOpts = suspendable
+                ? { toolCallStreamingMode: 'stop-on-client-tool' }
+                : undefined;
+            const { stream, response } = innerAgent.stream(userPrompt, streamOpts);
+            for await (const chunk of stream) {
+                const update = project(chunk);
+                if (update)
+                    yield update;
+            }
+            const result = await response;
+            if (suspendable &&
+                result.finishReason === 'client_tool_calls' &&
+                result.pendingClientToolCalls?.length) {
+                const subRunId = generateSubRunId();
+                const snapshot = {
+                    messages: buildSubAgentSnapshotMessages(userPrompt, result),
+                    pendingToolCallIds: result.pendingClientToolCalls.map((tc) => tc.id),
+                    stepsSoFar: result.steps.length,
+                    tokensSoFar: result.usage?.totalTokens ?? 0,
+                };
+                await suspendable.runStore.store(subRunId, snapshot);
+                yield { kind: 'subagent_paused', subRunId, pendingToolCallIds: snapshot.pendingToolCallIds };
+                yield pauseForClientTools(result.pendingClientToolCalls, subRunId);
+                // Unreachable — the parent loop halts iteration after the pause chunk.
+                return undefined;
+            }
+            yield {
+                kind: 'agent_done',
+                steps: result.steps.length,
+                tokens: result.usage?.totalTokens ?? 0,
+            };
+            return result;
+        };
         return toolDefinition({
             name: options.name,
             description: options.description,
             inputSchema: schema,
         })
-            .server((input) => this.prompt(promptOf(input)))
+            .server(generatorExecute)
             .modelOutput(modelOutput);
     }
+    /**
+     * Resume a sub-agent run that previously paused with
+     * `pauseForClientTools` (typically from {@link Agent.asTool} with
+     * `suspendable: { runStore }` set). Loads the snapshot, validates the
+     * incoming tool-result ids against the pending set, and re-runs the
+     * inner loop with those results appended.
+     *
+     * Returns either a `'completed'` result (the inner agent finished) or
+     * a `'paused'` continuation pointing at a fresh `subRunId` for the
+     * next round-trip.
+     *
+     * @example
+     * const r = await Agent.resumeAsTool(subRunId, browserResults, { runStore, agent: subAgent })
+     * if (r.kind === 'completed') {
+     *   feedToolResultBackToParent(r.response.text)
+     * } else {
+     *   emitPendingClientToolsSse(r.subRunId, r.pendingToolCallIds)
+     * }
+     */
+    static async resumeAsTool(subRunId, clientToolResults, options) {
+        const snapshot = await options.runStore.consume(subRunId);
+        if (!snapshot) {
+            throw new Error(`[RudderJS AI] resumeAsTool: subRunId "${subRunId}" expired or never existed.`);
+        }
+        // Forgery guard — every incoming tool-result id must be in the pending set.
+        const pending = new Set(snapshot.pendingToolCallIds);
+        const seen = new Set();
+        for (const r of clientToolResults) {
+            if (!pending.has(r.toolCallId)) {
+                throw new Error(`[RudderJS AI] resumeAsTool: toolCallId "${r.toolCallId}" was not in the pending set.`);
+            }
+            if (seen.has(r.toolCallId)) {
+                throw new Error(`[RudderJS AI] resumeAsTool: duplicate result for toolCallId "${r.toolCallId}".`);
+            }
+            seen.add(r.toolCallId);
+        }
+        // Append client tool-result messages to the snapshot, in incoming order.
+        const messages = [...snapshot.messages];
+        for (const r of clientToolResults) {
+            messages.push({
+                role: 'tool',
+                content: typeof r.result === 'string' ? r.result : JSON.stringify(r.result),
+                toolCallId: r.toolCallId,
+            });
+        }
+        const result = await options.agent.prompt('', {
+            messages,
+            toolCallStreamingMode: 'stop-on-client-tool',
+        });
+        if (result.finishReason === 'client_tool_calls' &&
+            result.pendingClientToolCalls?.length) {
+            const newSubRunId = generateSubRunId();
+            const newSnapshot = {
+                messages: buildResumeSnapshotMessages(messages, result),
+                pendingToolCallIds: result.pendingClientToolCalls.map((tc) => tc.id),
+                stepsSoFar: snapshot.stepsSoFar + result.steps.length,
+                tokensSoFar: snapshot.tokensSoFar + (result.usage?.totalTokens ?? 0),
+                ...(snapshot.meta !== undefined ? { meta: snapshot.meta } : {}),
+            };
+            await options.runStore.store(newSubRunId, newSnapshot);
+            return {
+                kind: 'paused',
+                subRunId: newSubRunId,
+                pendingToolCallIds: newSnapshot.pendingToolCallIds,
+            };
+        }
+        return { kind: 'completed', response: result };
+    }
+}
+/**
+ * Default projection from inner-agent stream chunks to {@link SubAgentUpdate}
+ * events. Emits one `tool_call` per inner `tool-call` chunk; everything
+ * else is suppressed (the wrapping execute emits the `agent_start` /
+ * `agent_done` bookends and the suspend path emits `subagent_paused`).
+ *
+ * Hosts wanting different cadence (e.g. surfacing `text-delta` previews
+ * or per-step usage) pass `streaming: chunk => …` and own the discriminator.
+ */
+function defaultSubAgentProjector(chunk) {
+    if (chunk.type === 'tool-call' && chunk.toolCall?.name) {
+        return {
+            kind: 'tool_call',
+            tool: chunk.toolCall.name,
+            ...(chunk.toolCall.arguments ? { args: chunk.toolCall.arguments } : {}),
+        };
+    }
+    return null;
+}
+/**
+ * Reconstruct the inner-agent message history at the point the loop
+ * paused, so a subsequent {@link Agent.resumeAsTool} can rerun the loop
+ * with the appended client tool results. The shape is `[user, …(message
+ * + serverToolResults)*]` — system messages are omitted because the
+ * `messages` mode of the agent loop prepends `system` itself.
+ *
+ * Each step's `message` includes ALL `toolCalls` (server + client).
+ * Server-side `toolResults` are interleaved; client-side calls remain
+ * unfulfilled until resume appends their results.
+ */
+function buildSubAgentSnapshotMessages(userPrompt, response) {
+    const out = [{ role: 'user', content: userPrompt }];
+    for (const step of response.steps) {
+        out.push(step.message);
+        for (const tr of step.toolResults) {
+            const resultStr = typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result);
+            out.push({ role: 'tool', content: resultStr, toolCallId: tr.toolCallId });
+        }
+    }
+    return out;
+}
+/**
+ * Snapshot reconstruction for a resume-time pause. The `priorMessages`
+ * already include the original user prompt + every step prior to the
+ * resume call. Append the freshly-completed steps' messages and any
+ * server-side tool results so the next resume sees the full history.
+ */
+function buildResumeSnapshotMessages(priorMessages, response) {
+    const out = [...priorMessages];
+    for (const step of response.steps) {
+        out.push(step.message);
+        for (const tr of step.toolResults) {
+            const resultStr = typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result);
+            out.push({ role: 'tool', content: resultStr, toolCallId: tr.toolCallId });
+        }
+    }
+    return out;
+}
+function generateSubRunId() {
+    if (typeof globalThis.crypto?.randomUUID === 'function')
+        return globalThis.crypto.randomUUID();
+    return `sub-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 12)}`;
 }
 // ─── Conversable Agent (conversation persistence) ───────
 /**
@@ -116,84 +359,35 @@ export class ConversableAgent {
         return this;
     }
     async prompt(input, options) {
-        const store = resolveConversationStore();
-        if (!store)
-            throw new Error('[RudderJS AI] No ConversationStore registered. Register one via the DI container with key "ai.conversations".');
-        // Load or create conversation
-        let history = options?.history ?? [];
-        if (this._conversationId) {
-            history = [...(await store.load(this._conversationId)), ...history];
-        }
-        else {
-            const meta = this._userId ? { userId: this._userId } : undefined;
-            this._conversationId = await store.create(undefined, meta);
-        }
-        const response = await runAgentLoop(this.agent, input, { ...options, history });
-        // Persist messages
-        const newMessages = [
-            { role: 'user', content: input },
-            ...response.steps.flatMap(s => {
-                const msgs = [s.message];
-                for (const tr of s.toolResults) {
-                    const resultStr = typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result);
-                    msgs.push({ role: 'tool', content: resultStr, toolCallId: tr.toolCallId });
-                }
-                return msgs;
-            }),
-        ];
-        await store.append(this._conversationId, newMessages);
-        return { text: response.text, steps: response.steps, usage: response.usage, conversationId: this._conversationId };
+        const spec = this.toSpec();
+        return runWithPersistence(spec, this.agent.constructor.name, resolveConversationStore, input, options, (effOptions) => runAgentLoop(this.agent, input, effOptions)).then((r) => {
+            // Track the resolved id back on the wrapper so a subsequent
+            // `wrapper.prompt()` call resumes the same thread.
+            if (r.conversationId)
+                this._conversationId = r.conversationId;
+            return r;
+        });
     }
     stream(input, options) {
-        const store = resolveConversationStore();
-        if (!store)
-            throw new Error('[RudderJS AI] No ConversationStore registered. Register one via the DI container with key "ai.conversations".');
-        // We need to handle async setup, so wrap the streaming
-        let resolveReady;
-        const ready = new Promise(r => { resolveReady = r; });
-        let loadedHistory = [];
-        let convId = this._conversationId;
-        // Kick off async setup
-        const setupPromise = (async () => {
-            if (convId) {
-                loadedHistory = await store.load(convId);
-            }
-            else {
-                const meta = this._userId ? { userId: this._userId } : undefined;
-                convId = await store.create(undefined, meta);
-                this._conversationId = convId;
-            }
-            resolveReady();
-        })();
-        let resolveResponse;
-        const responsePromise = new Promise(r => { resolveResponse = r; });
-        const self = this; // eslint-disable-line @typescript-eslint/no-this-alias
-        const storeRef = store;
-        async function* generateStream() {
-            await setupPromise;
-            const history = [...loadedHistory, ...(options?.history ?? [])];
-            const inner = runAgentLoopStreaming(self.agent, input, { ...options, history });
-            for await (const chunk of inner.stream) {
-                yield chunk;
-            }
-            const response = await inner.response;
-            // Persist messages
-            const newMessages = [
-                { role: 'user', content: input },
-                ...response.steps.flatMap(s => {
-                    const msgs = [s.message];
-                    for (const tr of s.toolResults) {
-                        const resultStr = typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result);
-                        msgs.push({ role: 'tool', content: resultStr, toolCallId: tr.toolCallId });
-                    }
-                    return msgs;
-                }),
-            ];
-            await storeRef.append(convId, newMessages);
-            const result = { text: response.text, steps: response.steps, usage: response.usage, conversationId: convId };
-            resolveResponse(result);
-        }
-        return { stream: generateStream(), response: responsePromise };
+        const spec = this.toSpec();
+        const persisted = runWithPersistenceStreaming(spec, this.agent.constructor.name, resolveConversationStore, input, options, (effOptions) => runAgentLoopStreaming(this.agent, input, effOptions));
+        // Update the wrapper's id once the run completes.
+        persisted.response.then((r) => { if (r.conversationId)
+            this._conversationId = r.conversationId; }, () => { });
+        return persisted;
+    }
+    /**
+     * Translate the wrapper's explicit-form state (`forUser` / `continue`)
+     * into a {@link ConversationalSpec}. The explicit chain bypasses the
+     * agent's `conversational()` declaration entirely — `forUser` always
+     * wins over class defaults.
+     */
+    toSpec() {
+        if (this._conversationId)
+            return { user: this._userId ?? '', id: this._conversationId };
+        if (this._userId)
+            return { user: this._userId };
+        throw new Error('[RudderJS AI] ConversableAgent requires forUser() or continue() to be called before prompt().');
     }
 }
 // ─── Anonymous Agent ─────────────────────────────────────
@@ -243,6 +437,76 @@ export function setConversationStore(store) {
 function resolveConversationStore() {
     return _conversationStore;
 }
+/**
+ * Streaming counterpart of `Agent.prompt`'s auto-persist branch. The spec
+ * resolution is async (since `conversational()` may return a Promise), so
+ * we defer the decision into the outer wrapper that handles the inner
+ * stream's setup the same way `runWithPersistenceStreaming` does for the
+ * persisted path.
+ */
+function runStreamWithMaybeAutoPersist(a, input, options) {
+    // Synchronous fast path — most agents don't override `conversational()`,
+    // so we'd pay an extra microtask boundary on every streaming call. Bail
+    // out cheaply when we can prove the call is stateless.
+    const declared = a.conversational();
+    const isFast = (options?.conversation === false ||
+        (declared === false && (options?.conversation === undefined)));
+    if (isFast) {
+        return runAgentLoopStreaming(a, input, options);
+    }
+    // Async path — resolve the spec, then dispatch to the persisted or plain stream.
+    let resolveResp;
+    let rejectResp;
+    const responsePromise = new Promise((res, rej) => { resolveResp = res; rejectResp = rej; });
+    async function* outer() {
+        let spec;
+        try {
+            spec = await resolveAutoPersistSpec(() => a.conversational(), options?.conversation);
+        }
+        catch (err) {
+            rejectResp(err);
+            throw err;
+        }
+        if (!spec) {
+            const inner = runAgentLoopStreaming(a, input, options);
+            try {
+                for await (const chunk of inner.stream)
+                    yield chunk;
+            }
+            catch (err) {
+                rejectResp(err);
+                throw err;
+            }
+            try {
+                const r = await inner.response;
+                resolveResp(r);
+            }
+            catch (err) {
+                rejectResp(err);
+                throw err;
+            }
+            return;
+        }
+        const persisted = runWithPersistenceStreaming(spec, a.constructor.name, resolveConversationStore, input, options, (effOptions) => runAgentLoopStreaming(a, input, effOptions));
+        try {
+            for await (const chunk of persisted.stream)
+                yield chunk;
+        }
+        catch (err) {
+            rejectResp(err);
+            throw err;
+        }
+        try {
+            const r = await persisted.response;
+            resolveResp(r);
+        }
+        catch (err) {
+            rejectResp(err);
+            throw err;
+        }
+    }
+    return { stream: outer(), response: responsePromise };
+}
 // ─── Helpers ─────────────────────────────────────────────
 function getTools(a) {
     return 'tools' in a && typeof a.tools === 'function'
@@ -328,6 +592,7 @@ async function runFailover(loopCtx, currentModel, call) {
                 temperature: loopCtx.agent.temperature(),
                 maxTokens: loopCtx.agent.maxTokens(),
                 signal: loopCtx.options?.signal,
+                cache: resolveCacheMarkers(loopCtx.agent, loopCtx.options),
             };
             return await call(adapter, modelId, reqOptions);
         }
@@ -344,6 +609,39 @@ async function runFailover(loopCtx, currentModel, call) {
     }
     throw lastError ?? new Error('No provider available');
 }
+/**
+ * Merge agent-level `cacheable()` declaration with per-call override.
+ *
+ * - Per-call `cache: false` → returns `undefined` (caching disabled).
+ * - Per-call `cache: { ... }` → replaces the agent default.
+ * - Per-call omitted → uses `agent.cacheable()` unchanged.
+ *
+ * Returns `undefined` when no markers are set so the provider request
+ * carries no `cache` field at all.
+ */
+function resolveCacheMarkers(agent, options) {
+    if (options && options.cache === false)
+        return undefined;
+    const perCall = options?.cache === false ? undefined : options?.cache;
+    const config = perCall ?? agent.cacheable();
+    if (!config)
+        return undefined;
+    const markers = {};
+    if (config.instructions)
+        markers.instructions = true;
+    if (config.tools)
+        markers.tools = true;
+    if (config.messages !== undefined && config.messages > 0) {
+        markers.messages = Math.floor(config.messages);
+    }
+    if (config.ttl)
+        markers.ttl = config.ttl;
+    // ttl alone with no region markers is meaningless — drop it.
+    const hasRegion = markers.instructions || markers.tools || (markers.messages && markers.messages > 0);
+    if (!hasRegion)
+        return undefined;
+    return markers;
+}
 /** Emit the `agent.failed` observer event from the shared loop state. */
 function emitObserverFailed(loopCtx, err, streaming) {
     const obs = _getAiObservers();