Update chatStream.js

server/chatStream.js

@@ -409,12 +409,54 @@ function contentToText(content, { preview = false } = {}) {
   return pieces.join(preview ? " " : "\n");
 }
 
+function getAllowedToolNames(toolDefs = []) {
+  return toolDefs
+    .map((tool) => tool?.function?.name)
+    .filter(Boolean);
+}
+
+function sanitizeToolName(rawName, allowedToolNames = []) {
+  if (!rawName) return null;
+
+  let name = String(rawName).trim();
+  if (!name) return null;
+
+  name = name.replace(/<\|[\s\S]*$/, "").trim();
+  name = name.replace(/^["'`]+|["'`]+$/g, "");
+
+  const identifierMatch = name.match(/^[A-Za-z0-9_-]+/);
+  if (identifierMatch) {
+    name = identifierMatch[0];
+  }
+
+  if (!allowedToolNames.length) {
+    return name || null;
+  }
+
+  if (allowedToolNames.includes(name)) {
+    return name;
+  }
+
+  const normalizedRaw = String(rawName).toLowerCase();
+  const prefixedMatch = allowedToolNames.find((allowedName) =>
+    normalizedRaw.startsWith(allowedName.toLowerCase())
+  );
+  if (prefixedMatch) {
+    return prefixedMatch;
+  }
+
+  const embeddedMatch = allowedToolNames.find((allowedName) =>
+    normalizedRaw.includes(allowedName.toLowerCase())
+  );
+  return embeddedMatch || null;
+}
+
 function normalizeStoredToolCalls(toolCalls = []) {
   return toolCalls.map((call) => ({
     id: call.id || `call_${crypto.randomUUID()}`,
     type: "function",
     function: {
-      name: call.name || call.function?.name || "unknown_tool",
+      name: sanitizeToolName(call.name || call.function?.name || "unknown_tool"),
       arguments: (() => {
         const rawArgs = call.args ?? call.function?.arguments ?? {};
         return typeof rawArgs === "string" ? rawArgs : JSON.stringify(rawArgs);
@@ -991,9 +1033,20 @@ async function websocketChatStreamWithRetry(body, headers, onToken, abortSignal)
   }
 }
 
+function serializeToolCalls(toolCallBuffer) {
+  return [...toolCallBuffer.values()]
+    .filter((toolCall) => toolCall?.name)
+    .map((toolCall) => ({
+      id: toolCall.id || `call_${crypto.randomUUID()}`,
+      type: "function",
+      function: { name: toolCall.name, arguments: toolCall.arguments },
+    }));
+}
+
 export async function websocketChatStream(body, headers, onToken, abortSignal) {
   const ws = await getSafeWebSocket();
   const currentRequestId = ++requestIdCounter;
+  const allowedToolNames = getAllowedToolNames(body?.tools);
   const safeParse = (str) => {
     try { return JSON.parse(str.startsWith("data: ") ? str.slice(6) : str); } catch { return null; }
   };
@@ -1007,11 +1060,7 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       if (!finished) {
         finished = true;
         cleanup();
-        const toolCalls = [...toolCallBuffer.values()].map((t) => ({
-          id: t.id || `call_${crypto.randomUUID()}`,
-          type: "function",
-          function: { name: t.name, arguments: t.arguments },
-        }));
+        const toolCalls = serializeToolCalls(toolCallBuffer);
         resolve({ assistantText, toolCalls });
       }
     }, 120000);
@@ -1103,7 +1152,9 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
         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?.name) {
+            entry.name = sanitizeToolName(call.function.name, allowedToolNames);
+          }
           if (call.function?.arguments) entry.arguments += call.function.arguments;
           toolCallBuffer.set(call.index, entry);
         }
@@ -1112,11 +1163,7 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       if (payload.choices?.[0]?.finish_reason && !finished) {
         finished = true;
         cleanup();
-        const toolCalls = [...toolCallBuffer.values()].map((t) => ({
-          id: t.id || `call_${crypto.randomUUID()}`,
-          type: "function",
-          function: { name: t.name, arguments: t.arguments },
-        }));
+        const toolCalls = serializeToolCalls(toolCallBuffer);
         resolve({ assistantText, toolCalls });
       }
     };
@@ -1227,6 +1274,37 @@ export async function streamChat({
       }
     }
 
+    if (!finished) {
+      const finalMessages = [
+        ...buildModelMessages(baseMessages, workingMessages, sessionId),
+        {
+          role: "system",
+          content: "Tool-use budget is exhausted for this response. Do not call tools. Answer directly using the information already gathered. If something is still missing, briefly say what is missing without calling tools.",
+        },
+      ];
+
+      const { assistantText: finalStepText } = await websocketChatStreamWithRetry(
+        {
+          model: model || "lightning",
+          messages: finalMessages,
+          stream: true,
+        },
+        headers,
+        onToken,
+        abortSignal
+      );
+
+      if (finalStepText) {
+        assistantText += finalStepText;
+        workingMessages.push({ role: "assistant", content: finalStepText });
+      }
+      finished = true;
+    }
+
+    if (!assistantText.trim()) {
+      assistantText = "I wasn’t able to finish that response cleanly. Please try again.";
+    }
+
     const sessionName = extractSessionName(assistantText);
 
     if (typeof onDone === "function") {
@@ -1252,10 +1330,20 @@ function normalizeMessage(msg) {
   if (!VALID_ROLES.has(msg.role)) return null;
 
   if (msg.role === "assistant" && (msg.tool_calls || msg.toolCalls)) {
+    const normalizedToolCalls = (msg.tool_calls || normalizeStoredToolCalls(msg.toolCalls))
+      .map((call) => ({
+        ...call,
+        function: {
+          ...call.function,
+          name: sanitizeToolName(call.function?.name || call.name || "unknown_tool"),
+        },
+      }))
+      .filter((call) => call.function?.name);
+
     return {
       role: "assistant",
       content: msg.content ?? "",
-      tool_calls: msg.tool_calls || normalizeStoredToolCalls(msg.toolCalls),
+      ...(normalizedToolCalls.length ? { tool_calls: normalizedToolCalls } : {}),
     };
   }
   if (Array.isArray(msg.content)) {
@@ -1452,6 +1540,7 @@ async function processToolCalls({
 }) {
   const nextMessages = [];
   const authHeaders = {};
+  const allowedToolNames = getAllowedToolNames(buildToolList(tools));
   if (accessToken) {
     authHeaders["Authorization"] = `Bearer ${accessToken}`;
   }
@@ -1462,18 +1551,26 @@ async function processToolCalls({
   for (const call of toolCalls) {
     let args;
     try { args = JSON.parse(call.function.arguments || "{}"); } catch { args = {}; }
+    const toolName = sanitizeToolName(call.function?.name, allowedToolNames);
+    if (toolName) {
+      call.function.name = toolName;
+    }
 
-    onToolCall({ id: call.id, name: call.function.name, state: "pending", args });
+    onToolCall({ id: call.id, name: toolName || call.function?.name, state: "pending", args });
 
     let result = "Tool completed.";
 
     try {
-      if (call.function.name === "list_prompt_resources") {
+      if (!toolName || !allowedToolNames.includes(toolName)) {
+        result = `Invalid tool name "${call.function?.name || "unknown"}".`;
+      }
+
+      else if (toolName === "list_prompt_resources") {
         const state = getPromptState(sessionId);
         result = JSON.stringify(buildPromptResourceManifest(state), null, 2);
       }
 
-      else if (call.function.name === "read_prompt_chunk") {
+      else if (toolName === "read_prompt_chunk") {
         const state = getPromptState(sessionId);
         const chunkIndex = Number.isInteger(Number(args.chunk_index))
           ? Number(args.chunk_index)
@@ -1500,7 +1597,7 @@ async function processToolCalls({
         }
       }
 
-      else if (call.function.name === "load_prompt_images") {
+      else if (toolName === "load_prompt_images") {
         const state = getPromptState(sessionId);
         const resource = findPromptResource(state, args.resource_id);
 
@@ -1533,22 +1630,22 @@ async function processToolCalls({
                 ...selectedImages,
               ],
             });
-            onToolCall({ id: call.id, name: call.function.name, state: "resolved", result });
+            onToolCall({ id: call.id, name: toolName, state: "resolved", result });
             continue;
           }
         }
       }
 
-      else if (call.function.name === "write_notes") {
+      else if (toolName === "write_notes") {
         appendAssistantNote(sessionId, args.note);
         result = "Note stored for the rest of this response.";
       }
 
-      else if (call.function.name === "ollama_search") {
+      else if (toolName === "ollama_search") {
         result = await gradioSearch(args.query);
       }
 
-      else if (call.function.name === "read_web_page") {
+      else if (toolName === "read_web_page") {
         const { convert } = await import("html-to-text");
         const res = await fetch(args.url, { signal: abortSignal });
         if (!res.ok) {
@@ -1563,7 +1660,7 @@ async function processToolCalls({
         }
       }
 
-      else if (call.function.name === "generate_image") {
+      else if (toolName === "generate_image") {
         const body = { prompt: args.prompt };
         if (args.mode) body.mode = args.mode;
         if (args.image_urls?.length) body.image_urls = args.image_urls;
@@ -1590,7 +1687,7 @@ async function processToolCalls({
         }
       }
 
-      else if (call.function.name === "generate_video") {
+      else if (toolName === "generate_video") {
         const body = { prompt: args.prompt };
         if (args.ratio) body.ratio = args.ratio;
         if (args.mode) body.mode = args.mode;
@@ -1618,7 +1715,7 @@ async function processToolCalls({
         }
       }
 
-      else if (call.function.name === "generate_audio") {
+      else if (toolName === "generate_audio") {
         const res = await fetch(`${LIGHTNING_BASE}/gen/sfx`, {
           method: "POST",
           headers: { "Content-Type": "application/json", ...authHeaders },
@@ -1641,7 +1738,7 @@ async function processToolCalls({
       result = `Tool error: ${String(err)}`;
     }
 
-    onToolCall({ id: call.id, name: call.function.name, state: "resolved", result });
+    onToolCall({ id: call.id, name: toolName || call.function?.name, state: "resolved", result });
 
     nextMessages.push({
       role: "tool",