Update server/chatStream.js

server/chatStream.js

@@ -1,3 +1,4 @@
+import OpenAI from "openai";
 import { LIGHTNING_BASE } from "./config.js";
 
 const SYSTEM_PROMPT =
@@ -15,6 +16,59 @@ const SYSTEM_PROMPT =
   "Never use single backslashes. You may use emojis where appropriate. " +
   "Use markdown for everything other than coloring your text. Use tables, lists, and other markdown elements.";
 
+/**
+ * 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 }                : {}),
+    },
+  });
+}
+
+/**
+ * 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();
+
+  for await (const chunk of stream) {
+    const delta = chunk.choices?.[0]?.delta;
+    if (!delta) continue;
+
+    if (delta.content) {
+      assistantText += delta.content;
+      onToken(delta.content);
+    }
+
+    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.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",
+    function: { name: t.name, arguments: t.arguments },
+  }));
+
+  return { assistantText, toolCalls };
+}
+
 export async function streamChat(ws, {
   sessionId,
   model,
@@ -30,12 +84,8 @@ export async function streamChat(ws, {
   onNewAsset,
   abortSignal,
 }) {
-  const headers = {
-    "Content-Type": "application/json",
-    "Accept": "text/event-stream",
-  };
-  if (accessToken) headers["Authorization"] = `Bearer ${accessToken}`;
-  if (clientId) headers["X-Client-ID"] = clientId;
+  const client       = makeClient(accessToken, clientId);
+  const enabledTools = buildToolList(tools);
 
   const messages = [
     { role: "system", content: SYSTEM_PROMPT },
@@ -43,128 +93,44 @@ export async function streamChat(ws, {
     { role: "user", content: userMessage },
   ];
 
-  const enabledTools = buildToolList(tools);
-
   try {
-    const response = await fetch(`${LIGHTNING_BASE}/gen/chat/completions`, {
-      method: "POST",
-      headers,
-      body: JSON.stringify({
-        model: model || "lightning",
-        messages,
-        tools: enabledTools.length > 0 ? enabledTools : undefined,
-        stream: true,
-      }),
-      signal: abortSignal,
-    });
-
-    if (!response.ok) {
-      const err = await response.text();
-      onError(err);
-      return;
-    }
+    // ── First stream ────────────────────────────────────────────────────────
+    const stream = await client.chat.completions.create({
+      model:  model || "lightning",
+      messages,
+      tools:  enabledTools.length > 0 ? enabledTools : undefined,
+      stream: true,
+    }, { signal: abortSignal });
 
-    const reader = response.body.getReader();
-    const decoder = new TextDecoder();
-    let assistantText = "";
-    const toolCallBuffer = new Map();
-    let buffer = "";
-
-    while (true) {
-      const { done, value } = await reader.read();
-      if (done) break;
-
-      buffer += decoder.decode(value, { stream: true });
-      const lines = buffer.split("\n");
-      buffer = lines.pop() || "";
-
-      for (const line of lines) {
-        if (!line.startsWith("data: ")) continue;
-        const data = line.slice(6).trim();
-        if (data === "[DONE]") continue;
-
-        let chunk;
-        try { chunk = JSON.parse(data); } catch { continue; }
-
-        const delta = chunk.choices?.[0]?.delta;
-        if (!delta) continue;
-
-        if (delta.content) {
-          assistantText += delta.content;
-          onToken(delta.content);
-        }
-
-        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.function?.arguments) entry.arguments += call.function.arguments;
-            toolCallBuffer.set(call.index, entry);
-          }
-        }
-      }
-    }
-
-    // Process tool calls
-    const toolCalls = [...toolCallBuffer.values()].map(t => ({
-      id: t.id || `call_${crypto.randomUUID()}`,
-      type: "function",
-      function: { name: t.name, arguments: t.arguments },
-    }));
+    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 toolResults = await processToolCalls(
+        ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset,
+      );
 
-      // Follow-up response after tool calls
       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 followUp = await fetch(`${LIGHTNING_BASE}/gen/chat/completions`, {
-        method: "POST",
-        headers: { ...headers, "Accept": "text/event-stream" },
-        body: JSON.stringify({
-          model: model || "lightning",
-          messages: followUpMessages,
-          stream: true,
-        }),
-        signal: abortSignal,
-      });
-
-      if (followUp.ok) {
-        const fuReader = followUp.body.getReader();
-        const fuDecoder = new TextDecoder(); // fresh decoder — reusing the first one corrupts the stream
-        let fuBuffer = "";
-        while (true) {
-          const { done, value } = await fuReader.read();
-          if (done) break;
-          fuBuffer += fuDecoder.decode(value, { stream: true });
-          const fuLines = fuBuffer.split("\n");
-          fuBuffer = fuLines.pop() || "";
-          for (const line of fuLines) {
-            if (!line.startsWith("data: ")) continue;
-            const data = line.slice(6).trim();
-            if (data === "[DONE]") continue;
-            let chunk;
-            try { chunk = JSON.parse(data); } catch { continue; }
-            const delta = chunk.choices?.[0]?.delta;
-            if (delta?.content) {
-              assistantText += delta.content;
-              onToken(delta.content);
-            }
-          }
-        }
-      }
+      const followUpStream = await client.chat.completions.create({
+        model:  model || "lightning",
+        messages: followUpMessages,
+        stream: true,
+      }, { signal: abortSignal });
+
+      const followUp = await consumeStream(followUpStream, onToken);
+      assistantText += followUp.assistantText;
     }
 
     onDone(assistantText, toolCalls);
   } catch (err) {
-    if (err.name === "AbortError") {
+    if (err.name === "AbortError" || err.constructor?.name === "APIUserAbortError") {
       onDone(null, null, true); // aborted
     } else {
       onError(String(err));