zeitlich 0.2.32 → 0.2.33

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "zeitlich",
-  "version": "0.2.32",
+  "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",
  *     ),
  *   };
  * }

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

@@ -3,6 +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 { getActivityContext } from "../../../lib/activity";
 
 export interface AnthropicModelInvokerConfig {
   redis: Redis;
@@ -27,8 +28,8 @@ 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`, 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
@@ -44,7 +45,7 @@ function toAnthropicTools(
  *   model: 'claude-sonnet-4-20250514',
  * });
  *
- * return { runAgent: createRunAgentActivity(client, invoker) };
+ * return { ...createRunAgentActivity(client, invoker, "myAgent") };
  * ```
  */
 export function createAnthropicModelInvoker({
@@ -58,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();
@@ -65,13 +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?.();
+    }
+
+    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",
  *     ),
  *   };
  * }

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

@@ -1,8 +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 { getActivityContext } from "../../../lib/activity";
 
 export interface GoogleGenAIModelInvokerConfig {
   redis: Redis;
@@ -25,8 +26,8 @@ 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`, 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
@@ -42,7 +43,7 @@ function toFunctionDeclarations(
  *   model: 'gemini-2.5-flash',
  * });
  *
- * return { runAgent: createRunAgentActivity(client, invoker) };
+ * return { ...createRunAgentActivity(client, invoker, "myAgent") };
  * ```
  */
 export function createGoogleGenAIModelInvoker({
@@ -55,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 } =
@@ -64,19 +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?.();
+    }
 
-    const functionCalls = response.functionCalls ?? [];
+    if (!lastChunk) {
+      throw new Error("Google GenAI stream ended without producing any chunks");
+    }
+
+    const modelContent: Content = { role: "model", parts: allParts };
+    const functionCalls = lastChunk.functionCalls ?? [];
 
     return {
       message: modelContent,
@@ -86,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"),
  *   };
  * }
  * ```

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,10 @@ 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,
- * 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
@@ -29,51 +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,
+      });
+
+      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);
+    }
   };
 }
 
@@ -83,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/workflow.ts

@@ -133,6 +133,9 @@ export type {
   ModelInvokerConfig,
 } from "./lib/model";
 
+// Model proxy (workflow-safe proxy with LLM-optimised defaults)
+export { proxyRunAgent } from "./lib/model/proxy";
+
 // Subagent types
 export type {
   SubagentConfig,

package/src/lib/.env

@@ -1 +0,0 @@
-E2B_API_KEY=e2b_39af116424059782e2aee6942fd70237cc2126c9

package/src/tools/bash/.env

@@ -1 +0,0 @@
-E2B_API_KEY=e2b_39af116424059782e2aee6942fd70237cc2126c9