fix(runtime): erase pi tool array type safely

This commit is contained in:
ZacharyZhang-NY committed 2026-04-22 18:37:03 -04:00
1 parent 710448a75f
commit 9f35dc7446
3 files changed
+23 -18

No files matched your search

+6 -12
View File
@@ -1,10 +1,10 @@
import type { Agent, AgentMessage, AgentTool } from "@mariozechner/pi-agent-core";
import type { Agent, AgentMessage } from "@mariozechner/pi-agent-core";
import { resolveModel, resolveExtraStreamHeaders } from "../../model-resolver";
import type { RunStatus } from "@shared/types";
import type { ChildMessage, RunJob } from "./helpers";
import { applyEnvOverrides, buildRedactor } from "./helpers";
import { createCodingTools } from "@mariozechner/pi-coding-agent";
import { createPiAgent, resolveThinkingLevel } from "./pi-agent-factory";
import { createPiAgent, resolveThinkingLevel, type ErasedAgentTool } from "./pi-agent-factory";
import type { WorkerRpcClient } from "../entry/rpc";
import { createHireAgentTool, HIRE_AGENT_TOOL_NAME } from "./internal-tools/hire-agent-tool";
import { createFileApprovalTool, FILE_APPROVAL_TOOL_NAME } from "./internal-tools/file-approval-tool";
@@ -149,20 +149,14 @@ export class WorkerRunner {
const cwd = job.workspacePath ?? process.cwd();
const codingTools = createCodingTools(cwd);
// pi's `AgentTool<TSchema>` is generic over the concrete TypeBox
// schema; two tools with different schemas cannot share a single
// `AgentTool<TSchema>[]` without widening to the base shape. The
// internal-tool factories return properly-typed instances; we
// widen via unknown only at the array boundary so each execute
// signature stays type-checked at definition time.
const internalTools: AgentTool[] = [];
const internalTools: ErasedAgentTool[] = [];
for (const toolName of job.internalTools ?? []) {
switch (toolName) {
case HIRE_AGENT_TOOL_NAME:
internalTools.push(createHireAgentTool(rpc, job.runId) as unknown as AgentTool);
internalTools.push(createHireAgentTool(rpc, job.runId));
break;
case FILE_APPROVAL_TOOL_NAME:
internalTools.push(createFileApprovalTool(rpc, job.runId) as unknown as AgentTool);
internalTools.push(createFileApprovalTool(rpc, job.runId));
break;
default: {
// Exhaustiveness guard: a new entry in the `internalTools`
@@ -173,7 +167,7 @@ export class WorkerRunner {
}
}
}
const tools = [...codingTools, ...internalTools] as never;
const tools: ErasedAgentTool[] = [...codingTools, ...internalTools];
const agent = createPiAgent({
runId: job.runId,
+4 -5
View File
@@ -16,6 +16,7 @@
import {
Agent,
type AgentMessage,
type AgentState,
type ThinkingLevel,
} from "@mariozechner/pi-agent-core";
import type { Api, Model } from "@mariozechner/pi-ai";
@@ -27,16 +28,14 @@ import { safeStreamSimple } from "./pi-safe-stream";
// graph so any drift in pi-agent-core hook shapes breaks the build here.
import "./pi-hook-contract";
export type ErasedAgentTool = AgentState["tools"][number];
export interface PiAgentFactoryInput {
runId: string;
taskSessionId: string;
systemPrompt: string;
model: Model<Api>;
// `never` preserves the pre-C4 contract where the caller owns tool
// typing (pi-coding-agent's union is too wide to mirror here and C4
// adds local internal tools with their own TSchema). The caller casts
// a concrete tools array at the `createPiAgent` site.
tools: never;
tools: ErasedAgentTool[];
priorMessages: AgentMessage[];
thinkingLevel: ThinkingLevel;
apiKey: string;
+13 -1
View File
@@ -1,3 +1,5 @@
import { readFileSync } from "node:fs";
import { join } from "node:path";
import type { Model } from "@mariozechner/pi-ai";
import { describe, expect, it } from "vitest";
import { createPiAgent } from "../../src/main/worker/runtime/pi-agent-factory";
@@ -25,7 +27,7 @@ describe("pi agent factory", () => {
taskSessionId: "session_1",
systemPrompt: "You are an engineer.",
model: createModel(),
tools: [] as never,
tools: [],
priorMessages: [],
thinkingLevel: "off",
apiKey: "test-key",
@@ -39,4 +41,14 @@ describe("pi agent factory", () => {
expect(agent.toolExecution).toBe("sequential");
});
it("keeps worker tool arrays on the erased pi tool contract", () => {
const runnerSource = readFileSync(join(process.cwd(), "src/main/worker/runtime/index.ts"), "utf8");
const factorySource = readFileSync(join(process.cwd(), "src/main/worker/runtime/pi-agent-factory.ts"), "utf8");
expect(runnerSource).not.toContain("as unknown as AgentTool");
expect(runnerSource).not.toContain("as never");
expect(factorySource).not.toContain("tools: never");
expect(factorySource).toContain('type ErasedAgentTool = AgentState["tools"][number]');
});
});