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.
9.6 kB · 282 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283import { mkdirSync, rmSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { getBundledSkillsDir } from "../src/config.js";import { type KernelSentAgentMessage, ReplKernelManager } from "../src/core/kernel/index.js";import type { PythonSkillRuntimeInfo } from "../src/core/skills.js";import { IpythonKernelProvisioner } from "../src/core/tools/ipython.js";
function bundledAgentMessageSkill(): PythonSkillRuntimeInfo { const packagePath = join(getBundledSkillsDir(), "agent-message"); return { name: "agent-message", importName: "agent_message", packagePath, pyprojectPath: join(packagePath, "pyproject.toml"), };}
type LateHandlerRetentionHost = { lateSentAgentMessageHandlers: Map<string, (message: KernelSentAgentMessage) => void>; registerLateSentAgentMessageHandler: ( requestMessageId: string, handler: (message: KernelSentAgentMessage) => void, ) => void;};
describe("agent-message skill over the kernel host bridge", () => { let tempDir: string; let provisioner: IpythonKernelProvisioner | undefined;
beforeEach(() => { tempDir = join(tmpdir(), `pi-agent-message-skill-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); });
afterEach(async () => { await provisioner?.dispose(); provisioner = undefined; rmSync(tempDir, { recursive: true, force: true }); });
it("lists agents and sends without exposing a spoofable sender", async () => { const requests: Array<{ type: string; payload: Record<string, unknown> }> = []; provisioner = new IpythonKernelProvisioner(tempDir, { pythonSkills: [bundledAgentMessageSkill()], hostHandlers: { "agent_message.list_agents": async (payload) => { requests.push({ type: "agent_message.list_agents", payload }); return { current: { name: "alpha", id: "session-alpha", depth: 0 }, entries: [{ relationship: "sibling", name: "Beta", id: "session-beta", depth: 0, status: "idle" }], }; }, "agent_message.send": async (payload) => { requests.push({ type: "agent_message.send", payload }); return { id: "agentmsg-test", source: "agent_message", target: { activeSessionId: payload.receiver_name, sessionId: "session-beta", sessionName: "Beta" }, from: { activeSessionId: "alpha", sessionId: "session-alpha" }, message: payload.message, deliveryStatus: "queued", queuedAt: "2026-06-16T00:00:00.000Z", }; }, }, });
const manager = await provisioner.ensure(); const result = await manager.execute(`import jsonagents = await agent_message.list_agents()receipt = await agent_message.send( "hello beta", receiver_role="sibling", receiver_name="beta")print(json.dumps({"agents": agents, "receipt": receipt}, sort_keys=True))`);
expect(result.status).toBe("ok"); const output = JSON.parse(result.stdout.trim()); expect(output.agents).toMatchObject({ current: { id: "session-alpha", depth: 0 }, entries: [{ relationship: "sibling", id: "session-beta", status: "idle" }], }); expect(output.receipt).toMatchObject({ id: "agentmsg-test", source: "agent_message", message: "hello beta", deliveryStatus: "queued", }); expect(result.sentAgentMessages).toEqual([ { id: "agentmsg-test", message: "hello beta", deliveryStatus: "queued", receiverRole: "sibling", target: { activeSessionId: "beta", sessionId: "session-beta", sessionName: "Beta" }, }, ]); expect(requests[0]).toMatchObject({ type: "agent_message.list_agents", payload: { type: "agent_message.list_agents" }, }); expect(requests[1]).toMatchObject({ type: "agent_message.send", payload: { type: "agent_message.send", message: "hello beta", receiver_role: "sibling", receiver_name: "beta", }, }); expect(requests[1].payload).not.toHaveProperty("from"); });
it("emits successful broadcast receipts and leaves short errors in the result", async () => { provisioner = new IpythonKernelProvisioner(tempDir, { pythonSkills: [bundledAgentMessageSkill()], hostHandlers: { "agent_message.send": async (payload) => ({ receipts: [ { id: "agentmsg-root", source: "agent_message", target: { activeSessionId: "root", sessionId: "session-root" }, message: payload.message, deliveryStatus: "delivered", deliveredAt: "2026-08-03T00:00:00.000Z", deliveryMode: payload.mode, }, { target: "sibling", error: "rate limited" }, ], }), }, });
const manager = await provisioner.ensure(); const result = await manager.execute(`import jsonreceipt = await agent_message.send("all", "status")print(json.dumps(receipt, sort_keys=True))`);
expect(result.status).toBe("ok"); expect(JSON.parse(result.stdout.trim())).toMatchObject({ receipts: [ { id: "agentmsg-root", deliveryStatus: "delivered" }, { target: "sibling", error: "rate limited" }, ], }); expect(result.sentAgentMessages).toEqual([ { id: "agentmsg-root", message: "status", deliveryStatus: "delivered", target: { activeSessionId: "root", sessionId: "session-root" }, }, ]); });
it("rejects broadcast combined with role selectors before reaching the host", async () => { provisioner = new IpythonKernelProvisioner(tempDir, { pythonSkills: [bundledAgentMessageSkill()], hostHandlers: { "agent_message.send": async () => { throw new Error("should not reach host"); }, }, });
const manager = await provisioner.ensure(); const result = await manager.execute(`try: await agent_message.send("all", "secret", receiver_role="sibling", receiver_name="beta")except TypeError as error: print(f"TypeError: {error}")`); expect(result.status).toBe("ok"); expect(result.stdout.trim()).toBe("TypeError: broadcast cannot be combined with receiver_role/receiver_name"); });
it("rejects a positional name target before reaching the host", async () => { provisioner = new IpythonKernelProvisioner(tempDir, { pythonSkills: [bundledAgentMessageSkill()], hostHandlers: { "agent_message.send": async () => { throw new Error("should not reach host"); }, }, });
const manager = await provisioner.ensure(); const result = await manager.execute(`try: await agent_message.send("beta", "done")except TypeError as error: print(f"TypeError: {error}")`); expect(result.status).toBe("ok"); expect(result.stdout.trim()).toBe( "TypeError: positional agent_message.send targets are not supported; use receiver_role and receiver_name", ); });
it("does not expose a queueable delivery mode", async () => { provisioner = new IpythonKernelProvisioner(tempDir, { pythonSkills: [bundledAgentMessageSkill()], hostHandlers: { "agent_message.send": async () => { throw new Error("should not reach host"); }, }, });
const manager = await provisioner.ensure(); const result = await manager.execute(`try: await agent_message.send("hello", receiver_role="sibling", receiver_name="beta", mode="broadcast")except TypeError as error: print(f"TypeError: {error}")`); expect(result.status).toBe("ok"); expect(result.stdout.trim()).toContain("TypeError: send() got an unexpected keyword argument 'mode'"); });
it("captures sent messages from detached tasks after the cell is idle", async () => { provisioner = new IpythonKernelProvisioner(tempDir, { pythonSkills: [bundledAgentMessageSkill()], hostHandlers: { "agent_message.send": async (payload) => ({ id: "agentmsg-background", source: "agent_message", target: { activeSessionId: payload.receiver_name, sessionId: "session-beta", sessionName: "Beta" }, message: payload.message, deliveryStatus: "delivered", deliveredAt: "2026-07-10T00:00:00.000Z", deliveryMode: payload.mode, }), }, });
const manager = await provisioner.ensure(); let resolveLateMessage!: (message: KernelSentAgentMessage) => void; const lateMessage = new Promise<KernelSentAgentMessage>((resolve) => { resolveLateMessage = resolve; }); const result = await manager.execute( `async def send_later(): await asyncio.sleep(0.05) await agent_message.send("background hello", receiver_role="sibling", receiver_name="beta")
background_send = asyncio.create_task(send_later())`, { onLateSentAgentMessage: (message) => resolveLateMessage(message) }, );
expect(result.status).toBe("ok"); expect(result.sentAgentMessages).toBeUndefined(); await expect(lateMessage).resolves.toEqual({ id: "agentmsg-background", message: "background hello", deliveryStatus: "delivered", receiverRole: "sibling", target: { activeSessionId: "beta", sessionId: "session-beta", sessionName: "Beta" }, }); });
it("bounds retained handlers for late sent messages", async () => { const manager = new ReplKernelManager({ cwd: tempDir }); const host = manager as unknown as LateHandlerRetentionHost; const handler = () => {};
for (let index = 0; index < 300; index += 1) { host.registerLateSentAgentMessageHandler(`request-${index}`, handler); }
expect(host.lateSentAgentMessageHandlers.size).toBe(256); expect(host.lateSentAgentMessageHandlers.has("request-0")).toBe(false); expect(host.lateSentAgentMessageHandlers.has("request-299")).toBe(true); await manager.shutdown({ snapshot: true, drainHostRequests: true }); });});