feat(api): harden embed server, implement missing endpoints needed for interrupts (#1520)

.changeset/red-hands-hunt.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph-api": patch
+---
+
+feat(api): harden embed server, implement missing endpoints needed for interrupts

libs/langgraph-api/src/api/threads.mts

@@ -95,20 +95,7 @@ api.get(
 api.post(
   "/threads/:thread_id/state",
   zValidator("param", z.object({ thread_id: z.string().uuid() })),
-  zValidator(
-    "json",
-    z.object({
-      values: z
-        .union([
-          z.record(z.string(), z.unknown()),
-          z.array(z.record(z.string(), z.unknown())),
-        ])
-        .nullish(),
-      as_node: z.string().optional(),
-      checkpoint_id: z.string().optional(),
-      checkpoint: schemas.CheckpointSchema.nullish(),
-    })
-  ),
+  zValidator("json", schemas.ThreadStateUpdate),
   async (c) => {
     // Update Thread State
     const { thread_id } = c.req.valid("param");

libs/langgraph-api/src/experimental/embed.mts

@@ -14,7 +14,7 @@ import * as schemas from "../schemas.mjs";
 
 import { z } from "zod";
 import { streamState } from "../stream.mjs";
-import { serialiseAsDict } from "../utils/serde.mjs";
+import { serialiseAsDict, serializeError } from "../utils/serde.mjs";
 import { getDisconnectAbortSignal, jsonExtra } from "../utils/hono.mjs";
 import { stateSnapshotToThreadState } from "../state.mjs";
 import { ensureContentType } from "../http/middleware.mjs";
@@ -26,10 +26,26 @@ interface Thread {
   metadata: Metadata;
 }
 
-interface ThreadSaver {
+/**
+ * Interface for storing and retrieving threads used by `createEmbedServer`.
+ * @experimental Does not follow semver.
+ */
+export interface ThreadSaver {
   get: (id: string) => Promise<Thread>;
-  put: (id: string, options: { metadata?: Metadata }) => Promise<void>;
+
+  set: (
+    id: string,
+    options: { kind: "put" | "patch"; metadata?: Metadata }
+  ) => Promise<Thread>;
   delete: (id: string) => Promise<void>;
+
+  search?: (options: {
+    metadata?: Metadata;
+    limit: number;
+    offset: number;
+    sortBy: "created_at" | "updated_at";
+    sortOrder: "asc" | "desc";
+  }) => AsyncGenerator<{ thread: Thread; total: number }>;
 }
 
 function createStubRun(
@@ -39,12 +55,13 @@ function createStubRun(
   const now = new Date();
   const runId = uuidv4();
 
-  const streamMode = Array.isArray(payload.stream_mode)
+  let streamMode = Array.isArray(payload.stream_mode)
     ? payload.stream_mode
     : payload.stream_mode
     ? [payload.stream_mode]
     : undefined;
 
+  if (streamMode == null || streamMode.length === 0) streamMode = ["values"];
   const config = Object.assign(
     {},
     payload.config ?? {},
@@ -78,6 +95,7 @@ function createStubRun(
       input: payload.input,
       command: payload.command,
       config,
+      context: payload.context,
       stream_mode: streamMode,
       interrupt_before: payload.interrupt_before,
       interrupt_after: payload.interrupt_after,
@@ -92,7 +110,8 @@ function createStubRun(
 }
 
 /**
- * Attach LangGraph Platform-esque routes to a given Hono instance.
+ * Create a Hono server with a subset of LangGraph Platform routes.
+ *
  * @experimental Does not follow semver.
  */
 export function createEmbedServer(options: {
@@ -113,12 +132,16 @@ export function createEmbedServer(options: {
   api.use(ensureContentType());
 
   api.post("/threads", zValidator("json", schemas.ThreadCreate), async (c) => {
-    // create a new threaad
+    // create a new thread
     const payload = c.req.valid("json");
     const threadId = payload.thread_id || uuidv4();
-
-    await options.threads.put(threadId, payload);
-    return jsonExtra(c, { thread_id: threadId });
+    return jsonExtra(
+      c,
+      await options.threads.set(threadId, {
+        kind: "put",
+        metadata: payload.metadata,
+      })
+    );
   });
 
   api.get(
@@ -127,8 +150,25 @@ export function createEmbedServer(options: {
     async (c) => {
       // Get Thread
       const { thread_id } = c.req.valid("param");
-      const thread = await options.threads.get(thread_id);
-      return jsonExtra(c, thread);
+      return jsonExtra(c, await options.threads.get(thread_id));
+    }
+  );
+
+  api.patch(
+    "/threads/:thread_id",
+    zValidator("param", z.object({ thread_id: z.string().uuid() })),
+    zValidator("json", schemas.ThreadCreate),
+    async (c) => {
+      // Update Thread
+      const { thread_id } = c.req.valid("param");
+      const payload = c.req.valid("json");
+      return jsonExtra(
+        c,
+        await options.threads.set(thread_id, {
+          kind: "patch",
+          metadata: payload.metadata,
+        })
+      );
     }
   );
 
@@ -136,12 +176,45 @@ export function createEmbedServer(options: {
     "/threads/:thread_id",
     zValidator("param", z.object({ thread_id: z.string().uuid() })),
     async (c) => {
+      // Delete Thread
       const { thread_id } = c.req.valid("param");
       await options.threads.delete(thread_id);
       return new Response(null, { status: 204 });
     }
   );
 
+  api.post(
+    "/threads/search",
+    zValidator("json", schemas.ThreadSearchRequest),
+    async (c) => {
+      const payload = c.req.valid("json");
+      const result: unknown[] = [];
+
+      if (!options.threads.search)
+        return c.json({ error: "Threads search not implemented" }, 422);
+
+      const sortBy =
+        payload.sort_by === "created_at" || payload.sort_by === "updated_at"
+          ? payload.sort_by
+          : "created_at";
+
+      let total = 0;
+      for await (const item of options.threads.search({
+        metadata: payload.metadata,
+        limit: payload.limit ?? 10,
+        offset: payload.offset ?? 0,
+        sortBy,
+        sortOrder: payload.sort_order ?? "desc",
+      })) {
+        result.push(item.thread);
+        // Only set total if it's the first item
+        if (total === 0) total = item.total;
+      }
+      c.res.headers.set("X-Pagination-Total", total.toString());
+      return jsonExtra(c, result);
+    }
+  );
+
   api.get(
     "/threads/:thread_id/state",
     zValidator("param", z.object({ thread_id: z.string().uuid() })),
@@ -179,6 +252,99 @@ export function createEmbedServer(options: {
     }
   );
 
+  api.post(
+    "/threads/:thread_id/state",
+    zValidator("param", z.object({ thread_id: z.string().uuid() })),
+    zValidator("json", schemas.ThreadStateUpdate),
+    async (c) => {
+      // Update Thread State
+      const { thread_id } = c.req.valid("param");
+      const payload = c.req.valid("json");
+      const config: RunnableConfig = { configurable: { thread_id } };
+      config.configurable ??= {};
+
+      if (payload.checkpoint_id) {
+        config.configurable.checkpoint_id = payload.checkpoint_id;
+      }
+
+      if (payload.checkpoint) {
+        Object.assign(config.configurable, payload.checkpoint);
+      }
+
+      const thread = await options.threads.get(thread_id);
+      const graphId = thread.metadata?.graph_id as string | undefined | null;
+      const graph = graphId ? await getGraph(graphId) : undefined;
+      if (graph == null) return c.json({ error: "Graph not found" }, 404);
+
+      const result = await graph.updateState(
+        config,
+        payload.values,
+        payload.as_node
+      );
+      return jsonExtra(c, { checkpoint: result.configurable });
+    }
+  );
+
+  // get thread state at checkpoint
+  api.get(
+    "/threads/:thread_id/state/:checkpoint_id",
+    zValidator(
+      "param",
+      z.object({
+        thread_id: z.string().uuid(),
+        checkpoint_id: z.string().uuid(),
+      })
+    ),
+    zValidator(
+      "query",
+      z.object({ subgraphs: schemas.coercedBoolean.optional() })
+    ),
+    async (c) => {
+      // Get Thread State At Checkpoint
+      const { thread_id, checkpoint_id } = c.req.valid("param");
+      const { subgraphs } = c.req.valid("query");
+
+      const thread = await options.threads.get(thread_id);
+      const graphId = thread.metadata?.graph_id as string | undefined | null;
+      const graph = graphId ? await getGraph(graphId) : undefined;
+      if (graph == null) return c.json({ error: "Graph not found" }, 404);
+
+      const result = await graph.getState(
+        { configurable: { thread_id, checkpoint_id } },
+        { subgraphs }
+      );
+      return jsonExtra(c, stateSnapshotToThreadState(result));
+    }
+  );
+
+  api.post(
+    "/threads/:thread_id/state/checkpoint",
+    zValidator("param", z.object({ thread_id: z.string().uuid() })),
+    zValidator(
+      "json",
+      z.object({
+        subgraphs: schemas.coercedBoolean.optional(),
+        checkpoint: schemas.CheckpointSchema.nullish(),
+      })
+    ),
+    async (c) => {
+      // Get Thread State At Checkpoint post
+      const { thread_id } = c.req.valid("param");
+      const { checkpoint, subgraphs } = c.req.valid("json");
+
+      const thread = await options.threads.get(thread_id);
+      const graphId = thread.metadata?.graph_id as string | undefined | null;
+      const graph = graphId ? await getGraph(graphId) : undefined;
+      if (graph == null) return c.json({ error: "Graph not found" }, 404);
+
+      const result = await graph.getState(
+        { configurable: { thread_id, ...checkpoint } },
+        { subgraphs }
+      );
+      return jsonExtra(c, stateSnapshotToThreadState(result));
+    }
+  );
+
   api.post(
     "/threads/:thread_id/history",
     zValidator("param", z.object({ thread_id: z.string().uuid() })),
@@ -218,29 +384,37 @@ export function createEmbedServer(options: {
     zValidator("json", schemas.RunCreate),
     async (c) => {
       // Stream Run
-      return streamSSE(c, async (stream) => {
-        const { thread_id } = c.req.valid("param");
-        const payload = c.req.valid("json");
+      const { thread_id } = c.req.valid("param");
+      const payload = c.req.valid("json");
 
+      const thread = await options.threads.get(thread_id);
+      if (thread == null) return c.json({ error: "Thread not found" }, 404);
+
+      return streamSSE(c, async (stream) => {
         const signal = getDisconnectAbortSignal(c, stream);
         const run = createStubRun(thread_id, payload);
 
-        // update thread with new graph_id
-        const thread = await options.threads.get(thread_id);
-        await options.threads.put(thread_id, {
+        await options.threads.set(thread_id, {
+          kind: "patch",
           metadata: {
-            ...thread.metadata,
             graph_id: payload.assistant_id,
             assistant_id: payload.assistant_id,
           },
         });
 
-        for await (const { event, data } of streamState(run, {
-          attempt: 1,
-          getGraph,
-          signal,
-        })) {
-          await stream.writeSSE({ data: serialiseAsDict(data), event });
+        try {
+          for await (const { event, data } of streamState(run, {
+            attempt: 1,
+            getGraph,
+            signal,
+          })) {
+            await stream.writeSSE({ data: serialiseAsDict(data), event });
+          }
+        } catch (error) {
+          await stream.writeSSE({
+            data: serialiseAsDict(serializeError(error)),
+            event: "error",
+          });
         }
       });
     }
@@ -253,7 +427,8 @@ export function createEmbedServer(options: {
       const signal = getDisconnectAbortSignal(c, stream);
       const threadId = uuidv4();
 
-      await options.threads.put(threadId, {
+      await options.threads.set(threadId, {
+        kind: "put",
         metadata: {
           graph_id: payload.assistant_id,
           assistant_id: payload.assistant_id,
@@ -262,12 +437,19 @@ export function createEmbedServer(options: {
 
       try {
         const run = createStubRun(threadId, payload);
-        for await (const { event, data } of streamState(run, {
-          attempt: 1,
-          getGraph,
-          signal,
-        })) {
-          await stream.writeSSE({ data: serialiseAsDict(data), event });
+        try {
+          for await (const { event, data } of streamState(run, {
+            attempt: 1,
+            getGraph,
+            signal,
+          })) {
+            await stream.writeSSE({ data: serialiseAsDict(data), event });
+          }
+        } catch (error) {
+          await stream.writeSSE({
+            data: serialiseAsDict(serializeError(error)),
+            event: "error",
+          });
         }
       } finally {
         await options.threads.delete(threadId);
@@ -275,5 +457,12 @@ export function createEmbedServer(options: {
     });
   });
 
+  api.notFound((c) => {
+    return c.json(
+      { error: `${c.req.method} ${c.req.path} not implemented` },
+      404
+    );
+  });
+
   return api;
 }

libs/langgraph-api/src/schemas.mts

@@ -421,13 +421,13 @@ export const ThreadStateUpdate = z
   .object({
     values: z
       .union([
-        z.array(z.object({}).catchall(z.any())),
-        z.object({}).catchall(z.any()),
-        z.null(),
+        z.record(z.string(), z.unknown()),
+        z.array(z.record(z.string(), z.unknown())),
       ])
-      .optional(),
-    checkpoint_id: z.string().optional(),
+      .nullish(),
     as_node: z.string().optional(),
+    checkpoint_id: z.string().optional(),
+    checkpoint: CheckpointSchema.nullish(),
   })
   .describe("Payload for adding state to a thread.");
 

libs/sdk-validation/src/tests/stream.test.tsx

@@ -12,27 +12,28 @@ import { StateGraph, MessagesAnnotation, START } from "@langchain/langgraph";
 import { MemorySaver } from "@langchain/langgraph-checkpoint";
 import { FakeStreamingChatModel } from "@langchain/core/utils/testing";
 import { AIMessage } from "@langchain/core/messages";
-import { createEmbedServer } from "@langchain/langgraph-api/experimental/embed";
+import {
+  createEmbedServer,
+  type ThreadSaver,
+} from "@langchain/langgraph-api/experimental/embed";
 import { randomUUID } from "node:crypto";
 import { useState } from "react";
 
-const threads = (() => {
+const threads: ThreadSaver = (() => {
   const THREADS: Record<
     string,
     { thread_id: string; metadata: Record<string, unknown> }
   > = {};
 
   return {
-    get: async (id: string) => THREADS[id],
-    put: async (
-      threadId: string,
-      { metadata }: { metadata?: Record<string, unknown> }
-    ) => {
-      THREADS[threadId] = { thread_id: threadId, metadata: metadata ?? {} };
-    },
-    delete: async (threadId: string) => {
-      delete THREADS[threadId];
+    get: async (id) => THREADS[id],
+    set: async (threadId, { metadata }) => {
+      return (THREADS[threadId] = {
+        thread_id: threadId,
+        metadata: { ...(THREADS[threadId]?.metadata ?? {}), ...metadata },
+      });
     },
+    delete: async (threadId) => void delete THREADS[threadId],
   };
 })();