zeitlich 0.2.31 → 0.2.33

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "zeitlich",
-  "version": "0.2.31",
+  "version": "0.2.33",
   "description": "[EXPERIMENTAL] An opinionated AI agent implementation for Temporal",
   "main": "./dist/index.js",
   "module": "./dist/index.js",

package/src/adapters/thread/anthropic/activities.ts

@@ -107,7 +107,7 @@ export interface AnthropicAdapter {
  * export function createActivities(temporalClient: WorkflowClient) {
  *   return {
  *     ...adapter.createActivities("codingAgent"),
- *     runCodingAgent: createRunAgentActivity(temporalClient, adapter.invoker),
+ *     ...createRunAgentActivity(temporalClient, adapter.invoker, "codingAgent"),
  *   };
  * }
  * ```
@@ -118,10 +118,11 @@ export interface AnthropicAdapter {
  *   return {
  *     ...adapter.createActivities("codingAgent"),
  *     ...adapter.createActivities("researchAgent"),
- *     runCodingAgent: createRunAgentActivity(temporalClient, adapter.invoker),
- *     runResearchAgent: createRunAgentActivity(
+ *     ...createRunAgentActivity(temporalClient, adapter.invoker, "codingAgent"),
+ *     ...createRunAgentActivity(
  *       temporalClient,
  *       adapter.createModelInvoker('claude-sonnet-4-20250514'),
+ *       "researchAgent",
  *     ),
  *   };
  * }
@@ -164,6 +165,16 @@ export function createAnthropicAdapter(
       await thread.appendToolResult(id, toolCallId, toolName, content);
     },
 
+    async appendAgentMessage(
+      threadId: string,
+      id: string,
+      message: Anthropic.Messages.Message,
+      threadKey?: string,
+    ): Promise<void> {
+      const thread = createAnthropicThreadManager({ redis, threadId, key: threadKey });
+      await thread.appendAssistantMessage(id, message.content);
+    },
+
     async forkThread(
       sourceThreadId: string,
       targetThreadId: string,

package/src/adapters/thread/anthropic/model-invoker.ts

@@ -3,7 +3,7 @@ import type Anthropic from "@anthropic-ai/sdk";
 import type { SerializableToolDefinition } from "../../../lib/types";
 import type { AgentResponse, ModelInvokerConfig } from "../../../lib/model";
 import { createAnthropicThreadManager, type AnthropicThreadManagerHooks } from "./thread-manager";
-import { v4 as uuidv4 } from "uuid";
+import { getActivityContext } from "../../../lib/activity";
 
 export interface AnthropicModelInvokerConfig {
   redis: Redis;
@@ -28,9 +28,9 @@ function toAnthropicTools(
  * Creates an Anthropic model invoker that satisfies the generic
  * `ModelInvoker<Anthropic.Messages.Message>` contract.
  *
- * Loads the conversation thread from Redis, invokes the Claude model via
- * `client.messages.create`, appends the AI response, and returns
- * a normalised AgentResponse.
+ * Internally streams the response and emits Temporal heartbeats on each
+ * event so that long-running LLM calls remain visible to the scheduler.
+ * The caller is responsible for appending the response to the thread.
  *
  * @example
  * ```typescript
@@ -45,7 +45,7 @@ function toAnthropicTools(
  *   model: 'claude-sonnet-4-20250514',
  * });
  *
- * return { runAgent: createRunAgentActivity(client, invoker) };
+ * return { ...createRunAgentActivity(client, invoker, "myAgent") };
  * ```
  */
 export function createAnthropicModelInvoker({
@@ -59,6 +59,7 @@ export function createAnthropicModelInvoker({
     config: ModelInvokerConfig,
   ): Promise<AgentResponse<Anthropic.Messages.Message>> {
     const { threadId, threadKey, state } = config;
+    const { heartbeat, signal } = getActivityContext();
 
     const thread = createAnthropicThreadManager({ redis, threadId, key: threadKey, hooks });
     const { messages, system } = await thread.prepareForInvocation();
@@ -66,15 +67,21 @@ export function createAnthropicModelInvoker({
     const anthropicTools = toAnthropicTools(state.tools);
     const tools = anthropicTools.length > 0 ? anthropicTools : undefined;
 
-    const response = await client.messages.create({
+    const params: Anthropic.MessageCreateParams = {
       model,
       max_tokens: maxTokens,
       messages,
       ...(system ? { system } : {}),
       ...(tools ? { tools } : {}),
-    });
+    };
+
+    const stream = client.messages.stream(params, { signal });
+
+    for await (const _event of stream) {
+      heartbeat?.();
+    }
 
-    await thread.appendAssistantMessage(uuidv4(), response.content);
+    const response: Anthropic.Messages.Message = await stream.finalMessage();
 
     const toolCalls = response.content.filter(
       (block): block is Anthropic.Messages.ToolUseBlock =>

package/src/adapters/thread/google-genai/activities.ts

@@ -113,7 +113,7 @@ export interface GoogleGenAIAdapter {
  * export function createActivities(temporalClient: WorkflowClient) {
  *   return {
  *     ...adapter.createActivities("codingAgent"),
- *     runCodingAgent: createRunAgentActivity(temporalClient, adapter.invoker),
+ *     ...createRunAgentActivity(temporalClient, adapter.invoker, "codingAgent"),
  *   };
  * }
  * ```
@@ -124,10 +124,11 @@ export interface GoogleGenAIAdapter {
  *   return {
  *     ...adapter.createActivities("codingAgent"),
  *     ...adapter.createActivities("researchAgent"),
- *     runCodingAgent: createRunAgentActivity(temporalClient, adapter.invoker),
- *     runResearchAgent: createRunAgentActivity(
+ *     ...createRunAgentActivity(temporalClient, adapter.invoker, "codingAgent"),
+ *     ...createRunAgentActivity(
  *       temporalClient,
  *       adapter.createModelInvoker('gemini-2.5-pro'),
+ *       "researchAgent",
  *     ),
  *   };
  * }
@@ -194,6 +195,20 @@ export function createGoogleGenAIAdapter(
       );
     },
 
+    async appendAgentMessage(
+      threadId: string,
+      id: string,
+      message: Content,
+      threadKey?: string,
+    ): Promise<void> {
+      const thread = createGoogleGenAIThreadManager({
+        redis,
+        threadId,
+        key: threadKey,
+      });
+      await thread.appendModelContent(id, message.parts ?? []);
+    },
+
     async forkThread(
       sourceThreadId: string,
       targetThreadId: string,

package/src/adapters/thread/google-genai/model-invoker.ts

@@ -1,9 +1,9 @@
 import type Redis from "ioredis";
-import type { GoogleGenAI, Content, FunctionDeclaration } from "@google/genai";
+import type { GoogleGenAI, Content, FunctionDeclaration, Part, GenerateContentResponse } from "@google/genai";
 import type { SerializableToolDefinition } from "../../../lib/types";
 import type { AgentResponse, ModelInvokerConfig } from "../../../lib/model";
 import { createGoogleGenAIThreadManager, type GoogleGenAIThreadManagerHooks } from "./thread-manager";
-import { v4 as uuidv4 } from "uuid";
+import { getActivityContext } from "../../../lib/activity";
 
 export interface GoogleGenAIModelInvokerConfig {
   redis: Redis;
@@ -26,9 +26,9 @@ function toFunctionDeclarations(
  * Creates a Google GenAI model invoker that satisfies the generic
  * `ModelInvoker<Content>` contract.
  *
- * Loads the conversation thread from Redis, invokes the Gemini model via
- * `client.models.generateContent`, appends the AI response, and returns
- * a normalised AgentResponse.
+ * Internally streams the response and emits Temporal heartbeats on each
+ * chunk so that long-running LLM calls remain visible to the scheduler.
+ * The caller is responsible for appending the response to the thread.
  *
  * @example
  * ```typescript
@@ -43,7 +43,7 @@ function toFunctionDeclarations(
  *   model: 'gemini-2.5-flash',
  * });
  *
- * return { runAgent: createRunAgentActivity(client, invoker) };
+ * return { ...createRunAgentActivity(client, invoker, "myAgent") };
  * ```
  */
 export function createGoogleGenAIModelInvoker({
@@ -56,6 +56,7 @@ export function createGoogleGenAIModelInvoker({
     config: ModelInvokerConfig,
   ): Promise<AgentResponse<Content>> {
     const { threadId, threadKey, state } = config;
+    const { heartbeat, signal } = getActivityContext();
 
     const thread = createGoogleGenAIThreadManager({ redis, threadId, key: threadKey, hooks });
     const { contents, systemInstruction } =
@@ -65,21 +66,30 @@ export function createGoogleGenAIModelInvoker({
     const tools =
       functionDeclarations.length > 0 ? [{ functionDeclarations }] : undefined;
 
-    const response = await client.models.generateContent({
+    const stream = await client.models.generateContentStream({
       model,
       contents,
       config: {
         ...(systemInstruction ? { systemInstruction } : {}),
         ...(tools ? { tools } : {}),
+        abortSignal: signal,
       },
     });
 
-    const responseParts = response.candidates?.[0]?.content?.parts ?? [];
-    const modelContent: Content = { role: "model", parts: responseParts };
+    const allParts: Part[] = [];
+    let lastChunk: GenerateContentResponse | undefined;
+    for await (const chunk of stream) {
+      lastChunk = chunk;
+      allParts.push(...(chunk.candidates?.[0]?.content?.parts ?? []));
+      heartbeat?.();
+    }
 
-    await thread.appendModelContent(uuidv4(), responseParts);
+    if (!lastChunk) {
+      throw new Error("Google GenAI stream ended without producing any chunks");
+    }
 
-    const functionCalls = response.functionCalls ?? [];
+    const modelContent: Content = { role: "model", parts: allParts };
+    const functionCalls = lastChunk.functionCalls ?? [];
 
     return {
       message: modelContent,
@@ -89,9 +99,9 @@ export function createGoogleGenAIModelInvoker({
         args: fc.args ?? {},
       })),
       usage: {
-        inputTokens: response.usageMetadata?.promptTokenCount,
-        outputTokens: response.usageMetadata?.candidatesTokenCount,
-        cachedReadTokens: response.usageMetadata?.cachedContentTokenCount,
+        inputTokens: lastChunk.usageMetadata?.promptTokenCount,
+        outputTokens: lastChunk.usageMetadata?.candidatesTokenCount,
+        cachedReadTokens: lastChunk.usageMetadata?.cachedContentTokenCount,
       },
     };
   };

package/src/adapters/thread/langchain/activities.ts

@@ -94,7 +94,7 @@ export interface LangChainAdapter {
  * export function createActivities(client: WorkflowClient) {
  *   return {
  *     ...adapter.createActivities("codingAgent"),
- *     runCodingAgent: createRunAgentActivity(client, adapter.invoker),
+ *     ...createRunAgentActivity(client, adapter.invoker, "codingAgent"),
  *   };
  * }
  * ```
@@ -105,8 +105,8 @@ export interface LangChainAdapter {
  *   return {
  *     ...adapter.createActivities("codingAgent"),
  *     ...adapter.createActivities("researchAgent"),
- *     runCodingAgent: createRunAgentActivity(client, adapter.invoker),
- *     runResearchAgent: createRunAgentActivity(client, adapter.createModelInvoker(claude)),
+ *     ...createRunAgentActivity(client, adapter.invoker, "codingAgent"),
+ *     ...createRunAgentActivity(client, adapter.createModelInvoker(claude), "researchAgent"),
  *   };
  * }
  * ```
@@ -148,6 +148,17 @@ export function createLangChainAdapter(
       await thread.appendToolResult(id, toolCallId, "", content);
     },
 
+    async appendAgentMessage(
+      threadId: string,
+      id: string,
+      message: StoredMessage,
+      threadKey?: string,
+    ): Promise<void> {
+      const thread = createLangChainThreadManager({ redis, threadId, key: threadKey });
+      const patched = { ...message, data: { ...message.data, id } };
+      await thread.append([patched]);
+    },
+
     async forkThread(
       sourceThreadId: string,
       targetThreadId: string,

package/src/adapters/thread/langchain/model-invoker.ts

@@ -3,10 +3,16 @@ import type { AgentResponse, ModelInvokerConfig } from "../../../lib/model";
 import type { StoredMessage } from "@langchain/core/messages";
 import { v4 as uuidv4 } from "uuid";
 import type { BaseChatModel } from "@langchain/core/language_models/chat_models";
-import { createLangChainThreadManager, type LangChainThreadManagerHooks } from "./thread-manager";
+import {
+  createLangChainThreadManager,
+  type LangChainThreadManagerHooks,
+} from "./thread-manager";
+import { getActivityContext } from "../../../lib/activity";
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export interface LangChainModelInvokerConfig<TModel extends BaseChatModel<any> = BaseChatModel<any>> {
+export interface LangChainModelInvokerConfig<
+  TModel extends BaseChatModel<any> = BaseChatModel<any>,
+> {
   redis: Redis;
   model: TModel;
   hooks?: LangChainThreadManagerHooks;
@@ -16,8 +22,11 @@ export interface LangChainModelInvokerConfig<TModel extends BaseChatModel<any> =
  * Creates a LangChain-based model invoker that satisfies the generic
  * `ModelInvoker<StoredMessage>` contract.
  *
- * Loads the conversation thread from Redis, invokes a LangChain chat model,
- * appends the AI response, and returns a normalised AgentResponse.
+ * Uses interval-based Temporal heartbeats during model.invoke() to keep
+ * long-running LLM calls visible to the scheduler. LangChain's streaming
+ * chunk accumulation is unreliable across providers (e.g. reasoning_content
+ * blocks don't merge correctly), so we use invoke() for correctness.
+ * The caller is responsible for appending the response to the thread.
  *
  * @example
  * ```typescript
@@ -28,53 +37,70 @@ export interface LangChainModelInvokerConfig<TModel extends BaseChatModel<any> =
  * const model = new ChatAnthropic({ model: "claude-sonnet-4-6" });
  * const invoker = createLangChainModelInvoker({ redis, model });
  *
- * return { runAgent: createRunAgentActivity(client, invoker) };
+ * return { ...createRunAgentActivity(client, invoker, "myAgent") };
  * ```
  */
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function createLangChainModelInvoker<TModel extends BaseChatModel<any> = BaseChatModel<any>>(
-  { redis, model, hooks }: LangChainModelInvokerConfig<TModel>,
-) {
+export function createLangChainModelInvoker<
+  TModel extends BaseChatModel<any> = BaseChatModel<any>,
+>({ redis, model, hooks }: LangChainModelInvokerConfig<TModel>) {
   return async function invokeLangChainModel(
-    config: ModelInvokerConfig,
+    config: ModelInvokerConfig
   ): Promise<AgentResponse<StoredMessage>> {
     const { threadId, threadKey, agentName, state, metadata } = config;
+    const { heartbeat, signal } = getActivityContext();
 
-    const thread = createLangChainThreadManager({ redis, threadId, key: threadKey, hooks });
+    const thread = createLangChainThreadManager({
+      redis,
+      threadId,
+      key: threadKey,
+      hooks,
+    });
     const runId = uuidv4();
 
     const { messages } = await thread.prepareForInvocation();
-    const response = await model.invoke(
-      messages,
-      {
+
+    const heartbeatInterval = heartbeat
+      ? setInterval(() => heartbeat(), 30_000)
+      : undefined;
+
+    try {
+      const response = await model.invoke(messages, {
         runName: agentName,
         runId,
         metadata: { thread_id: `${agentName}-${threadId}`, ...metadata },
         tools: state.tools,
-      },
-    );
+        signal,
+      });
 
-    await thread.append([response.toDict()]);
+      const toolCalls = response.tool_calls ?? [];
 
-    const toolCalls = response.tool_calls ?? [];
+      const providerUsage =
+        (response.response_metadata?.usage as Record<string, unknown>) ?? {};
 
-    return {
-      message: response.toDict(),
-      rawToolCalls: toolCalls.map((tc) => ({
-        id: tc.id,
-        name: tc.name,
-        args: tc.args,
-      })),
-      usage: {
-        inputTokens: response.usage_metadata?.input_tokens,
-        outputTokens: response.usage_metadata?.output_tokens,
-        reasonTokens: response.usage_metadata?.output_token_details?.reasoning,
-        cachedWriteTokens:
-          response.usage_metadata?.input_token_details?.cache_creation,
-        cachedReadTokens:
-          response.usage_metadata?.input_token_details?.cache_read,
-      },
-    };
+      return {
+        message: response.toDict(),
+        rawToolCalls: toolCalls.map((tc) => ({
+          id: tc.id,
+          name: tc.name,
+          args: tc.args,
+        })),
+        usage: {
+          inputTokens: response.usage_metadata?.input_tokens,
+          outputTokens: response.usage_metadata?.output_tokens,
+          reasonTokens:
+            response.usage_metadata?.output_token_details?.reasoning,
+          cachedWriteTokens:
+            response.usage_metadata?.input_token_details?.cache_creation ||
+            (providerUsage.cacheWriteInputTokens as number | undefined),
+          cachedReadTokens:
+            response.usage_metadata?.input_token_details?.cache_read ||
+            (providerUsage.cacheReadInputTokens as number | undefined),
+        },
+      };
+    } finally {
+      if (heartbeatInterval) clearInterval(heartbeatInterval);
+    }
   };
 }
 
@@ -84,7 +110,9 @@ export function createLangChainModelInvoker<TModel extends BaseChatModel<any> =
  * you don't need to reuse the invoker.
  */
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export async function invokeLangChainModel<TModel extends BaseChatModel<any> = BaseChatModel<any>>({
+export async function invokeLangChainModel<
+  TModel extends BaseChatModel<any> = BaseChatModel<any>,
+>({
   redis,
   model,
   hooks,

package/src/index.ts

@@ -52,6 +52,7 @@ export {
   queryParentWorkflowState,
   createRunAgentActivity,
   withParentWorkflowState,
+  getActivityContext,
 } from "./lib/activity";
 export type { AgentStateContext } from "./lib/activity";
 

package/src/lib/activity.ts

@@ -8,6 +8,22 @@ import type {
   ToolHandlerResponse,
 } from "./tool-router/types";
 
+/**
+ * Safely retrieve Temporal activity heartbeat and cancellation signal.
+ * Returns empty object when called outside a Temporal activity (e.g. tests).
+ */
+export function getActivityContext(): {
+  heartbeat?: () => void;
+  signal?: AbortSignal;
+} {
+  try {
+    const ctx = Context.current();
+    return { heartbeat: () => ctx.heartbeat(), signal: ctx.cancellationSignal };
+  } catch {
+    return {};
+  }
+}
+
 /**
  * Query the parent workflow's state from within an activity.
  * Resolves the workflow handle from the current activity context.
@@ -24,25 +40,36 @@ export async function queryParentWorkflowState<T>(
 }
 
 /**
- * Wraps a handler into a `RunAgentActivity` by auto-fetching the parent
- * workflow's agent state before each invocation.
+ * Wraps a handler into a scope-prefixed `RunAgentActivity` by auto-fetching
+ * the parent workflow's agent state before each invocation.
+ *
+ * Returns a `Record` with a single key `run<Scope>` so it can be spread
+ * into the activities object alongside adapter activities.
+ *
+ * @param scope - Workflow scope used to derive the activity name.
+ *   `"myAgentWorkflow"` produces `{ runMyAgentWorkflow: fn }`.
  *
  * @example
  * ```typescript
  * import { createRunAgentActivity } from 'zeitlich';
- * import { createLangChainModelInvoker } from 'zeitlich/adapters/thread/langchain';
  *
- * const invoker = createLangChainModelInvoker({ redis, model });
- * return { runAgent: createRunAgentActivity(client, invoker) };
+ * return {
+ *   ...adapter.createActivities("myAgentWorkflow"),
+ *   ...createRunAgentActivity(client, adapter.invoker, "myAgentWorkflow"),
+ * };
  * ```
  */
 export function createRunAgentActivity<R, S extends BaseAgentState = BaseAgentState>(
   client: WorkflowClient,
   handler: (config: RunAgentConfig & { state: S }) => Promise<R>,
-): (config: RunAgentConfig) => Promise<R> {
-  return async (config: RunAgentConfig) => {
-    const state = await queryParentWorkflowState<S>(client);
-    return handler({ ...config, state });
+  scope: string,
+): Record<string, (config: RunAgentConfig) => Promise<R>> {
+  const name = `run${scope.charAt(0).toUpperCase()}${scope.slice(1)}`;
+  return {
+    [name]: async (config: RunAgentConfig) => {
+      const state = await queryParentWorkflowState<S>(client);
+      return handler({ ...config, state });
+    },
   };
 }
 

package/src/lib/model/helpers.ts

@@ -2,5 +2,6 @@ export {
   queryParentWorkflowState,
   createRunAgentActivity,
   withParentWorkflowState,
+  getActivityContext,
 } from "../activity";
 export type { AgentStateContext } from "../activity";

package/src/lib/model/index.ts

@@ -2,6 +2,7 @@ export {
   queryParentWorkflowState,
   createRunAgentActivity,
   withParentWorkflowState,
+  getActivityContext,
 } from "./helpers";
 export type { AgentStateContext } from "./helpers";
 

package/src/lib/model/proxy.ts

@@ -0,0 +1,50 @@
+/**
+ * Workflow-safe proxy for runAgent activities with LLM-optimised defaults.
+ *
+ * Resolves the activity name from the scope using the same convention as
+ * {@link createRunAgentActivity}: `run<Scope>`.
+ * When no scope is provided, defaults to `workflowInfo().workflowType`.
+ *
+ * Import this from `zeitlich/workflow` in your Temporal workflow files.
+ *
+ * @typeParam M - SDK-native message type (e.g. `StoredMessage` for LangChain,
+ *   `Anthropic.Messages.Message` for Anthropic, `Content` for Google GenAI).
+ *   Must be provided for `SessionResult.finalMessage` to be correctly typed.
+ *
+ * @example
+ * ```typescript
+ * import { proxyRunAgent } from 'zeitlich/workflow';
+ * import type { StoredMessage } from '@langchain/core/messages';
+ *
+ * // Auto-scoped to the current workflow name
+ * const runAgent = proxyRunAgent<StoredMessage>();
+ *
+ * // Explicit scope for subagents
+ * const runResearcher = proxyRunAgent<StoredMessage>("Researcher");
+ * ```
+ */
+import { proxyActivities, workflowInfo } from "@temporalio/workflow";
+import type { AgentResponse } from "./types";
+import type { RunAgentConfig } from "../types";
+
+export function proxyRunAgent<M = unknown>(
+  scope?: string,
+  options?: Parameters<typeof proxyActivities>[0],
+): (config: RunAgentConfig) => Promise<AgentResponse<M>> {
+  const resolvedScope = scope ?? workflowInfo().workflowType;
+  const name = `run${resolvedScope.charAt(0).toUpperCase()}${resolvedScope.slice(1)}`;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const acts = proxyActivities<Record<string, (...args: any[]) => any>>(
+    options ?? {
+      startToCloseTimeout: "10m",
+      heartbeatTimeout: "1m",
+      retry: {
+        maximumAttempts: 3,
+        initialInterval: "10s",
+        maximumInterval: "2m",
+        backoffCoefficient: 3,
+      },
+    },
+  );
+  return acts[name] as (config: RunAgentConfig) => Promise<AgentResponse<M>>;
+}

package/src/lib/model/types.ts

@@ -33,8 +33,9 @@ export interface ModelInvokerConfig {
 
 /**
  * Generic model invocation contract.
- * Implementations load the thread, call the LLM, append the response,
- * and return a normalised AgentResponse.
+ * Implementations load the thread, call the LLM, and return a normalised
+ * AgentResponse. The caller (workflow) is responsible for appending the
+ * response to the thread with a deterministic ID.
  *
  * Framework adapters (e.g. `zeitlich/langchain`) provide concrete
  * implementations of this type.

package/src/lib/session/session-edge-cases.integration.test.ts

@@ -88,6 +88,9 @@ function createMockThreadOps() {
     appendSystemMessage: async (threadId, id, content) => {
       log.push({ op: "appendSystemMessage", args: [threadId, id, content] });
     },
+    appendAgentMessage: async (threadId, id, message) => {
+      log.push({ op: "appendAgentMessage", args: [threadId, id, message] });
+    },
     forkThread: async (source, target) => {
       log.push({ op: "forkThread", args: [source, target] });
     },
@@ -752,6 +755,9 @@ describe("createSession edge cases", () => {
       appendSystemMessage: async (threadId, id, content) => {
         log.push({ op: "appendSystemMessage", args: [threadId, id, content] });
       },
+      appendAgentMessage: async (threadId, id, message) => {
+        log.push({ op: "appendAgentMessage", args: [threadId, id, message] });
+      },
       forkThread: async (source, target) => {
         log.push({ op: "forkThread", args: [source, target] });
       },

package/src/lib/session/session.integration.test.ts

@@ -99,6 +99,9 @@ function createMockThreadOps() {
     appendSystemMessage: async (threadId, id, content) => {
       log.push({ op: "appendSystemMessage", args: [threadId, id, content] });
     },
+    appendAgentMessage: async (threadId, id, message) => {
+      log.push({ op: "appendAgentMessage", args: [threadId, id, message] });
+    },
     forkThread: async (source, target) => {
       log.push({ op: "forkThread", args: [source, target] });
     },

package/src/lib/session/session.ts

@@ -139,6 +139,7 @@ export async function createSession<
     appendHumanMessage,
     initializeThread,
     appendSystemMessage,
+    appendAgentMessage,
     forkThread,
   } = threadOps;
 
@@ -285,6 +286,7 @@ export async function createSession<
           : result.fileTree;
         stateManager.mergeUpdate({
           fileTree,
+          virtualFsCtx: virtualFsConfig.ctx,
           ...(skillFiles && { inlineFiles: skillFiles }),
         } as Partial<AgentState<TState>>);
       }
@@ -355,6 +357,8 @@ export async function createSession<
             metadata,
           });
 
+          await appendAgentMessage(threadId, uuid4(), message, threadKey);
+
           if (usage) {
             stateManager.updateUsage(usage);
           }

package/src/lib/session/types.ts

@@ -39,6 +39,13 @@ export interface ThreadOps<TContent = string> {
   ): Promise<void>;
   /** Append a tool result to the thread */
   appendToolResult(id: string, config: ToolResultConfig): Promise<void>;
+  /** Append the model's response to the thread */
+  appendAgentMessage(
+    threadId: string,
+    id: string,
+    message: unknown,
+    threadKey?: string
+  ): Promise<void>;
   /** Append a system message to the thread */
   appendSystemMessage(
     threadId: string,