zeitlich 0.2.7 → 0.2.8

package/src/lib/session.ts

@@ -24,7 +24,11 @@ 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"]>;
+  }>;
 }
 
 /**
@@ -40,6 +44,7 @@ export interface SessionLifecycleHooks {
 export const createSession = async <T extends ToolMap, M = unknown>({
   threadId,
   agentName,
+  description,
   maxTurns = 50,
   metadata = {},
   runAgent,
@@ -86,7 +91,13 @@ 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) => {
@@ -115,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);
@@ -134,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
@@ -165,9 +186,18 @@ 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") {
             const conditionMet = await condition(
@@ -195,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

@@ -114,6 +114,25 @@ export interface AgentStateManager<TCustom extends JsonSerializable<TCustom>> {
 
   /** Set the tools (converts Zod schemas to JSON Schema for serialization) */
   setTools(newTools: ToolDefinition[]): void;
+
+  /** 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;
+  };
 }
 
 /**
@@ -140,6 +159,11 @@ export function createAgentStateManager<
   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);
@@ -283,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

@@ -32,6 +32,10 @@ export interface BaseAgentState {
   turns: number;
   tasks: Map<string, WorkflowTask>;
   systemPrompt: string;
+  totalInputTokens: number;
+  totalOutputTokens: number;
+  cachedWriteTokens: number;
+  cachedReadtTokens: number;
 }
 
 /**
@@ -50,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;
 }
 
 /**
@@ -143,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>;
 }
 
@@ -155,6 +164,7 @@ export interface RunAgentConfig {
 export type RunAgentActivity<M = StoredMessage> = (
   config: RunAgentConfig
 ) => Promise<AgentResponse<M>>;
+
 /**
  * Configuration for appending a tool result
  */
@@ -186,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 */

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 }),
     };
   };
 }