Spaces:
Paused
Paused
| import { afterEach, beforeAll, describe, expect, it, vi } from "vitest"; | |
| import { ObjectId } from "mongodb"; | |
| import { randomUUID } from "node:crypto"; | |
| import { collections, ready } from "$lib/server/database"; | |
| import { | |
| MessageUpdateStatus, | |
| MessageUpdateType, | |
| type MessageUpdate, | |
| } from "$lib/types/MessageUpdate"; | |
| import type { AgentPendingBatch, Generation } from "$lib/types/Generation"; | |
| import type { Workspace } from "$lib/types/Workspace"; | |
| import { | |
| claimWorkspaceOperation, | |
| releaseWorkspaceOperation, | |
| WORKSPACE_OPERATION_LEASE_MS, | |
| } from "./WorkspaceOperation"; | |
| import { recordGenerationStop } from "$lib/server/generation/stopGeneration"; | |
| import { | |
| AgentGenerationIncarnationMismatchError, | |
| AgentGenerationStore, | |
| AgentRelayLeaseLostError, | |
| } from "./AgentGenerationStore"; | |
| import { reapStaleGenerations } from "$lib/server/generation/reaper"; | |
| beforeAll(async () => { | |
| await ready; | |
| }); | |
| afterEach(async () => { | |
| vi.restoreAllMocks(); | |
| await Promise.all([ | |
| collections.conversations.deleteMany({}), | |
| collections.workspaces.deleteMany({}), | |
| collections.generations.deleteMany({}), | |
| collections.generationEvents.deleteMany({}), | |
| collections.abortedGenerations.deleteMany({}), | |
| ]); | |
| }); | |
| describe.sequential("AgentGenerationStore", () => { | |
| it("commits a replayed bridge sequence exactly once", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| const update: MessageUpdate = { type: MessageUpdateType.Stream, token: "once" }; | |
| const bound = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(bound?.agent).toMatchObject(sandboxIncarnation(seeded.workspace)); | |
| const first = await store.commitBridgeBatch(41, [update]); | |
| const replay = await store.commitBridgeBatch(41, [update]); | |
| expect(first.inserted).toBe(true); | |
| expect(replay.inserted).toBe(false); | |
| expect( | |
| await collections.generationEvents.countDocuments({ generationId: seeded.generationId }) | |
| ).toBe(1); | |
| const conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| const message = conversation?.messages.find((item) => item.id === seeded.messageId); | |
| expect(message?.content).toBe("once"); | |
| expect(message?.agentSequence).toBe(41); | |
| expect(message?.materializedSeq).toBe(1); | |
| expect( | |
| (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.piSession?.lastSequence | |
| ).toBe(41); | |
| }); | |
| it("binds once and refuses to acquire the same prompt for a replacement sandbox", async () => { | |
| const seeded = await seedAgentRun(); | |
| const original = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await original.release(); | |
| await expect( | |
| AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| sandboxId: "sandbox-replacement", | |
| sandboxGeneration: "generation-replacement", | |
| }) | |
| ).rejects.toBeInstanceOf(AgentGenerationIncarnationMismatchError); | |
| const generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation?.agent).toMatchObject(sandboxIncarnation(seeded.workspace)); | |
| expect(generation?.agent?.relay).toBeUndefined(); | |
| expect(generation?.status).toBe("running"); | |
| }); | |
| it("continues after a durable provisioning event even when its heartbeat cursor lagged", async () => { | |
| const seeded = await seedAgentRun(); | |
| await collections.generationEvents.insertOne({ | |
| _id: new ObjectId(), | |
| generationId: seeded.generationId, | |
| seq: 1, | |
| event: { type: MessageUpdateType.AgentSandbox, state: "ready" }, | |
| createdAt: new Date(), | |
| }); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| const committed = await store.commitBridgeBatch(1, [ | |
| { type: MessageUpdateType.Stream, token: "after provisioning" }, | |
| ]); | |
| expect(committed.endSeq).toBe(2); | |
| expect( | |
| ( | |
| await collections.generationEvents | |
| .find({ generationId: seeded.generationId }) | |
| .sort({ seq: 1 }) | |
| .toArray() | |
| ).map((event) => event.event.type) | |
| ).toEqual([MessageUpdateType.AgentSandbox, MessageUpdateType.Stream]); | |
| const conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| const message = conversation?.messages.find((item) => item.id === seeded.messageId); | |
| expect(message?.materializedSeq).toBe(2); | |
| expect(message?.updates?.map((event) => event.type)).toEqual([ | |
| MessageUpdateType.AgentSandbox, | |
| MessageUpdateType.Stream, | |
| ]); | |
| }); | |
| it("tracks native approval requests as live generation state and clears them at terminal", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await store.commitBridgeBatch(1, [ | |
| { | |
| type: MessageUpdateType.AgentApproval, | |
| requestId: "approval-live", | |
| method: "confirm", | |
| title: "Allow command?", | |
| }, | |
| ]); | |
| expect( | |
| (await collections.generations.findOne({ generationId: seeded.generationId }))?.agent | |
| ?.pendingApprovalIds | |
| ).toEqual(["approval-live"]); | |
| expect(await store.finish("completed")).toBe(true); | |
| expect( | |
| (await collections.generations.findOne({ generationId: seeded.generationId }))?.agent | |
| ?.pendingApprovalIds | |
| ).toBeUndefined(); | |
| }); | |
| it("keeps one exact relay authoritative through every login ownership handoff phase", async () => { | |
| const seeded = await seedAgentRun(); | |
| const userId = new ObjectId(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await collections.workspaces.updateOne({ _id: seeded.workspace._id }, { $set: { userId } }); | |
| expect(await store.authoritativeOwner()).toEqual({ sessionId: seeded.workspace.sessionId }); | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $set: { userId } } | |
| ); | |
| expect(await store.authoritativeOwner()).toEqual({ sessionId: seeded.workspace.sessionId }); | |
| await collections.conversations.updateOne( | |
| { _id: seeded.conversationId }, | |
| { $set: { userId }, $unset: { sessionId: "" } } | |
| ); | |
| expect(await store.authoritativeOwner()).toEqual({ userId }); | |
| await collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { $unset: { sessionId: "" } } | |
| ); | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $unset: { sessionId: "" } } | |
| ); | |
| await store.assertPromptDispatch(seeded.workspace); | |
| expect(await store.authoritativeOwner()).toEqual({ userId }); | |
| await store.release(); | |
| }); | |
| it("cannot complete after an exact stop intent wins the relay terminal barrier", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await store.commitBridgeBatch(1, [ | |
| { type: MessageUpdateType.Stream, token: "stopped-late-output" }, | |
| ]); | |
| expect( | |
| await recordGenerationStop({ | |
| conversationId: seeded.conversationId, | |
| generationId: seeded.generationId, | |
| seenContentLength: 4, | |
| }) | |
| ).toBe(true); | |
| expect(await store.finish("completed")).toBe(true); | |
| expect( | |
| await collections.generations.findOne({ generationId: seeded.generationId }) | |
| ).toMatchObject({ | |
| status: "interrupted", | |
| stopRequested: { seenContentLength: 4 }, | |
| }); | |
| const message = (await collections.conversations.findOne({ _id: seeded.conversationId })) | |
| ?.messages[0]; | |
| expect(message?.content).toBe("stop"); | |
| expect(message?.interrupted).toBe(true); | |
| }); | |
| it("finishes a crash-partial write-ahead batch without gaps or duplication", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| const events: MessageUpdate[] = [ | |
| { type: MessageUpdateType.Stream, token: "alpha " }, | |
| { type: MessageUpdateType.Stream, token: "beta" }, | |
| ]; | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { | |
| $set: { | |
| seq: 2, | |
| "agent.pending": { | |
| kind: "bridge", | |
| sourceSequence: 7, | |
| startSeq: 1, | |
| events, | |
| }, | |
| }, | |
| } | |
| ); | |
| await collections.generationEvents.insertOne({ | |
| _id: new ObjectId(), | |
| generationId: seeded.generationId, | |
| seq: 1, | |
| event: events[0], | |
| sourceSequence: 7, | |
| sourceIndex: 0, | |
| createdAt: new Date(), | |
| }); | |
| await store.resumePending(); | |
| const persisted = await collections.generationEvents | |
| .find({ generationId: seeded.generationId }) | |
| .sort({ seq: 1 }) | |
| .toArray(); | |
| expect(persisted.map((item) => item.seq)).toEqual([1, 2]); | |
| expect(persisted.map((item) => item.sourceIndex)).toEqual([0, 1]); | |
| const conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| expect(conversation?.messages.find((item) => item.id === seeded.messageId)?.content).toBe( | |
| "alpha beta" | |
| ); | |
| const generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation?.agent?.pending).toBeUndefined(); | |
| expect(generation?.agent?.sourceSequence).toBe(7); | |
| }); | |
| it("keeps the WAL pending when a relay crashes after message materialization", async () => { | |
| const seeded = await seedAgentRun(); | |
| let timestamp = 0; | |
| const now = () => new Date(timestamp); | |
| const first = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| owner: "crashing-relay", | |
| leaseMs: 10, | |
| now, | |
| }) | |
| ); | |
| timestamp = 5; | |
| const restoreMaterialization = failAfterNextPendingMaterialization(first); | |
| try { | |
| await expect( | |
| first.commitBridgeBatch(1, [ | |
| { type: MessageUpdateType.Stream, token: "materialized exactly once" }, | |
| ]) | |
| ).rejects.toThrow("simulated process crash after Agent WAL materialization"); | |
| } finally { | |
| restoreMaterialization(); | |
| } | |
| let generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation?.agent?.pending).toMatchObject({ | |
| kind: "bridge", | |
| sourceSequence: 1, | |
| }); | |
| expect(generation?.agent?.sourceSequence).toBe(0); | |
| expect(generation?.agent?.relay?.expiresAt).toEqual(new Date(15)); | |
| let workspace = await collections.workspaces.findOne({ _id: seeded.workspace._id }); | |
| expect(workspace?.execution?.expiresAt).toEqual(new Date(15)); | |
| let conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| expect(conversation?.messages[0]?.content).toBe("materialized exactly once"); | |
| await first.release(); | |
| const replacement = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| owner: "replacement-relay", | |
| leaseMs: 10, | |
| now, | |
| }) | |
| ); | |
| await replacement.resumePending(); | |
| generation = await collections.generations.findOne({ generationId: seeded.generationId }); | |
| expect(generation?.agent?.pending).toBeUndefined(); | |
| expect(generation?.agent?.sourceSequence).toBe(1); | |
| workspace = await collections.workspaces.findOne({ _id: seeded.workspace._id }); | |
| expect(workspace?.piSession?.lastSequence).toBe(1); | |
| conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| expect(conversation?.messages[0]?.content).toBe("materialized exactly once"); | |
| expect( | |
| await collections.generationEvents.countDocuments({ generationId: seeded.generationId }) | |
| ).toBe(1); | |
| }); | |
| it("recovers a sandbox-change crash after pending message materialization", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| const pending: AgentPendingBatch = { | |
| kind: "bridge", | |
| sourceSequence: 1, | |
| startSeq: 1, | |
| events: [{ type: MessageUpdateType.Stream, token: "terminal batch exactly once" }], | |
| }; | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $set: { seq: 1, "agent.pending": pending, "agent.controlsOpen": true } } | |
| ); | |
| const originalUpdateOne = collections.conversations.updateOne.bind(collections.conversations); | |
| let crashed = false; | |
| const updateSpy = vi | |
| .spyOn(collections.conversations, "updateOne") | |
| .mockImplementation(async (filter, update, options) => { | |
| const result = await originalUpdateOne(filter, update, options); | |
| if (!crashed && "$set" in update && update.$set && "messages.$" in update.$set) { | |
| crashed = true; | |
| throw new Error("simulated sandbox terminalizer crash after materialization"); | |
| } | |
| return result; | |
| }); | |
| await expect( | |
| AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| sandboxId: "replacement-sandbox", | |
| sandboxGeneration: "replacement-generation", | |
| }) | |
| ).rejects.toThrow("simulated sandbox terminalizer crash after materialization"); | |
| let generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation).toMatchObject({ | |
| status: "running", | |
| agent: { | |
| controlsOpen: false, | |
| pending, | |
| terminalization: { kind: "sandbox-change" }, | |
| }, | |
| }); | |
| expect(generation?.agent?.relay).toBeUndefined(); | |
| expect( | |
| (await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0] | |
| ?.content | |
| ).toBe("terminal batch exactly once"); | |
| expect( | |
| await collections.generationEvents.countDocuments({ generationId: seeded.generationId }) | |
| ).toBe(1); | |
| await expect(store.commitBridgeBatch(2, [])).rejects.toBeInstanceOf(AgentRelayLeaseLostError); | |
| expect( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ).toBeNull(); | |
| updateSpy.mockRestore(); | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $set: { "agent.terminalization.expiresAt": new Date(0) } } | |
| ); | |
| expect( | |
| await AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| sandboxId: "replacement-sandbox", | |
| sandboxGeneration: "replacement-generation", | |
| }) | |
| ).toBe(true); | |
| generation = await collections.generations.findOne({ generationId: seeded.generationId }); | |
| expect(generation).toMatchObject({ | |
| status: "interrupted", | |
| agent: { sourceSequence: 1, controlsOpen: false }, | |
| }); | |
| expect(generation?.agent?.pending).toBeUndefined(); | |
| expect(generation?.agent?.terminalization).toBeUndefined(); | |
| expect( | |
| (await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0] | |
| ).toMatchObject({ content: "terminal batch exactly once", interrupted: true }); | |
| expect( | |
| await collections.generationEvents.countDocuments({ generationId: seeded.generationId }) | |
| ).toBe(1); | |
| }); | |
| it("reaper resumes an expired post-terminal marker before clearing side effects", async () => { | |
| const seeded = await seedAgentRun(); | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }); | |
| const originalFindOne = collections.conversations.findOne.bind(collections.conversations); | |
| let crashed = false; | |
| const findSpy = vi | |
| .spyOn(collections.conversations, "findOne") | |
| .mockImplementation(async (filter, options) => { | |
| const generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| if (!crashed && generation?.status === "interrupted" && generation.agent?.terminalization) { | |
| crashed = true; | |
| throw new Error("simulated crash after terminal CAS"); | |
| } | |
| return originalFindOne(filter, options); | |
| }); | |
| await expect( | |
| AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| sandboxId: "replacement-sandbox", | |
| sandboxGeneration: "replacement-generation", | |
| }) | |
| ).rejects.toThrow("simulated crash after terminal CAS"); | |
| let generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation).toMatchObject({ | |
| status: "interrupted", | |
| agent: { terminalization: { kind: "sandbox-change" } }, | |
| }); | |
| expect( | |
| (await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0] | |
| ?.interrupted | |
| ).not.toBe(true); | |
| expect( | |
| (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.execution | |
| ).toBeDefined(); | |
| findSpy.mockRestore(); | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $set: { "agent.terminalization.expiresAt": new Date(0) } } | |
| ); | |
| await reapStaleGenerations(); | |
| generation = await collections.generations.findOne({ generationId: seeded.generationId }); | |
| expect(generation?.agent?.terminalization).toBeUndefined(); | |
| expect( | |
| (await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0] | |
| ?.interrupted | |
| ).toBe(true); | |
| expect( | |
| (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.execution | |
| ).toBeUndefined(); | |
| }); | |
| it("does not materialize a captured WAL batch after its relay fence moves", async () => { | |
| const seeded = await seedAgentRun(); | |
| let timestamp = 0; | |
| const now = () => new Date(timestamp); | |
| const first = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| owner: "stale-relay", | |
| leaseMs: 10, | |
| now, | |
| }) | |
| ); | |
| const pending: AgentPendingBatch = { | |
| kind: "bridge", | |
| sourceSequence: 1, | |
| startSeq: 1, | |
| events: [{ type: MessageUpdateType.Stream, token: "authoritative output" }], | |
| }; | |
| await collections.generations.updateOne( | |
| { generationId: seeded.generationId }, | |
| { $set: { seq: 1, "agent.pending": pending } } | |
| ); | |
| const captured = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| if (!captured) throw new Error("Missing captured generation WAL"); | |
| timestamp = 11; | |
| const replacement = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| owner: "replacement-relay", | |
| leaseMs: 10, | |
| now, | |
| }) | |
| ); | |
| const capturedFlush = first as unknown as { | |
| flushPending(generation: Generation, batch: AgentPendingBatch): Promise<unknown>; | |
| }; | |
| await expect(capturedFlush.flushPending(captured, pending)).rejects.toBeInstanceOf( | |
| AgentRelayLeaseLostError | |
| ); | |
| let conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| expect(conversation?.messages[0]?.content).toBe(""); | |
| await replacement.resumePending(); | |
| conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| expect(conversation?.messages[0]?.content).toBe("authoritative output"); | |
| expect( | |
| await collections.generationEvents.countDocuments({ generationId: seeded.generationId }) | |
| ).toBe(1); | |
| }); | |
| it("fences an expired producer and lets only the replacement finish", async () => { | |
| const seeded = await seedAgentRun(); | |
| let timestamp = 0; | |
| const now = () => new Date(timestamp); | |
| const first = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| owner: "first", | |
| leaseMs: 10, | |
| now, | |
| }) | |
| ); | |
| timestamp = 11; | |
| const replacement = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| owner: "replacement", | |
| leaseMs: 10, | |
| now, | |
| }) | |
| ); | |
| await expect(first.renew()).rejects.toBeInstanceOf(AgentRelayLeaseLostError); | |
| await replacement.commitBridgeBatch(1, [ | |
| { type: MessageUpdateType.Stream, token: "replacement" }, | |
| ]); | |
| expect(await replacement.finish("completed")).toBe(true); | |
| expect( | |
| (await collections.generations.findOne({ generationId: seeded.generationId }))?.status | |
| ).toBe("completed"); | |
| }); | |
| it("never lets advisory pi stats lower authoritative proxy or sandbox spend", async () => { | |
| const seeded = await seedAgentRun(); | |
| await collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { | |
| $set: { | |
| spend: { | |
| accruedUsd: 12, | |
| inferenceUsd: 7, | |
| sandboxUsd: 5, | |
| lastCalculatedAt: new Date(), | |
| }, | |
| }, | |
| } | |
| ); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await store.commitBridgeBatch(1, [ | |
| { | |
| type: MessageUpdateType.AgentStats, | |
| scope: "session", | |
| tokens: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, total: 2 }, | |
| costUsd: 0.01, | |
| }, | |
| ]); | |
| const spend = (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.spend; | |
| expect(spend).toMatchObject({ accruedUsd: 12, inferenceUsd: 7, sandboxUsd: 5 }); | |
| }); | |
| it("does not advance the pi cursor on a replacement sandbox incarnation", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { | |
| $set: { | |
| "sandbox.id": "sandbox-replacement", | |
| "sandbox.generation": "generation-replacement", | |
| "piSession.lastSequence": 0, | |
| }, | |
| } | |
| ); | |
| await expect( | |
| store.commitBridgeBatch(41, [{ type: MessageUpdateType.Stream, token: "old incarnation" }]) | |
| ).rejects.toBeInstanceOf(AgentRelayLeaseLostError); | |
| const workspace = await collections.workspaces.findOne({ _id: seeded.workspace._id }); | |
| expect(workspace?.sandbox.id).toBe("sandbox-replacement"); | |
| expect(workspace?.piSession?.lastSequence).toBe(0); | |
| }); | |
| it("treats an exact workspace operation as contention instead of a sandbox mismatch", async () => { | |
| const seeded = await seedAgentRun(); | |
| const claimed = await claimWorkspaceOperation({ | |
| conversationId: seeded.conversationId, | |
| workspace: seeded.workspace, | |
| owner: { sessionId: seeded.workspace.sessionId }, | |
| kind: "configuration", | |
| action: "compact", | |
| }); | |
| expect( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ).toBeNull(); | |
| const generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation?.status).toBe("running"); | |
| expect(generation?.agent?.relay).toBeUndefined(); | |
| await releaseWorkspaceOperation( | |
| seeded.workspace._id, | |
| { sessionId: seeded.workspace.sessionId }, | |
| claimed.lease.id | |
| ); | |
| }); | |
| it("atomically reclaims orphaned operation leases on relay acquire and renew", async () => { | |
| const seeded = await seedAgentRun(); | |
| const orphanedOperation = (id: string) => ({ | |
| id, | |
| kind: "configuration" as const, | |
| action: "set-model", | |
| claimedAt: new Date(Date.now() - WORKSPACE_OPERATION_LEASE_MS - 1), | |
| workspaceUpdatedAt: seeded.workspace.updatedAt, | |
| }); | |
| await collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { $set: { operation: orphanedOperation("orphaned-acquire") } } | |
| ); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| expect( | |
| (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.operation | |
| ).toBeUndefined(); | |
| await collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { $set: { operation: orphanedOperation("orphaned-renew") } } | |
| ); | |
| await expect(store.renew()).resolves.toBeUndefined(); | |
| expect( | |
| (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.operation | |
| ).toBeUndefined(); | |
| await store.finish("completed"); | |
| }); | |
| it("keeps a live operation lease blocking relay renewal", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await collections.workspaces.updateOne( | |
| { _id: seeded.workspace._id }, | |
| { | |
| $set: { | |
| operation: { | |
| id: "live-renew-operation", | |
| kind: "configuration", | |
| action: "set-policy", | |
| claimedAt: new Date(), | |
| workspaceUpdatedAt: seeded.workspace.updatedAt, | |
| }, | |
| }, | |
| } | |
| ); | |
| await expect(store.renew()).rejects.toBeInstanceOf(AgentRelayLeaseLostError); | |
| expect( | |
| (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.operation?.id | |
| ).toBe("live-renew-operation"); | |
| }); | |
| it("clears an unrecoverable pending WAL while preserving its terminal error", async () => { | |
| const seeded = await seedAgentRun(); | |
| const store = requireStore( | |
| await AgentGenerationStore.acquire({ | |
| generationId: seeded.generationId, | |
| workspaceId: seeded.workspace._id, | |
| ...sandboxIncarnation(seeded.workspace), | |
| }) | |
| ); | |
| await collections.generations.updateOne( | |
| { | |
| generationId: seeded.generationId, | |
| "agent.relay.owner": store.owner, | |
| "agent.relay.fence": store.fence, | |
| }, | |
| { | |
| $set: { | |
| "agent.controlsOpen": true, | |
| "agent.control": { | |
| id: "undeliverable-control", | |
| fingerprint: "fingerprint", | |
| expectedPromptId: "prompt", | |
| payload: { kind: "compact" }, | |
| requestedAt: new Date(), | |
| state: "pending", | |
| }, | |
| }, | |
| } | |
| ); | |
| await collections.workspaces.deleteOne({ _id: seeded.workspace._id }); | |
| await expect(store.finish("interrupted")).rejects.toBeInstanceOf(AgentRelayLeaseLostError); | |
| await collections.generations.updateOne( | |
| { | |
| generationId: seeded.generationId, | |
| "agent.relay.owner": store.owner, | |
| "agent.relay.fence": store.fence, | |
| }, | |
| { | |
| $set: { | |
| seq: 1, | |
| "agent.pending": { | |
| kind: "system", | |
| key: "unrecoverable-pending", | |
| startSeq: 1, | |
| events: [ | |
| { | |
| type: MessageUpdateType.Status, | |
| status: MessageUpdateStatus.Error, | |
| message: "historical terminal failure", | |
| }, | |
| ], | |
| }, | |
| }, | |
| } | |
| ); | |
| const restoreMaterialization = failAfterNextPendingMaterialization(store); | |
| try { | |
| await expect( | |
| store.finishUnrecoverable("interrupted", "Agent workspace is no longer available") | |
| ).rejects.toThrow("simulated process crash after Agent WAL materialization"); | |
| } finally { | |
| restoreMaterialization(); | |
| } | |
| let generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation?.status).toBe("running"); | |
| expect(generation?.agent?.pending).toMatchObject({ | |
| kind: "system", | |
| key: "unrecoverable-pending", | |
| }); | |
| let conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| expect(conversation?.messages[0]?.updates).toContainEqual({ | |
| type: MessageUpdateType.Status, | |
| status: MessageUpdateStatus.Error, | |
| message: "historical terminal failure", | |
| }); | |
| expect( | |
| await store.finishUnrecoverable("interrupted", "Agent workspace is no longer available") | |
| ).toBe(true); | |
| generation = await collections.generations.findOne({ | |
| generationId: seeded.generationId, | |
| }); | |
| expect(generation).toMatchObject({ | |
| status: "interrupted", | |
| error: "Agent workspace is no longer available", | |
| agent: { controlsOpen: false }, | |
| }); | |
| expect(generation?.agent?.relay).toBeUndefined(); | |
| expect(generation?.agent?.control).toBeUndefined(); | |
| expect(generation?.agent?.pending).toBeUndefined(); | |
| expect(generation?.agent?.systemKeys).toContain("unrecoverable-pending"); | |
| expect( | |
| await collections.generationEvents.countDocuments({ generationId: seeded.generationId }) | |
| ).toBe(1); | |
| conversation = await collections.conversations.findOne({ _id: seeded.conversationId }); | |
| const message = conversation?.messages.find((item) => item.id === seeded.messageId); | |
| expect(message?.interrupted).toBe(true); | |
| expect(message?.updates).toContainEqual({ | |
| type: MessageUpdateType.Status, | |
| status: MessageUpdateStatus.Error, | |
| message: "historical terminal failure", | |
| }); | |
| }); | |
| }); | |
| function requireStore(store: AgentGenerationStore | null): AgentGenerationStore { | |
| if (!store) throw new Error("Expected to acquire the Agent generation relay lease"); | |
| return store; | |
| } | |
| function failAfterNextPendingMaterialization(store: AgentGenerationStore): () => void { | |
| const internals = store as unknown as { | |
| applyPendingToMessage( | |
| generation: Generation, | |
| pending: AgentPendingBatch, | |
| endSeq: number | |
| ): Promise<void>; | |
| }; | |
| const applyPendingToMessage = internals.applyPendingToMessage.bind(store); | |
| internals.applyPendingToMessage = async (generation, pending, endSeq) => { | |
| await applyPendingToMessage(generation, pending, endSeq); | |
| throw new Error("simulated process crash after Agent WAL materialization"); | |
| }; | |
| return () => { | |
| internals.applyPendingToMessage = applyPendingToMessage; | |
| }; | |
| } | |
| function sandboxIncarnation(workspace: Workspace): { | |
| sandboxId: string; | |
| sandboxGeneration: string; | |
| } { | |
| if (!workspace.sandbox.id || !workspace.sandbox.generation) { | |
| throw new Error("Expected a complete sandbox incarnation"); | |
| } | |
| return { | |
| sandboxId: workspace.sandbox.id, | |
| sandboxGeneration: workspace.sandbox.generation, | |
| }; | |
| } | |
| async function seedAgentRun() { | |
| const now = new Date(); | |
| const conversationId = new ObjectId(); | |
| const workspaceId = new ObjectId(); | |
| const messageId = randomUUID(); | |
| const generationId = randomUUID(); | |
| const sandboxId = `sandbox-${randomUUID()}`; | |
| const workspace: Workspace = { | |
| _id: workspaceId, | |
| conversationId, | |
| sessionId: "agent-store-owner", | |
| repo: { kind: "blank" }, | |
| flavor: "cpu-basic", | |
| model: "fake/model", | |
| thinkingLevel: "medium", | |
| policy: "auto", | |
| sandbox: { | |
| status: "ready", | |
| generation: `generation-${randomUUID()}`, | |
| id: sandboxId, | |
| namespace: "tester", | |
| baseUrl: `https://${sandboxId}--49983.hf.jobs`, | |
| nonce: "nonce", | |
| image: "node:22", | |
| flavor: "cpu-basic", | |
| createdAt: now, | |
| }, | |
| piSession: { id: "pi-session", lastSequence: 0 }, | |
| spend: { accruedUsd: 0 }, | |
| createdAt: now, | |
| updatedAt: now, | |
| }; | |
| await collections.conversations.insertOne({ | |
| _id: conversationId, | |
| sessionId: workspace.sessionId, | |
| mode: "agent", | |
| workspaceId, | |
| messages: [ | |
| { | |
| id: messageId, | |
| from: "assistant", | |
| content: "", | |
| updates: [], | |
| generationId, | |
| createdAt: now, | |
| updatedAt: now, | |
| }, | |
| ], | |
| createdAt: now, | |
| updatedAt: now, | |
| } as never); | |
| await collections.workspaces.insertOne(workspace); | |
| await collections.generations.insertOne({ | |
| _id: new ObjectId(), | |
| generationId, | |
| conversationId, | |
| messageId, | |
| sessionId: workspace.sessionId, | |
| mode: "agent", | |
| workspaceId, | |
| agent: { sourceSequence: 0 }, | |
| status: "running", | |
| seq: 0, | |
| startedAt: now, | |
| lastHeartbeatAt: now, | |
| createdAt: now, | |
| updatedAt: now, | |
| } satisfies Generation); | |
| return { conversationId, messageId, generationId, workspace }; | |
| } | |