abort

2 files changed, 62 insertions, 1 deletions

diff --git a/js/src/server/server.ts b/js/src/server/server.ts
index ea70cd5ef..11262f4fd 100644
--- a/js/src/server/server.ts
+++ b/js/src/server/server.ts
@@ -159,7 +159,46 @@ export namespace Server {
         },
       )
       .post(
+        "/session_abort",
+        describeRoute({
+          description: "Abort a session",
+          responses: {
+            200: {
+              description: "Aborted session",
+              content: {
+                "application/json": {
+                  schema: resolver(z.boolean()),
+                },
+              },
+            },
+          },
+        }),
+        zValidator(
+          "json",
+          z.object({
+            sessionID: z.string(),
+          }),
+        ),
+        async (c) => {
+          const body = c.req.valid("json");
+          return c.json(Session.abort(body.sessionID));
+        },
+      )
+      .post(
         "/session_chat",
+        describeRoute({
+          description: "Chat with a model",
+          responses: {
+            200: {
+              description: "Chat with a model",
+              content: {
+                "application/json": {
+                  schema: resolver(SessionMessage),
+                },
+              },
+            },
+          },
+        }),
         zValidator(
           "json",
           z.object({
diff --git a/js/src/session/session.ts b/js/src/session/session.ts
index 406ab27f8..011519048 100644
--- a/js/src/session/session.ts
+++ b/js/src/session/session.ts
@@ -129,6 +129,16 @@ export namespace Session {
     }
   }
 
+  const pending = new Map<string, AbortController>();
+
+  export function abort(sessionID: string) {
+    const controller = pending.get(sessionID);
+    if (!controller) return false;
+    controller.abort();
+    pending.delete(sessionID);
+    return true;
+  }
+
   export async function chat(input: {
     sessionID: string;
     providerID: string;
@@ -225,6 +235,8 @@ export namespace Session {
         tool: {},
       },
     };
+    const controller = new AbortController();
+    pending.set(input.sessionID, controller);
     const result = streamText({
       onStepFinish: (step) => {
         update(input.sessionID, (draft) => {
@@ -240,6 +252,8 @@ export namespace Session {
             .toNumber();
         });
       },
+      abortSignal: controller.signal,
+      maxRetries: 6,
       stopWhen: stepCountIs(1000),
       messages: convertToModelMessages(msgs),
       temperature: 0,
@@ -251,7 +265,14 @@ export namespace Session {
     let text: TextUIPart | undefined;
     const reader = result.toUIMessageStream().getReader();
     while (true) {
-      const { done, value } = await reader.read();
+      const result = await reader.read().catch((e) => {
+        if (e instanceof DOMException && e.name === "AbortError") {
+          return;
+        }
+        throw e;
+      });
+      if (!result) break;
+      const { done, value } = result;
       if (done) break;
       l.info("part", {
         type: value.type,
@@ -316,6 +337,7 @@ export namespace Session {
       }
       await write(next);
     }
+    pending.delete(input.sessionID);
     next.metadata!.time.completed = Date.now();
     await write(next);
     return next;