zeitlich 0.2.6 → 0.2.8

package/src/lib/model-invoker.ts

@@ -1,6 +1,6 @@
 import type Redis from "ioredis";
 import { createThreadManager } from "./thread-manager";
-import type { AgentResponse } from "./types";
+import type { AgentResponse, BaseAgentState } from "./types";
 import { Context } from "@temporalio/activity";
 import type { WorkflowClient } from "@temporalio/client";
 import { mapStoredMessagesToChatMessages } from "@langchain/core/messages";
@@ -10,7 +10,6 @@ import type {
   BaseChatModelCallOptions,
   BindToolsInput,
 } from "@langchain/core/language_models/chat_models";
-import { getStateQuery } from "./state-manager";
 
 /**
  * Configuration for invoking the model
@@ -49,7 +48,7 @@ export async function invokeModel({
   const parentRunId = info.workflowExecution.runId;
 
   const handle = client.getHandle(parentWorkflowId, parentRunId);
-  const { tools } = await handle.query(getStateQuery);
+  const { tools } = await handle.query<BaseAgentState>(`get${agentName}State`);
 
   const messages = await thread.load();
   const response = await model.invoke(
@@ -74,9 +73,13 @@ export async function invokeModel({
       args: tc.args,
     })),
     usage: {
-      input_tokens: response.usage_metadata?.input_tokens,
-      output_tokens: response.usage_metadata?.output_tokens,
-      total_tokens: response.usage_metadata?.total_tokens,
+      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,
     },
   };
 }

package/src/lib/session.ts

@@ -1,11 +1,17 @@
-import { proxyActivities } from "@temporalio/workflow";
+import {
+  proxyActivities,
+  condition,
+  defineUpdate,
+  setHandler,
+} from "@temporalio/workflow";
 import type { ZeitlichSharedActivities } from "../activities";
 import type {
   ThreadOps,
-  ZeitlichAgentConfig,
+  AgentConfig,
   SessionStartHook,
   SessionEndHook,
   SessionExitReason,
+  SessionConfig,
 } from "./types";
 import { type AgentStateManager, type JsonSerializable } from "./state-manager";
 import {
@@ -13,11 +19,16 @@ import {
   type ParsedToolCallUnion,
   type ToolMap,
 } from "./tool-router";
+import type { MessageContent } from "@langchain/core/messages";
 
 export interface ZeitlichSession<M = unknown> {
   runSession<T extends JsonSerializable<T>>(args: {
     stateManager: AgentStateManager<T>;
-  }): Promise<M | null>;
+  }): Promise<{
+    finalMessage: M | null;
+    exitReason: SessionExitReason;
+    usage: ReturnType<AgentStateManager<T>["getTotalUsage"]>;
+  }>;
 }
 
 /**
@@ -33,6 +44,7 @@ export interface SessionLifecycleHooks {
 export const createSession = async <T extends ToolMap, M = unknown>({
   threadId,
   agentName,
+  description,
   maxTurns = 50,
   metadata = {},
   runAgent,
@@ -44,13 +56,15 @@ export const createSession = async <T extends ToolMap, M = unknown>({
   hooks = {},
   appendSystemPrompt = true,
   systemPrompt,
-}: ZeitlichAgentConfig<T, M>): Promise<ZeitlichSession<M>> => {
+  waitForInputTimeout = "48h",
+}: SessionConfig<T, M> & AgentConfig): Promise<ZeitlichSession<M>> => {
   const {
     appendToolResult,
     appendHumanMessage,
     initializeThread,
     appendSystemMessage,
   } = threadOps ?? proxyDefaultThreadOps();
+
   const toolRouter = createToolRouter({
     tools,
     appendToolResult,
@@ -77,7 +91,33 @@ export const createSession = async <T extends ToolMap, M = unknown>({
   };
 
   return {
-    runSession: async ({ stateManager }): Promise<M | null> => {
+    runSession: async ({
+      stateManager,
+    }): Promise<{
+      finalMessage: M | null;
+      exitReason: SessionExitReason;
+      usage: ReturnType<typeof stateManager.getTotalUsage>;
+    }> => {
+      setHandler(
+        defineUpdate<unknown, [MessageContent]>(`add${agentName}Message`),
+        async (message: MessageContent) => {
+          if (hooks.onPreHumanMessageAppend) {
+            await hooks.onPreHumanMessageAppend({
+              message,
+              threadId,
+            });
+          }
+          await appendHumanMessage(threadId, message);
+          if (hooks.onPostHumanMessageAppend) {
+            await hooks.onPostHumanMessageAppend({
+              message,
+              threadId,
+            });
+          }
+          stateManager.run();
+        }
+      );
+
       if (hooks.onSessionStart) {
         await hooks.onSessionStart({
           threadId,
@@ -86,8 +126,6 @@ export const createSession = async <T extends ToolMap, M = unknown>({
         });
       }
 
-      stateManager.setTools(toolRouter.getToolDefinitions());
-
       await initializeThread(threadId);
       if (appendSystemPrompt && systemPrompt && systemPrompt.trim() !== "") {
         await appendSystemMessage(threadId, systemPrompt);
@@ -105,17 +143,29 @@ export const createSession = async <T extends ToolMap, M = unknown>({
           stateManager.incrementTurns();
           const currentTurn = stateManager.getTurns();
 
-          const { message, rawToolCalls } = await runAgent({
+          stateManager.setTools(toolRouter.getToolDefinitions());
+
+          const { message, rawToolCalls, usage } = await runAgent({
             threadId,
             agentName,
             metadata,
+            systemPrompt,
+            description,
           });
 
+          if (usage) {
+            stateManager.updateUsage(usage);
+          }
+
           // No tools configured - treat any non-end_turn as completed
           if (!toolRouter.hasTools() || rawToolCalls.length === 0) {
             stateManager.complete();
             exitReason = "completed";
-            return message;
+            return {
+              finalMessage: message,
+              exitReason,
+              usage: stateManager.getTotalUsage(),
+            };
           }
 
           // Parse all tool calls uniformly through the router
@@ -136,13 +186,30 @@ export const createSession = async <T extends ToolMap, M = unknown>({
           }
 
           // Hooks can call stateManager.waitForInput() to pause the session
-          await toolRouter.processToolCalls(parsedToolCalls, {
-            turn: currentTurn,
-          });
+          const toolCallResults = await toolRouter.processToolCalls(
+            parsedToolCalls,
+            {
+              turn: currentTurn,
+            }
+          );
+
+          for (const result of toolCallResults) {
+            if (result.usage) {
+              stateManager.updateUsage(result.usage);
+            }
+          }
 
           if (stateManager.getStatus() === "WAITING_FOR_INPUT") {
-            exitReason = "waiting_for_input";
-            break;
+            const conditionMet = await condition(
+              () => stateManager.getStatus() === "RUNNING",
+              waitForInputTimeout
+            );
+            if (!conditionMet) {
+              stateManager.cancel();
+              // Wait briefly to allow pending waitForStateChange handlers to complete
+              await condition(() => false, "2s");
+              break;
+            }
           }
         }
 
@@ -158,7 +225,11 @@ export const createSession = async <T extends ToolMap, M = unknown>({
         await callSessionEnd(exitReason, stateManager.getTurns());
       }
 
-      return null;
+      return {
+        finalMessage: null,
+        exitReason,
+        usage: stateManager.getTotalUsage(),
+      };
     },
   };
 };

package/src/lib/state-manager.ts

@@ -1,5 +1,11 @@
-import { defineQuery, setHandler } from "@temporalio/workflow";
 import {
+  condition,
+  defineQuery,
+  defineUpdate,
+  setHandler,
+} from "@temporalio/workflow";
+import {
+  type AgentConfig,
   type AgentStatus,
   type BaseAgentState,
   type WorkflowTask,
@@ -108,9 +114,26 @@ export interface AgentStateManager<TCustom extends JsonSerializable<TCustom>> {
 
   /** Set the tools (converts Zod schemas to JSON Schema for serialization) */
   setTools(newTools: ToolDefinition[]): void;
-}
 
-export const getStateQuery = defineQuery<BaseAgentState>("getState");
+  /** Update the usage */
+  updateUsage(usage: {
+    inputTokens?: number;
+    outputTokens?: number;
+    cachedWriteTokens?: number;
+    cachedReadTokens?: number;
+    reasonTokens?: number;
+  }): void;
+
+  /** Get the total usage */
+  getTotalUsage(): {
+    totalInputTokens: number;
+    totalOutputTokens: number;
+    totalCachedWriteTokens: number;
+    totalCachedReadTokens: number;
+    totalReasonTokens: number;
+    turns: number;
+  };
+}
 
 /**
  * Creates an agent state manager for tracking workflow state.
@@ -124,14 +147,23 @@ export const getStateQuery = defineQuery<BaseAgentState>("getState");
  */
 export function createAgentStateManager<
   TCustom extends JsonSerializable<TCustom> = Record<string, never>,
->(
-  initialState?: Partial<BaseAgentState> & TCustom
-): AgentStateManager<TCustom> {
+>({
+  initialState,
+  agentConfig,
+}: {
+  initialState?: Partial<BaseAgentState> & TCustom;
+  agentConfig: AgentConfig;
+}): AgentStateManager<TCustom> {
   // Default state (BaseAgentState fields)
   let status: AgentStatus = initialState?.status ?? "RUNNING";
   let version = initialState?.version ?? 0;
   let turns = initialState?.turns ?? 0;
   let tools = initialState?.tools ?? [];
+  let totalInputTokens = 0;
+  let totalOutputTokens = 0;
+  let totalCachedWriteTokens = 0;
+  let totalCachedReadTokens = 0;
+  let totalReasonTokens = 0;
 
   // Tasks state
   const tasks = new Map<string, WorkflowTask>(initialState?.tasks);
@@ -157,10 +189,23 @@ export function createAgentStateManager<
     } as AgentState<TCustom>;
   }
 
-  setHandler(getStateQuery, () => {
+  setHandler(defineQuery(`get${agentConfig.agentName}State`), () => {
     return buildState();
   });
 
+  setHandler(
+    defineUpdate<AgentState<TCustom>, [number]>(
+      `waitFor${agentConfig.agentName}StateChange`
+    ),
+    async (lastKnownVersion: number) => {
+      await condition(
+        () => version > lastKnownVersion || isTerminalStatus(status),
+        "55s"
+      );
+      return buildState();
+    }
+  );
+
   return {
     getStatus(): AgentStatus {
       return status;
@@ -262,6 +307,38 @@ export function createAgentStateManager<
       }
       return deleted;
     },
+
+    updateUsage(usage: {
+      inputTokens?: number;
+      outputTokens?: number;
+      cachedWriteTokens?: number;
+      cachedReadTokens?: number;
+      reasonTokens?: number;
+    }): void {
+      totalInputTokens += usage.inputTokens ?? 0;
+      totalOutputTokens += usage.outputTokens ?? 0;
+      totalCachedWriteTokens += usage.cachedWriteTokens ?? 0;
+      totalCachedReadTokens += usage.cachedReadTokens ?? 0;
+      totalReasonTokens += usage.reasonTokens ?? 0;
+    },
+
+    getTotalUsage(): {
+      totalInputTokens: number;
+      totalOutputTokens: number;
+      totalCachedWriteTokens: number;
+      totalCachedReadTokens: number;
+      totalReasonTokens: number;
+      turns: number;
+    } {
+      return {
+        totalInputTokens,
+        totalOutputTokens,
+        totalCachedWriteTokens,
+        totalCachedReadTokens,
+        totalReasonTokens,
+        turns,
+      };
+    },
   };
 }
 

package/src/lib/tool-router.ts

@@ -6,6 +6,7 @@ import type {
   PreToolUseHookResult,
   SubagentConfig,
   SubagentHooks,
+  TokenUsage,
   ToolHooks,
   ToolResultConfig,
 } from "./types";
@@ -53,7 +54,7 @@ export interface ToolWithHandler<
   strict?: boolean;
   max_uses?: number;
   /** Whether this tool is available to the agent (default: true). Disabled tools are excluded from definitions and rejected at parse time. */
-  enabled?: boolean;
+  enabled?: () => boolean;
   /** Per-tool lifecycle hooks (run in addition to global hooks) */
   hooks?: ToolHooks<z.infer<TSchema>, TResult>;
 }
@@ -76,7 +77,7 @@ export type ToolMap = Record<
     handler: ToolHandler<any, any, any>;
     strict?: boolean;
     max_uses?: number;
-    enabled?: boolean;
+    enabled?: () => boolean;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     hooks?: ToolHooks<any, any>;
   }
@@ -153,6 +154,8 @@ export interface ToolHandlerResponse<TResult = null> {
    * payloads through Temporal's activity payload limit.
    */
   resultAppended?: boolean;
+  /** Token usage from the tool execution (e.g. child agent invocations) */
+  usage?: TokenUsage;
 }
 
 /**
@@ -228,6 +231,7 @@ export interface ToolCallResult<
   toolCallId: string;
   name: TName;
   data: TResult;
+  usage?: TokenUsage;
 }
 
 /**
@@ -412,16 +416,14 @@ export function createToolRouter<T extends ToolMap>(
   }
 
   /** Check if a tool is enabled (defaults to true when not specified) */
-  const isEnabled = (tool: ToolMap[string]): boolean => tool.enabled !== false;
+  const isEnabled = (tool: ToolMap[string]): boolean =>
+    tool.enabled?.() ?? true;
 
   if (options.subagents) {
-    const enabledSubagents = options.subagents.filter(
-      (s) => s.enabled !== false
-    );
-    if (enabledSubagents.length > 0) {
+    if (options.subagents.length > 0) {
       // Build per-subagent hook dispatcher keyed by subagent name
       const subagentHooksMap = new Map<string, SubagentHooks>();
-      for (const s of enabledSubagents) {
+      for (const s of options.subagents) {
         if (s.hooks) subagentHooksMap.set(s.agentName, s.hooks);
       }
 
@@ -429,8 +431,8 @@ export function createToolRouter<T extends ToolMap>(
         (args as SubagentArgs).subagent;
 
       toolMap.set("Subagent", {
-        ...createSubagentTool(enabledSubagents),
-        handler: createSubagentHandler(enabledSubagents),
+        ...createSubagentTool(options.subagents),
+        handler: createSubagentHandler(options.subagents),
         ...(subagentHooksMap.size > 0 && {
           hooks: {
             onPreToolUse: async (ctx): Promise<PreToolUseHookResult> => {

package/src/lib/types.ts

@@ -8,6 +8,7 @@ import type {
 } from "./tool-router";
 
 import type { MessageContent, StoredMessage } from "@langchain/core/messages";
+import type { Duration } from "@temporalio/common";
 import type { Workflow } from "@temporalio/workflow";
 import type { z } from "zod";
 
@@ -31,6 +32,10 @@ export interface BaseAgentState {
   turns: number;
   tasks: Map<string, WorkflowTask>;
   systemPrompt: string;
+  totalInputTokens: number;
+  totalOutputTokens: number;
+  cachedWriteTokens: number;
+  cachedReadtTokens: number;
 }
 
 /**
@@ -49,17 +54,21 @@ export interface AgentFile {
   mimeType?: string;
 }
 
+export interface TokenUsage {
+  inputTokens?: number;
+  outputTokens?: number;
+  cachedWriteTokens?: number;
+  cachedReadTokens?: number;
+  reasonTokens?: number;
+}
+
 /**
  * Agent response from LLM invocation
  */
 export interface AgentResponse<M = StoredMessage> {
   message: M;
   rawToolCalls: RawToolCall[];
-  usage?: {
-    input_tokens?: number;
-    output_tokens?: number;
-    total_tokens?: number;
-  };
+  usage?: TokenUsage;
 }
 
 /**
@@ -82,16 +91,28 @@ export interface ThreadOps {
 }
 
 /**
- * Configuration for a Zeitlich agent session
+ * Configuration for a Zeitlich agent
  */
-export interface ZeitlichAgentConfig<T extends ToolMap, M = StoredMessage> {
-  threadId: string;
+export interface AgentConfig {
+  /** The name of the agent, should be unique within the workflows, ideally Pascal Case */
   agentName: string;
   /** Description, used for sub agents */
   description?: string;
+  /** The system prompt to append to the thread */
   systemPrompt?: string;
+}
+
+/**
+ * Configuration for a Zeitlich agent session
+ */
+export interface SessionConfig<T extends ToolMap, M = StoredMessage> {
+  /** The thread ID to use for the session */
+  threadId: string;
+  /** Metadata for the session */
   metadata?: Record<string, unknown>;
+  /** Whether to append the system prompt as message to the thread */
   appendSystemPrompt?: boolean;
+  /** How many turns to run the session for */
   maxTurns?: number;
   /** Workflow-specific runAgent activity (with tools pre-bound) */
   runAgent: RunAgentActivity<M>;
@@ -110,6 +131,8 @@ export interface ZeitlichAgentConfig<T extends ToolMap, M = StoredMessage> {
    * Returns MessageContent array for the initial HumanMessage.
    */
   buildContextMessage: () => MessageContent | Promise<MessageContent>;
+  /** How long to wait for input before cancelling the workflow */
+  waitForInputTimeout?: Duration;
 }
 
 /**
@@ -128,9 +151,10 @@ export interface SerializableToolDefinition {
 /**
  * Configuration passed to runAgent activity
  */
-export interface RunAgentConfig {
+export interface RunAgentConfig extends AgentConfig {
+  /** The thread ID to use for the session */
   threadId: string;
-  agentName: string;
+  /** Metadata for the session */
   metadata?: Record<string, unknown>;
 }
 
@@ -140,6 +164,7 @@ export interface RunAgentConfig {
 export type RunAgentActivity<M = StoredMessage> = (
   config: RunAgentConfig
 ) => Promise<AgentResponse<M>>;
+
 /**
  * Configuration for appending a tool result
  */
@@ -171,7 +196,7 @@ export interface SubagentConfig<TResult extends z.ZodType = z.ZodType> {
   /** Description shown to the parent agent explaining what this subagent does */
   description: string;
   /** Whether this subagent is available (default: true). Disabled subagents are excluded from the Subagent tool. */
-  enabled?: boolean;
+  enabled?: () => boolean;
   /** Temporal workflow function or type name (used with executeChild) */
   workflow: string | Workflow;
   /** Optional task queue - defaults to parent's queue if not specified */
@@ -369,6 +394,40 @@ export type SessionStartHook = (
   ctx: SessionStartHookContext
 ) => void | Promise<void>;
 
+/**
+ * Context for PreHumanMessageAppend hook - called before each human message is appended to the thread
+ */
+export interface PreHumanMessageAppendHookContext {
+  /** The message about to be appended */
+  message: MessageContent;
+  /** Thread identifier */
+  threadId: string;
+}
+
+/**
+ * PreHumanMessageAppend hook - called before each human message is appended to the thread
+ */
+export type PreHumanMessageAppendHook = (
+  ctx: PreHumanMessageAppendHookContext
+) => void | Promise<void>;
+
+/**
+ * PostHumanMessageAppend hook - called after each human message is appended to the thread
+ */
+export type PostHumanMessageAppendHook = (
+  ctx: PostHumanMessageAppendHookContext
+) => void | Promise<void>;
+
+/**
+ * Context for PostHumanMessageAppend hook - called after each human message is appended to the thread
+ */
+export interface PostHumanMessageAppendHookContext {
+  /** The message that was appended */
+  message: MessageContent;
+  /** Thread identifier */
+  threadId: string;
+}
+
 /**
  * Context for SessionEnd hook - called when session ends
  */
@@ -424,6 +483,10 @@ export interface ToolHooks<TArgs = unknown, TResult = unknown> {
  * Combined hooks interface for session lifecycle
  */
 export interface Hooks<T extends ToolMap, TResult = unknown> {
+  /** Called before each human message is appended to the thread */
+  onPreHumanMessageAppend?: PreHumanMessageAppendHook;
+  /** Called after each human message is appended to the thread */
+  onPostHumanMessageAppend?: PostHumanMessageAppendHook;
   /** Called before each tool execution - can block or modify */
   onPreToolUse?: PreToolUseHook<T>;
   /** Called after each successful tool execution */

package/src/tools/ask-user-question/handler.ts

@@ -1,25 +1,24 @@
-import { AIMessage, type StoredMessage } from "@langchain/core/messages";
-import type { ActivityToolHandler } from "../../lib/tool-router";
+import type { ToolHandler } from "../../lib/tool-router";
 import type { AskUserQuestionArgs } from "./tool";
 
 /**
  * Creates handler for user interaction tool - creates AI messages for display.
  */
-export const createAskUserQuestionHandler = (): ActivityToolHandler<
-  AskUserQuestionArgs,
-  { chatMessages: StoredMessage[] }
-> => async (args) => {
-  const messages = args.questions.map(
-    ({ question, header, options, multiSelect }) =>
-      new AIMessage({
-        content: question,
-        additional_kwargs: {
-          header,
-          options,
-          multiSelect,
-        },
-      }).toDict()
-  );
-
-  return { toolResponse: "Question submitted", data: { chatMessages: messages } };
-};
+export const createAskUserQuestionHandler =
+  (): ToolHandler<
+    AskUserQuestionArgs,
+    {
+      questions: {
+        question: string;
+        header: string;
+        options: { label: string; description: string }[];
+        multiSelect: boolean;
+      }[];
+    }
+  > =>
+  (args) => {
+    return {
+      toolResponse: "Question submitted",
+      data: { questions: args.questions },
+    };
+  };

package/src/tools/subagent/handler.ts

@@ -53,7 +53,7 @@ export function createSubagentHandler<
       taskQueue: config.taskQueue ?? parentTaskQueue,
     };
 
-    const { toolResponse, data } =
+    const { toolResponse, data, usage } =
       typeof config.workflow === "string"
         ? await executeChild(config.workflow, childOpts)
         : await executeChild(config.workflow, childOpts);
@@ -66,6 +66,7 @@ export function createSubagentHandler<
     return {
       toolResponse,
       data: validated,
+      ...(usage && { usage }),
     };
   };
 }

package/src/workflow.ts

@@ -74,7 +74,7 @@ export type {
   AgentFile,
   AgentResponse,
   ThreadOps,
-  ZeitlichAgentConfig,
+  AgentConfig,
   RunAgentConfig,
   RunAgentActivity,
   ToolResultConfig,
@@ -109,8 +109,6 @@ export type { SubagentArgs } from "./tools/subagent/tool";
 export type { ZeitlichSharedActivities } from "./activities";
 
 // Tool definitions (schemas only - no handlers)
-export { askUserQuestionTool } from "./tools/ask-user-question/tool";
-export type { AskUserQuestionArgs } from "./tools/ask-user-question/tool";
 export { globTool } from "./tools/glob/tool";
 export type { GlobArgs } from "./tools/glob/tool";
 export { grepTool } from "./tools/grep/tool";
@@ -141,3 +139,7 @@ export { createTaskUpdateHandler } from "./tools/task-update/handler";
 
 export { bashTool, createBashToolDescription } from "./tools/bash/tool";
 export type { BashArgs } from "./tools/bash/tool";
+
+export { askUserQuestionTool } from "./tools/ask-user-question/tool";
+export type { AskUserQuestionArgs } from "./tools/ask-user-question/tool";
+export { createAskUserQuestionHandler } from "./tools/ask-user-question/handler";