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.
8.4 kB · 274 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275import * as fs from "node:fs";import * as os from "node:os";import * as path from "node:path";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { AuthStorage } from "../src/core/auth-storage.js";import { discoverAndLoadExtensions } from "../src/core/extensions/loader.js";import { ExtensionRunner } from "../src/core/extensions/runner.js";import { ModelRegistry } from "../src/core/model-registry.js";import { SessionManager } from "../src/core/session-manager.js";
describe("Before Agent Message Send Event", () => { let tempDir: string; let extensionsDir: string;
beforeEach(() => { tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-agent-msg-ext-test-")); extensionsDir = path.join(tempDir, "extensions"); fs.mkdirSync(extensionsDir); });
afterEach(() => fs.rmSync(tempDir, { recursive: true, force: true }));
async function createRunner(...extensions: string[]) { fs.rmSync(extensionsDir, { recursive: true, force: true }); fs.mkdirSync(extensionsDir); for (let i = 0; i < extensions.length; i++) { fs.writeFileSync(path.join(extensionsDir, `e${i}.ts`), extensions[i]!); } const result = await discoverAndLoadExtensions([], tempDir, tempDir); const sm = SessionManager.inMemory(); const mr = ModelRegistry.create(AuthStorage.create(path.join(tempDir, "auth.json"))); return new ExtensionRunner(result.extensions, result.runtime, tempDir, sm, mr); }
it("returns undefined when no handlers registered", async () => { const runner = await createRunner(); const result = await runner.emitBeforeAgentMessageSend({ type: "before_agent_message_send", target: "parent-session-123", message: "hello parent", receiverRole: "parent", }); expect(result).toBeUndefined(); });
it("allows an extension to append context usage when sending to parent", async () => { const ext = `export default function(pi) { pi.on("before_agent_message_send", async (event, ctx) => { if (event.receiverRole === "parent") { const usage = ctx.getContextUsage(); const tokenText = usage ? \`[\${usage.tokens}/\${usage.contextWindow}]\` : "[no-usage]"; return { message: \`\${event.message} \${tokenText}\` }; } });}`; const runner = await createRunner(ext); runner.bindCore( {} as any, { getContextUsage: () => ({ tokens: 5000, contextWindow: 128000, percent: 3.90625 }), } as any, );
const result = await runner.emitBeforeAgentMessageSend({ type: "before_agent_message_send", target: "parent-1", message: "work finished", receiverRole: "parent", }); expect(result).toEqual({ message: "work finished [5000/128000]" }); });
it("does not modify message when receiverRole is not parent", async () => { const ext = `export default function(pi) { pi.on("before_agent_message_send", async (event, ctx) => { if (event.receiverRole === "parent") { return { message: event.message + " (to-parent)" }; } });}`; const runner = await createRunner(ext); const result = await runner.emitBeforeAgentMessageSend({ type: "before_agent_message_send", target: "sibling-1", message: "hey sibling", receiverRole: "sibling", receiverName: "sibling-1", }); expect(result).toBeUndefined(); });
it("chains modifications across multiple extensions", async () => { const ext1 = `export default function(pi) { pi.on("before_agent_message_send", async (event) => { return { message: event.message + " [first]" }; });}`; const ext2 = `export default function(pi) { pi.on("before_agent_message_send", async (event) => { return { message: event.message + " [second]" }; });}`; const runner = await createRunner(ext1, ext2); const result = await runner.emitBeforeAgentMessageSend({ type: "before_agent_message_send", target: "target-1", message: "base", }); expect(result).toEqual({ message: "base [first] [second]" }); });
it("cancels message delivery when handler returns cancel: true", async () => { const ext = `export default function(pi) { pi.on("before_agent_message_send", async (event) => { if (event.message.includes("secret")) { return { cancel: true, reason: "secrets cannot be sent" }; } });}`; const runner = await createRunner(ext); const result = await runner.emitBeforeAgentMessageSend({ type: "before_agent_message_send", target: "parent-1", message: "contains secret data", }); expect(result).toEqual({ cancel: true, reason: "secrets cannot be sent" }); });
it("handles extension errors gracefully without crashing", async () => { const failingExt = `export default function(pi) { pi.on("before_agent_message_send", async () => { throw new Error("handler exploded"); });}`; const okExt = `export default function(pi) { pi.on("before_agent_message_send", async (event) => { return { message: event.message + " ok" }; });}`; const runner = await createRunner(failingExt, okExt); const errors: any[] = []; runner.onError((e) => errors.push(e));
const result = await runner.emitBeforeAgentMessageSend({ type: "before_agent_message_send", target: "parent-1", message: "test", }); expect(result).toEqual({ message: "test ok" }); expect(errors.length).toBe(1); expect(errors[0].event).toBe("before_agent_message_send"); });});
import { createHarness, type Harness } from "./suite/harness.js";
describe("AgentSession agent message host request with extensions and context usage", () => { let harness: Harness | undefined;
afterEach(() => { if (harness) { harness.cleanup(); harness = undefined; } });
it("intercepts agent_message.send and appends context usage via extension", async () => { let sentMessage = ""; let sentRole: any; const mockController = { listAgents: () => ({ parent: null, siblings: [], children: [] }), roster: () => ({ entries: [ { id: "parent-id", name: "parent", relationship: "parent" as const, depth: 0, status: "running" as const, }, ], }), sendAgentMessage: async (input: any) => { sentMessage = input.message; sentRole = input.receiverRole; return { id: "msg-1", source: "agent_message" as const, target: { activeSessionId: input.target, sessionId: "sess-1" }, message: input.message, deliveryStatus: "delivered" as const, }; }, };
const extFactory = (pi: any) => { pi.on("before_agent_message_send", async (event: any, ctx: any) => { if (event.receiverRole === "parent") { const usage = ctx.getContextUsage(); const tokenPart = usage ? `[tokens: ${usage.tokens}/${usage.contextWindow}]` : "[no-usage]"; return { message: `${event.message} ${tokenPart}` }; } }); };
harness = await createHarness({ agentMessageController: mockController as any, extensionFactories: [extFactory], });
await harness.session.handleAgentMessageHostRequest("agent_message.send", { target: "parent-id", message: "task complete", receiverRole: "parent", });
expect(sentMessage).toContain("task complete [tokens:"); expect(sentRole).toBe("parent"); });
it("throws when extension cancels agent_message.send", async () => { const mockController = { listAgents: () => ({ parent: null, siblings: [], children: [] }), sendAgentMessage: async () => { throw new Error("should not be called"); }, };
const extFactory = (pi: any) => { pi.on("before_agent_message_send", async () => { return { cancel: true, reason: "blocked by policy" }; }); };
harness = await createHarness({ agentMessageController: mockController as any, extensionFactories: [extFactory], });
await expect( harness.session.handleAgentMessageHostRequest("agent_message.send", { target: "parent-id", message: "forbidden text", receiverRole: "parent", }), ).rejects.toThrow("blocked by policy"); });
it("exposes rlm.context_usage in kernel host handlers", async () => { harness = await createHarness(); const handlers = (harness.session as any)._createKernelHostHandlers(); expect(handlers["rlm.context_usage"]).toBeDefined(); expect(handlers["rlm.get_context_usage"]).toBeDefined();
const res = await handlers["rlm.context_usage"](); expect(res).toBeDefined(); expect(res.usage).toBeDefined(); expect(res.usage.context_window).toBeGreaterThan(0); });});