zeitlich 0.2.45 → 0.2.46

package/dist/index.js

@@ -3,6 +3,8 @@ 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 { Context } from '@temporalio/activity';
 import { promises } from 'fs';
 
@@ -1032,7 +1034,9 @@ async function createSession(config) {
     appendAgentMessage,
     forkThread,
     loadThreadState,
-    saveThreadState
+    saveThreadState,
+    hydrateThread,
+    flushThread
   } = threadOps;
   const plugins = [];
   let destroySubagentSandboxes;
@@ -1176,10 +1180,12 @@ async function createSession(config) {
         });
       };
       if (threadMode === "fork" && sourceThreadId) {
+        await hydrateThread(sourceThreadId, threadKey);
         await forkThread(sourceThreadId, threadId, threadKey);
         const forkedSlice = await loadThreadState(threadId, threadKey);
         if (forkedSlice) rehydrateFromSlice(forkedSlice);
       } else if (threadMode === "continue") {
+        await hydrateThread(threadId, threadKey);
         const continuedSlice = await loadThreadState(threadId, threadKey);
         if (continuedSlice) rehydrateFromSlice(continuedSlice);
       } else {
@@ -1361,6 +1367,15 @@ async function createSession(config) {
             error: persistError instanceof Error ? persistError.message : String(persistError)
           });
         }
+        try {
+          await flushThread(threadId, threadKey);
+        } catch (flushError) {
+          log.warn("failed to flush thread to cold tier", {
+            agentName,
+            threadId,
+            error: flushError instanceof Error ? flushError.message : String(flushError)
+          });
+        }
         await callSessionEnd(exitReason, stateManager.getTurns());
         if (sandboxOwned && sandboxId && sandboxOps) {
           switch (resolvedShutdown) {
@@ -1444,6 +1459,9 @@ function getThreadMetaKey(threadKey, threadId) {
 function getThreadStateKey(threadKey, threadId) {
   return `${threadKey}:state:thread:${threadId}`;
 }
+function getThreadDedupKey(threadId, dedupId) {
+  return `dedup:${dedupId}:thread:${threadId}`;
+}
 
 // src/lib/types.ts
 function isTerminalStatus(status) {
@@ -2588,9 +2606,6 @@ redis.call('EXPIRE', KEYS[2], tonumber(ARGV[1]))
 redis.call('SET', KEYS[1], '1', 'EX', tonumber(ARGV[1]))
 return 1
 `;
-function getDedupKey(threadId, id) {
-  return `dedup:${id}:thread:${threadId}`;
-}
 function createThreadManager(config) {
   const {
     redis,
@@ -2598,11 +2613,13 @@ function createThreadManager(config) {
     key = "messages",
     serialize = (m) => JSON.stringify(m),
     deserialize = (raw) => JSON.parse(raw),
-    idOf
+    idOf,
+    ttlSeconds = THREAD_TTL_SECONDS
   } = config;
   const redisKey = getThreadListKey(key, threadId);
   const metaKey = getThreadMetaKey(key, threadId);
   const stateKey = getThreadStateKey(key, threadId);
+  const dedupKey = (id) => getThreadDedupKey(threadId, id);
   async function assertThreadExists() {
     const exists = await redis.exists(metaKey);
     if (!exists) {
@@ -2612,7 +2629,7 @@ function createThreadManager(config) {
   return {
     async initialize() {
       await redis.del(redisKey);
-      await redis.set(metaKey, "1", "EX", THREAD_TTL_SECONDS);
+      await redis.set(metaKey, "1", "EX", ttlSeconds);
     },
     async load() {
       await assertThreadExists();
@@ -2624,18 +2641,17 @@ function createThreadManager(config) {
       await assertThreadExists();
       if (idOf) {
         const dedupId = messages.map(idOf).join(":");
-        const dedupKey = getDedupKey(threadId, dedupId);
         await redis.eval(
           APPEND_IDEMPOTENT_SCRIPT,
           2,
-          dedupKey,
+          dedupKey(dedupId),
           redisKey,
-          String(THREAD_TTL_SECONDS),
+          String(ttlSeconds),
           ...messages.map(serialize)
         );
       } else {
         await redis.rpush(redisKey, ...messages.map(serialize));
-        await redis.expire(redisKey, THREAD_TTL_SECONDS);
+        await redis.expire(redisKey, ttlSeconds);
       }
     },
     async fork(newThreadId) {
@@ -2650,11 +2666,11 @@ function createThreadManager(config) {
       if (data.length > 0) {
         const newKey = getThreadListKey(key, newThreadId);
         await redis.rpush(newKey, ...data);
-        await redis.expire(newKey, THREAD_TTL_SECONDS);
+        await redis.expire(newKey, ttlSeconds);
       }
       if (stateRaw != null) {
         const newStateKey = getThreadStateKey(key, newThreadId);
-        await redis.set(newStateKey, stateRaw, "EX", THREAD_TTL_SECONDS);
+        await redis.set(newStateKey, stateRaw, "EX", ttlSeconds);
       }
       return forked;
     },
@@ -2669,15 +2685,13 @@ function createThreadManager(config) {
       const existingIds = existing.map((raw) => idOf(deserialize(raw))).filter((id) => typeof id === "string");
       await redis.del(redisKey);
       if (existingIds.length > 0) {
-        await redis.del(
-          ...existingIds.map((id) => getDedupKey(threadId, id))
-        );
+        await redis.del(...existingIds.map(dedupKey));
       }
       if (messages.length > 0) {
         await redis.rpush(redisKey, ...messages.map(serialize));
-        await redis.expire(redisKey, THREAD_TTL_SECONDS);
+        await redis.expire(redisKey, ttlSeconds);
       }
-      await redis.expire(metaKey, THREAD_TTL_SECONDS);
+      await redis.expire(metaKey, ttlSeconds);
     },
     async delete() {
       await redis.del(redisKey, metaKey, stateKey);
@@ -2689,12 +2703,7 @@ function createThreadManager(config) {
     },
     async saveState(state) {
       await assertThreadExists();
-      await redis.set(
-        stateKey,
-        JSON.stringify(state),
-        "EX",
-        THREAD_TTL_SECONDS
-      );
+      await redis.set(stateKey, JSON.stringify(state), "EX", ttlSeconds);
     },
     async deleteState() {
       await redis.del(stateKey);
@@ -2723,19 +2732,202 @@ function createThreadManager(config) {
       if (idx === -1) return;
       if (idx === 0) {
         await redis.del(redisKey);
-        await redis.expire(metaKey, THREAD_TTL_SECONDS);
+        await redis.expire(metaKey, ttlSeconds);
       } else {
         await redis.ltrim(redisKey, 0, idx - 1);
-        await redis.expire(redisKey, THREAD_TTL_SECONDS);
+        await redis.expire(redisKey, ttlSeconds);
       }
       if (removedIds.length > 0) {
-        await redis.del(
-          ...removedIds.map((id) => getDedupKey(threadId, id))
+        await redis.del(...removedIds.map(dedupKey));
+      }
+    }
+  };
+}
+function joinKey(parts) {
+  return parts.map((p) => p.replace(/^\/+|\/+$/g, "")).filter((p) => p.length > 0).join("/");
+}
+function buildKey(prefix, threadKey, threadId, gzip) {
+  const ext = gzip ? "json.gz" : "json";
+  return joinKey([
+    prefix ?? "threads",
+    threadKey,
+    `${threadId}.${ext}`
+  ]);
+}
+async function streamToBuffer(body) {
+  if (body == null) return Buffer.alloc(0);
+  if (body instanceof Uint8Array) return Buffer.from(body);
+  if (typeof body.transformToByteArray === "function") {
+    const bytes = await body.transformToByteArray();
+    return Buffer.from(bytes);
+  }
+  if (typeof body.arrayBuffer === "function") {
+    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);
+}
+function createS3ColdStore(config) {
+  const { s3, bucket, prefix, gzip = true } = config;
+  const contentType = config.contentType ?? (gzip ? "application/gzip" : "application/json");
+  return {
+    async read(threadKey, threadId) {
+      const Key = buildKey(prefix, threadKey, threadId, gzip);
+      try {
+        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");
+        return JSON.parse(json);
+      } catch (err) {
+        if (isNotFound(err)) return null;
+        throw err;
       }
+    },
+    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
+        })
+      );
+    },
+    async delete(threadKey, threadId) {
+      const Key = buildKey(prefix, threadKey, threadId, gzip);
+      await s3.send(new DeleteObjectCommand({ Bucket: bucket, Key }));
     }
   };
 }
+function isNotFound(err) {
+  if (typeof err !== "object" || err === null) return false;
+  const e = err;
+  return e.name === "NoSuchKey" || e.Code === "NoSuchKey" || e.code === "NoSuchKey" || e.name === "NotFound" || e.$metadata?.httpStatusCode === 404;
+}
+
+// src/lib/thread/snapshot.ts
+async function encodeSnapshot(config) {
+  const { redis, threadKey, threadId, idOf } = config;
+  const metaKey = getThreadMetaKey(threadKey, threadId);
+  if (await redis.exists(metaKey) === 0) {
+    return null;
+  }
+  const listKey = getThreadListKey(threadKey, threadId);
+  const stateKey = getThreadStateKey(threadKey, threadId);
+  const messages = await redis.lrange(listKey, 0, -1);
+  const stateRaw = await redis.get(stateKey);
+  const state = stateRaw == null ? null : JSON.parse(stateRaw);
+  const dedupIds = idOf ? messages.map(idOf) : [];
+  return { v: 1, messages, state, dedupIds };
+}
+async function applySnapshot(config) {
+  const {
+    redis,
+    threadKey,
+    threadId,
+    snapshot,
+    ttlSeconds = THREAD_TTL_SECONDS
+  } = config;
+  const metaKey = getThreadMetaKey(threadKey, threadId);
+  if (await redis.exists(metaKey) === 1) {
+    return;
+  }
+  const listKey = getThreadListKey(threadKey, threadId);
+  const stateKey = getThreadStateKey(threadKey, threadId);
+  await redis.del(listKey, stateKey);
+  const pipeline = redis.pipeline();
+  if (snapshot.messages.length > 0) {
+    pipeline.rpush(listKey, ...snapshot.messages);
+    pipeline.expire(listKey, ttlSeconds);
+  }
+  if (snapshot.state != null) {
+    pipeline.set(stateKey, JSON.stringify(snapshot.state), "EX", ttlSeconds);
+  }
+  for (const id of snapshot.dedupIds) {
+    pipeline.set(getThreadDedupKey(threadId, id), "1", "EX", ttlSeconds);
+  }
+  const results = await pipeline.exec();
+  if (results) {
+    const firstErr = results.find(([err]) => err)?.[0] ?? null;
+    if (firstErr) {
+      await redis.del(
+        listKey,
+        stateKey,
+        ...snapshot.dedupIds.map((id) => getThreadDedupKey(threadId, id))
+      ).catch(() => void 0);
+      throw firstErr;
+    }
+  }
+  await redis.set(metaKey, "1", "EX", ttlSeconds);
+}
+async function clearHotTier(config) {
+  const { redis, threadKey, threadId, dedupIds = [] } = config;
+  const keys = [
+    getThreadListKey(threadKey, threadId),
+    getThreadMetaKey(threadKey, threadId),
+    getThreadStateKey(threadKey, threadId),
+    ...dedupIds.map((id) => getThreadDedupKey(threadId, id))
+  ];
+  await redis.del(...keys);
+}
+
+// src/lib/thread/tiered.ts
+function createTieredThreadManager(config) {
+  const {
+    redis,
+    threadId,
+    key = "messages",
+    coldStore,
+    idOf,
+    deserialize = (raw) => JSON.parse(raw),
+    ttlSeconds = THREAD_TTL_SECONDS
+  } = config;
+  const base = createThreadManager(config);
+  const rawIdOf = idOf ? (raw) => idOf(deserialize(raw)) : void 0;
+  return Object.assign(base, {
+    async hydrate() {
+      if (!coldStore) return;
+      const snapshot = await coldStore.read(key, threadId);
+      if (!snapshot) return;
+      await applySnapshot({
+        redis,
+        threadKey: key,
+        threadId,
+        snapshot,
+        ttlSeconds
+      });
+    },
+    async flush(opts) {
+      if (!coldStore) return;
+      const snapshot = await encodeSnapshot({
+        redis,
+        threadKey: key,
+        threadId,
+        ...rawIdOf ? { idOf: rawIdOf } : {}
+      });
+      if (!snapshot) return;
+      await coldStore.write(key, threadId, snapshot);
+      const deleteHot = opts?.deleteHot ?? true;
+      if (deleteHot) {
+        await clearHotTier({
+          redis,
+          threadKey: key,
+          threadId,
+          dedupIds: snapshot.dedupIds
+        });
+      }
+    }
+  });
+}
 function getActivityContext() {
   try {
     const ctx = Context.current();
@@ -3671,6 +3863,6 @@ var toTree = async (fs, opts = {}) => {
   return base + subtree;
 };
 
-export { DEFAULT_SUBAGENT_WORKFLOW_RUN_TIMEOUT, FileSystemSkillProvider, NodeFsSandboxFileSystem, SandboxManager, SandboxNotFoundError, SandboxNotSupportedError, THREAD_TTL_SECONDS, VirtualFileSystem, applyVirtualTreeMutations, askUserQuestionTool, bashHandler, bashTool, composeHooks, createAgentStateManager, createAskUserQuestionHandler, createBashToolDescription, createObservabilityHooks, createReadSkillHandler, createReadSkillTool, createRunAgentActivity, createSession, createTaskCreateHandler, createTaskGetHandler, createTaskListHandler, createTaskUpdateHandler, createThreadManager, createToolRouter, createVirtualFsActivities, defineSubagent, defineSubagentWorkflow, defineTool, defineWorkflow, editHandler, editTool, filesWithMimeType, formatVirtualFileTree, getActivityContext, getShortId, getThreadListKey, getThreadMetaKey, 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, 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