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.
22 kB · 605 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { registerFauxProvider } from "@earendil-works/pi-ai";import { afterEach, describe, expect, it, vi } from "vitest";import { ENV_AGENT_DIR } from "../src/config.js";import { AGENT_MESSAGE_SKILL_NAME, type AgentSessionMessageController } from "../src/core/agent-messages.js";import { AGENT_OBSERVE_SKILL_NAME, type AgentObserveController } from "../src/core/agent-observe.js";import { createAgentSessionFromServices, createAgentSessionServices } from "../src/core/agent-session-services.js";import { AuthStorage } from "../src/core/auth-storage.js";import { McpConnectionStore } from "../src/core/mcp/connection-store.js";import { McpManager } from "../src/core/mcp/mcp-manager.js";import { ModelRegistry } from "../src/core/model-registry.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import { createSyntheticSourceInfo } from "../src/core/source-info.js";
describe("createAgentSessionFromServices", () => { const cleanupPaths: string[] = []; const unregisters: Array<() => void> = [];
afterEach(() => { vi.unstubAllEnvs(); while (unregisters.length > 0) { unregisters.pop()?.(); } while (cleanupPaths.length > 0) { const path = cleanupPaths.pop(); if (path && existsSync(path)) { rmSync(path, { recursive: true, force: true }); } } });
it("enables CLI login reuse only for default services storage", async () => { const tempDir = join(tmpdir(), `pi-default-services-auth-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); vi.stubEnv("HOME", tempDir); vi.stubEnv(ENV_AGENT_DIR, ""); const injected = AuthStorage.inMemory(); for (const options of [{}, { agentDir: join(tempDir, "custom") }, { authStorage: injected }]) { const services = await createAgentSessionServices({ cwd: tempDir, ...options, telemetryDisabled: true, resourceLoaderOptions: { noExtensions: true, noSkills: true, noPromptTemplates: true, noThemes: true, noContextFiles: true, }, }); expect(services.modelRegistry.authStorage).toBe(services.authStorage); expect(services.authStorage.getPrimeCliConfigPath()).toBe( "agentDir" in options || "authStorage" in options ? undefined : join(tempDir, ".prime", "config.json"), ); if ("authStorage" in options) expect(services.authStorage).toBe(injected); } });
it("shows the telemetry disclosure independently of the Herdr reporter", async () => { vi.stubEnv("DO_NOT_TRACK", "0"); vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const tempDir = join(tmpdir(), `pi-session-telemetry-notice-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); const settingsManager = SettingsManager.inMemory(); settingsManager.setOnboardingShown(true);
const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, settingsManager, noBuiltinHerdrReporter: true, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true }, });
expect(services.diagnostics).toContainEqual( expect.objectContaining({ type: "info", message: expect.stringContaining("pseudonymous usage") }), ); expect(settingsManager.getTelemetryNoticeShown()).toBe(true); });
it("defers the telemetry disclosure on a first interactive launch", async () => { vi.stubEnv("DO_NOT_TRACK", "0"); vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const tempDir = join(tmpdir(), `pi-session-telemetry-first-run-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); // A settings profile that has never seen onboarding: the notice would land // on the welcome screen, so it waits for the next launch. const settingsManager = SettingsManager.inMemory();
const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, settingsManager, noBuiltinHerdrReporter: true, deferTelemetryNoticeForOnboarding: true, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true }, });
expect(services.diagnostics).not.toContainEqual( expect.objectContaining({ message: expect.stringContaining("pseudonymous usage") }), ); expect(settingsManager.getTelemetryNoticeShown()).toBe(false); });
it("discloses telemetry immediately for sessions that never onboard", async () => { vi.stubEnv("DO_NOT_TRACK", "0"); vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const tempDir = join(tmpdir(), `pi-session-telemetry-headless-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); // No onboarding will run here, so holding the notice back would hide it forever. const settingsManager = SettingsManager.inMemory();
const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, settingsManager, noBuiltinHerdrReporter: true, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true }, });
expect(services.diagnostics).toContainEqual( expect.objectContaining({ message: expect.stringContaining("pseudonymous usage") }), ); expect(settingsManager.getTelemetryNoticeShown()).toBe(true); });
it("honors an explicit daemon-carried telemetry opt-out", async () => { vi.stubEnv("DO_NOT_TRACK", "0"); vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const tempDir = join(tmpdir(), `pi-session-daemon-telemetry-opt-out-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); const settingsManager = SettingsManager.inMemory(); const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, settingsManager, telemetryDisabled: true, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true }, });
expect(services.diagnostics).not.toContainEqual( expect.objectContaining({ message: expect.stringContaining("pseudonymous usage") }), ); expect(settingsManager.getTelemetryNoticeShown()).toBe(false);
const { session } = await createAgentSessionFromServices({ services, sessionManager: SessionManager.create(tempDir, join(tempDir, "sessions")), telemetryDisabled: true, }); try { expect(existsSync(join(tempDir, "telemetry.json"))).toBe(false); } finally { session.dispose(); } });
it("does not install top-level telemetry for a resumed child session", async () => { vi.stubEnv("DO_NOT_TRACK", "0"); vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const tempDir = join(tmpdir(), `pi-session-child-telemetry-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, settingsManager: SettingsManager.inMemory({ telemetry: { noticeShown: true } }), resourceLoaderOptions: { noPromptTemplates: true, noThemes: true }, }); const sessionManager = SessionManager.create(tempDir, join(tempDir, "sessions")); sessionManager.newSession({ rlmDepth: 1 });
const { session } = await createAgentSessionFromServices({ services, sessionManager }); try { expect(session.rlmDepth).toBe(1); expect(existsSync(join(tempDir, "telemetry.json"))).toBe(false); } finally { session.dispose(); } });
it("advertises enabled generic MCP servers and refreshes the prompt on reload", async () => { /** * The enabled-generic-servers line: user-declared servers plus per-account * records; a catalog entry appears only when it is explicitly public * no-auth AND setup-ready (credential-free dispatch — api_key and * requires-setup rows fail closed). The exact catalog contents evolve, so * the contract pins WHICH servers appear, not the full list. */ const enabledServersLine = (prompt: string): string => { const match = prompt.match(/Enabled generic MCP servers: ([^\n]*)\./); return match?.[1] ?? ""; };
const tempDir = join(tmpdir(), `pi-session-mcp-prompt-${Date.now()}-${Math.random().toString(36).slice(2)}`); const projectDir = join(tempDir, "project"); const agentDir = join(tempDir, "agent"); mkdirSync(join(projectDir, ".prime", "agent"), { recursive: true }); mkdirSync(agentDir, { recursive: true }); cleanupPaths.push(tempDir);
writeFileSync( join(agentDir, "settings.json"), JSON.stringify({ mcpServers: { zebra: { type: "http", url: "https://secret.example/mcp", headers: { Authorization: "secret" } }, filesystem: { type: "stdio", command: "/secret/bin/filesystem", args: ["/private/data"], cwd: "/secret/cwd", env: { TOKEN: { env: "FILESYSTEM_SECRET" } }, }, disabled: { type: "stdio", command: "disabled-secret", enabled: false }, linear: { type: "stdio", command: "reserved-secret" }, }, }), ); writeFileSync( join(projectDir, ".prime", "agent", "settings.json"), JSON.stringify({ mcpServers: { projectOnly: { type: "stdio", command: "project-secret" } } }), );
const settingsManager = SettingsManager.create(projectDir, agentDir); const services = await createAgentSessionServices({ cwd: projectDir, agentDir, settingsManager, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true }, }); const { session } = await createAgentSessionFromServices({ services, sessionManager: SessionManager.create(projectDir, join(tempDir, "sessions")), });
try { const initialPrompt = session.systemPrompt; expect(initialPrompt).toContain( "Generic MCP connections are accessed through the pre-imported Python `mcp` object in the Python REPL, not as top-level native tool namespaces or installed Python skills.", ); // Zero none+ready catalog rows exist today, so the enabled line is exactly the user-declared servers. expect(enabledServersLine(initialPrompt)).toBe("`filesystem`, `zebra`"); expect(initialPrompt).toContain('await mcp.list_tools("filesystem")'); expect(initialPrompt).toContain('await mcp.call_tool("filesystem", "<tool>", arguments)'); for (const hidden of [ "disabled", "projectOnly", "https://secret.example/mcp", "Authorization", "/secret/bin/filesystem", "/private/data", "/secret/cwd", "FILESYSTEM_SECRET", "reserved-secret", ]) { expect(initialPrompt).not.toContain(hidden); } expect(initialPrompt).not.toContain("Enabled generic MCP servers: `linear`");
const rebuildRuntime = vi.spyOn( session as unknown as { _rebuildRuntimeForAcpMcpServers(): void }, "_rebuildRuntimeForAcpMcpServers", ); session.replaceAcpMcpServers( [ { name: "task", type: "http", url: "https://task-secret.example/mcp", headers: { Authorization: "Bearer task-secret" }, }, ], "owner-a", ); expect(enabledServersLine(session.systemPrompt)).toBe("`filesystem`, `zebra`"); expect(session.systemPrompt).not.toContain('await mcp.list_tools("task")'); expect(session.getActiveToolNames()).toEqual(expect.arrayContaining(["mcp_list_tools_task", "mcp_call_task"])); expect(session.systemPrompt).not.toContain("task-secret"); rebuildRuntime.mockClear(); const waitForIdle = vi.spyOn(session.agent, "waitForIdle"); await session.releaseAcpMcpServers("unknown-owner", ["task"]); expect(waitForIdle).not.toHaveBeenCalled(); const originalProvisioner = Reflect.get(session, "_ipythonKernelProvisioner"); const execute = vi.fn(async (_code: string) => ({ status: "ok" })); Reflect.set(session, "_ipythonKernelProvisioner", { manager: { isRunning: true, execute } }); await session.releaseAcpMcpServers("owner-a", ["task"]); Reflect.set(session, "_ipythonKernelProvisioner", originalProvisioner); expect(rebuildRuntime).not.toHaveBeenCalled(); expect(execute).toHaveBeenCalledOnce(); expect(execute.mock.calls[0]?.[0]).toContain("await _prime_mcp.reload(_prime_mcp_name)"); expect(execute.mock.calls[0]?.[0]).toContain('["task"]'); expect(enabledServersLine(session.systemPrompt)).toBe("`filesystem`, `zebra`"); expect(session.getAllTools().map((tool) => tool.name)).not.toContain("mcp_call_task"); expect(session.getActiveToolNames()).not.toContain("mcp_call_task");
settingsManager.setGlobalMcpServer("added", { type: "stdio", command: "new-secret" }); settingsManager.removeGlobalMcpServer("filesystem"); await settingsManager.flush(); await session.reload();
expect(enabledServersLine(session.systemPrompt)).toBe("`added`, `zebra`"); expect(session.systemPrompt).toContain('await mcp.list_tools("added")'); expect(session.systemPrompt).not.toContain('await mcp.list_tools("filesystem")'); expect(session.systemPrompt).not.toContain("new-secret"); } finally { session.dispose(); } });
it("ENG-6108: activates a verified catalog connection in the live conversation after /plugins login", async () => { // A /plugins login must become usable in the current (daemon-backed) conversation without a restart: credentials // land in the shared credential store, the connection is verified with a real MCP handshake (token presence alone // is never Connected), and the session reload the picker triggers rebuilds the system prompt and serves the catalog // service through the generic mcp route. const tempDir = join( tmpdir(), `pi-session-catalog-activation-${Date.now()}-${Math.random().toString(36).slice(2)}`, ); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir); const faux = registerFauxProvider({ provider: "faux-eng6108" }); unregisters.push(() => faux.unregister()); const authStorage = AuthStorage.inMemory(); const store = McpConnectionStore.open(join(tempDir, "mcp-connections.json")); // Pin a minimal catalog: this exercises connect-then-activate, not the merged catalog's breadth. const mcpManager = new McpManager({ authStorage, connectionStore: store, getServiceCatalog: () => [ { serviceId: "notion", label: "Notion", aliases: [], transport: { type: "http", url: "https://mcp.notion.com/mcp" }, authStrategy: "oauth", setup: { status: "ready" }, metadataReviewed: true, legacyBuiltin: true, }, ], noBackgroundVerification: true, probeConnection: async () => ({ ok: true, toolCount: 3 }), }); const model = faux.getModel(); const modelRegistry = ModelRegistry.inMemory(authStorage); modelRegistry.registerProvider(model.provider, { baseUrl: model.baseUrl, apiKey: "faux-key", api: faux.api, models: faux.models.map((registeredModel) => ({ id: registeredModel.id, name: registeredModel.name, api: registeredModel.api, reasoning: registeredModel.reasoning, input: registeredModel.input, cost: registeredModel.cost, contextWindow: registeredModel.contextWindow, maxTokens: registeredModel.maxTokens, baseUrl: registeredModel.baseUrl, })), }); const services = await createAgentSessionServices({ cwd: tempDir, agentDir: join(tempDir, "agent"), authStorage, settingsManager: SettingsManager.inMemory(), modelRegistry, mcpManager, telemetryDisabled: true, noBuiltinHerdrReporter: true, resourceLoaderOptions: { noExtensions: true }, }); const { session } = await createAgentSessionFromServices({ services, sessionManager: SessionManager.inMemory(), model, }); try { expect(session.agent.state.systemPrompt).not.toContain("Enabled generic MCP servers"); expect(services.mcpManager.getDisabledBuiltinSkillOverrides()).toContain("-notion/SKILL.md");
// The /plugins connect flow, minus the UI: browser OAuth writes the // shared credential, then the host verifies with a real handshake. authStorage.set("mcp:notion", { type: "oauth", access: "notion-access", refresh: "notion-refresh", expires: Date.now() + 3600_000, endpoint: "https://mcp.notion.com/mcp", }); const record = await services.mcpManager.verifyConnection("notion"); expect(record.status).toBe("connected");
// The picker triggers the same session reload the interactive client performs: the prompt now // advertises the service and the authored skill override clears. await session.reload(); const promptAfter = session.agent.state.systemPrompt; expect(promptAfter).toContain("Enabled generic MCP servers"); expect(promptAfter).toContain('await mcp.list_tools("notion")'); expect(services.mcpManager.getDisabledBuiltinSkillOverrides()).not.toContain("-notion/SKILL.md"); } finally { session.dispose(); } });
it("forwards daemon-backed agent message controllers into AgentSession", async () => { const tempDir = join(tmpdir(), `pi-session-services-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir);
const faux = registerFauxProvider(); unregisters.push(() => faux.unregister());
const authStorage = AuthStorage.inMemory(); authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, authStorage, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true, skillsOverride: () => ({ skills: [ { name: AGENT_MESSAGE_SKILL_NAME, description: "hidden agent message skill", filePath: "<test:agent-message>", baseDir: tempDir, sourceInfo: createSyntheticSourceInfo("<test:agent-message>", { source: "test" }), disableModelInvocation: true, kind: "python" as const, python: { importName: "agent_message", packagePath: tempDir, pyprojectPath: join(tempDir, "pyproject.toml"), }, }, ], diagnostics: [], }), }, }); services.modelRegistry.registerProvider(faux.getModel().provider, { baseUrl: faux.getModel().baseUrl, apiKey: "faux-key", api: faux.api, models: faux.models, });
const agentMessageController: AgentSessionMessageController = { listAgents: () => ({ current: { activeSessionId: "current", sessionId: "session-current", runtimeKind: "top-level" }, agents: [ { activeSessionId: "worker", sessionId: "session-worker", runtimeKind: "top-level", cwd: tempDir, isStreaming: false, unfinishedActionCount: 0, }, ], }), sendAgentMessage: async () => { throw new Error("not used"); }, };
const { session } = await createAgentSessionFromServices({ services, sessionManager: SessionManager.create(tempDir, join(tempDir, "sessions")), model: faux.getModel(), agentMessageController, });
try { await expect(session.handleAgentMessageHostRequest("agent_message.list")).rejects.toThrow( "unknown agent message request", ); expect( ( session as unknown as { _createKernelHostHandlers(): Record<string, unknown>; } )._createKernelHostHandlers(), ).not.toHaveProperty("agent_message.send"); } finally { session.dispose(); } });
it("hides daemon-backed orchestration skills unless their host bridges are available", async () => { const tempDir = join(tmpdir(), `pi-session-skills-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); cleanupPaths.push(tempDir);
const authStorage = AuthStorage.inMemory(); const services = await createAgentSessionServices({ cwd: tempDir, agentDir: tempDir, authStorage, resourceLoaderOptions: { noPromptTemplates: true, noThemes: true, }, });
const createSession = async (options: Parameters<typeof createAgentSessionFromServices>[0]) => { const { session } = await createAgentSessionFromServices(options); return session; }; const visibleSkillNames = (session: unknown) => ( session as { _modelVisibleSkills(): Array<{ name: string }>; } ) ._modelVisibleSkills() .map((skill) => skill.name); const kernelHostHandlers = (session: unknown) => ( session as { _createKernelHostHandlers(): Record<string, unknown>; } )._createKernelHostHandlers();
const withoutControllers = await createSession({ services, sessionManager: SessionManager.create(tempDir, join(tempDir, "sessions-without")), }); try { expect(visibleSkillNames(withoutControllers)).not.toContain(AGENT_MESSAGE_SKILL_NAME); expect(visibleSkillNames(withoutControllers)).not.toContain(AGENT_OBSERVE_SKILL_NAME); } finally { withoutControllers.dispose(); }
const agentObserveController: AgentObserveController = { listAgents: () => ({ current: { activeSessionId: "current", sessionId: "session-current", runtimeKind: "top-level", cwd: tempDir, status: "idle", isCurrent: true, isStreaming: false, isCompacting: false, attachedClients: 1, messageCount: 0, queuedCount: 0, isSessionActive: false, }, agents: [], }), getAgent: () => { throw new Error("not used"); }, recentMessages: () => { throw new Error("not used"); }, }; const withControllers = await createSession({ services, sessionManager: SessionManager.create(tempDir, join(tempDir, "sessions-with")), agentObserveController, }); try { expect(visibleSkillNames(withControllers)).toContain(AGENT_OBSERVE_SKILL_NAME); expect(visibleSkillNames(withControllers)).not.toContain(AGENT_MESSAGE_SKILL_NAME); } finally { withControllers.dispose(); }
const agentMessageController: AgentSessionMessageController = { listAgents: () => ({ current: { activeSessionId: "current", sessionId: "session-current" }, agents: [], }), sendAgentMessage: async () => { throw new Error("not used"); }, }; const withMessageController = await createSession({ services, sessionManager: SessionManager.create(tempDir, join(tempDir, "sessions-with-message")), agentMessageController, }); try { expect(visibleSkillNames(withMessageController)).toContain(AGENT_MESSAGE_SKILL_NAME); expect(kernelHostHandlers(withMessageController)).toHaveProperty("agent_message.send"); } finally { withMessageController.dispose(); } });});