feat(app): allow stopping requests

src/app/api/chat/cancel/route.ts

@@ -0,0 +1,50 @@
+import { NextRequest } from 'next/server';
+
+// In-memory map to store cancel tokens by messageId
+const cancelTokens: Record<string, AbortController> = {};
+
+// Export for use in chat/route.ts
+export function registerCancelToken(
+  messageId: string,
+  controller: AbortController,
+) {
+  cancelTokens[messageId] = controller;
+}
+
+export function cleanupCancelToken(messageId: string) {
+  var cancelled = false;
+  if (messageId in cancelTokens) {
+    delete cancelTokens[messageId];
+    cancelled = true;
+  }
+  return cancelled;
+}
+
+export function cancelRequest(messageId: string) {
+  const controller = cancelTokens[messageId];
+  if (controller) {
+    try {
+      controller.abort();
+    } catch (e) {
+      console.error(`Error aborting request for messageId ${messageId}:`, e);
+    }
+    return true;
+  }
+  return false;
+}
+
+export async function POST(req: NextRequest) {
+  const { messageId } = await req.json();
+  if (!messageId) {
+    return Response.json({ error: 'Missing messageId' }, { status: 400 });
+  }
+  const cancelled = cancelRequest(messageId);
+  if (cancelled) {
+    return Response.json({ success: true });
+  } else {
+    return Response.json(
+      { error: 'No in-progress request for this messageId' },
+      { status: 404 },
+    );
+  }
+}

src/app/api/chat/route.ts

@@ -18,6 +18,10 @@ import { ChatOpenAI } from '@langchain/openai';
 import crypto from 'crypto';
 import { and, eq, gt } from 'drizzle-orm';
 import { EventEmitter } from 'stream';
+import {
+  registerCancelToken,
+  cleanupCancelToken,
+} from './cancel/route';
 
 export const runtime = 'nodejs';
 export const dynamic = 'force-dynamic';
@@ -62,6 +66,7 @@ const handleEmitterEvents = async (
   aiMessageId: string,
   chatId: string,
   startTime: number,
+  userMessageId: string,
 ) => {
   let recievedMessage = '';
   let sources: any[] = [];
@@ -139,6 +144,9 @@ const handleEmitterEvents = async (
     );
     writer.close();
 
+    // Clean up the abort controller reference
+    cleanupCancelToken(userMessageId);
+
     db.insert(messagesSchema)
       .values({
         content: recievedMessage,
@@ -329,6 +337,28 @@ export const POST = async (req: Request) => {
       );
     }
 
+    const responseStream = new TransformStream();
+    const writer = responseStream.writable.getWriter();
+    const encoder = new TextEncoder();
+
+    // --- Cancellation logic ---
+    const abortController = new AbortController();
+    registerCancelToken(message.messageId, abortController);
+
+    abortController.signal.addEventListener('abort', () => {
+      console.log('Stream aborted, sending cancel event');
+      writer.write(
+        encoder.encode(
+          JSON.stringify({
+            type: 'error',
+            data: 'Request cancelled by user',
+          }),
+        ),
+      );
+      cleanupCancelToken(message.messageId);
+    });
+
+    // Pass the abort signal to the search handler
     const stream = await handler.searchAndAnswer(
       message.content,
       history,
@@ -337,12 +367,9 @@ export const POST = async (req: Request) => {
       body.optimizationMode,
       body.files,
       body.systemInstructions,
+      abortController.signal,
     );
 
-    const responseStream = new TransformStream();
-    const writer = responseStream.writable.getWriter();
-    const encoder = new TextEncoder();
-
     handleEmitterEvents(
       stream,
       writer,
@@ -350,7 +377,9 @@ export const POST = async (req: Request) => {
       aiMessageId,
       message.chatId,
       startTime,
+      message.messageId,
     );
+
     handleHistorySave(message, humanMessageId, body.focusMode, body.files);
 
     return new Response(responseStream.readable, {

src/app/api/search/route.ts

@@ -124,6 +124,8 @@ export const POST = async (req: Request) => {
     if (!searchHandler) {
       return Response.json({ message: 'Invalid focus mode' }, { status: 400 });
     }
+    const abortController = new AbortController();
+    const { signal } = abortController;
 
     const emitter = await searchHandler.searchAndAnswer(
       body.query,
@@ -133,6 +135,7 @@ export const POST = async (req: Request) => {
       body.optimizationMode,
       [],
       body.systemInstructions || '',
+      signal,
     );
 
     if (!body.stream) {
@@ -180,9 +183,6 @@ export const POST = async (req: Request) => {
 
     const encoder = new TextEncoder();
 
-    const abortController = new AbortController();
-    const { signal } = abortController;
-
     const stream = new ReadableStream({
       start(controller) {
         let sources: any[] = [];