fix(FN-2316): harden pi session message normalization

2026-04-23 23:38:05 -07:00
parent 6feb697041
commit 49a5605454
3 changed files with 404 additions and 1 deletions
--- a/packages/engine/src/pi-create-fn-agent.test.ts
+++ b/packages/engine/src/pi-create-fn-agent.test.ts
@@ -678,6 +678,124 @@ describe("createFnAgent", () => {
    );
  });

+  it("preserves tool result content when extension hooks only modify metadata", async () => {
+    const originalAfterToolCall = vi.fn().mockResolvedValue({ isError: false });
+    const session = {
+      agent: {
+        afterToolCall: originalAfterToolCall,
+      },
+      prompt: vi.fn(),
+      subscribe: vi.fn(),
+      dispose: vi.fn(),
+      setThinkingLevel: vi.fn(),
+    };
+    createAgentSessionMock.mockResolvedValueOnce({ session });
+
+    const { createFnAgent } = await import("./pi.js");
+    const { session: guardedSession } = await createFnAgent({
+      cwd: "/tmp",
+      systemPrompt: "test",
+      tools: "readonly",
+    });
+
+    const result = await (guardedSession as any).agent.afterToolCall({
+      toolCall: { id: "tool-1", name: "read" },
+      args: { path: "file.txt" },
+      result: { content: [{ type: "text", text: "ok" }], details: { source: "tool" } },
+      isError: false,
+    });
+
+    expect(result).toEqual({
+      content: [{ type: "text", text: "ok" }],
+      details: { source: "tool" },
+      isError: false,
+    });
+  });
+
+  it("repairs malformed persisted session messages missing content", async () => {
+    const rewriteFile = vi.fn();
+    const sessionManager = {
+      fileEntries: [
+        { type: "message", message: { role: "toolResult", toolName: "read" } },
+        { type: "message", message: { role: "assistant", stopReason: "error" } },
+      ],
+      _rewriteFile: rewriteFile,
+    };
+
+    const { createFnAgent } = await import("./pi.js");
+    await createFnAgent({
+      cwd: "/tmp",
+      systemPrompt: "test",
+      tools: "readonly",
+      sessionManager: sessionManager as any,
+    });
+
+    expect(sessionManager.fileEntries[0]?.message).toMatchObject({
+      role: "toolResult",
+      content: [],
+    });
+    expect(sessionManager.fileEntries[1]?.message).toMatchObject({
+      role: "assistant",
+      content: [],
+    });
+    expect(rewriteFile).toHaveBeenCalledTimes(1);
+  });
+
+  it("normalizes malformed live tool results before persistence and replay", async () => {
+    const listeners: Array<(event: unknown) => void> = [];
+    const stateMessages = [
+      { role: "toolResult", toolCallId: "call-1", toolName: "read", timestamp: 123 },
+    ];
+    const originalAppendMessage = vi.fn();
+    const sessionManager = {
+      fileEntries: [],
+      appendMessage: originalAppendMessage,
+    };
+    const session = {
+      agent: {
+        afterToolCall: vi.fn().mockResolvedValue(undefined),
+        state: {
+          messages: stateMessages,
+        },
+      },
+      prompt: vi.fn(),
+      subscribe: vi.fn((listener: (event: unknown) => void) => {
+        listeners.push(listener);
+        return vi.fn();
+      }),
+      dispose: vi.fn(),
+      setThinkingLevel: vi.fn(),
+    };
+    createAgentSessionMock.mockResolvedValueOnce({ session });
+
+    const { createFnAgent } = await import("./pi.js");
+    await createFnAgent({
+      cwd: "/tmp",
+      systemPrompt: "test",
+      tools: "readonly",
+      sessionManager: sessionManager as any,
+    });
+
+    const liveMessage = stateMessages[0]!;
+    listeners[0]?.({ type: "message_end", message: liveMessage });
+    expect(liveMessage).toMatchObject({
+      role: "toolResult",
+      content: [],
+    });
+
+    const persistedMessage = {
+      role: "toolResult",
+      toolCallId: "call-1",
+      toolName: "read",
+      timestamp: 123,
+    };
+    sessionManager.appendMessage(persistedMessage as any);
+    expect(originalAppendMessage).toHaveBeenCalledWith(expect.objectContaining({
+      role: "toolResult",
+      content: [],
+    }));
+  });
+
  describe("skill selection", () => {
    beforeEach(() => {
      // Reset modules to ensure fresh imports for each test
--- a/packages/engine/src/pi.js
+++ b/packages/engine/src/pi.js
@@ -286,6 +286,117 @@ function readJsonObject(path) {
        return {};
    }
 }
+function normalizeSessionHistoryEntries(sessionManager) {
+    const entries = sessionManager.fileEntries;
+    if (!Array.isArray(entries) || entries.length === 0) {
+        return;
+    }
+    let changed = false;
+    for (const entry of entries) {
+        if (entry?.type !== "message" || !entry.message || typeof entry.message !== "object") {
+            continue;
+        }
+        const role = entry.message.role;
+        if (role !== "assistant" && role !== "toolResult") {
+            continue;
+        }
+        if (!("content" in entry.message)) {
+            entry.message.content = [];
+            changed = true;
+        }
+    }
+    if (changed) {
+        sessionManager._rewriteFile?.();
+    }
+}
+function normalizeAssistantOrToolResultMessage(message) {
+    if (!message || typeof message !== "object") {
+        return false;
+    }
+    const role = message.role;
+    if (role !== "assistant" && role !== "toolResult") {
+        return false;
+    }
+    if (!Array.isArray(message.content)) {
+        message.content = [];
+    }
+    return true;
+}
+function syncNormalizedMessageIntoAgentState(session, message) {
+    const messages = session.agent?.state?.messages;
+    if (!Array.isArray(messages) || messages.length === 0) {
+        return;
+    }
+    for (let i = messages.length - 1; i >= 0; i--) {
+        const candidate = messages[i];
+        if (!candidate || typeof candidate !== "object") {
+            continue;
+        }
+        if (candidate === message) {
+            normalizeAssistantOrToolResultMessage(candidate);
+            return;
+        }
+        if (candidate.role !== message.role) {
+            continue;
+        }
+        if (candidate.role === "toolResult") {
+            if (candidate.toolCallId === message.toolCallId && candidate.toolName === message.toolName) {
+                normalizeAssistantOrToolResultMessage(candidate);
+                return;
+            }
+            continue;
+        }
+        if (candidate.timestamp === message.timestamp) {
+            normalizeAssistantOrToolResultMessage(candidate);
+            return;
+        }
+    }
+}
+function installToolResultContentGuard(session) {
+    if (session.__fusionToolResultGuardInstalled || !session.agent?.afterToolCall) {
+        return;
+    }
+    const originalAfterToolCall = session.agent.afterToolCall.bind(session.agent);
+    session.agent.afterToolCall = async (payload) => {
+        const hookResult = await originalAfterToolCall(payload);
+        if (!hookResult || typeof hookResult !== "object") {
+            return hookResult;
+        }
+        const content = hookResult.content ?? payload.result?.content ?? [];
+        return {
+            content: Array.isArray(content) ? content : [],
+            details: hookResult.details ?? payload.result?.details,
+            isError: hookResult.isError ?? payload.isError,
+        };
+    };
+    session.__fusionToolResultGuardInstalled = true;
+}
+function installMessageContentGuard(session, sessionManager) {
+    if (session.__fusionMessageContentGuardInstalled) {
+        return;
+    }
+    if (typeof session.subscribe === "function") {
+        session.subscribe((event) => {
+            if (!event || typeof event !== "object" || event.type !== "message_end") {
+                return;
+            }
+            const message = event.message;
+            if (!normalizeAssistantOrToolResultMessage(message)) {
+                return;
+            }
+            syncNormalizedMessageIntoAgentState(session, message);
+        });
+    }
+    if (typeof sessionManager.appendMessage === "function") {
+        const originalAppendMessage = sessionManager.appendMessage.bind(sessionManager);
+        sessionManager.appendMessage = (message) => {
+            normalizeAssistantOrToolResultMessage(message);
+            syncNormalizedMessageIntoAgentState(session, message);
+            return originalAppendMessage(message);
+        };
+    }
+    session.__fusionMessageContentGuardInstalled = true;
+}
 function hasPackageManagerSettings(settings) {
    return Array.isArray(settings.packages) || Array.isArray(settings.npmCommand);
 }
@@ -572,6 +683,7 @@ export async function createFnAgent(options) {
    });
    await resourceLoader.reload();
    const sessionManager = options.sessionManager ?? SessionManager.inMemory();
+    normalizeSessionHistoryEntries(sessionManager);
    const createSessionWithModel = async (modelOverride) => {
        return createAgentSession({
            cwd: options.cwd,
@@ -602,6 +714,8 @@ export async function createFnAgent(options) {
        piLog.log("Fallback session created successfully");
    }
    const { session } = sessionResult;
+    installToolResultContentGuard(session);
+    installMessageContentGuard(session, sessionManager);
    session.__fusionMemoryAppendAvailable = options.customTools?.some((tool) => tool.name === "memory_append") === true;
    const promptableSession = session;
    promptableSession.promptWithFallback = async (prompt, promptOptions) => {
@@ -657,6 +771,8 @@ export async function createFnAgent(options) {
            }
            const fallbackSessionResult = await createSessionWithModel(fallbackModel);
            const fallbackSession = fallbackSessionResult.session;
+            installToolResultContentGuard(fallbackSession);
+            installMessageContentGuard(fallbackSession, sessionManager);
            fallbackSession.__fusionMemoryAppendAvailable = options.customTools?.some((tool) => tool.name === "memory_append") === true;
            if (options.defaultThinkingLevel) {
                fallbackSession.setThinkingLevel(options.defaultThinkingLevel);
@@ -747,4 +863,4 @@ export async function createFnAgent(options) {
    });
    return { session: promptableSession, sessionFile: promptableSession.sessionFile };
 }
-//# sourceMappingURL=pi.js.map
+//# sourceMappingURL=pi.js.map
--- a/packages/engine/src/pi.ts
+++ b/packages/engine/src/pi.ts
@@ -46,6 +46,36 @@ export interface PromptableSession extends AgentSession {
  promptWithFallback: (prompt: string, options?: unknown) => Promise<void>;
 }

+interface SessionManagerLike {
+  fileEntries?: Array<{ type?: string; message?: Record<string, unknown> }>;
+  appendMessage?: (message: Record<string, unknown>) => void;
+  _rewriteFile?: () => void;
+}
+
+interface ToolHookPayload {
+  toolCall: unknown;
+  args: unknown;
+  result: { content?: unknown; details?: unknown };
+  isError: boolean;
+}
+
+interface ToolHookResult {
+  content?: unknown;
+  details?: unknown;
+  isError?: boolean;
+}
+
+type AgentToolHookSession = AgentSession & {
+  agent?: {
+    afterToolCall?: (payload: ToolHookPayload) => Promise<ToolHookResult | undefined>;
+    state?: {
+      messages?: Array<Record<string, unknown>>;
+    };
+  };
+  __fusionToolResultGuardInstalled?: boolean;
+  __fusionMessageContentGuardInstalled?: boolean;
+};
+
 function getSessionStateError(session: AgentSession): string {
  const state = (session as any).state;
  const error = state?.errorMessage ?? state?.error;
@@ -414,6 +444,137 @@ function readJsonObject(path: string): Record<string, any> {
  }
 }

+function normalizeSessionHistoryEntries(sessionManager: SessionManagerLike): void {
+  const entries = sessionManager.fileEntries;
+  if (!Array.isArray(entries) || entries.length === 0) {
+    return;
+  }
+
+  let changed = false;
+  for (const entry of entries) {
+    if (entry?.type !== "message" || !entry.message || typeof entry.message !== "object") {
+      continue;
+    }
+    const role = entry.message.role;
+    if (role !== "assistant" && role !== "toolResult") {
+      continue;
+    }
+    if (!("content" in entry.message)) {
+      entry.message.content = [];
+      changed = true;
+    }
+  }
+
+  if (changed) {
+    sessionManager._rewriteFile?.();
+  }
+}
+
+function normalizeAssistantOrToolResultMessage(message: unknown): message is Record<string, unknown> {
+  if (!message || typeof message !== "object") {
+    return false;
+  }
+
+  const role = (message as Record<string, unknown>).role;
+  if (role !== "assistant" && role !== "toolResult") {
+    return false;
+  }
+
+  if (!Array.isArray((message as Record<string, unknown>).content)) {
+    (message as Record<string, unknown>).content = [];
+  }
+  return true;
+}
+
+function syncNormalizedMessageIntoAgentState(session: AgentToolHookSession, message: Record<string, unknown>): void {
+  const messages = session.agent?.state?.messages;
+  if (!Array.isArray(messages) || messages.length === 0) {
+    return;
+  }
+
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const candidate = messages[i];
+    if (!candidate || typeof candidate !== "object") {
+      continue;
+    }
+    if (candidate === message) {
+      normalizeAssistantOrToolResultMessage(candidate);
+      return;
+    }
+
+    if (candidate.role !== message.role) {
+      continue;
+    }
+
+    if (candidate.role === "toolResult") {
+      if (candidate.toolCallId === message.toolCallId && candidate.toolName === message.toolName) {
+        normalizeAssistantOrToolResultMessage(candidate);
+        return;
+      }
+      continue;
+    }
+
+    if (candidate.timestamp === message.timestamp) {
+      normalizeAssistantOrToolResultMessage(candidate);
+      return;
+    }
+  }
+}
+
+function installToolResultContentGuard(session: AgentToolHookSession): void {
+  if (session.__fusionToolResultGuardInstalled || !session.agent?.afterToolCall) {
+    return;
+  }
+
+  const originalAfterToolCall = session.agent.afterToolCall.bind(session.agent) as any;
+  (session.agent as any).afterToolCall = async (payload: ToolHookPayload) => {
+    const hookResult = await originalAfterToolCall(payload);
+    if (!hookResult || typeof hookResult !== "object") {
+      return hookResult;
+    }
+
+    const content = hookResult.content ?? payload.result?.content ?? [];
+    return {
+      content: Array.isArray(content) ? content : [],
+      details: hookResult.details ?? payload.result?.details,
+      isError: hookResult.isError ?? payload.isError,
+    };
+  };
+  session.__fusionToolResultGuardInstalled = true;
+}
+
+function installMessageContentGuard(session: AgentToolHookSession, sessionManager: SessionManagerLike): void {
+  if (session.__fusionMessageContentGuardInstalled) {
+    return;
+  }
+
+  if (typeof session.subscribe === "function") {
+    session.subscribe((event: unknown) => {
+      if (!event || typeof event !== "object" || (event as { type?: string }).type !== "message_end") {
+        return;
+      }
+
+      const message = (event as { message?: unknown }).message;
+      if (!normalizeAssistantOrToolResultMessage(message)) {
+        return;
+      }
+
+      syncNormalizedMessageIntoAgentState(session, message);
+    });
+  }
+
+  if (typeof sessionManager.appendMessage === "function") {
+    const originalAppendMessage = sessionManager.appendMessage.bind(sessionManager);
+    sessionManager.appendMessage = (message: Record<string, unknown>) => {
+      normalizeAssistantOrToolResultMessage(message);
+      syncNormalizedMessageIntoAgentState(session, message);
+      return originalAppendMessage(message);
+    };
+  }
+
+  session.__fusionMessageContentGuardInstalled = true;
+}
+
 function hasPackageManagerSettings(settings: Record<string, any>): boolean {
  return Array.isArray(settings.packages) || Array.isArray(settings.npmCommand);
 }
@@ -760,6 +921,7 @@ export async function createFnAgent(options: AgentOptions): Promise<AgentResult>
  await resourceLoader.reload();

  const sessionManager = options.sessionManager ?? SessionManager.inMemory();
+  normalizeSessionHistoryEntries(sessionManager as unknown as SessionManagerLike);

  const createSessionWithModel = async (modelOverride?: typeof selectedModel) => {
    // pi-coding-agent 0.68+: `tools` is a string[] allowlist of tool names, not
@@ -801,6 +963,8 @@ export async function createFnAgent(options: AgentOptions): Promise<AgentResult>
  }

  const { session } = sessionResult;
+  installToolResultContentGuard(session as AgentToolHookSession);
+  installMessageContentGuard(session as AgentToolHookSession, sessionManager as unknown as SessionManagerLike);
  (session as any).__fusionMemoryAppendAvailable = options.customTools?.some((tool) => tool.name === "memory_append") === true;
  const promptableSession = session as PromptableSession;

@@ -857,6 +1021,11 @@ export async function createFnAgent(options: AgentOptions): Promise<AgentResult>

      const fallbackSessionResult = await createSessionWithModel(fallbackModel);
      const fallbackSession = fallbackSessionResult.session as PromptableSession;
+      installToolResultContentGuard(fallbackSession as unknown as AgentToolHookSession);
+      installMessageContentGuard(
+        fallbackSession as unknown as AgentToolHookSession,
+        sessionManager as unknown as SessionManagerLike,
+      );
      (fallbackSession as any).__fusionMemoryAppendAvailable = options.customTools?.some((tool) => tool.name === "memory_append") === true;

      if (options.defaultThinkingLevel) {