fix: guard session tool results

src/agents/session-tool-result-guard.test.ts

@@ -0,0 +1,152 @@
+import { describe, expect, it } from "vitest";
+
+import { SessionManager } from "@mariozechner/pi-coding-agent";
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+
+import { installSessionToolResultGuard } from "./session-tool-result-guard.js";
+
+const toolCallMessage = {
+  role: "assistant",
+  content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }],
+} satisfies AgentMessage;
+
+describe("installSessionToolResultGuard", () => {
+  it("inserts synthetic toolResult before non-tool message when pending", () => {
+    const sm = SessionManager.inMemory();
+    installSessionToolResultGuard(sm);
+
+    sm.appendMessage(toolCallMessage);
+    sm.appendMessage({
+      role: "assistant",
+      content: [{ type: "text", text: "error" }],
+      stopReason: "error",
+    } as AgentMessage);
+
+    const entries = sm
+      .getEntries()
+      .filter((e) => e.type === "message")
+      .map((e) => (e as { message: AgentMessage }).message);
+
+    expect(entries.map((m) => m.role)).toEqual([
+      "assistant",
+      "toolResult",
+      "assistant",
+    ]);
+    const synthetic = entries[1] as {
+      toolCallId?: string;
+      isError?: boolean;
+      content?: Array<{ type?: string; text?: string }>;
+    };
+    expect(synthetic.toolCallId).toBe("call_1");
+    expect(synthetic.isError).toBe(true);
+    expect(synthetic.content?.[0]?.text).toContain("missing tool result");
+  });
+
+  it("flushes pending tool calls when asked explicitly", () => {
+    const sm = SessionManager.inMemory();
+    const guard = installSessionToolResultGuard(sm);
+
+    sm.appendMessage(toolCallMessage);
+    guard.flushPendingToolResults();
+
+    const messages = sm
+      .getEntries()
+      .filter((e) => e.type === "message")
+      .map((e) => (e as { message: AgentMessage }).message);
+
+    expect(messages.map((m) => m.role)).toEqual(["assistant", "toolResult"]);
+  });
+
+  it("does not add synthetic toolResult when a matching one exists", () => {
+    const sm = SessionManager.inMemory();
+    installSessionToolResultGuard(sm);
+
+    sm.appendMessage(toolCallMessage);
+    sm.appendMessage({
+      role: "toolResult",
+      toolCallId: "call_1",
+      content: [{ type: "text", text: "ok" }],
+      isError: false,
+    } as AgentMessage);
+
+    const messages = sm
+      .getEntries()
+      .filter((e) => e.type === "message")
+      .map((e) => (e as { message: AgentMessage }).message);
+
+    expect(messages.map((m) => m.role)).toEqual(["assistant", "toolResult"]);
+  });
+
+  it("preserves ordering with multiple tool calls and partial results", () => {
+    const sm = SessionManager.inMemory();
+    const guard = installSessionToolResultGuard(sm);
+
+    sm.appendMessage({
+      role: "assistant",
+      content: [
+        { type: "toolCall", id: "call_a", name: "one", arguments: {} },
+        { type: "toolUse", id: "call_b", name: "two", arguments: {} },
+      ],
+    } as AgentMessage);
+    sm.appendMessage({
+      role: "toolResult",
+      toolUseId: "call_a",
+      content: [{ type: "text", text: "a" }],
+      isError: false,
+    } as AgentMessage);
+    sm.appendMessage({
+      role: "assistant",
+      content: [{ type: "text", text: "after tools" }],
+    } as AgentMessage);
+
+    const messages = sm
+      .getEntries()
+      .filter((e) => e.type === "message")
+      .map((e) => (e as { message: AgentMessage }).message);
+
+    expect(messages.map((m) => m.role)).toEqual([
+      "assistant", // tool calls
+      "toolResult", // call_a real
+      "toolResult", // synthetic for call_b
+      "assistant", // text
+    ]);
+    expect(
+      (messages[2] as { toolCallId?: string }).toolCallId,
+    ).toBe("call_b");
+    expect(guard.getPendingIds()).toEqual([]);
+  });
+
+  it("flushes pending on guard when no toolResult arrived", () => {
+    const sm = SessionManager.inMemory();
+    const guard = installSessionToolResultGuard(sm);
+
+    sm.appendMessage(toolCallMessage);
+    sm.appendMessage({
+      role: "assistant",
+      content: [{ type: "text", text: "hard error" }],
+      stopReason: "error",
+    } as AgentMessage);
+    expect(guard.getPendingIds()).toEqual([]);
+  });
+
+  it("handles toolUseId on toolResult", () => {
+    const sm = SessionManager.inMemory();
+    installSessionToolResultGuard(sm);
+
+    sm.appendMessage({
+      role: "assistant",
+      content: [{ type: "toolUse", id: "use_1", name: "f", arguments: {} }],
+    } as AgentMessage);
+    sm.appendMessage({
+      role: "toolResult",
+      toolUseId: "use_1",
+      content: [{ type: "text", text: "ok" }],
+    } as AgentMessage);
+
+    const messages = sm
+      .getEntries()
+      .filter((e) => e.type === "message")
+      .map((e) => (e as { message: AgentMessage }).message);
+    expect(messages.map((m) => m.role)).toEqual(["assistant", "toolResult"]);
+  });
+});