feat: carry assistant reasoning on the generation output too

src/index.ts

@@ -189,11 +189,11 @@ export function extractText(content: unknown): string {
     .join("");
 }
 
-export function extractToolCalls(content: unknown): Array<{ id: string; name: string }> {
+export function extractToolCalls(content: unknown): ChatMlToolCall[] {
   if (!Array.isArray(content)) return [];
   return content
     .filter((p): p is { type: string; id: string; name: string } => !!p && typeof p === "object" && (p as { type?: string }).type === "toolCall")
-    .map((p) => ({ id: p.id, name: p.name }));
+    .map((p) => ({ id: p.id, type: "function" as const, function: { name: p.name } }));
 }
 
 export interface PiImagePart {
@@ -828,13 +828,15 @@ export default function (pi: ExtensionAPI) {
 
     const text = extractText(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;
 
     gen.obs.update({
       output: {
         role: "assistant",
         ...(text ? { content: text } : {}),
+        ...(thinking.length ? { thinking } : {}),
         ...(tools.length ? { tool_calls: tools } : {}),
       },
       model: message.responseModel || message.model,

test/integration.test.ts

@@ -63,6 +63,11 @@ function contentsOf(span: CapturedSpan): string[] {
   return historyOf(span).map((m) => m.content ?? "");
 }
 
+function outputOf(span: CapturedSpan): TracedMessage {
+  const raw = span.attrs["langfuse.observation.output"];
+  return (typeof raw === "string" ? JSON.parse(raw) : raw) as TracedMessage;
+}
+
 function wireText(content: unknown): string {
   if (typeof content === "string") return content;
   if (!Array.isArray(content)) return "";
@@ -679,11 +684,23 @@ describe("integration: pi -> extension -> Langfuse export", () => {
       await waitForRequests(capture, 2);
 
       const roots = byStart(findSpansByName(capture.spans(), "Conversational Turn"));
+      const firstTurn = byStart(
+        findSpansByName(capture.spans(), "LLM Call").filter((s) => s.traceId === roots[0]!.traceId),
+      );
       const secondTurn = byStart(
         findSpansByName(capture.spans(), "LLM Call").filter((s) => s.traceId === roots[1]!.traceId),
       );
       const history = historyOf(secondTurn[0]!);
 
+      const answer = outputOf(firstTurn.at(-1)!);
+      assert.deepEqual(answer.thinking, [{ type: "thinking", content: FINAL_ANSWER_THINKING }]);
+      assert.equal(answer.content, "This is the test workspace. Done.");
+      const toolStep = outputOf(firstTurn[0]!);
+      assert.equal(toolStep.thinking, undefined, "a step that did not reason carries no thinking block");
+      assert.deepEqual(toolStep.tool_calls, [
+        { id: toolStep.tool_calls![0]!.id, type: "function", function: { name: "bash" } },
+      ]);
+
       const reasoning = history.filter((m) => m.thinking);
       assert.equal(reasoning.length, 1, "only the step that reasoned carries a thinking block");
       assert.deepEqual(reasoning[0]!.thinking, [{ type: "thinking", content: FINAL_ANSWER_THINKING }]);

test/unit.test.ts

@@ -100,8 +100,11 @@ describe("extractText / extractToolCalls", () => {
     assert.equal(extractText({}), "");
   });
 
-  it("extracts tool calls as {id, name} only (no arguments)", () => {
-    assert.deepEqual(extractToolCalls(content), [{ id: "call_1", name: "bash" }]);
+  it("extracts tool calls in the nested shape, without the arguments", () => {
+    assert.deepEqual(extractToolCalls(content), [
+      { id: "call_1", type: "function", function: { name: "bash" } },
+    ]);
+    assert.ok(!JSON.stringify(extractToolCalls(content)).includes("ls"));
   });
 });