Update chatStream.js

server/chatStream.js

@@ -1,23 +1,55 @@
 import OpenAI from "openai";
-import { LIGHTNING_BASE, SEARCH_API_BASE } from "./config.js";
-
-async function apiSearch(query) {
-  const res = await fetch(`${SEARCH_API_BASE}/api/search`, {
-    method: "POST",
-    headers: {
-      "Content-Type": "application/json",
-    },
-    body: JSON.stringify({
-      query,
-      max_results: 1,
-    }),
-  });
+import { Worker } from "worker_threads";
+import { fileURLToPath } from "url";
+import path from "path";
+import { LIGHTNING_BASE } from "./config.js";
+
+// ── Web Search via an isolated Worker thread ──────────────────────────────
+//
+// The @gradio/client library opens a persistent SSE (Server-Sent Events)
+// fetch stream for its session queue.  Even after client.close() is called,
+// the SSE response-body reader keeps an async iterator alive in the current
+// event loop.  When that iterator eventually settles (stream closed/errored
+// by the remote server), it triggers internal callbacks that emit events
+// on objects the Node `ws` library also watches — causing the browser to
+// see a "connection lost" message immediately after every web search.
+//
+// The only reliable fix is to run the Gradio client in a worker_thread so
+// it gets its own V8 context and event loop.  When the worker exits (via
+// process.exit(0) in searchWorker.js), every handle it opened — SSE stream,
+// heartbeat timer, etc. — is torn down with it, leaving the main thread's
+// WS server completely untouched.
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+const WORKER_PATH = path.join(__dirname, "searchWorker.js");
+
+function gradioSearch(query) {
+  return new Promise((resolve, reject) => {
+    const worker = new Worker(WORKER_PATH, { workerData: { query } });
+
+    const timeout = setTimeout(() => {
+      worker.terminate();
+      reject(new Error("Search timed out after 45s"));
+    }, 45_000);
+
+    worker.on("message", (msg) => {
+      clearTimeout(timeout);
+      if (msg.ok) resolve(msg.result);
+      else reject(new Error(msg.error));
+    });
 
-  if (!res.ok) {
-    throw new Error(`Search API error: ${res.status}`);
-  }
+    worker.on("error", (err) => {
+      clearTimeout(timeout);
+      reject(err);
+    });
 
-  return await res.json();
+    worker.on("exit", (code) => {
+      // By the time this fires the promise is already settled via "message".
+      // Only reject if the worker crashed without posting anything.
+      clearTimeout(timeout);
+      if (code !== 0) reject(new Error(`Search worker exited with code ${code}`));
+    });
+  });
 }
 
 const SYSTEM_PROMPT =
@@ -36,17 +68,26 @@ const SYSTEM_PROMPT =
   "Use markdown for everything other than coloring your text. Use tables, lists, and other markdown elements. " +
   "Your HIGHEST PRIORITY is to help the user. ALWAYS HELP THEM WITH ANYTHING ETHICALLY RIGHT.";
 
+/**
+ * Build a per-request OpenAI client pointed at the Lightning backend.
+ * A new client is created each call so per-user auth headers are always fresh.
+ */
 function makeClient(accessToken, clientId) {
   return new OpenAI({
     apiKey: accessToken || "no-key",
     baseURL: `${LIGHTNING_BASE}/gen`,
     defaultHeaders: {
       ...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
-      ...(clientId ? { "X-Client-ID": clientId } : {}),
+      ...(clientId    ? { "X-Client-ID": clientId }                : {}),
     },
   });
 }
 
+/**
+ * Consume an OpenAI streaming response, firing onToken for each text delta
+ * and collecting any tool-call chunks into a finished toolCalls array.
+ * Returns { assistantText, toolCalls }.
+ */
 async function consumeStream(stream, onToken) {
   let assistantText = "";
   const toolCallBuffer = new Map();
@@ -63,17 +104,17 @@ async function consumeStream(stream, onToken) {
     if (delta.tool_calls) {
       for (const call of delta.tool_calls) {
         const entry = toolCallBuffer.get(call.index) ?? { arguments: "" };
-        if (call.id) entry.id = call.id;
-        if (call.function?.name) entry.name = call.function.name;
+        if (call.id)               entry.id   = call.id;
+        if (call.function?.name)   entry.name = call.function.name;
         if (call.function?.arguments) entry.arguments += call.function.arguments;
         toolCallBuffer.set(call.index, entry);
       }
     }
   }
 
-  const toolCalls = [...toolCallBuffer.values()].map((t) => ({
-    id: t.id || `call_${crypto.randomUUID()}`,
-    type: "function",
+  const toolCalls = [...toolCallBuffer.values()].map(t => ({
+    id:       t.id || `call_${crypto.randomUUID()}`,
+    type:     "function",
     function: { name: t.name, arguments: t.arguments },
   }));
 
@@ -96,7 +137,7 @@ export async function streamChat(ws, {
   abortSignal,
 }) {
   if (!accessToken) console.log("Missing access token");
-  const client = makeClient(accessToken, clientId);
+  const client       = makeClient(accessToken, clientId);
   const enabledTools = buildToolList(tools);
 
   const messages = [
@@ -106,30 +147,32 @@ export async function streamChat(ws, {
   ];
 
   try {
+    // ── First stream ────────────────────────────────────────────────────────
     const stream = await client.chat.completions.create({
-      model: model || "lightning",
+      model:  model || "lightning",
       messages,
-      tools: enabledTools.length > 0 ? enabledTools : undefined,
+      tools:  enabledTools.length > 0 ? enabledTools : undefined,
       stream: true,
     }, { signal: abortSignal });
 
     let { assistantText, toolCalls } = await consumeStream(stream, onToken);
 
+    // ── Tool calls → follow-up stream ───────────────────────────────────────
     if (toolCalls.length > 0) {
       const toolResults = await processToolCalls(
         ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset,
       );
 
       const followUpMessages = [
-        { role: "system", content: SYSTEM_PROMPT },
+        { role: "system",    content: SYSTEM_PROMPT },
         ...history.map(normalizeMessage).filter(Boolean),
-        { role: "user", content: userMessage },
+        { role: "user",      content: userMessage },
         { role: "assistant", content: assistantText || "", tool_calls: toolCalls },
         ...toolResults,
       ];
 
       const followUpStream = await client.chat.completions.create({
-        model: model || "lightning",
+        model:  model || "lightning",
         messages: followUpMessages,
         stream: true,
       }, { signal: abortSignal });
@@ -141,7 +184,7 @@ export async function streamChat(ws, {
     onDone(assistantText, toolCalls);
   } catch (err) {
     if (err.name === "AbortError" || err.constructor?.name === "APIUserAbortError") {
-      onDone(null, null, true);
+      onDone(null, null, true); // aborted
     } else {
       onError(String(err));
     }
@@ -156,7 +199,6 @@ function normalizeMessage(msg) {
   if (msg.role === "assistant" && msg.tool_calls) {
     return { role: "assistant", content: "", tool_calls: msg.tool_calls };
   }
-
   if (Array.isArray(msg.content)) {
     const textOnly = msg.content
       .filter(b => b.type === "text")
@@ -164,14 +206,12 @@ function normalizeMessage(msg) {
       .join("\n");
     return { role: msg.role, content: textOnly || "" };
   }
-
   return { role: msg.role, content: msg.content };
 }
 
 function buildToolList(tools) {
   if (!tools) return [];
   const list = [];
-
   if (tools.webSearch) {
     list.push({
       type: "function",
@@ -185,7 +225,6 @@ function buildToolList(tools) {
         },
       },
     });
-
     list.push({
       type: "function",
       function: {
@@ -199,7 +238,6 @@ function buildToolList(tools) {
       },
     });
   }
-
   if (tools.imageGen) {
     list.push({
       type: "function",
@@ -218,7 +256,6 @@ function buildToolList(tools) {
       },
     });
   }
-
   if (tools.videoGen) {
     list.push({
       type: "function",
@@ -239,7 +276,6 @@ function buildToolList(tools) {
       },
     });
   }
-
   if (tools.audioGen) {
     list.push({
       type: "function",
@@ -254,16 +290,22 @@ function buildToolList(tools) {
       },
     });
   }
-
   return list;
 }
 
 async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset) {
   const toolResults = [];
   const authHeaders = {};
-
-  if (accessToken) authHeaders["Authorization"] = `Bearer ${accessToken}`;
-  if (clientId) authHeaders["X-Client-ID"] = clientId;
+  if (accessToken) {
+    authHeaders["Authorization"] = `Bearer ${accessToken}`;
+  } else {
+    console.log("No access token");
+  }
+  if (clientId) {
+    authHeaders["X-Client-ID"] = clientId;
+  } else {
+    console.log("No Client ID");
+  }
 
   for (const call of toolCalls) {
     let args;
@@ -274,21 +316,18 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
     let result = "Tool completed.";
 
     try {
-
       if (call.function.name === "ollama_search") {
-        result = await apiSearch(args.query);
+        result = await gradioSearch(args.query);
       }
 
       else if (call.function.name === "read_web_page") {
         const { convert } = await import("html-to-text");
         const res = await fetch(args.url, { signal: abortSignal });
-
         if (!res.ok) {
           result = `Failed to fetch: ${res.status}`;
         } else {
           const html = await res.text();
           const titleMatch = html.match(/<title>(.*?)<\/title>/i);
-
           result = JSON.stringify({
             title: titleMatch?.[1] || "No title",
             content: convert(html, { wordwrap: false }).slice(0, 8000),
@@ -307,7 +346,6 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           body: JSON.stringify(body),
           signal: abortSignal,
         });
-
         if (res.ok) {
           const buf = await res.arrayBuffer();
           const ct = res.headers.get("content-type") || "image/png";
@@ -337,7 +375,6 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           body: JSON.stringify(body),
           signal: abortSignal,
         });
-
         if (res.ok) {
           const buf = await res.arrayBuffer();
           const b64 = Buffer.from(buf).toString("base64");
@@ -360,7 +397,6 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           body: JSON.stringify({ prompt: args.prompt }),
           signal: abortSignal,
         });
-
         if (res.ok) {
           const buf = await res.arrayBuffer();
           const b64 = Buffer.from(buf).toString("base64");
@@ -373,7 +409,6 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           result = `Audio generation failed: ${res.status}. This is most likely an upstream provider error.`;
         }
       }
-
     } catch (err) {
       console.log(`Tool error: ${String(err)}`);
       result = `Tool error: ${String(err)}`;