From a0acdfdcb9cc995be31079f079973b762042d94d Mon Sep 17 00:00:00 2001 From: diegosouzapw Date: Sat, 28 Mar 2026 21:04:47 -0300 Subject: [PATCH] fix: context pinning bypass during tool-call responses (#721) Non-streaming: Fixed json.messages check to use json.choices[0].message (OpenAI format). Streaming: inject pin tag before finish_reason chunk for tool-call-only streams. injectModelTag now appends synthetic assistant message when content is null/array (tool_calls). --- open-sse/services/combo.ts | 45 +++++- open-sse/services/comboAgentMiddleware.ts | 12 +- .../unit/context-pinning-tool-calls.test.mjs | 132 ++++++++++++++++++ 3 files changed, 181 insertions(+), 8 deletions(-) create mode 100644 tests/unit/context-pinning-tool-calls.test.mjs diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 044679c61a..c9bf705d0f 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -464,14 +464,23 @@ export async function handleComboChat({ const res = await handleSingleModel(b, modelStr); if (!res.ok) return res; - // Non-streaming: inject tag into JSON response (existing logic) + // Non-streaming: inject tag into JSON response + // Fix #721: Use OpenAI choices format (json.choices[0].message) not json.messages if (!b.stream) { try { const json = await res.clone().json(); - const msgs = Array.isArray(json?.messages) ? json.messages : []; - if (msgs.length > 0) { - const tagged = injectModelTag(msgs, modelStr); - return new Response(JSON.stringify({ ...json, messages: tagged }), { + const choice = json?.choices?.[0]; + if (choice?.message) { + // Wrap single message in array for injectModelTag, then unwrap + const tagged = injectModelTag([choice.message], modelStr); + // If the message had tool_calls but no string content, injectModelTag + // appends a synthetic assistant message — use the last one + const taggedMsg = tagged[tagged.length - 1]; + const updatedJson = { + ...json, + choices: [{ ...choice, message: taggedMsg }, ...(json.choices?.slice(1) || [])], + }; + return new Response(JSON.stringify(updatedJson), { status: res.status, headers: res.headers, }); @@ -502,8 +511,9 @@ export async function handleComboChat({ const text = decoder.decode(chunk, { stream: true }); - // Look for the first SSE data line with non-empty content - // Pattern: "content":"" — we inject tag at the start + // Fix #721: Look for either non-empty content OR tool_calls in the + // SSE data. Tool-call-only responses have content:null, so we inject + // the tag when we see a finish_reason approaching, or on first content. const contentMatch = text.match(/"content":"([^"]+)/); if (contentMatch) { // Inject tag at the beginning of the first content value @@ -516,6 +526,27 @@ export async function handleComboChat({ return; } + // Fix #721: For tool-call-only streams, inject the tag when we see + // the finish_reason chunk (before it reaches the client SDK which + // would close the connection). This ensures the tag roundtrips + // through the conversation history even when there's no text content. + if (text.includes('"finish_reason"') && !text.includes('"finish_reason":null')) { + // Inject a content chunk with the tag just before this finish chunk + const tagChunk = `data: ${JSON.stringify({ + choices: [ + { + delta: { content: tagContent }, + index: 0, + finish_reason: null, + }, + ], + })}\n\n`; + tagInjected = true; + controller.enqueue(encoder.encode(tagChunk)); + controller.enqueue(chunk); + return; + } + // No content yet — passthrough controller.enqueue(chunk); }, diff --git a/open-sse/services/comboAgentMiddleware.ts b/open-sse/services/comboAgentMiddleware.ts index aa06211c8f..4903c9bd99 100644 --- a/open-sse/services/comboAgentMiddleware.ts +++ b/open-sse/services/comboAgentMiddleware.ts @@ -67,7 +67,17 @@ export function injectModelTag(messages: Message[], providerModel: string): Mess } const msg = cleaned[lastAssistantIdx]; - if (typeof msg.content !== "string") return cleaned; + // Fix #721: Handle messages where content is not a string (tool_calls responses). + // In this case, append a synthetic assistant message with the tag so the pin + // roundtrips through the conversation history. + if (typeof msg.content !== "string") { + // If the message has tool_calls but no string content, append a new assistant + // message with the tag rather than silently failing. + return [ + ...cleaned, + { role: "assistant", content: `\n${providerModel}` }, + ]; + } const tagged = [...cleaned]; tagged[lastAssistantIdx] = { diff --git a/tests/unit/context-pinning-tool-calls.test.mjs b/tests/unit/context-pinning-tool-calls.test.mjs new file mode 100644 index 0000000000..045e6d2f24 --- /dev/null +++ b/tests/unit/context-pinning-tool-calls.test.mjs @@ -0,0 +1,132 @@ +import { describe, test } from "node:test"; +import assert from "node:assert/strict"; +import { + injectModelTag, + extractPinnedModel, +} from "../../open-sse/services/comboAgentMiddleware.ts"; + +describe("Context pinning — tool call responses (#721)", () => { + test("injectModelTag appends synthetic tag when last assistant has null content (tool_calls)", () => { + const messages = [ + { role: "user", content: "List the files" }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_abc123", + type: "function", + function: { name: "read", arguments: '{"filePath":"/mnt/e/deer-flow"}' }, + }, + ], + }, + ]; + + const result = injectModelTag(messages, "ollamacloud/glm-5"); + + // Should append a synthetic assistant message with the pin tag + assert.equal(result.length, 3, "Should have 3 messages (original 2 + synthetic)"); + assert.equal(result[2].role, "assistant"); + assert.ok( + result[2].content.includes("ollamacloud/glm-5"), + "Synthetic message should contain the pin tag" + ); + }); + + test("injectModelTag appends synthetic tag when last assistant has array content", () => { + const messages = [ + { role: "user", content: "Explain the code" }, + { + role: "assistant", + content: [ + { type: "text", text: "Here is the analysis" }, + { type: "text", text: "And here is part 2" }, + ], + }, + ]; + + const result = injectModelTag(messages, "nvidia/llama-3.4-70b"); + + // Array content → should append synthetic message + assert.equal(result.length, 3); + assert.equal(result[2].role, "assistant"); + assert.ok(result[2].content.includes("nvidia/llama-3.4-70b")); + }); + + test("extractPinnedModel finds tag in synthetic message after tool_calls", () => { + const messages = [ + { role: "user", content: "List the files" }, + { + role: "assistant", + content: null, + tool_calls: [ + { id: "call_abc", type: "function", function: { name: "read", arguments: "{}" } }, + ], + }, + { role: "assistant", content: "\nollamacloud/glm-5" }, + ]; + + const pinned = extractPinnedModel(messages); + assert.equal(pinned, "ollamacloud/glm-5"); + }); + + test("injectModelTag still works for normal string content", () => { + const messages = [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi there!" }, + ]; + + const result = injectModelTag(messages, "openai/gpt-4o"); + + assert.equal(result.length, 2, "Should not add a new message"); + assert.ok(result[1].content.includes("openai/gpt-4o")); + assert.ok(result[1].content.startsWith("Hi there!")); + }); + + test("roundtrip: inject → extract works for tool-call messages", () => { + const messages = [ + { role: "user", content: "List the files" }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_abc123", + type: "function", + function: { name: "read", arguments: '{"filePath":"/home"}' }, + }, + ], + }, + ]; + + const tagged = injectModelTag(messages, "qwen/coder-model"); + const pinned = extractPinnedModel(tagged); + + assert.equal(pinned, "qwen/coder-model", "Should roundtrip the pinned model"); + }); + + test("re-injection clears old pin and sets new one", () => { + const messages = [ + { role: "user", content: "Follow up" }, + { role: "assistant", content: "Previous answer\nold/model" }, + { role: "user", content: "Continue" }, + { + role: "assistant", + content: null, + tool_calls: [ + { id: "call_xyz", type: "function", function: { name: "exec", arguments: "{}" } }, + ], + }, + ]; + + const tagged = injectModelTag(messages, "new/model"); + const pinned = extractPinnedModel(tagged); + + assert.equal(pinned, "new/model", "Should return new pinned model, not old one"); + // Verify old tag was cleaned + const oldTagPresent = tagged.some( + (m) => typeof m.content === "string" && m.content.includes("old/model") + ); + assert.equal(oldTagPresent, false, "Old pin tag should be cleaned"); + }); +});