@axiom-lattice/core 1.0.50 → 2.0.0

package/dist/index.mjs

@@ -2205,13 +2205,16 @@ var ThreadStatus = /* @__PURE__ */ ((ThreadStatus2) => {
 })(ThreadStatus || {});
 
 // src/chunk_buffer_lattice/InMemoryChunkBuffer.ts
+import { ReplaySubject } from "rxjs";
 var InMemoryChunkBuffer = class extends ChunkBuffer {
   constructor(config) {
     super();
     this.buffers = /* @__PURE__ */ new Map();
     this.config = {
       ttl: config?.ttl ?? 60 * 60 * 1e3,
-      cleanupInterval: config?.cleanupInterval ?? 0
+      cleanupInterval: config?.cleanupInterval ?? 0,
+      maxChunks: config?.maxChunks ?? 1e3
+      // Default max chunks per thread
     };
     if (this.config.cleanupInterval > 0) {
       this.startCleanupTimer();
@@ -2260,11 +2263,14 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
    */
   getOrCreateBuffer(threadId) {
     let buffer = this.getBufferIfValid(threadId);
+    4;
     if (!buffer) {
       const now = Date.now();
       buffer = {
         threadId,
-        chunks: [],
+        chunks$: new ReplaySubject(
+          this.config.maxChunks ?? 1e4
+        ),
         status: "active" /* ACTIVE */,
         createdAt: now,
         updatedAt: now,
@@ -2272,34 +2278,30 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
       };
       this.buffers.set(threadId, buffer);
     }
+    if (buffer.status !== "active" /* ACTIVE */) {
+      buffer.status = "active" /* ACTIVE */;
+      buffer.chunks$ = new ReplaySubject(
+        this.config.maxChunks ?? 1e4
+      );
+      buffer.updatedAt = Date.now();
+      buffer.expiresAt = Date.now() + this.config.ttl;
+    }
     return buffer;
   }
   async addChunk(threadId, content) {
     const buffer = this.getOrCreateBuffer(threadId);
     const chunk = content;
-    buffer.chunks.push(chunk);
+    buffer.chunks$.next(chunk);
     buffer.updatedAt = Date.now();
     buffer.expiresAt = Date.now() + this.config.ttl;
-  }
-  async getChunks(threadId) {
-    const buffer = this.getBufferIfValid(threadId);
-    return buffer ? [...buffer.chunks] : [];
-  }
-  async getAccumulatedContent(threadId) {
-    const buffer = this.getBufferIfValid(threadId);
-    if (!buffer) return "";
-    return buffer.chunks.map((chunk) => chunk.data?.content).join("");
-  }
-  async getChunksByMessageId(threadId, messageId) {
-    const buffer = this.getBufferIfValid(threadId);
-    if (!buffer) return [];
-    return buffer.chunks.filter((chunk) => chunk.data?.id === messageId);
+    buffer.status = "active" /* ACTIVE */;
   }
   async completeThread(threadId) {
     const buffer = this.getBufferIfValid(threadId);
     if (buffer) {
       buffer.status = "completed" /* COMPLETED */;
       buffer.updatedAt = Date.now();
+      buffer.chunks$.complete();
     }
   }
   async abortThread(threadId) {
@@ -2307,6 +2309,7 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
     if (buffer) {
       buffer.status = "aborted" /* ABORTED */;
       buffer.updatedAt = Date.now();
+      buffer.chunks$.error(new Error("Thread aborted"));
     }
   }
   async isThreadActive(threadId) {
@@ -2319,10 +2322,7 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
   async getThreadBuffer(threadId) {
     const buffer = this.getBufferIfValid(threadId);
     if (!buffer) return void 0;
-    return {
-      ...buffer,
-      chunks: [...buffer.chunks]
-    };
+    return buffer;
   }
   async clearThread(threadId) {
     this.buffers.delete(threadId);
@@ -2369,9 +2369,6 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
     }
     return cleanedCount;
   }
-  /**
-   * Extend thread TTL
-   */
   async extendThreadTTL(threadId, additionalMs) {
     const buffer = this.getBufferIfValid(threadId);
     if (buffer) {
@@ -2380,62 +2377,69 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
       buffer.updatedAt = Date.now();
     }
   }
-  /**
-   * Get new chunks since known content
-   * Used for resuming streams from a known position
-   * Matches the known content and returns chunks after that position
-   * Continues to yield new chunks as they are added until thread completes/aborts
-   */
-  async *getNewChunksSinceContent(threadId, messageId, knownContent) {
-    let buffer = this.getBufferIfValid(threadId);
+  async *getNewChunksSinceContentIterator(threadId, messageId, knownContent) {
+    const buffer = this.getBufferIfValid(threadId);
     if (!buffer) return;
-    let lastYieldedIndex = -1;
     let accumulatedContent = "";
-    for (let i = 0; i < buffer.chunks.length; i++) {
-      const chunk = buffer.chunks[i];
-      if (chunk.data?.id === messageId) {
-        accumulatedContent += chunk.data?.content || "";
-        if (accumulatedContent === knownContent) {
-          lastYieldedIndex = i;
-          break;
+    const queue = [];
+    let resolveNext = null;
+    let errorNext = null;
+    let isCompleted = false;
+    const subscription = buffer.chunks$.subscribe({
+      next: (chunk) => {
+        queue.push(chunk);
+        if (resolveNext) {
+          const resolve = resolveNext;
+          resolveNext = null;
+          resolve();
         }
-        if (accumulatedContent.length > knownContent.length) {
-          if (accumulatedContent.startsWith(knownContent)) {
-            lastYieldedIndex = i;
-            break;
-          }
+      },
+      error: (err) => {
+        if (errorNext) {
+          const reject = errorNext;
+          errorNext = null;
+          reject(err);
         }
-      }
-    }
-    const pollingInterval = 100;
-    while (true) {
-      buffer = this.getBufferIfValid(threadId);
-      if (!buffer) break;
-      let hasNewChunks = false;
-      for (let i = lastYieldedIndex + 1; i < buffer.chunks.length; i++) {
-        const chunk = buffer.chunks[i];
-        if (chunk.data?.id === messageId) {
-          yield chunk;
-          lastYieldedIndex = i;
-          hasNewChunks = true;
+      },
+      complete: () => {
+        isCompleted = true;
+        if (resolveNext) {
+          const resolve = resolveNext;
+          resolveNext = null;
+          resolve();
         }
       }
-      if (buffer.status === "completed" /* COMPLETED */ || buffer.status === "aborted" /* ABORTED */) {
-        break;
-      }
-      if (!hasNewChunks) {
-        await new Promise((resolve) => setTimeout(resolve, pollingInterval));
+    });
+    try {
+      while (true) {
+        if (queue.length === 0) {
+          if (isCompleted) break;
+          await new Promise((resolve, reject) => {
+            resolveNext = resolve;
+            errorNext = reject;
+          });
+        }
+        while (queue.length > 0) {
+          const chunk = queue.shift();
+          if (!chunk) continue;
+          if (chunk.data?.id === messageId) {
+            accumulatedContent += chunk.data?.content || "";
+          }
+          if (accumulatedContent.length > knownContent.length) {
+            if (accumulatedContent.startsWith(knownContent)) {
+              yield chunk;
+            }
+          }
+        }
       }
+    } finally {
+      subscription.unsubscribe();
     }
   }
-  /**
-   * Get statistics about the buffer
-   */
   getStats() {
     let activeCount = 0;
     let completedCount = 0;
     let abortedCount = 0;
-    let totalChunks = 0;
     const validBuffers = [];
     for (const [threadId, buffer] of this.buffers.entries()) {
       if (this.isExpired(buffer)) {
@@ -2445,7 +2449,6 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
       }
     }
     for (const buffer of validBuffers) {
-      totalChunks += buffer.chunks.length;
       switch (buffer.status) {
         case "active" /* ACTIVE */:
           activeCount++;
@@ -2463,13 +2466,9 @@ var InMemoryChunkBuffer = class extends ChunkBuffer {
       activeThreads: activeCount,
       completedThreads: completedCount,
       abortedThreads: abortedCount,
-      totalChunks,
       config: this.config
     };
   }
-  /**
-   * Cleanup method for graceful shutdown
-   */
   dispose() {
     this.stopCleanupTimer();
     this.buffers.clear();