feat: trace inline <think> reasoning and time thinking tokens

src/index.ts

@@ -445,6 +445,31 @@ export function extractThinking(content: unknown): ChatMlThinkingPart[] {
   return parts;
 }
 
+const THINK_TAG_PATTERN = /<think>([\s\S]*?)(?:<\/think>|$)/g;
+
+export function splitInlineThinking(text: string): {
+  text: string;
+  thinking: ChatMlThinkingPart[];
+} {
+  if (!text.trimStart().startsWith("<think>")) 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;
 

test/helpers.ts

@@ -121,6 +121,8 @@ export function startMockProvider(): Promise<MockProvider> {
       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<MockProvider> {
         return;
       }
 
-      if (canDelegate && stage === 0) {
+      if (inlineThinkMode) {
+        streamChunks(res, model, {
+          text: `<think>${FINAL_ANSWER_THINKING}</think>\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" } },

test/integration.test.ts

@@ -726,4 +726,39 @@ describe("integration: pi -> extension -> Langfuse export", () => {
       capture.close();
     }
   });
+
+  it("splits reasoning a server streams inline as <think> 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();
+    }
+  });
 });

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 <think> block out of the answer text", () => {
+    assert.deepEqual(splitInlineThinking("<think>hm, which files?</think>\nThe answer."), {
+      text: "The answer.",
+      thinking: [{ type: "thinking", content: "hm, which files?" }],
+    });
+  });
+
+  it("keeps a <think> tag the answer only talks about", () => {
+    const answer = "Strip the <think> 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: "<think>inline</think>answer" },
+    ]);
+    assert.deepEqual(result.thinking, [{ type: "thinking", content: "structured" }]);
+    assert.equal(result.text, "<think>inline</think>answer");
+  });
+});