zeitlich 0.2.46 → 0.2.48

package/dist/index.js

@@ -3,8 +3,10 @@ import z14, { z } from 'zod';
 import { randomUUID, randomFillSync } from 'crypto';
 import { ApplicationFailure as ApplicationFailure$1 } from '@temporalio/common';
 import { join, resolve, posix } from 'path';
-import { gzipSync, gunzipSync } from 'zlib';
-import { DeleteObjectCommand, PutObjectCommand, GetObjectCommand } from '@aws-sdk/client-s3';
+import { gzip, gunzip } from 'zlib';
+import { promisify } from 'util';
+import { DeleteObjectCommand, GetObjectCommand } from '@aws-sdk/client-s3';
+import { Upload } from '@aws-sdk/lib-storage';
 import { Context } from '@temporalio/activity';
 import { promises } from 'fs';
 
@@ -123,7 +125,7 @@ function createToolRouter(options) {
       });
     }
   }
-  async function processToolCall(toolCall, turn, sandboxId, onRewindRequested) {
+  async function processToolCall(toolCall, turn, sandboxId, onRewindRequested, assistantMessageId) {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
     const preResult = await runPreHooks(toolCall, tool, turn);
@@ -158,7 +160,8 @@ function createToolRouter(options) {
           ...options.threadKey && { threadKey: options.threadKey },
           toolCallId: toolCall.id,
           toolName: toolCall.name,
-          ...sandboxId !== void 0 && { sandboxId }
+          ...sandboxId !== void 0 && { sandboxId },
+          ...assistantMessageId !== void 0 && { assistantMessageId }
         };
         const response = await tool.handler(
           effectiveArgs,
@@ -285,6 +288,7 @@ function createToolRouter(options) {
       }
       const turn = context?.turn ?? 0;
       const sandboxId = context?.sandboxId;
+      const assistantMessageId = context?.assistantMessageId;
       let rewindSignal;
       if (options.parallel) {
         const scope = new CancellationScope({ cancellable: true });
@@ -297,7 +301,13 @@ function createToolRouter(options) {
         const outcomes = await scope.run(
           async () => Promise.allSettled(
             toolCalls.map(
-              (tc) => processToolCall(tc, turn, sandboxId, onRewindRequested)
+              (tc) => processToolCall(
+                tc,
+                turn,
+                sandboxId,
+                onRewindRequested,
+                assistantMessageId
+              )
             )
           )
         );
@@ -317,7 +327,13 @@ function createToolRouter(options) {
       }
       const results = [];
       for (const toolCall of toolCalls) {
-        const outcome = await processToolCall(toolCall, turn, sandboxId);
+        const outcome = await processToolCall(
+          toolCall,
+          turn,
+          sandboxId,
+          void 0,
+          assistantMessageId
+        );
         if (outcome.kind === "rewind") {
           rewindSignal = outcome.signal;
           break;
@@ -341,6 +357,9 @@ function createToolRouter(options) {
           toolName: toolCall.name,
           ...context?.sandboxId !== void 0 && {
             sandboxId: context.sandboxId
+          },
+          ...context?.assistantMessageId !== void 0 && {
+            assistantMessageId: context.assistantMessageId
           }
         };
         const response = await handler(
@@ -566,13 +585,27 @@ function createSubagentHandler(subagents) {
     }
     const threadMode = config.thread ?? "new";
     const allowsContinuation = threadMode !== "new";
-    const continuationThreadId = args.threadId && allowsContinuation ? args.threadId : void 0;
+    const newThreadSource = config.newThreadSource ?? "new";
+    const usingParentFallback = allowsContinuation && !args.threadId && newThreadSource === "from-parent";
+    const continuationThreadId = !allowsContinuation ? void 0 : args.threadId ?? (usingParentFallback ? context.threadId : void 0);
     let thread;
     if (continuationThreadId) {
-      thread = {
-        mode: threadMode,
-        threadId: continuationThreadId
-      };
+      if (threadMode === "fork") {
+        thread = {
+          mode: "fork",
+          threadId: continuationThreadId,
+          ...usingParentFallback && context.assistantMessageId ? {
+            truncateAfterFork: {
+              fromMessageId: context.assistantMessageId
+            }
+          } : {}
+        };
+      } else {
+        thread = {
+          mode: "continue",
+          threadId: continuationThreadId
+        };
+      }
     }
     let sandbox;
     let sandboxShutdownOverride;
@@ -1033,6 +1066,7 @@ async function createSession(config) {
     appendSystemMessage,
     appendAgentMessage,
     forkThread,
+    truncateThread,
     loadThreadState,
     saveThreadState,
     hydrateThread,
@@ -1182,6 +1216,10 @@ async function createSession(config) {
       if (threadMode === "fork" && sourceThreadId) {
         await hydrateThread(sourceThreadId, threadKey);
         await forkThread(sourceThreadId, threadId, threadKey);
+        const truncate = threadInit.truncateAfterFork;
+        if (truncate?.fromMessageId) {
+          await truncateThread(threadId, truncate.fromMessageId, threadKey);
+        }
         const forkedSlice = await loadThreadState(threadId, threadKey);
         if (forkedSlice) rehydrateFromSlice(forkedSlice);
       } else if (threadMode === "continue") {
@@ -1311,7 +1349,10 @@ async function createSession(config) {
             parsedToolCalls,
             {
               turn: currentTurn,
-              ...sandboxId !== void 0 && { sandboxId }
+              ...sandboxId !== void 0 && { sandboxId },
+              ...assistantId !== void 0 && {
+                assistantMessageId: assistantId
+              }
             }
           );
           for (const result of toolCallResults) {
@@ -2214,6 +2255,13 @@ IMPORTANT:
   }),
   strict: true
 };
+var textEditSchema = z.object({
+  old_string: z.string().describe("The exact text to replace"),
+  new_string: z.string().describe("The text to replace it with"),
+  replace_all: z.boolean().optional().describe(
+    "If true, replace all occurrences of old_string for this edit (default: false)"
+  )
+});
 var editTool = {
   name: "FileEdit",
   description: `Edit specific sections of a file by replacing text.
@@ -2242,6 +2290,27 @@ IMPORTANT:
   }),
   strict: true
 };
+var multiEditTool = {
+  name: "FileMultiEdit",
+  description: `Apply multiple exact text replacements to one file in order.
+
+Usage:
+- Use this when a task needs several related edits in the same file
+- Each edit is applied to the file content produced by the prior edit
+- The operation is atomic: if any edit fails, the file is left unchanged
+
+IMPORTANT:
+- You must read the file first (in this session) before editing it
+- Each old_string must match exactly (whitespace-sensitive)
+- Each old_string must be unique unless that edit uses replace_all: true
+- old_string and new_string must be different for every edit
+`,
+  schema: z.object({
+    file_path: z.string().describe("The absolute virtual path to the file to modify"),
+    edits: z.array(textEditSchema).min(1).describe("Exact replacements to apply sequentially to the file")
+  }),
+  strict: true
+};
 var taskCreateTool = {
   name: "TaskCreate",
   description: `Use this tool to create a structured task list. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.
@@ -2312,7 +2381,7 @@ function createTaskCreateHandler(stateManager) {
     };
     stateManager.setTask(task);
     return {
-      toolResponse: JSON.stringify(task, null, 2),
+      toolResponse: `Task ${task.id} created`,
       data: task
     };
   };
@@ -2411,7 +2480,7 @@ function createTaskUpdateHandler(stateManager) {
     }
     stateManager.setTask(task);
     return {
-      toolResponse: JSON.stringify(task, null, 2),
+      toolResponse: `Task ${task.id} updated`,
       data: task
     };
   };
@@ -2743,6 +2812,44 @@ function createThreadManager(config) {
     }
   };
 }
+function getActivityContext() {
+  try {
+    const ctx = Context.current();
+    return { heartbeat: () => ctx.heartbeat(), signal: ctx.cancellationSignal };
+  } catch {
+    return {};
+  }
+}
+async function queryParentWorkflowState(client) {
+  const { workflowExecution } = Context.current().info;
+  if (!workflowExecution) {
+    throw new Error("No workflow execution found");
+  }
+  const handle = client.getHandle(
+    workflowExecution.workflowId,
+    workflowExecution.runId
+  );
+  return handle.query("getAgentState");
+}
+function createRunAgentActivity(client, handler, scope) {
+  const name = `run${scope.charAt(0).toUpperCase()}${scope.slice(1)}`;
+  return {
+    [name]: async (config) => {
+      const state = await queryParentWorkflowState(client);
+      return handler({ ...config, state });
+    }
+  };
+}
+function withParentWorkflowState(client, handler) {
+  return async (args, context) => {
+    const state = await queryParentWorkflowState(client);
+    return handler(args, { ...context, state });
+  };
+}
+
+// src/lib/thread/cold-store.ts
+var gzipAsync = promisify(gzip);
+var gunzipAsync = promisify(gunzip);
 function joinKey(parts) {
   return parts.map((p) => p.replace(/^\/+|\/+$/g, "")).filter((p) => p.length > 0).join("/");
 }
@@ -2754,9 +2861,17 @@ function buildKey(prefix, threadKey, threadId, gzip) {
     `${threadId}.${ext}`
   ]);
 }
-async function streamToBuffer(body) {
+async function streamToBuffer(body, onChunk) {
   if (body == null) return Buffer.alloc(0);
   if (body instanceof Uint8Array) return Buffer.from(body);
+  if (typeof body[Symbol.asyncIterator] === "function") {
+    const chunks = [];
+    for await (const chunk of body) {
+      chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
+      onChunk?.();
+    }
+    return Buffer.concat(chunks);
+  }
   if (typeof body.transformToByteArray === "function") {
     const bytes = await body.transformToByteArray();
     return Buffer.from(bytes);
@@ -2765,11 +2880,7 @@ async function streamToBuffer(body) {
     const ab = await body.arrayBuffer();
     return Buffer.from(ab);
   }
-  const chunks = [];
-  for await (const chunk of body) {
-    chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
-  }
-  return Buffer.concat(chunks);
+  return Buffer.alloc(0);
 }
 function createS3ColdStore(config) {
   const { s3, bucket, prefix, gzip = true } = config;
@@ -2781,8 +2892,9 @@ function createS3ColdStore(config) {
         const resp = await s3.send(
           new GetObjectCommand({ Bucket: bucket, Key })
         );
-        const buf = await streamToBuffer(resp.Body);
-        const json = gzip ? gunzipSync(buf).toString("utf8") : buf.toString("utf8");
+        const { heartbeat } = getActivityContext();
+        const buf = await streamToBuffer(resp.Body, heartbeat);
+        const json = gzip ? (await gunzipAsync(buf)).toString("utf8") : buf.toString("utf8");
         return JSON.parse(json);
       } catch (err) {
         if (isNotFound(err)) return null;
@@ -2792,15 +2904,14 @@ function createS3ColdStore(config) {
     async write(threadKey, threadId, snapshot) {
       const Key = buildKey(prefix, threadKey, threadId, gzip);
       const json = JSON.stringify(snapshot);
-      const body = gzip ? gzipSync(Buffer.from(json, "utf8")) : json;
-      await s3.send(
-        new PutObjectCommand({
-          Bucket: bucket,
-          Key,
-          Body: body,
-          ContentType: contentType
-        })
-      );
+      const body = gzip ? await gzipAsync(Buffer.from(json, "utf8")) : json;
+      const upload = new Upload({
+        client: s3,
+        params: { Bucket: bucket, Key, Body: body, ContentType: contentType }
+      });
+      const { heartbeat } = getActivityContext();
+      if (heartbeat) upload.on("httpUploadProgress", heartbeat);
+      await upload.done();
     },
     async delete(threadKey, threadId) {
       const Key = buildKey(prefix, threadKey, threadId, gzip);
@@ -2928,40 +3039,6 @@ function createTieredThreadManager(config) {
     }
   });
 }
-function getActivityContext() {
-  try {
-    const ctx = Context.current();
-    return { heartbeat: () => ctx.heartbeat(), signal: ctx.cancellationSignal };
-  } catch {
-    return {};
-  }
-}
-async function queryParentWorkflowState(client) {
-  const { workflowExecution } = Context.current().info;
-  if (!workflowExecution) {
-    throw new Error("No workflow execution found");
-  }
-  const handle = client.getHandle(
-    workflowExecution.workflowId,
-    workflowExecution.runId
-  );
-  return handle.query("getAgentState");
-}
-function createRunAgentActivity(client, handler, scope) {
-  const name = `run${scope.charAt(0).toUpperCase()}${scope.slice(1)}`;
-  return {
-    [name]: async (config) => {
-      const state = await queryParentWorkflowState(client);
-      return handler({ ...config, state });
-    }
-  };
-}
-function withParentWorkflowState(client, handler) {
-  return async (args, context) => {
-    const state = await queryParentWorkflowState(client);
-    return handler(args, { ...context, state });
-  };
-}
 
 // src/lib/sandbox/manager.ts
 var CAP_METHOD_TO_CAPABILITY = [
@@ -3633,57 +3710,190 @@ ${result.stderr}`,
 };
 
 // src/tools/edit/handler.ts
-function escapeRegExp(str) {
-  return str.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+function splitLines(text) {
+  if (text.length === 0) return [];
+  return text.replace(/\r\n/g, "\n").split("\n");
+}
+function lineNumberAt(content, index) {
+  let line = 1;
+  for (let i = 0; i < index; i++) {
+    if (content.charCodeAt(i) === 10) line++;
+  }
+  return line;
+}
+function lineEnd(startLine, lines) {
+  return lines.length === 0 ? startLine : startLine + lines.length - 1;
+}
+function indicesOf(content, needle) {
+  const indices = [];
+  let cursor = 0;
+  while (cursor <= content.length) {
+    const index = content.indexOf(needle, cursor);
+    if (index === -1) break;
+    indices.push(index);
+    cursor = index + needle.length;
+  }
+  return indices;
+}
+function makeHunk(editIndex, beforeContent, replacementIndex, oldString, newString) {
+  const oldStartLine = lineNumberAt(beforeContent, replacementIndex);
+  const oldLines = splitLines(oldString);
+  const newLines = splitLines(newString);
+  return {
+    editIndex,
+    oldStartLine,
+    oldEndLine: lineEnd(oldStartLine, oldLines),
+    newStartLine: oldStartLine,
+    newEndLine: lineEnd(oldStartLine, newLines),
+    oldLines,
+    newLines
+  };
+}
+function applyOneEdit(content, edit, editIndex) {
+  const { old_string, new_string, replace_all = false } = edit;
+  if (old_string.length === 0) {
+    return {
+      ok: false,
+      editIndex,
+      message: `Error: old_string for edit ${editIndex} must not be empty.`
+    };
+  }
+  if (old_string === new_string) {
+    return {
+      ok: false,
+      editIndex,
+      message: `Error: old_string and new_string must be different for edit ${editIndex}.`
+    };
+  }
+  const matches = indicesOf(content, old_string);
+  if (matches.length === 0) {
+    return {
+      ok: false,
+      editIndex,
+      message: `Error: Could not find old_string for edit ${editIndex}. Make sure it matches exactly (whitespace-sensitive).`
+    };
+  }
+  if (!replace_all && matches.length > 1) {
+    return {
+      ok: false,
+      editIndex,
+      message: `Error: old_string for edit ${editIndex} appears ${matches.length} times. Provide more context to make it unique, or use replace_all: true for that edit.`
+    };
+  }
+  if (replace_all) {
+    const hunks = matches.map(
+      (index2) => makeHunk(editIndex, content, index2, old_string, new_string)
+    );
+    return {
+      ok: true,
+      content: content.split(old_string).join(new_string),
+      replacements: matches.length,
+      hunks
+    };
+  }
+  const index = matches[0];
+  if (index === void 0) {
+    return {
+      ok: false,
+      editIndex,
+      message: `Error: Could not find old_string for edit ${editIndex}.`
+    };
+  }
+  return {
+    ok: true,
+    content: content.slice(0, index) + new_string + content.slice(index + old_string.length),
+    replacements: 1,
+    hunks: [makeHunk(editIndex, content, index, old_string, new_string)]
+  };
+}
+function applyEditPlan(content, edits) {
+  if (edits.length === 0) {
+    return {
+      ok: false,
+      message: "Error: edits must contain at least one edit."
+    };
+  }
+  let current = content;
+  let replacements = 0;
+  const hunks = [];
+  for (const [index, edit] of edits.entries()) {
+    const result = applyOneEdit(current, edit, index);
+    if (!result.ok) return result;
+    current = result.content;
+    replacements += result.replacements;
+    hunks.push(...result.hunks);
+  }
+  return { ok: true, content: current, replacements, hunks };
+}
+function editFailureResult(filePath, message) {
+  return {
+    toolResponse: message,
+    data: { path: filePath, success: false, replacements: 0 }
+  };
 }
 var editHandler = async (args, { sandbox }) => {
   const { fs } = sandbox;
   const { file_path, old_string, new_string, replace_all = false } = args;
-  if (old_string === new_string) {
+  try {
+    const exists = await fs.exists(file_path);
+    if (!exists) {
+      return editFailureResult(
+        file_path,
+        `Error: File "${file_path}" does not exist.`
+      );
+    }
+    const content = await fs.readFile(file_path);
+    const result = applyEditPlan(content, [
+      { old_string, new_string, replace_all }
+    ]);
+    if (!result.ok) {
+      return editFailureResult(file_path, result.message);
+    }
+    await fs.writeFile(file_path, result.content);
+    const summary = replace_all ? `Replaced ${result.replacements} occurrence(s)` : `Replaced 1 occurrence`;
     return {
-      toolResponse: `Error: old_string and new_string must be different.`,
+      toolResponse: `${summary} in ${file_path}`,
+      data: {
+        path: file_path,
+        success: true,
+        replacements: result.replacements,
+        hunks: result.hunks
+      }
+    };
+  } catch (error) {
+    const message = error instanceof Error ? error.message : "Unknown error";
+    return {
+      toolResponse: `Error editing file "${file_path}": ${message}`,
       data: { path: file_path, success: false, replacements: 0 }
     };
   }
+};
+var multiEditHandler = async (args, { sandbox }) => {
+  const { fs } = sandbox;
+  const { file_path, edits } = args;
   try {
     const exists = await fs.exists(file_path);
     if (!exists) {
-      return {
-        toolResponse: `Error: File "${file_path}" does not exist.`,
-        data: { path: file_path, success: false, replacements: 0 }
-      };
+      return editFailureResult(
+        file_path,
+        `Error: File "${file_path}" does not exist.`
+      );
     }
     const content = await fs.readFile(file_path);
-    if (!content.includes(old_string)) {
-      return {
-        toolResponse: `Error: Could not find the specified text in "${file_path}". Make sure old_string matches exactly (whitespace-sensitive).`,
-        data: { path: file_path, success: false, replacements: 0 }
-      };
-    }
-    const escapedOldString = escapeRegExp(old_string);
-    const globalRegex = new RegExp(escapedOldString, "g");
-    const occurrences = (content.match(globalRegex) || []).length;
-    if (!replace_all && occurrences > 1) {
-      return {
-        toolResponse: `Error: old_string appears ${occurrences} times in "${file_path}". Either provide more context to make it unique, or use replace_all: true.`,
-        data: { path: file_path, success: false, replacements: 0 }
-      };
-    }
-    let newContent;
-    let replacements;
-    if (replace_all) {
-      newContent = content.split(old_string).join(new_string);
-      replacements = occurrences;
-    } else {
-      const index = content.indexOf(old_string);
-      newContent = content.slice(0, index) + new_string + content.slice(index + old_string.length);
-      replacements = 1;
+    const result = applyEditPlan(content, edits);
+    if (!result.ok) {
+      const suffix = result.editIndex === void 0 ? "" : ` in ${file_path}`;
+      return editFailureResult(file_path, `${result.message}${suffix}`);
     }
-    await fs.writeFile(file_path, newContent);
-    const summary = replace_all ? `Replaced ${replacements} occurrence(s)` : `Replaced 1 occurrence`;
+    await fs.writeFile(file_path, result.content);
     return {
-      toolResponse: `${summary} in ${file_path}`,
-      data: { path: file_path, success: true, replacements }
+      toolResponse: `Applied ${edits.length} edit(s), ${result.replacements} replacement(s) in ${file_path}`,
+      data: {
+        path: file_path,
+        success: true,
+        replacements: result.replacements,
+        hunks: result.hunks
+      }
     };
   } catch (error) {
     const message = error instanceof Error ? error.message : "Unknown error";
@@ -3863,6 +4073,6 @@ var toTree = async (fs, opts = {}) => {
   return base + subtree;
 };
 
-export { DEFAULT_SUBAGENT_WORKFLOW_RUN_TIMEOUT, FileSystemSkillProvider, NodeFsSandboxFileSystem, SandboxManager, SandboxNotFoundError, SandboxNotSupportedError, THREAD_TTL_SECONDS, VirtualFileSystem, applySnapshot, applyVirtualTreeMutations, askUserQuestionTool, bashHandler, bashTool, clearHotTier, composeHooks, createAgentStateManager, createAskUserQuestionHandler, createBashToolDescription, createObservabilityHooks, createReadSkillHandler, createReadSkillTool, createRunAgentActivity, createS3ColdStore, createSession, createTaskCreateHandler, createTaskGetHandler, createTaskListHandler, createTaskUpdateHandler, createThreadManager, createTieredThreadManager, createToolRouter, createVirtualFsActivities, defineSubagent, defineSubagentWorkflow, defineTool, defineWorkflow, editHandler, editTool, encodeSnapshot, filesWithMimeType, formatVirtualFileTree, getActivityContext, getShortId, getThreadDedupKey, getThreadListKey, getThreadMetaKey, getThreadStateKey, globHandler, globTool, grepTool, hasDirectory, hasFileWithMimeType, hasNoOtherToolCalls, isTerminalStatus, parseSkillFile, proxyRunAgent, proxyVirtualFsOps, queryParentWorkflowState, readFileHandler, readFileTool, taskCreateTool, taskGetTool, taskListTool, taskUpdateTool, toTree, withAutoAppend, withParentWorkflowState, withSandbox, withVirtualFs, writeFileHandler, writeFileTool };
+export { DEFAULT_SUBAGENT_WORKFLOW_RUN_TIMEOUT, FileSystemSkillProvider, NodeFsSandboxFileSystem, SandboxManager, SandboxNotFoundError, SandboxNotSupportedError, THREAD_TTL_SECONDS, VirtualFileSystem, applySnapshot, applyVirtualTreeMutations, askUserQuestionTool, bashHandler, bashTool, clearHotTier, composeHooks, createAgentStateManager, createAskUserQuestionHandler, createBashToolDescription, createObservabilityHooks, createReadSkillHandler, createReadSkillTool, createRunAgentActivity, createS3ColdStore, createSession, createTaskCreateHandler, createTaskGetHandler, createTaskListHandler, createTaskUpdateHandler, createThreadManager, createTieredThreadManager, createToolRouter, createVirtualFsActivities, defineSubagent, defineSubagentWorkflow, defineTool, defineWorkflow, editHandler, editTool, encodeSnapshot, filesWithMimeType, formatVirtualFileTree, getActivityContext, getShortId, getThreadDedupKey, getThreadListKey, getThreadMetaKey, getThreadStateKey, globHandler, globTool, grepTool, hasDirectory, hasFileWithMimeType, hasNoOtherToolCalls, isTerminalStatus, multiEditHandler, multiEditTool, parseSkillFile, proxyRunAgent, proxyVirtualFsOps, queryParentWorkflowState, readFileHandler, readFileTool, taskCreateTool, taskGetTool, taskListTool, taskUpdateTool, toTree, withAutoAppend, withParentWorkflowState, withSandbox, withVirtualFs, writeFileHandler, writeFileTool };
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map