@axiom-lattice/gateway 1.0.46 → 2.0.0

package/RESUME_STREAM_README.md

@@ -0,0 +1,388 @@
+# Resume Stream Feature
+
+## Overview
+
+The `resume_stream` feature allows you to continue receiving streaming chunks from a known position. This is particularly useful for scenarios where the client connection is interrupted (e.g., page refresh, network issues) and you want to seamlessly continue from where you left off.
+
+## Architecture
+
+### ChunkBuffer Integration
+
+The feature leverages the `ChunkBuffer` module from `@axiom-lattice/core` to:
+
+- Store streaming chunks in memory
+- Track thread status (active/completed/aborted)
+- Provide TTL-based automatic cleanup
+- Calculate and return only new chunks since a known position
+
+### How It Works
+
+```
+┌─────────────────┐
+│  Original       │
+│  Streaming      │ ──► Chunks stored in ChunkBuffer
+│  (agent_stream) │     (thread_id, message_id, content)
+└─────────────────┘
+         │
+         │ Connection lost / Page refresh
+         │
+         ▼
+┌─────────────────┐
+│  Resume Stream  │
+│  (from known    │ ──► Polls for new chunks
+│  position)      │     Returns only new content
+└─────────────────┘     Ends when thread completes
+```
+
+## API Reference
+
+### `resume_stream(options)`
+
+Creates an async iterator that yields new chunks as they arrive.
+
+#### Parameters
+
+```typescript
+{
+  thread_id: string;        // Thread identifier
+  message_id: string;       // Message identifier (usually run_id)
+  known_content: string;    // Content already received (used to find resume position)
+  poll_interval?: number;   // Polling interval in ms (default: 100)
+}
+```
+
+#### Returns
+
+An async iterable object with `Symbol.asyncIterator`:
+
+```typescript
+{
+  content: string; // Chunk content
+  timestamp: number; // When chunk was added
+  messageId: string; // Message identifier
+}
+```
+
+#### Behavior
+
+- **Polling**: Checks for new chunks every `poll_interval` milliseconds
+- **Status Check**: Monitors thread status (active/completed/aborted)
+- **Timeout**: Automatically stops after 30 seconds of no new data
+- **Completion**: Exits when thread is no longer active
+
+## Usage Examples
+
+### Example 1: Basic Resume After Page Refresh
+
+```typescript
+import {
+  resume_stream,
+  get_accumulated_content,
+} from "./services/agent_service";
+
+// Client already received content before refresh (e.g., from localStorage or state)
+const knownContent = "Hello world! This is content I already received...";
+
+const stream = await resume_stream({
+  thread_id: "thread-123",
+  message_id: "msg-abc-456",
+  known_content: knownContent,
+});
+
+// Consume new chunks
+for await (const chunk of stream) {
+  console.log("New content:", chunk.content);
+  displayInUI(chunk.content); // Update your UI
+}
+
+console.log("Stream completed!");
+```
+
+### Example 2: Check Status Before Resuming
+
+```typescript
+import { get_thread_status, resume_stream } from "./services/agent_service";
+
+// Check thread status first
+const status = await get_thread_status("thread-123");
+
+if (!status.exists) {
+  console.log("Thread not found or expired");
+} else if (status.status === "completed") {
+  // Thread already finished, get full content
+  const content = await get_accumulated_content("thread-123");
+  displayFullContent(content);
+} else if (status.status === "active") {
+  // Thread still streaming, resume from current position
+  const currentContent = await get_accumulated_content("thread-123");
+
+  const stream = await resume_stream({
+    thread_id: "thread-123",
+    message_id: "msg-123",
+    known_content: currentContent,
+  });
+
+  for await (const chunk of stream) {
+    appendToUI(chunk.content);
+  }
+}
+```
+
+### Example 3: Express/Fastify Endpoint (SSE)
+
+```typescript
+import express from "express";
+import { resume_stream } from "./services/agent_service";
+
+const app = express();
+
+app.post("/api/resume-stream", async (req, res) => {
+  const { thread_id, message_id, known_content } = req.body;
+
+  // Set up Server-Sent Events
+  res.setHeader("Content-Type", "text/event-stream");
+  res.setHeader("Cache-Control", "no-cache");
+  res.setHeader("Connection", "keep-alive");
+
+  try {
+    const stream = await resume_stream({
+      thread_id,
+      message_id,
+      known_content,
+    });
+
+    for await (const chunk of stream) {
+      res.write(`data: ${JSON.stringify(chunk)}\n\n`);
+    }
+
+    res.write("event: complete\ndata: {}\n\n");
+    res.end();
+  } catch (error) {
+    res.write(
+      `event: error\ndata: ${JSON.stringify({ error: error.message })}\n\n`
+    );
+    res.end();
+  }
+});
+```
+
+### Example 4: React Hook
+
+```typescript
+import { useState, useEffect } from 'react';
+
+function useResumeStream(threadId: string, messageId: string, knownContent: string) {
+  const [content, setContent] = useState('');
+  const [isStreaming, setIsStreaming] = useState(false);
+  const [error, setError] = useState<Error | null>(null);
+
+  useEffect(() => {
+    let isCancelled = false;
+
+    async function startStream() {
+      setIsStreaming(true);
+
+      try {
+        const stream = await resume_stream({
+          thread_id: threadId,
+          message_id: messageId,
+          known_content: knownContent,
+        });
+
+        for await (const chunk of stream) {
+          if (isCancelled) break;
+          setContent(prev => prev + chunk.content);
+        }
+      } catch (err) {
+        if (!isCancelled) {
+          setError(err as Error);
+        }
+      } finally {
+        if (!isCancelled) {
+          setIsStreaming(false);
+        }
+      }
+    }
+
+    startStream();
+
+    return () => {
+      isCancelled = true;
+    };
+  }, [threadId, messageId, knownContent]);
+
+  return { content, isStreaming, error };
+}
+
+// Usage in component
+function ChatMessage({ threadId, messageId, initialContent }) {
+  const { content, isStreaming } = useResumeStream(
+    threadId,
+    messageId,
+    initialContent
+  );
+
+  return (
+    <div>
+      {initialContent + content}
+      {isStreaming && <span className="cursor">▋</span>}
+    </div>
+  );
+}
+```
+
+## Related Functions
+
+### `get_accumulated_content(thread_id)`
+
+Get all accumulated content for a thread.
+
+```typescript
+const content = await get_accumulated_content("thread-123");
+```
+
+### `get_thread_status(thread_id)`
+
+Get thread status and metadata.
+
+```typescript
+const status = await get_thread_status("thread-123");
+// {
+//   exists: true,
+//   status: 'active' | 'completed' | 'aborted',
+//   chunkCount: 42,
+//   createdAt: 1234567890,
+//   updatedAt: 1234567900
+// }
+```
+
+### `get_active_threads()`
+
+Get all currently active threads.
+
+```typescript
+const activeThreads = await get_active_threads();
+// ['thread-1', 'thread-2', 'thread-3']
+```
+
+### `clear_thread_buffer(thread_id)`
+
+Manually clear a thread's buffer.
+
+```typescript
+await clear_thread_buffer("thread-123");
+```
+
+## Configuration
+
+### TTL (Time-To-Live)
+
+Threads are automatically cleaned up after 1 hour of inactivity:
+
+```typescript
+// In agent_service.ts
+const buffer = new InMemoryChunkBuffer({
+  ttl: 60 * 60 * 1000, // 1 hour
+  cleanupInterval: 5 * 60 * 1000, // Clean every 5 minutes
+});
+```
+
+### Polling Interval
+
+Adjust polling frequency based on your needs:
+
+```typescript
+// Fast polling for real-time updates
+const stream = await resume_stream({
+  thread_id: "thread-123",
+  message_id: "msg-123",
+  known_content: knownContent,
+  poll_interval: 50, // Check every 50ms
+});
+
+// Slower polling to reduce server load
+const stream = await resume_stream({
+  thread_id: "thread-123",
+  message_id: "msg-123",
+  known_content: knownContent,
+  poll_interval: 500, // Check every 500ms
+});
+```
+
+### Timeout
+
+The resume stream automatically times out after 30 seconds of no new data:
+
+```typescript
+// In agent_service.ts - resume_stream function
+const maxIdleTime = 30000; // 30 seconds
+```
+
+## Error Handling
+
+### Thread Not Found
+
+```typescript
+const stream = await resume_stream({
+  thread_id: "non-existent",
+  message_id: "msg-123",
+  known_content: "",
+});
+
+// Stream will exit immediately if thread doesn't exist
+for await (const chunk of stream) {
+  // Won't execute
+}
+```
+
+### Network Errors
+
+```typescript
+try {
+  const stream = await resume_stream({
+    thread_id: "thread-123",
+    message_id: "msg-123",
+    known_content: previousContent,
+  });
+
+  for await (const chunk of stream) {
+    displayChunk(chunk);
+  }
+} catch (error) {
+  console.error("Stream error:", error);
+  showErrorToUser("Failed to resume stream");
+}
+```
+
+## Best Practices
+
+1. **Always Check Status First**: Check thread status before attempting to resume
+2. **Store Message ID**: Keep track of the `message_id` (usually `run_id`) with your content
+3. **Preserve Content Exactly**: Store the exact content as received (no formatting/processing)
+4. **Handle Completion**: Be prepared for the stream to end at any time
+5. **Implement Retry Logic**: Add retry mechanism for transient failures
+6. **Set Appropriate Poll Interval**: Balance between responsiveness and server load
+7. **Content Matching**: The algorithm matches content intelligently - exact match, prefix, or suffix
+
+## Performance Considerations
+
+- **Memory Usage**: Chunks are stored in memory (consider TTL for long-running applications)
+- **Polling Overhead**: Lower `poll_interval` increases server load but improves responsiveness
+- **Concurrent Streams**: Multiple resume streams can run concurrently
+- **Cleanup**: Old threads are automatically cleaned up based on TTL
+
+## Limitations
+
+- **In-Memory Only**: Current implementation stores chunks in memory (lost on server restart)
+- **Single Message Per Thread**: Designed for one active message per thread
+- **Content Matching**: Relies on exact content matching (whitespace, encoding must match)
+- **No Persistence**: Chunks are not persisted to disk or database
+- **Performance**: Content matching is O(n) where n is number of chunks
+
+## Future Enhancements
+
+- **Persistent Storage**: Add Redis or database backend for persistence across restarts
+- **WebSocket Support**: Direct WebSocket integration for lower latency
+- **Event-Based Notifications**: Use event emitters instead of polling
+- **Compression**: Compress stored chunks to reduce memory usage
+- **Multi-Message Support**: Handle multiple concurrent messages per thread

package/dist/index.js

@@ -42,6 +42,18 @@ var import_messages = require("@langchain/core/messages");
 var import_langgraph = require("@langchain/langgraph");
 var import_uuid = require("uuid");
 var import_core = require("@axiom-lattice/core");
+function getOrCreateChunkBuffer() {
+  if (!(0, import_core.hasChunkBuffer)("default")) {
+    const buffer = new import_core.InMemoryChunkBuffer({
+      ttl: 60 * 60 * 1e3,
+      // 1 hour TTL
+      cleanupInterval: 5 * 60 * 1e3
+      // Clean every 5 minutes
+    });
+    (0, import_core.registerChunkBuffer)("default", buffer);
+  }
+  return (0, import_core.getChunkBuffer)("default");
+}
 async function agent_invoke({
   input,
   thread_id,
@@ -64,11 +76,11 @@ async function agent_invoke({
       configurable: {
         thread_id,
         run_id: run_id || (0, import_uuid.v4)(),
-        recursionLimit: 200,
         "x-tenant-id": tenant_id,
         "x-request-id": run_id,
         "x-thread-id": thread_id
-      }
+      },
+      recursionLimit: 200
     }
   );
   const data = result.messages.map((message2) => {
@@ -96,6 +108,7 @@ async function agent_stream({
     humanMessage.additional_kwargs = { files };
     messages = [humanMessage];
   }
+  const chunkBuffer = getOrCreateChunkBuffer();
   try {
     if (!runnable_agent) {
       throw new Error(`Agent ${assistant_id} not found`);
@@ -115,7 +128,8 @@ async function agent_stream({
           "x-thread-id": thread_id
         },
         streamMode: ["updates", "messages"],
-        subgraphs: false
+        subgraphs: false,
+        recursionLimit: 200
       }
     );
     return {
@@ -123,6 +137,7 @@ async function agent_stream({
         try {
           for await (const chunk of agentStream) {
             let data;
+            let chunkContent = "";
             if (chunk[0] === "updates") {
               const update = chunk[1];
               const values = Object.values(update);
@@ -141,16 +156,20 @@ async function agent_stream({
               };
             }
             if (data) {
+              await chunkBuffer.addChunk(thread_id, data);
               yield data;
             }
           }
+          await chunkBuffer.completeThread(thread_id);
         } catch (error) {
           console.error("Stream error:", error);
+          await chunkBuffer.abortThread(thread_id);
           throw error;
         }
       }
     };
   } catch (error) {
+    await chunkBuffer.abortThread(thread_id);
     throw error;
   }
 }
@@ -212,6 +231,32 @@ async function draw_graph(assistant_id) {
   const image = await drawableGraph.drawMermaid();
   return image;
 }
+async function resume_stream({
+  thread_id,
+  message_id,
+  known_content,
+  poll_interval = 100
+}) {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  const stream = await chunkBuffer.getNewChunksSinceContentIterator(
+    thread_id,
+    message_id,
+    known_content
+  );
+  return {
+    [Symbol.asyncIterator]: async function* () {
+      try {
+        for await (const chunk of stream) {
+          console.log(chunk.data?.content);
+          yield chunk;
+        }
+      } catch (error) {
+        console.error("Resume stream error:", error);
+        throw error;
+      }
+    }
+  };
+}
 
 // src/controllers/run.ts
 var import_uuid2 = require("uuid");
@@ -243,6 +288,7 @@ var createRun = async (request, reply) => {
         tenant_id,
         run_id: x_request_id
       });
+      reply.hijack();
       reply.raw.writeHead(200, {
         "Content-Type": "text/event-stream",
         "Cache-Control": "no-cache",
@@ -256,9 +302,9 @@ var createRun = async (request, reply) => {
 `);
         }
       } catch (error) {
+        console.error("Stream processing error:", error);
       } finally {
         reply.raw.end();
-        return reply.hijack();
       }
     } else {
       const result = await agent_invoke({
@@ -278,6 +324,47 @@ var createRun = async (request, reply) => {
     });
   }
 };
+var resumeStream = async (request, reply) => {
+  try {
+    const { thread_id, message_id, known_content, poll_interval } = request.body;
+    if (!thread_id || !message_id || known_content === void 0) {
+      reply.status(400).send({
+        success: false,
+        error: "thread_id, message_id, and known_content are required"
+      });
+      return;
+    }
+    reply.hijack();
+    reply.raw.writeHead(200, {
+      "Content-Type": "text/event-stream",
+      "Cache-Control": "no-cache",
+      Connection: "keep-alive",
+      "Access-Control-Allow-Origin": "*"
+    });
+    try {
+      const stream = await resume_stream({
+        thread_id,
+        message_id,
+        known_content: "",
+        poll_interval: poll_interval || 100
+      });
+      for await (const chunk of stream) {
+        reply.raw.write(`data: ${JSON.stringify(chunk)}
+
+`);
+      }
+    } catch (error) {
+      console.error("Resume stream processing error:", error);
+    } finally {
+      reply.raw.end();
+    }
+  } catch (error) {
+    reply.status(500).send({
+      success: false,
+      error: `Error resuming stream: ${error.message}`
+    });
+  }
+};
 
 // src/controllers/memory.ts
 var setMemoryItem = async (request, reply) => {
@@ -455,39 +542,6 @@ var getAgentGraph = async (request, reply) => {
 };
 
 // src/schemas/index.ts
-var createRunSchema = {
-  description: "Create a new agent run",
-  tags: ["Runs"],
-  summary: "Create Agent Run",
-  body: {
-    type: "object",
-    properties: {
-      thread_id: { type: "string", description: "Thread ID" },
-      assistant_id: { type: "string", description: "Assistant ID" },
-      message: { type: "string", description: "Message data for the run" },
-      command: {
-        type: "object",
-        description: "Command data for the run",
-        nullable: true
-      },
-      streaming: {
-        type: "boolean",
-        description: "Whether to stream the response",
-        nullable: true
-      },
-      background: {
-        type: "boolean",
-        description: "Whether to run in background",
-        nullable: true
-      }
-    },
-    required: ["thread_id", "assistant_id", "message"]
-  },
-  response: {
-    200: {},
-    400: {}
-  }
-};
 var getAllMemoryItemsSchema = {
   description: "Get all memory items for an assistant thread",
   tags: ["Memory"],
@@ -607,7 +661,8 @@ var getAgentGraphSchema = {
 
 // src/routes/index.ts
 var registerLatticeRoutes = (app2) => {
-  app2.post("/api/runs", { schema: createRunSchema }, createRun);
+  app2.post("/api/runs", createRun);
+  app2.post("/api/resume_stream", resumeStream);
   app2.get(
     "/api/assistants/:assistantId/:thread_id/memory",
     { schema: getAllMemoryItemsSchema },