From ede2956d76173d6f93691093270297a3ee8151aa Mon Sep 17 00:00:00 2001 From: Milana Gurbanova Date: Thu, 17 Sep 2026 16:58:13 +0200 Subject: [PATCH] feat: trace inline reasoning and time thinking tokens --- src/index.ts | 35 ++++++++++++++++++++++++++++++----- test/helpers.ts | 23 ++++++++++++++++++++++- test/integration.test.ts | 35 +++++++++++++++++++++++++++++++++++ test/unit.test.ts | 27 +++++++++++++++++++++++++++ 4 files changed, 114 insertions(+), 6 deletions(-) diff --git a/src/index.ts b/src/index.ts index 3814614..6bbaaf9 100644 --- a/src/index.ts +++ b/src/index.ts @@ -445,6 +445,31 @@ export function extractThinking(content: unknown): ChatMlThinkingPart[] { return parts; } +const THINK_TAG_PATTERN = /([\s\S]*?)(?:<\/think>|$)/g; + +export function splitInlineThinking(text: string): { + text: string; + thinking: ChatMlThinkingPart[]; +} { + if (!text.trimStart().startsWith("")) return { text, thinking: [] }; + const blocks: Array<{ type: "thinking"; thinking: string }> = []; + const answer = text.replace(THINK_TAG_PATTERN, (_match, inner: string) => { + if (inner.trim()) blocks.push({ type: "thinking", thinking: inner.trim() }); + return ""; + }); + return { text: answer.trim(), thinking: extractThinking(blocks) }; +} + +export function extractAnswerAndThinking(content: unknown): { + text: string; + thinking: ChatMlThinkingPart[]; +} { + const text = extractText(content); + const thinking = extractThinking(content); + if (thinking.length) return { text, thinking }; + return splitInlineThinking(text); +} + export function toChatMlMessage(message: unknown): ChatMlMessage | undefined { if (!message || typeof message !== "object") return undefined; const msg = message as { @@ -458,8 +483,8 @@ export function toChatMlMessage(message: unknown): ChatMlMessage | undefined { return { role: "user", content: markDataUris(renderHistoryContent(msg.content)) }; } if (msg.role === "assistant") { - const content = markDataUris(extractText(msg.content)); - const thinking = extractThinking(msg.content); + const { text, thinking } = extractAnswerAndThinking(msg.content); + const content = markDataUris(text); const toolCalls = historyToolCalls(msg.content); if (!content && !thinking.length && !toolCalls.length) return undefined; return { @@ -909,7 +934,8 @@ export default function (pi: ExtensionAPI) { pi.on("message_update", (event) => { const gen = state?.openGeneration; if (!gen || gen.finished || gen.sawFirstToken) return; - if (extractText((event.message as { content?: unknown })?.content).length > 0) { + const content = (event.message as { content?: unknown })?.content; + if (extractText(content).length > 0 || extractThinking(content).length > 0) { gen.sawFirstToken = true; gen.obs.update({ completionStartTime: new Date() }); } @@ -933,9 +959,8 @@ export default function (pi: ExtensionAPI) { const gen = state.openGeneration; if (!gen || gen.finished) return; - const text = extractText(message.content); + const { text, thinking } = extractAnswerAndThinking(message.content); const tools = extractToolCalls(message.content); - const thinking = extractThinking(message.content); const isError = message.stopReason === "error" || message.stopReason === "aborted"; if (message.stopReason === "error") state.sawError = true; diff --git a/test/helpers.ts b/test/helpers.ts index 4b8a298..42102fa 100644 --- a/test/helpers.ts +++ b/test/helpers.ts @@ -121,6 +121,8 @@ export function startMockProvider(): Promise { const stage = messages.slice(lastUserIdx + 1).filter((m) => m.role === "tool").length; const lastUser = messages[lastUserIdx]; const failMode = JSON.stringify(lastUser?.content ?? "").includes("[fail]"); + const inlineThinkMode = JSON.stringify(lastUser?.content ?? "").includes("[inline-think]"); + const thinkThenToolMode = JSON.stringify(lastUser?.content ?? "").includes("[think-tool]"); // True only when a test loads the subagent fixture. The default script // does not change for the other tests. const canDelegate = (payload.tools ?? []).some((t) => t.function?.name === "subagent"); @@ -147,7 +149,26 @@ export function startMockProvider(): Promise { return; } - if (canDelegate && stage === 0) { + if (inlineThinkMode) { + streamChunks(res, model, { + text: `${FINAL_ANSWER_THINKING}\nThis is the test workspace. Done.`, + finish: "stop", + usage: usage(800, 60, 0, 30), + }); + } else if (thinkThenToolMode && stage === 0) { + streamChunks(res, model, { + thinking: FINAL_ANSWER_THINKING, + tool: { name: "bash", args: { command: "ls" } }, + finish: "tool_calls", + usage: usage(800, 44, 0, 30), + }); + } else if (thinkThenToolMode) { + streamChunks(res, model, { + text: "This is the test workspace. Done.", + finish: "stop", + usage: usage(900, 30, 0), + }); + } else if (canDelegate && stage === 0) { streamChunks(res, model, { text: "Delegating to a subagent. ", tool: { name: "subagent", args: { task: "inspect the repository" } }, diff --git a/test/integration.test.ts b/test/integration.test.ts index 2cc5365..faf59c9 100644 --- a/test/integration.test.ts +++ b/test/integration.test.ts @@ -726,4 +726,39 @@ describe("integration: pi -> extension -> Langfuse export", () => { capture.close(); } }); + + it("splits reasoning a server streams inline as tags", async () => { + const capture = await startCaptureServer(); + try { + const sandbox = createSandbox(mock.port); + const env = buildLangfuseEnv(capture); + assert.equal((await runPi(sandbox, "[inline-think] Summarize this", { env })).status, 0); + await waitForRequests(capture, 1); + + const answer = outputOf(byStart(findSpansByName(capture.spans(), "LLM Call")).at(-1)!); + assert.deepEqual(answer.thinking, [{ type: "thinking", content: FINAL_ANSWER_THINKING }]); + assert.equal(answer.content, "This is the test workspace. Done."); + } finally { + capture.close(); + } + }); + + it("times the first token from the first thinking token", async () => { + const capture = await startCaptureServer(); + try { + const sandbox = createSandbox(mock.port); + const env = buildLangfuseEnv(capture); + assert.equal((await runPi(sandbox, "[think-tool] Summarize this", { env })).status, 0); + await waitForRequests(capture, 1); + + const reasoned = byStart(findSpansByName(capture.spans(), "LLM Call"))[0]!; + assert.equal(outputOf(reasoned).content, undefined, "this step streamed no text"); + assert.ok( + reasoned.attrs["langfuse.observation.completion_start_time"], + "a thinking-only step must still report a time to first token", + ); + } finally { + capture.close(); + } + }); }); diff --git a/test/unit.test.ts b/test/unit.test.ts index c3ec0fd..eceae3e 100644 --- a/test/unit.test.ts +++ b/test/unit.test.ts @@ -14,6 +14,8 @@ import { buildHistoryInput, markDataUris, extractThinking, + splitInlineThinking, + extractAnswerAndThinking, toChatMlMessage, type ChatMlMessage, type PiUsage, @@ -685,3 +687,28 @@ describe("extractThinking", () => { assert.equal(parts[0]!.content.length, 120_000); }); }); + +describe("splitInlineThinking", () => { + it("splits a leading block out of the answer text", () => { + assert.deepEqual(splitInlineThinking("hm, which files?\nThe answer."), { + text: "The answer.", + thinking: [{ type: "thinking", content: "hm, which files?" }], + }); + }); + + it("keeps a tag the answer only talks about", () => { + const answer = "Strip the tags with a regex before parsing."; + assert.deepEqual(splitInlineThinking(answer), { text: answer, thinking: [] }); + }); +}); + +describe("extractAnswerAndThinking", () => { + it("prefers structured thinking parts over inline tags", () => { + const result = extractAnswerAndThinking([ + { type: "thinking", thinking: "structured" }, + { type: "text", text: "inlineanswer" }, + ]); + assert.deepEqual(result.thinking, [{ type: "thinking", content: "structured" }]); + assert.equal(result.text, "inlineanswer"); + }); +});