Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
5.6 kB · 182 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183import { existsSync, mkdirSync, rmSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { getModel } from "@earendil-works/pi-ai";import { Type } from "typebox";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { DefaultResourceLoader } from "../src/core/resource-loader.js";import { createAgentSession } from "../src/core/sdk.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";
describe("AgentSession dynamic tool registration", () => { let tempDir: string; let agentDir: string;
beforeEach(() => { tempDir = join(tmpdir(), `pi-dynamic-tool-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); agentDir = join(tempDir, "agent"); mkdirSync(agentDir, { recursive: true }); });
afterEach(() => { if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true, force: true }); } });
it("refreshes tool registry when tools are registered after initialization", async () => { const settingsManager = SettingsManager.create(tempDir, agentDir); const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir, settingsManager, extensionFactories: [ (pi) => { pi.on("session_start", () => { pi.registerTool({ name: "dynamic_tool", label: "Dynamic Tool", description: "Tool registered from session_start", promptSnippet: "Run dynamic test behavior", promptGuidelines: ["Use dynamic_tool when the user asks for dynamic behavior tests."], parameters: Type.Object({}), execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {}, }), }); }); }, ], }); await resourceLoader.reload();
const { session } = await createAgentSession({ cwd: tempDir, agentDir, model: getModel("anthropic", "claude-sonnet-4-5")!, settingsManager, sessionManager, resourceLoader, });
expect(session.getAllTools().map((tool) => tool.name)).not.toContain("dynamic_tool");
await session.bindExtensions({});
const allTools = session.getAllTools(); const dynamicTool = allTools.find((tool) => tool.name === "dynamic_tool"); const ipythonTool = allTools.find((tool) => tool.name === "ipython");
expect(allTools.map((tool) => tool.name)).toContain("dynamic_tool"); expect(dynamicTool?.sourceInfo).toMatchObject({ path: "<inline:1>", source: "inline", scope: "temporary", origin: "top-level", }); expect(ipythonTool?.sourceInfo).toMatchObject({ path: "<builtin:ipython>", source: "builtin", scope: "temporary", origin: "top-level", }); expect(session.getActiveToolNames()).toContain("dynamic_tool"); expect(session.systemPrompt).not.toContain("- dynamic_tool: Run dynamic test behavior"); expect(session.systemPrompt).toContain("- Use dynamic_tool when the user asks for dynamic behavior tests.");
session.dispose(); });
it("returns source metadata for SDK custom tools", async () => { const settingsManager = SettingsManager.create(tempDir, agentDir); const sessionManager = SessionManager.inMemory(); const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir, settingsManager, }); await resourceLoader.reload();
const { session } = await createAgentSession({ cwd: tempDir, agentDir, model: getModel("anthropic", "claude-sonnet-4-5")!, settingsManager, sessionManager, resourceLoader, customTools: [ { name: "sdk_tool", label: "SDK Tool", description: "Tool registered through createAgentSession", parameters: Type.Object({}), execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {}, }), }, ], });
const sdkTool = session.getAllTools().find((tool) => tool.name === "sdk_tool"); expect(sdkTool?.sourceInfo).toMatchObject({ path: "<sdk:sdk_tool>", source: "sdk", scope: "temporary", origin: "top-level", }); expect(session.getActiveToolNames()).toContain("sdk_tool");
session.dispose(); });
it("keeps custom tools active but omits them from the system prompt when prompt guidance is not provided", async () => { const settingsManager = SettingsManager.create(tempDir, agentDir); const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir, settingsManager, extensionFactories: [ (pi) => { pi.on("session_start", () => { pi.registerTool({ name: "hidden_tool", label: "Hidden Tool", description: "Description should not appear in available tools", parameters: Type.Object({}), execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {}, }), }); }); }, ], }); await resourceLoader.reload();
const { session } = await createAgentSession({ cwd: tempDir, agentDir, model: getModel("anthropic", "claude-sonnet-4-5")!, settingsManager, sessionManager, resourceLoader, });
await session.bindExtensions({});
expect(session.getAllTools().map((tool) => tool.name)).toContain("hidden_tool"); expect(session.getActiveToolNames()).toContain("hidden_tool"); expect(session.systemPrompt).not.toContain("hidden_tool"); expect(session.systemPrompt).not.toContain("Description should not appear in available tools");
session.dispose(); });});