evOo-local/packages/adapters/src/pi-runtime-computer.test.ts

import type { ConnectorTool } from "@rakazo/adapter-kit";
import { beforeEach, describe, expect, it, vi } from "vitest";

const fakeAgentState = vi.hoisted(() => ({
  result: undefined as unknown,
  systemPrompt: "",
}));

vi.mock("@earendil-works/pi-agent-core", () => ({
  Agent: class {
    state = { errorMessage: undefined, messages: [] };
    private readonly tool: {
      execute: (toolCallId: string, params: unknown) => Promise<unknown>;
    };

    constructor(options: {
      initialState: {
        systemPrompt: string;
        tools: Array<{
          name: string;
          execute: (toolCallId: string, params: unknown) => Promise<unknown>;
        }>;
      };
    }) {
      fakeAgentState.systemPrompt = options.initialState.systemPrompt;
      const tool = options.initialState.tools.find(
        (candidate) => candidate.name === "computer_observe",
      );
      if (!tool) throw new Error("computer_observe was not exposed");
      this.tool = tool;
    }

    subscribe(_listener: unknown) {}

    async prompt() {
      fakeAgentState.result = await this.tool.execute("observe-1", {});
    }

    async waitForIdle() {}

    abort() {}
  },
}));

vi.mock("@earendil-works/pi-ai/providers/all", () => ({
  builtinModels: () => ({
    getModel: (_provider: string, modelId: string) =>
      modelId === "computer-test-model" ? { provider: "test", id: modelId } : undefined,
    streamSimple: () => {
      throw new Error("the fake agent must not call a provider");
    },
  }),
}));

vi.mock("./pi-local-provider.js", () => ({
  registerLocalProvider: (models: unknown) => models,
}));

vi.mock("./pi-openai-compatible-provider.js", () => ({
  OPENAI_COMPATIBLE_PROVIDER_ID: "openai-compatible",
  registerOpenAiCompatibleCatalog: (models: unknown) => models,
  registerOpenAiCompatibleRuntime: (models: unknown) => models,
}));

import { COMPUTER_SCREEN_UNAVAILABLE } from "./computer-screens.js";
import { PiAgentRuntime, pruneComputerScreenshotContext } from "./pi-runtime.js";

const computerObserve: ConnectorTool = {
  name: "computer_observe",
  description: "Observe the computer",
  inputSchema: { type: "object", properties: {} },
};

describe("Pi computer tool dispatch", () => {
  beforeEach(() => {
    fakeAgentState.result = undefined;
    fakeAgentState.systemPrompt = "";
  });

  it("forwards screenshots as image content for any Pi model provider", async () => {
    const runtime = new PiAgentRuntime();
    for await (const _event of runtime.run(
      {
        botId: "bot",
        threadId: "thread",
        runId: "run",
        prompt: "look at the screen",
        instructions: "Follow the user's instructions.",
        history: [],
        tools: [computerObserve],
        model: { provider: "test", id: "computer-test-model" },
        executeTool: async () => ({
          kind: "agent_tool_result",
          content: [
            { type: "text", text: "computer observed" },
            { type: "image", data: "iVBORw0KGgo=", mimeType: "image/png" },
          ],
          details: { frameId: "frame-1" },
        }),
      },
      {
        operationId: "computer-test",
        traceId: "computer-test",
        spaceId: "workspace",
        userId: "user",
        signal: new AbortController().signal,
      },
    )) {
      // Exhaust the runtime so the fake agent executes the tool.
    }

    expect(fakeAgentState.result).toMatchObject({
      content: [{ type: "text" }, { type: "image", mimeType: "image/png", data: "iVBORw0KGgo=" }],
    });
    expect(fakeAgentState.systemPrompt).toBe("Follow the user's instructions.");
  });

  it("keeps the run alive when a graphical tool returns an error object", async () => {
    const runtime = new PiAgentRuntime();
    const events: Array<{ type: string; text?: string }> = [];
    for await (const event of runtime.run(
      {
        botId: "bot",
        threadId: "thread",
        runId: "run-screen-error",
        prompt: "look at the screen",
        instructions: "Follow the user's instructions.",
        history: [],
        tools: [computerObserve],
        model: { provider: "test", id: "computer-test-model" },
        executeTool: async () => ({ error: COMPUTER_SCREEN_UNAVAILABLE }),
      },
      {
        operationId: "computer-test",
        traceId: "computer-test",
        spaceId: "workspace",
        userId: "user",
        signal: new AbortController().signal,
      },
    )) {
      events.push(event);
    }

    expect(fakeAgentState.result).toMatchObject({
      content: [{ type: "text" }],
      details: { error: expect.stringMatching(/temporarily busy/) },
    });
    expect(events.some((event) => event.text?.includes("I hit a problem"))).toBe(false);
    expect(events.at(-1)?.type).toBe("done");
  });

  it("keeps the two latest computer screenshots by default", () => {
    const messages = ["frame-1", "frame-2", "frame-3"].map((frameId) => ({
      role: "toolResult" as const,
      toolCallId: frameId,
      toolName: "computer_observe",
      content: [
        { type: "text" as const, text: frameId },
        { type: "image" as const, data: frameId, mimeType: "image/png" as const },
      ],
      details: { frameId },
      isError: false,
      timestamp: 1,
    }));

    const pruned = pruneComputerScreenshotContext(messages);
    expect(
      pruned.map((message) =>
        (message as (typeof messages)[number]).content.some((part) => part.type === "image"),
      ),
    ).toEqual([false, true, true]);
  });

  it("honors a model-specific one-image limit", () => {
    const messages = ["frame-1", "frame-2", "frame-3"].map((frameId) => ({
      role: "toolResult" as const,
      toolCallId: frameId,
      toolName: "computer_observe",
      content: [
        { type: "text" as const, text: frameId },
        { type: "image" as const, data: frameId, mimeType: "image/png" as const },
      ],
      details: { frameId },
      isError: false,
      timestamp: 1,
    }));

    const pruned = pruneComputerScreenshotContext(messages, 1);
    expect(
      pruned.map((message) =>
        (message as (typeof messages)[number]).content.some((part) => part.type === "image"),
      ),
    ).toEqual([false, false, true]);
  });

  it("reserves the image budget for user attachments", () => {
    const messages = [
      {
        role: "user" as const,
        content: [{ type: "image" as const, data: "user-image", mimeType: "image/png" as const }],
        timestamp: 1,
      },
      ...["frame-1", "frame-2"].map((frameId) => ({
        role: "toolResult" as const,
        toolCallId: frameId,
        toolName: "computer_observe",
        content: [
          { type: "text" as const, text: frameId },
          { type: "image" as const, data: frameId, mimeType: "image/png" as const },
        ],
        details: { frameId },
        isError: false,
        timestamp: 1,
      })),
    ];

    const pruned = pruneComputerScreenshotContext(messages, 2);
    expect(
      pruned.map((message) =>
        (message as (typeof messages)[number]).content.some((part) => part.type === "image"),
      ),
    ).toEqual([true, false, true]);
  });

  it("rejects user images that exceed a configured model limit", () => {
    const messages = [
      {
        role: "user" as const,
        content: [
          { type: "image" as const, data: "user-image-1", mimeType: "image/png" as const },
          { type: "image" as const, data: "user-image-2", mimeType: "image/png" as const },
        ],
        timestamp: 1,
      },
    ];

    expect(() => pruneComputerScreenshotContext(messages, 1)).toThrow(
      "the prompt contains 2 non-screenshot images",
    );
  });

  it("does not apply the default screenshot retention to user images", () => {
    const messages = [
      {
        role: "user" as const,
        content: [
          { type: "image" as const, data: "user-image-1", mimeType: "image/png" as const },
          { type: "image" as const, data: "user-image-2", mimeType: "image/png" as const },
          { type: "image" as const, data: "user-image-3", mimeType: "image/png" as const },
        ],
        timestamp: 1,
      },
      ...["frame-1", "frame-2", "frame-3"].map((frameId) => ({
        role: "toolResult" as const,
        toolCallId: frameId,
        toolName: "computer_observe",
        content: [
          { type: "text" as const, text: frameId },
          { type: "image" as const, data: frameId, mimeType: "image/png" as const },
        ],
        details: { frameId },
        isError: false,
        timestamp: 1,
      })),
    ];

    const pruned = pruneComputerScreenshotContext(messages);
    expect(pruned[0]).toBe(messages[0]);
    expect((pruned[0] as (typeof messages)[number]).content).toHaveLength(3);
    expect(
      pruned
        .slice(1)
        .map((message) =>
          (message as (typeof messages)[number]).content.some((part) => part.type === "image"),
        ),
    ).toEqual([false, true, true]);
  });

  it("reuses the message array when no screenshot needs pruning", () => {
    const messages = [
      {
        role: "user" as const,
        content: [{ type: "text" as const, text: "hello" }],
        timestamp: 1,
      },
    ];

    expect(pruneComputerScreenshotContext(messages)).toBe(messages);
  });
});