Spaces:
Paused
Paused
| import { randomUUID } from "node:crypto"; | |
| import { ObjectId } from "mongodb"; | |
| import { afterEach, beforeAll, describe, expect, it, vi } from "vitest"; | |
| import { collections, ready } from "$lib/server/database"; | |
| import type { Generation } from "$lib/types/Generation"; | |
| import { MessageUpdateType } from "$lib/types/MessageUpdate"; | |
| import type { Workspace } from "$lib/types/Workspace"; | |
| import { enqueueAgentControl, type AgentControlTarget } from "./AgentControlMailbox"; | |
| import { AgentControlPump } from "./AgentControlPump"; | |
| import { AgentGenerationStore, AgentRelayLeaseLostError } from "./AgentGenerationStore"; | |
| import type { AgentSessionManager } from "./AgentSessionManager"; | |
| import { PiRpcCommandRejectedError } from "./PiRpcClient"; | |
| beforeAll(async () => { | |
| await ready; | |
| }); | |
| afterEach(async () => { | |
| await Promise.all([ | |
| collections.conversations.deleteMany({ sessionId: /^agent-pump-/ }), | |
| collections.workspaces.deleteMany({ sessionId: /^agent-pump-/ }), | |
| collections.generations.deleteMany({ sessionId: /^agent-pump-/ }), | |
| ]); | |
| }); | |
| describe.sequential("AgentControlPump", () => { | |
| it("dispatches, receipts, and seals a command before terminalization", async () => { | |
| const seeded = await seedPumpRun(); | |
| const queued = await enqueueAgentControl({ | |
| target: seeded.target, | |
| id: randomUUID(), | |
| payload: { kind: "steer", message: "take the safe path" }, | |
| }); | |
| expect(queued.status).toBe("queued"); | |
| const commands: Record<string, unknown>[] = []; | |
| const manager = fakeManager({ | |
| command: async (command) => { | |
| commands.push(command); | |
| return { type: "response", success: true }; | |
| }, | |
| }); | |
| const pump = new AgentControlPump({ | |
| store: seeded.store, | |
| manager, | |
| workspace: seeded.workspace, | |
| pollIntervalMs: 1, | |
| }); | |
| await pump.sealAndStop(); | |
| expect(commands).toHaveLength(1); | |
| expect(commands[0]).toMatchObject({ | |
| type: "steer", | |
| _chatUiExpectedPromptId: seeded.promptId, | |
| }); | |
| const generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation?.agent?.control).toBeUndefined(); | |
| expect(generation?.agent?.controlsOpen).toBe(false); | |
| expect(generation?.agent?.controlReceipts).toHaveLength(1); | |
| expect(await seeded.store.finish("completed")).toBe(true); | |
| }); | |
| it("persists always policy only after the guarded bridge acknowledgement", async () => { | |
| const seeded = await seedPumpRun({ approvalRequestId: "approval-1", policy: "ask" }); | |
| await enqueueAgentControl({ | |
| target: seeded.target, | |
| id: randomUUID(), | |
| payload: { | |
| kind: "approval", | |
| requestId: "approval-1", | |
| response: { confirmed: true, always: true }, | |
| }, | |
| }); | |
| const order: string[] = []; | |
| const manager = fakeManager({ | |
| command: async () => { | |
| order.push("ack"); | |
| return { type: "response", success: true }; | |
| }, | |
| writeFile: async () => { | |
| expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe( | |
| "auto" | |
| ); | |
| order.push("policy-file"); | |
| }, | |
| }); | |
| const pump = new AgentControlPump({ | |
| store: seeded.store, | |
| manager, | |
| workspace: seeded.workspace, | |
| pollIntervalMs: 1, | |
| }); | |
| await pump.sealAndStop(); | |
| expect(order).toEqual(["ack", "policy-file"]); | |
| expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe( | |
| "auto" | |
| ); | |
| }); | |
| it("commits durable auto before a runtime-file write can outlive the exact relay lease", async () => { | |
| const seeded = await seedPumpRun({ approvalRequestId: "approval-race", policy: "ask" }); | |
| await enqueueAgentControl({ | |
| target: seeded.target, | |
| id: randomUUID(), | |
| payload: { | |
| kind: "approval", | |
| requestId: "approval-race", | |
| response: { confirmed: true, always: true }, | |
| }, | |
| }); | |
| let policyAtRuntimeWrite: Workspace["policy"] | undefined; | |
| let reportFailure!: (cause: Error) => void; | |
| const failed = new Promise<Error>((resolve) => (reportFailure = resolve)); | |
| const pump = new AgentControlPump({ | |
| store: seeded.store, | |
| manager: fakeManager({ | |
| command: async () => ({ type: "response", success: true }), | |
| writeFile: async () => { | |
| policyAtRuntimeWrite = ( | |
| await collections.workspaces.findOne({ _id: seeded.workspace._id }) | |
| )?.policy; | |
| const expired = new Date(0); | |
| await Promise.all([ | |
| collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { $set: { "execution.expiresAt": expired } } | |
| ), | |
| collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $set: { "agent.relay.expiresAt": expired } } | |
| ), | |
| ]); | |
| }, | |
| }), | |
| workspace: seeded.workspace, | |
| pollIntervalMs: 1, | |
| onFailure: reportFailure, | |
| }); | |
| expect(await failed).toBeInstanceOf(AgentRelayLeaseLostError); | |
| await expect(pump.waitForStop()).rejects.toBeInstanceOf(AgentRelayLeaseLostError); | |
| expect(policyAtRuntimeWrite).toBe("auto"); | |
| expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe( | |
| "auto" | |
| ); | |
| }); | |
| it("receipts semantic rejection but leaves ambiguous transport failure for a replacement", async () => { | |
| const rejected = await seedPumpRun(); | |
| await enqueueAgentControl({ | |
| target: rejected.target, | |
| id: randomUUID(), | |
| payload: { kind: "compact" }, | |
| }); | |
| const rejectedPump = new AgentControlPump({ | |
| store: rejected.store, | |
| manager: fakeManager({ | |
| command: async () => { | |
| throw new PiRpcCommandRejectedError({ | |
| type: "response", | |
| success: false, | |
| error: "Agent control no longer belongs to the active prompt", | |
| }); | |
| }, | |
| }), | |
| workspace: rejected.workspace, | |
| pollIntervalMs: 1, | |
| }); | |
| await rejectedPump.sealAndStop(); | |
| const rejectedGeneration = await collections.generations.findOne({ | |
| generationId: rejected.generationId, | |
| }); | |
| expect(rejectedGeneration?.agent?.controlReceipts?.[0]?.outcome).toBe("rejected"); | |
| const ambiguous = await seedPumpRun(); | |
| await enqueueAgentControl({ | |
| target: ambiguous.target, | |
| id: randomUUID(), | |
| payload: { kind: "steer", message: "may have crossed the socket" }, | |
| }); | |
| let reportFailure!: (cause: Error) => void; | |
| const failed = new Promise<Error>((resolve) => (reportFailure = resolve)); | |
| const ambiguousPump = new AgentControlPump({ | |
| store: ambiguous.store, | |
| manager: fakeManager({ | |
| command: async () => { | |
| throw new Error("socket closed after write"); | |
| }, | |
| }), | |
| workspace: ambiguous.workspace, | |
| pollIntervalMs: 1, | |
| onFailure: reportFailure, | |
| }); | |
| expect((await failed).message).toContain("socket closed"); | |
| await expect(ambiguousPump.waitForStop()).rejects.toThrow("socket closed"); | |
| const ambiguousGeneration = await collections.generations.findOne({ | |
| generationId: ambiguous.generationId, | |
| }); | |
| expect(ambiguousGeneration?.agent?.control?.state).toBe("dispatching"); | |
| expect(ambiguousGeneration?.agent?.controlsOpen).not.toBe(false); | |
| expect(ambiguousGeneration?.agent?.controlReceipts).toBeUndefined(); | |
| }); | |
| it("never persists always policy when Pi says the native approval is no longer pending", async () => { | |
| const seeded = await seedPumpRun({ approvalRequestId: "approval-expired", policy: "ask" }); | |
| await enqueueAgentControl({ | |
| target: seeded.target, | |
| id: randomUUID(), | |
| payload: { | |
| kind: "approval", | |
| requestId: "approval-expired", | |
| response: { confirmed: true, always: true }, | |
| }, | |
| }); | |
| const writeFile = vi.fn(async () => undefined); | |
| const pump = new AgentControlPump({ | |
| store: seeded.store, | |
| manager: fakeManager({ | |
| command: async (command) => { | |
| expect(command).toMatchObject({ | |
| type: "extension_ui_response", | |
| requestId: "approval-expired", | |
| }); | |
| throw new PiRpcCommandRejectedError({ | |
| type: "response", | |
| success: false, | |
| error: "Extension UI request is no longer pending", | |
| }); | |
| }, | |
| writeFile, | |
| }), | |
| workspace: seeded.workspace, | |
| pollIntervalMs: 1, | |
| }); | |
| await pump.sealAndStop(); | |
| expect(writeFile).not.toHaveBeenCalled(); | |
| expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe( | |
| "ask" | |
| ); | |
| expect( | |
| (await collections.generations.findOne({ generationId: seeded.generationId }))?.agent | |
| ?.controlReceipts?.[0]?.outcome | |
| ).toBe("rejected"); | |
| }); | |
| }); | |
| function fakeManager(options: { | |
| command: (command: Record<string, unknown> & { type: string }) => Promise<unknown>; | |
| writeFile?: () => Promise<void>; | |
| }): AgentSessionManager { | |
| return { | |
| rpc: { command: options.command }, | |
| sandbox: { writeFile: options.writeFile ?? (async () => undefined) }, | |
| } as unknown as AgentSessionManager; | |
| } | |
| async function seedPumpRun( | |
| options: { approvalRequestId?: string; policy?: Workspace["policy"] } = {} | |
| ) { | |
| const now = new Date(); | |
| const conversationId = new ObjectId(); | |
| const workspaceId = new ObjectId(); | |
| const generationId = randomUUID(); | |
| const messageId = randomUUID(); | |
| const promptId = randomUUID(); | |
| const sessionId = `agent-pump-${randomUUID()}`; | |
| const sandboxId = `sandbox-${randomUUID()}`; | |
| const sandboxGeneration = `sandbox-generation-${randomUUID()}`; | |
| const workspace: Workspace = { | |
| _id: workspaceId, | |
| conversationId, | |
| sessionId, | |
| repo: { kind: "blank" }, | |
| flavor: "cpu-basic", | |
| model: "fake/model", | |
| thinkingLevel: "medium", | |
| policy: options.policy ?? "auto", | |
| sandbox: { | |
| status: "ready", | |
| id: sandboxId, | |
| generation: sandboxGeneration, | |
| namespace: "tester", | |
| baseUrl: `https://${sandboxId}.example.test`, | |
| nonce: "nonce", | |
| image: "node:22", | |
| flavor: "cpu-basic", | |
| createdAt: now, | |
| }, | |
| spend: { accruedUsd: 0 }, | |
| createdAt: now, | |
| updatedAt: now, | |
| }; | |
| await collections.conversations.insertOne({ | |
| _id: conversationId, | |
| sessionId, | |
| mode: "agent", | |
| workspaceId, | |
| messages: [ | |
| { | |
| id: messageId, | |
| from: "assistant", | |
| content: "", | |
| generationId, | |
| updates: options.approvalRequestId | |
| ? [ | |
| { | |
| type: MessageUpdateType.AgentApproval, | |
| requestId: options.approvalRequestId, | |
| method: "confirm", | |
| title: "Allow command?", | |
| }, | |
| ] | |
| : [], | |
| createdAt: now, | |
| updatedAt: now, | |
| }, | |
| ], | |
| createdAt: now, | |
| updatedAt: now, | |
| } as never); | |
| await collections.workspaces.insertOne(workspace); | |
| await collections.generations.insertOne({ | |
| _id: new ObjectId(), | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| messageId, | |
| sessionId, | |
| mode: "agent", | |
| status: "running", | |
| seq: 0, | |
| lastHeartbeatAt: now, | |
| startedAt: now, | |
| createdAt: now, | |
| updatedAt: now, | |
| agent: { | |
| sourceSequence: 0, | |
| controlsOpen: true, | |
| prompt: { id: promptId, message: "do work" }, | |
| ...(options.approvalRequestId ? { pendingApprovalIds: [options.approvalRequestId] } : {}), | |
| }, | |
| } satisfies Generation); | |
| const store = await AgentGenerationStore.acquire({ | |
| generationId, | |
| workspaceId, | |
| sandboxId, | |
| sandboxGeneration, | |
| sandboxNonce: workspace.sandbox.nonce, | |
| }); | |
| if (!store) throw new Error("Expected Agent relay lease"); | |
| const target: AgentControlTarget = { | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| messageId, | |
| owner: { sessionId }, | |
| }; | |
| return { generationId, promptId, workspace, store, target }; | |
| } | |