can1357
2867e1f4e3
- Added canonical `pi.zod` schema API exports and removed TypeBox package exports/imports. - Migrated Tool schema typing from TypeBox to shared `TSchema`/Zod flow with legacy TypeBox compatibility. - Updated AI provider adapters and MCP/agent builders to convert tool params through `toolWireSchema()`. - Reworked schema validation from AJV to Zod-safe parsing with `fromTypeBox`, `toolWireSchema`, and meta schema checks.
93 lines
3.1 KiB
TypeScript
93 lines
3.1 KiB
TypeScript
import { afterEach, beforeEach, expect, it } from "bun:test";
|
|
import * as path from "node:path";
|
|
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
|
|
import { getBundledModel } from "@oh-my-pi/pi-ai";
|
|
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
|
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
|
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
|
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
|
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
|
import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
|
|
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
|
import { TempDir } from "@oh-my-pi/pi-utils";
|
|
import * as z from "zod/v4";
|
|
|
|
let tempDir: TempDir;
|
|
let authStorage: AuthStorage | undefined;
|
|
let session: AgentSession;
|
|
|
|
beforeEach(async () => {
|
|
tempDir = TempDir.createSync("@pi-agent-session-force-tool-");
|
|
const model = getBundledModel("anthropic", "claude-sonnet-4-5");
|
|
if (!model) throw new Error("Expected claude-sonnet-4-5 model to exist");
|
|
|
|
authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db"));
|
|
authStorage.setRuntimeApiKey("anthropic", "test-key");
|
|
const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml"));
|
|
const settings = Settings.isolated({ "compaction.enabled": false });
|
|
const sessionManager = SessionManager.inMemory(tempDir.path());
|
|
|
|
const bashTool: AgentTool = {
|
|
name: "bash",
|
|
label: "Bash",
|
|
description: "Mock bash tool",
|
|
parameters: z.object({}),
|
|
execute: async () => ({ content: [{ type: "text" as const, text: "ok" }] }),
|
|
};
|
|
const writeTool: AgentTool = {
|
|
name: "write",
|
|
label: "Write",
|
|
description: "Mock write tool",
|
|
parameters: z.object({}),
|
|
execute: async () => ({ content: [{ type: "text" as const, text: "ok" }] }),
|
|
};
|
|
|
|
const agent = new Agent({
|
|
getApiKey: () => "test-key",
|
|
initialState: {
|
|
model,
|
|
systemPrompt: ["Test"],
|
|
tools: [bashTool, writeTool],
|
|
messages: [],
|
|
},
|
|
convertToLlm,
|
|
streamFn: () => new AssistantMessageEventStream(),
|
|
});
|
|
|
|
session = new AgentSession({
|
|
agent,
|
|
sessionManager,
|
|
settings,
|
|
modelRegistry,
|
|
toolRegistry: new Map([
|
|
[bashTool.name, bashTool],
|
|
[writeTool.name, writeTool],
|
|
]),
|
|
});
|
|
});
|
|
|
|
afterEach(async () => {
|
|
await session.dispose();
|
|
authStorage?.close();
|
|
authStorage = undefined;
|
|
tempDir.removeSync();
|
|
});
|
|
|
|
it("forces specific tool, then transitions to none, then clears", () => {
|
|
session.setForcedToolChoice("write");
|
|
|
|
const first = session.nextToolChoice();
|
|
const second = session.nextToolChoice();
|
|
const third = session.nextToolChoice();
|
|
|
|
expect(first).toEqual({ type: "tool", name: "write" });
|
|
// After the forced call, "none" prevents the loop from making more tool calls
|
|
expect(second).toBe("none");
|
|
// After "none" is consumed, override clears entirely
|
|
expect(third).toBeUndefined();
|
|
});
|
|
|
|
it("throws when forcing a non-active tool", () => {
|
|
expect(() => session.setForcedToolChoice("read")).toThrow('Tool "read" is not currently active.');
|
|
});
|