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.
24 kB · 713 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714import { mkdirSync, readFileSync, rmSync, symlinkSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { AuthStorage } from "../src/core/auth-storage.js";import { ExtensionRunner } from "../src/core/extensions/runner.js";import { ModelRegistry } from "../src/core/model-registry.js";import { DefaultResourceLoader } from "../src/core/resource-loader.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import type { Skill } from "../src/core/skills.js";import { createSyntheticSourceInfo } from "../src/core/source-info.js";
describe("DefaultResourceLoader", () => { let tempDir: string; let agentDir: string; let cwd: string; let previousSerperApiKey: string | undefined;
beforeEach(() => { tempDir = join(tmpdir(), `rl-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); agentDir = join(tempDir, "agent"); cwd = join(tempDir, "project"); previousSerperApiKey = process.env.SERPER_API_KEY; delete process.env.SERPER_API_KEY; mkdirSync(agentDir, { recursive: true }); mkdirSync(cwd, { recursive: true }); });
afterEach(() => { if (previousSerperApiKey === undefined) { delete process.env.SERPER_API_KEY; } else { process.env.SERPER_API_KEY = previousSerperApiKey; } rmSync(tempDir, { recursive: true, force: true }); });
describe("reload", () => { it("should initialize with empty results before reload", () => { const loader = new DefaultResourceLoader({ cwd, agentDir });
expect(loader.getExtensions().extensions).toEqual([]); expect(loader.getSkills().skills).toEqual([]); expect(loader.getPrompts().prompts).toEqual([]); expect(loader.getThemes().themes).toEqual([]); });
it("should discover skills from agentDir", async () => { const skillsDir = join(agentDir, "skills"); mkdirSync(skillsDir, { recursive: true }); writeFileSync( join(skillsDir, "test-skill.md"), `---name: test-skilldescription: A test skill---Skill content here.`, );
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { skills } = loader.getSkills(); expect(skills.some((s) => s.name === "test-skill")).toBe(true); });
it("should ignore extra markdown files in auto-discovered skill dirs", async () => { const skillDir = join(agentDir, "skills", "pi-skills", "browser-tools"); mkdirSync(skillDir, { recursive: true }); writeFileSync( join(skillDir, "SKILL.md"), `---name: browser-toolsdescription: Browser tools---Skill content here.`, ); writeFileSync(join(skillDir, "EFFICIENCY.md"), "No frontmatter here");
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { skills, diagnostics } = loader.getSkills(); expect(skills.some((s) => s.name === "browser-tools")).toBe(true); expect(diagnostics.some((d) => d.path?.endsWith("EFFICIENCY.md"))).toBe(false); });
it("should discover prompts from agentDir", async () => { const promptsDir = join(agentDir, "prompts"); mkdirSync(promptsDir, { recursive: true }); writeFileSync( join(promptsDir, "test-prompt.md"), `---description: A test prompt---Prompt content.`, );
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { prompts } = loader.getPrompts(); expect(prompts.some((p) => p.name === "test-prompt")).toBe(true); });
it("should prefer project resources over user on name collisions", async () => { const userPromptsDir = join(agentDir, "prompts"); const projectPromptsDir = join(cwd, ".prime", "agent", "prompts"); mkdirSync(userPromptsDir, { recursive: true }); mkdirSync(projectPromptsDir, { recursive: true }); const userPromptPath = join(userPromptsDir, "commit.md"); const projectPromptPath = join(projectPromptsDir, "commit.md"); writeFileSync(userPromptPath, "User prompt"); writeFileSync(projectPromptPath, "Project prompt");
const userSkillDir = join(agentDir, "skills", "collision-skill"); const projectSkillDir = join(cwd, ".prime", "agent", "skills", "collision-skill"); mkdirSync(userSkillDir, { recursive: true }); mkdirSync(projectSkillDir, { recursive: true }); const userSkillPath = join(userSkillDir, "SKILL.md"); const projectSkillPath = join(projectSkillDir, "SKILL.md"); writeFileSync( userSkillPath, `---name: collision-skilldescription: user---User skill`, ); writeFileSync( projectSkillPath, `---name: collision-skilldescription: project---Project skill`, );
const baseTheme = JSON.parse( readFileSync(join(process.cwd(), "src", "modes", "interactive", "theme", "dark.json"), "utf-8"), ) as { name: string; vars?: Record<string, string> }; baseTheme.name = "collision-theme"; const userThemePath = join(agentDir, "themes", "collision.json"); const projectThemePath = join(cwd, ".prime", "agent", "themes", "collision.json"); mkdirSync(join(agentDir, "themes"), { recursive: true }); mkdirSync(join(cwd, ".prime", "agent", "themes"), { recursive: true }); writeFileSync(userThemePath, JSON.stringify(baseTheme, null, 2)); if (baseTheme.vars) { baseTheme.vars.accent = "#ff00ff"; } writeFileSync(projectThemePath, JSON.stringify(baseTheme, null, 2));
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const prompt = loader.getPrompts().prompts.find((p) => p.name === "commit"); expect(prompt?.filePath).toBe(projectPromptPath);
const skill = loader.getSkills().skills.find((s) => s.name === "collision-skill"); expect(skill?.filePath).toBe(projectSkillPath);
const theme = loader.getThemes().themes.find((t) => t.name === "collision-theme"); expect(theme?.sourcePath).toBe(projectThemePath); });
it("should load symlinked user and project extensions once", async () => { const sharedExtDir = join(tempDir, "shared-extensions"); mkdirSync(sharedExtDir, { recursive: true }); writeFileSync( join(sharedExtDir, "shared.ts"), `export default function(pi) { pi.registerCommand("shared", { description: "shared command", handler: async () => {}, });}`, );
mkdirSync(agentDir, { recursive: true }); mkdirSync(join(cwd, ".prime", "agent"), { recursive: true }); symlinkSync(sharedExtDir, join(agentDir, "extensions"), "dir"); symlinkSync(sharedExtDir, join(cwd, ".prime", "agent", "extensions"), "dir");
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const extensionsResult = loader.getExtensions(); expect(extensionsResult.extensions).toHaveLength(1); expect(extensionsResult.errors).toEqual([]);
// mergePaths processes project paths before user paths, so the project // alias is the canonical survivor. expect(extensionsResult.extensions[0].path).toBe(join(cwd, ".prime", "agent", "extensions", "shared.ts")); });
it("should keep both extensions loaded when command names collide", async () => { const userExtDir = join(agentDir, "extensions"); const projectExtDir = join(cwd, ".prime", "agent", "extensions"); mkdirSync(userExtDir, { recursive: true }); mkdirSync(projectExtDir, { recursive: true });
writeFileSync( join(projectExtDir, "project.ts"), `export default function(pi) { pi.registerCommand("deploy", { description: "project deploy", handler: async () => {}, }); pi.registerCommand("project-only", { description: "project only", handler: async () => {}, });}`, );
writeFileSync( join(userExtDir, "user.ts"), `export default function(pi) { pi.registerCommand("deploy", { description: "user deploy", handler: async () => {}, }); pi.registerCommand("user-only", { description: "user only", handler: async () => {}, });}`, );
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const extensionsResult = loader.getExtensions(); expect(extensionsResult.extensions).toHaveLength(2); expect(extensionsResult.errors.some((e) => e.error.includes('Command "/deploy" conflicts'))).toBe(false);
const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage); const runner = new ExtensionRunner( extensionsResult.extensions, extensionsResult.runtime, cwd, sessionManager, modelRegistry, );
expect(runner.getCommand("deploy:1")?.description).toBe("project deploy"); expect(runner.getCommand("deploy:2")?.description).toBe("user deploy"); expect(runner.getCommand("project-only")?.description).toBe("project only"); expect(runner.getCommand("user-only")?.description).toBe("user only");
const commands = runner.getRegisteredCommands(); expect(commands.map((command) => command.invocationName)).toEqual([ "deploy:1", "project-only", "deploy:2", "user-only", ]); });
it("should honor overrides for auto-discovered resources", async () => { const settingsManager = SettingsManager.inMemory(); settingsManager.setExtensionPaths(["-extensions/disabled.ts"]); settingsManager.setSkillPaths(["-skills/skip-skill"]); settingsManager.setPromptTemplatePaths(["-prompts/skip.md"]); settingsManager.setThemePaths(["-themes/skip.json"]);
const extensionsDir = join(agentDir, "extensions"); mkdirSync(extensionsDir, { recursive: true }); writeFileSync(join(extensionsDir, "disabled.ts"), "export default function() {}");
const skillDir = join(agentDir, "skills", "skip-skill"); mkdirSync(skillDir, { recursive: true }); writeFileSync( join(skillDir, "SKILL.md"), `---name: skip-skilldescription: Skip me---Content`, );
const promptsDir = join(agentDir, "prompts"); mkdirSync(promptsDir, { recursive: true }); writeFileSync(join(promptsDir, "skip.md"), "Skip prompt");
const themesDir = join(agentDir, "themes"); mkdirSync(themesDir, { recursive: true }); writeFileSync(join(themesDir, "skip.json"), "{}");
const loader = new DefaultResourceLoader({ cwd, agentDir, settingsManager }); await loader.reload();
const { extensions } = loader.getExtensions(); const { skills } = loader.getSkills(); const { prompts } = loader.getPrompts(); const { themes } = loader.getThemes();
expect(extensions.some((e) => e.path.endsWith("disabled.ts"))).toBe(false); expect(skills.some((s) => s.name === "skip-skill")).toBe(false); expect(prompts.some((p) => p.name === "skip")).toBe(false); expect(themes.some((t) => t.sourcePath?.endsWith("skip.json"))).toBe(false); });
it("should discover AGENTS.md context files", async () => { writeFileSync(join(cwd, "AGENTS.md"), "# Project Guidelines\n\nBe helpful.");
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { agentsFiles } = loader.getAgentsFiles(); expect(agentsFiles.some((f) => f.path.includes("AGENTS.md"))).toBe(true); });
it("should skip AGENTS.md and CLAUDE.md discovery when noContextFiles is true", async () => { writeFileSync(join(cwd, "AGENTS.md"), "# Project Guidelines\n\nBe helpful."); writeFileSync(join(cwd, "CLAUDE.md"), "# Claude Guidelines\n\nBe helpful.");
const loader = new DefaultResourceLoader({ cwd, agentDir, noContextFiles: true }); await loader.reload();
const { agentsFiles } = loader.getAgentsFiles(); expect(agentsFiles).toEqual([]); });
it("should discover SYSTEM.md from cwd/.pi", async () => { const piDir = join(cwd, ".prime", "agent"); mkdirSync(piDir, { recursive: true }); writeFileSync(join(piDir, "SYSTEM.md"), "You are a helpful assistant.");
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
expect(loader.getSystemPrompt()).toBe("You are a helpful assistant."); });
it("should discover APPEND_SYSTEM.md", async () => { const piDir = join(cwd, ".prime", "agent"); mkdirSync(piDir, { recursive: true }); writeFileSync(join(piDir, "APPEND_SYSTEM.md"), "Additional instructions.");
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
expect(loader.getAppendSystemPrompt()).toContain("Additional instructions."); }); });
describe("extendResources", () => { it("should load skills and prompts with extension metadata", async () => { const extraSkillDir = join(tempDir, "extra-skills", "extra-skill"); mkdirSync(extraSkillDir, { recursive: true }); const skillPath = join(extraSkillDir, "SKILL.md"); writeFileSync( skillPath, `---name: extra-skilldescription: Extra skill---Extra content`, );
const extraPromptDir = join(tempDir, "extra-prompts"); mkdirSync(extraPromptDir, { recursive: true }); const promptPath = join(extraPromptDir, "extra.md"); writeFileSync( promptPath, `---description: Extra prompt---Extra prompt content`, );
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
loader.extendResources({ skillPaths: [ { path: extraSkillDir, metadata: { source: "extension:extra", scope: "temporary", origin: "top-level", baseDir: extraSkillDir, }, }, ], promptPaths: [ { path: promptPath, metadata: { source: "extension:extra", scope: "temporary", origin: "top-level", baseDir: extraPromptDir, }, }, ], });
const { skills } = loader.getSkills(); const loadedSkill = skills.find((skill) => skill.name === "extra-skill"); expect(loadedSkill).toBeDefined(); expect(loadedSkill?.sourceInfo?.source).toBe("extension:extra"); expect(loadedSkill?.sourceInfo?.path).toBe(skillPath);
const { prompts } = loader.getPrompts(); const loadedPrompt = prompts.find((prompt) => prompt.name === "extra"); expect(loadedPrompt).toBeDefined(); expect(loadedPrompt?.sourceInfo?.source).toBe("extension:extra"); expect(loadedPrompt?.sourceInfo?.path).toBe(promptPath); }); });
describe("noSkills option", () => { it("should skip skill discovery when noSkills is true", async () => { const skillsDir = join(agentDir, "skills"); mkdirSync(skillsDir, { recursive: true }); writeFileSync( join(skillsDir, "test-skill.md"), `---name: test-skilldescription: A test skill---Content`, );
const loader = new DefaultResourceLoader({ cwd, agentDir, noSkills: true }); await loader.reload();
const { skills } = loader.getSkills(); expect(skills).toEqual([]); });
it("should still load additional skill paths when noSkills is true", async () => { const customSkillDir = join(tempDir, "custom-skills"); mkdirSync(customSkillDir, { recursive: true }); writeFileSync( join(customSkillDir, "custom.md"), `---name: customdescription: Custom skill---Content`, );
const loader = new DefaultResourceLoader({ cwd, agentDir, noSkills: true, additionalSkillPaths: [customSkillDir], }); await loader.reload();
const { skills } = loader.getSkills(); expect(skills.some((s) => s.name === "custom")).toBe(true); }); });
describe("bundled skills", () => { it("should load the bundled websearch skill by default", async () => { const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { skills } = loader.getSkills(); const websearch = skills.find((s) => s.name === "websearch"); expect(websearch).toBeDefined(); expect(websearch?.description).toContain("/login"); expect(websearch?.description).toContain("MCP Connections"); expect(websearch?.description).toContain("Serper (web search)"); expect(websearch?.kind).toBe("python"); if (websearch?.kind === "python") { expect(websearch.python.importName).toBe("websearch"); expect(websearch.python.pyprojectPath.endsWith("pyproject.toml")).toBe(true); } });
it("should not emit a SERPER_API_KEY warning when the key is unset", async () => { const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { diagnostics } = loader.getSkills(); expect(diagnostics.some((d) => d.type === "warning" && d.message.includes("SERPER_API_KEY"))).toBe(false); });
it("should not load the bundled websearch skill when disabled in settings", async () => { const settingsManager = SettingsManager.inMemory({ bundledSkills: { websearch: false } }); const loader = new DefaultResourceLoader({ cwd, agentDir, settingsManager }); await loader.reload();
const { skills, diagnostics } = loader.getSkills(); expect(skills.some((s) => s.name === "websearch")).toBe(false); expect(diagnostics.some((d) => d.type === "warning" && d.message.includes("SERPER_API_KEY is not set"))).toBe( false, ); });
it("should not load bundled skills when noSkills is true", async () => { const loader = new DefaultResourceLoader({ cwd, agentDir, noSkills: true }); await loader.reload();
const { skills } = loader.getSkills(); expect(skills.some((s) => s.name === "websearch")).toBe(false); });
it("should let a project skill override the bundled websearch skill", async () => { const projectSkillDir = join(cwd, ".prime", "agent", "skills", "websearch"); mkdirSync(projectSkillDir, { recursive: true }); writeFileSync( join(projectSkillDir, "SKILL.md"), `---name: websearchdescription: Project-specific web search override.---Project override.`, );
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { skills, diagnostics } = loader.getSkills(); const websearch = skills.find((s) => s.name === "websearch"); expect(websearch).toBeDefined(); expect(websearch?.filePath).toBe(join(projectSkillDir, "SKILL.md")); expect(websearch?.kind).toBe("markdown"); expect(diagnostics.some((d) => d.type === "collision" && d.collision?.name === "websearch")).toBe(true); expect(diagnostics.some((d) => d.type === "warning" && d.message.includes("SERPER_API_KEY is not set"))).toBe( false, ); });
it("should let an explicit --skill path override the bundled websearch skill", async () => { const customSkillDir = join(tempDir, "custom-websearch", "websearch"); mkdirSync(customSkillDir, { recursive: true }); writeFileSync( join(customSkillDir, "SKILL.md"), `---name: websearchdescription: Explicit web search override.---Explicit override.`, );
const loader = new DefaultResourceLoader({ cwd, agentDir, additionalSkillPaths: [customSkillDir], }); await loader.reload();
const { skills } = loader.getSkills(); const websearch = skills.find((s) => s.name === "websearch"); expect(websearch).toBeDefined(); expect(websearch?.filePath).toBe(join(customSkillDir, "SKILL.md")); expect(websearch?.kind).toBe("markdown"); }); });
describe("override functions", () => { it("should apply skillsOverride", async () => { const injectedSkill: Skill = { name: "injected", description: "Injected skill", filePath: "/fake/path", baseDir: "/fake", sourceInfo: createSyntheticSourceInfo("/fake/path", { source: "custom" }), disableModelInvocation: false, kind: "markdown", }; const loader = new DefaultResourceLoader({ cwd, agentDir, skillsOverride: () => ({ skills: [injectedSkill], diagnostics: [], }), }); await loader.reload();
const { skills } = loader.getSkills(); expect(skills).toHaveLength(1); expect(skills[0].name).toBe("injected"); });
it("should apply systemPromptOverride", async () => { const loader = new DefaultResourceLoader({ cwd, agentDir, systemPromptOverride: () => "Custom system prompt", }); await loader.reload();
expect(loader.getSystemPrompt()).toBe("Custom system prompt"); }); });
describe("extension conflict detection", () => { it("should detect tool conflicts between extensions", async () => { const ext1Dir = join(agentDir, "extensions", "ext1"); const ext2Dir = join(agentDir, "extensions", "ext2"); mkdirSync(ext1Dir, { recursive: true }); mkdirSync(ext2Dir, { recursive: true });
writeFileSync( join(ext1Dir, "index.ts"), `import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { Type } from "typebox";export default function(pi: ExtensionAPI) { pi.registerTool({ name: "duplicate-tool", description: "First", parameters: Type.Object({}), execute: async () => ({ result: "1" }), });}`, );
writeFileSync( join(ext2Dir, "index.ts"), `import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { Type } from "typebox";export default function(pi: ExtensionAPI) { pi.registerTool({ name: "duplicate-tool", description: "Second", parameters: Type.Object({}), execute: async () => ({ result: "2" }), });}`, );
const loader = new DefaultResourceLoader({ cwd, agentDir }); await loader.reload();
const { errors } = loader.getExtensions(); expect(errors.some((e) => e.error.includes("duplicate-tool") && e.error.includes("conflicts"))).toBe(true); });
it("should prefer explicit CLI extensions over discovered extensions when commands and tools conflict", async () => { const globalExtDir = join(agentDir, "extensions"); mkdirSync(globalExtDir, { recursive: true }); const explicitExtPath = join(tempDir, "explicit-extension.ts");
writeFileSync( join(globalExtDir, "global.ts"), `import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { Type } from "typebox";export default function(pi: ExtensionAPI) { pi.registerTool({ name: "duplicate-tool", description: "global tool", parameters: Type.Object({}), execute: async () => ({ result: "global" }), }); pi.registerCommand("deploy", { description: "global command", handler: async () => {}, });}`, );
writeFileSync( explicitExtPath, `import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { Type } from "typebox";export default function(pi: ExtensionAPI) { pi.registerTool({ name: "duplicate-tool", description: "explicit tool", parameters: Type.Object({}), execute: async () => ({ result: "explicit" }), }); pi.registerCommand("deploy", { description: "explicit command", handler: async () => {}, });}`, );
const loader = new DefaultResourceLoader({ cwd, agentDir, additionalExtensionPaths: [explicitExtPath], }); await loader.reload();
const extensionsResult = loader.getExtensions(); expect(extensionsResult.extensions[0]?.path).toBe(explicitExtPath);
const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.create(join(tempDir, "auth-explicit.json")); const modelRegistry = ModelRegistry.create(authStorage); const runner = new ExtensionRunner( extensionsResult.extensions, extensionsResult.runtime, cwd, sessionManager, modelRegistry, );
expect(runner.getCommand("deploy:1")?.description).toBe("explicit command"); expect(runner.getCommand("deploy:2")?.description).toBe("global command"); expect(runner.getToolDefinition("duplicate-tool")?.description).toBe("explicit tool"); }); });});