Spaces:
Paused
Paused
| import { afterEach, beforeAll, describe, expect, it } from "vitest"; | |
| import { ObjectId } from "mongodb"; | |
| import { collections, ready } from "$lib/server/database"; | |
| import { | |
| prepareDetachedAgentPromptForSandbox, | |
| reseedAgentPromptForReplacement, | |
| type AgentPromptRecovery, | |
| } from "$lib/server/agent/agentPromptRecovery"; | |
| import { | |
| AgentGenerationIncarnationMismatchError, | |
| AgentGenerationStore, | |
| } from "$lib/server/agent/AgentGenerationStore"; | |
| beforeAll(async () => { | |
| await ready; | |
| }); | |
| afterEach(async () => { | |
| await Promise.all([ | |
| collections.generations.deleteMany({ generationId: /^prompt-recovery-/ }), | |
| collections.conversations.deleteMany({ sessionId: /^prompt-recovery-session-/ }), | |
| collections.workspaces.deleteMany({ sessionId: /^prompt-recovery-session-/ }), | |
| ]); | |
| }); | |
| describe.sequential("Agent replacement prompt recovery", () => { | |
| it("atomically seeds a ready-session prompt when reconnect returns a replacement sandbox", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const { conversationId, workspaceId } = await seedGeneration(generationId, { | |
| sandboxId: "job-before-dead-bridge", | |
| sandboxGeneration: "generation-before", | |
| }); | |
| const message = await reseedAgentPromptForReplacement({ | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| prompt: { id: "command-1", message: "continue" }, | |
| recovery: { | |
| sandboxId: "job-before-dead-bridge", | |
| sandboxGeneration: "generation-before", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }, | |
| sandbox: { id: "job-replacement", generation: "generation-after" }, | |
| }); | |
| expect(message).toBe("prior transcript\n\ncontinue"); | |
| expect((await collections.generations.findOne({ generationId }))?.agent?.prompt?.message).toBe( | |
| "prior transcript\n\ncontinue" | |
| ); | |
| }); | |
| it("keeps an unseeded prompt when the ready sandbox incarnation reconnects", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const recovery: AgentPromptRecovery = { | |
| sandboxId: "job-stable", | |
| sandboxGeneration: "generation-stable", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }; | |
| const { conversationId, workspaceId } = await seedGeneration(generationId, { | |
| sandboxId: "job-stable", | |
| sandboxGeneration: "generation-stable", | |
| recovery, | |
| }); | |
| const message = await reseedAgentPromptForReplacement({ | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| prompt: { id: "command-1", message: "continue" }, | |
| recovery, | |
| sandbox: { id: "job-stable", generation: "generation-stable" }, | |
| }); | |
| expect(message).toBe("continue"); | |
| const lease = await AgentGenerationStore.acquire({ | |
| generationId, | |
| workspaceId, | |
| sandboxId: "job-stable", | |
| sandboxGeneration: "generation-stable", | |
| }); | |
| expect(lease).not.toBeNull(); | |
| expect( | |
| (await collections.generations.findOne({ generationId }))?.agent?.prompt?.recovery | |
| ).toBeUndefined(); | |
| await lease?.release(); | |
| }); | |
| it("refuses to rewrite a prompt after a relay may have dispatched it", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const { conversationId, workspaceId } = await seedGeneration(generationId, { | |
| leased: true, | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| }); | |
| await expect( | |
| reseedAgentPromptForReplacement({ | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| prompt: { id: "command-1", message: "continue" }, | |
| recovery: { | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }, | |
| sandbox: { id: "job-after", generation: "generation-after" }, | |
| }) | |
| ).rejects.toThrow("changed before replacement context was restored"); | |
| expect((await collections.generations.findOne({ generationId }))?.agent?.prompt?.message).toBe( | |
| "continue" | |
| ); | |
| }); | |
| it("fails replacement relay acquisition closed until prompt and binding move together", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const { conversationId, workspaceId } = await seedGeneration(generationId, { | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| }); | |
| await expect( | |
| AgentGenerationStore.acquire({ | |
| generationId, | |
| workspaceId, | |
| sandboxId: "job-after", | |
| sandboxGeneration: "generation-after", | |
| }) | |
| ).rejects.toBeInstanceOf(AgentGenerationIncarnationMismatchError); | |
| await reseedAgentPromptForReplacement({ | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| prompt: { id: "command-1", message: "continue" }, | |
| recovery: { | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }, | |
| sandbox: { id: "job-after", generation: "generation-after" }, | |
| }); | |
| await collections.workspaces.updateOne( | |
| { _id: workspaceId }, | |
| { | |
| $set: { | |
| "sandbox.id": "job-after", | |
| "sandbox.generation": "generation-after", | |
| }, | |
| } | |
| ); | |
| const lease = await AgentGenerationStore.acquire({ | |
| generationId, | |
| workspaceId, | |
| sandboxId: "job-after", | |
| sandboxGeneration: "generation-after", | |
| }); | |
| expect(lease).not.toBeNull(); | |
| expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({ | |
| sandboxId: "job-after", | |
| sandboxGeneration: "generation-after", | |
| prompt: { message: "prior transcript\n\ncontinue" }, | |
| }); | |
| await lease?.release(); | |
| }); | |
| it("keeps a detached relay out while the replacement producer is still recovering", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const recovery: AgentPromptRecovery = { | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }; | |
| const { workspaceId } = await seedGeneration(generationId, { | |
| sandboxId: recovery.sandboxId, | |
| sandboxGeneration: recovery.sandboxGeneration, | |
| recovery, | |
| producer: { owner: "origin", expiresAt: new Date(Date.now() + 60_000) }, | |
| }); | |
| expect( | |
| await prepareDetachedAgentPromptForSandbox({ | |
| generationId, | |
| workspaceId, | |
| sandbox: { id: "job-after", generation: "generation-after" }, | |
| }) | |
| ).toBe("busy"); | |
| expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({ | |
| sandboxId: recovery.sandboxId, | |
| sandboxGeneration: recovery.sandboxGeneration, | |
| prompt: { message: "continue", recovery }, | |
| producer: { owner: "origin" }, | |
| }); | |
| }); | |
| it("lets the exact producer reseed a replacement while its lease is live", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const recovery: AgentPromptRecovery = { | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }; | |
| const { conversationId, workspaceId } = await seedGeneration(generationId, { | |
| sandboxId: recovery.sandboxId, | |
| sandboxGeneration: recovery.sandboxGeneration, | |
| recovery, | |
| producer: { owner: "origin", expiresAt: new Date(Date.now() + 60_000) }, | |
| }); | |
| expect( | |
| await reseedAgentPromptForReplacement({ | |
| generationId, | |
| conversationId, | |
| workspaceId, | |
| prompt: { id: "command-1", message: "continue" }, | |
| recovery, | |
| sandbox: { id: "job-after", generation: "generation-after" }, | |
| producerOwner: "origin", | |
| }) | |
| ).toBe(recovery.seededMessage); | |
| expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({ | |
| sandboxId: "job-after", | |
| sandboxGeneration: "generation-after", | |
| prompt: { message: recovery.seededMessage }, | |
| }); | |
| expect( | |
| (await collections.generations.findOne({ generationId }))?.agent?.prompt?.recovery | |
| ).toBeUndefined(); | |
| }); | |
| it("durably reseeds after producer expiry before a detached relay acquires", async () => { | |
| const generationId = `prompt-recovery-${new ObjectId().toString()}`; | |
| const recovery: AgentPromptRecovery = { | |
| sandboxId: "job-before", | |
| sandboxGeneration: "generation-before", | |
| seededMessage: "prior transcript\n\ncontinue", | |
| }; | |
| const { workspaceId } = await seedGeneration(generationId, { | |
| sandboxId: recovery.sandboxId, | |
| sandboxGeneration: recovery.sandboxGeneration, | |
| recovery, | |
| producer: { owner: "dead-origin", expiresAt: new Date(0) }, | |
| }); | |
| expect( | |
| await prepareDetachedAgentPromptForSandbox({ | |
| generationId, | |
| workspaceId, | |
| sandbox: { id: "job-after", generation: "generation-after" }, | |
| }) | |
| ).toBe("ready"); | |
| await collections.workspaces.updateOne( | |
| { _id: workspaceId }, | |
| { | |
| $set: { | |
| "sandbox.id": "job-after", | |
| "sandbox.generation": "generation-after", | |
| }, | |
| } | |
| ); | |
| const lease = await AgentGenerationStore.acquire({ | |
| generationId, | |
| workspaceId, | |
| sandboxId: "job-after", | |
| sandboxGeneration: "generation-after", | |
| }); | |
| expect(lease).not.toBeNull(); | |
| expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({ | |
| sandboxId: "job-after", | |
| sandboxGeneration: "generation-after", | |
| prompt: { message: recovery.seededMessage }, | |
| }); | |
| expect( | |
| (await collections.generations.findOne({ generationId }))?.agent?.prompt?.recovery | |
| ).toBeUndefined(); | |
| expect( | |
| (await collections.generations.findOne({ generationId }))?.agent?.producer | |
| ).toBeUndefined(); | |
| await lease?.release(); | |
| }); | |
| }); | |
| async function seedGeneration( | |
| generationId: string, | |
| options: { | |
| leased?: boolean; | |
| sandboxId?: string; | |
| sandboxGeneration?: string; | |
| recovery?: AgentPromptRecovery; | |
| producer?: { owner: string; expiresAt: Date }; | |
| } = {} | |
| ): Promise<{ conversationId: ObjectId; workspaceId: ObjectId }> { | |
| const now = new Date(); | |
| const workspaceId = new ObjectId(); | |
| const conversationId = new ObjectId(); | |
| const messageId = new ObjectId().toString(); | |
| const sessionId = `prompt-recovery-session-${generationId}`; | |
| await collections.conversations.insertOne({ | |
| _id: conversationId, | |
| sessionId, | |
| mode: "agent", | |
| workspaceId, | |
| messages: [{ id: messageId, from: "assistant", content: "", updates: [], generationId }], | |
| createdAt: now, | |
| updatedAt: now, | |
| } as never); | |
| await collections.workspaces.insertOne({ | |
| _id: workspaceId, | |
| conversationId, | |
| sessionId, | |
| flavor: "cpu-basic", | |
| model: "test-model", | |
| thinkingLevel: "medium", | |
| policy: "safe", | |
| sandbox: { | |
| status: "ready", | |
| id: options.sandboxId, | |
| generation: options.sandboxGeneration, | |
| nonce: "nonce", | |
| }, | |
| ...(options.producer | |
| ? { | |
| execution: { | |
| generationId, | |
| phase: "producer", | |
| owner: options.producer.owner, | |
| expiresAt: options.producer.expiresAt, | |
| }, | |
| } | |
| : {}), | |
| createdAt: now, | |
| updatedAt: now, | |
| } as never); | |
| await collections.generations.insertOne({ | |
| _id: new ObjectId(), | |
| generationId, | |
| conversationId, | |
| messageId, | |
| sessionId, | |
| status: "running", | |
| seq: 0, | |
| lastHeartbeatAt: now, | |
| startedAt: now, | |
| createdAt: now, | |
| updatedAt: now, | |
| mode: "agent", | |
| workspaceId, | |
| agent: { | |
| sourceSequence: 0, | |
| prompt: { | |
| id: "command-1", | |
| message: "continue", | |
| ...(options.recovery ? { recovery: options.recovery } : {}), | |
| }, | |
| ...(options.sandboxId ? { sandboxId: options.sandboxId } : {}), | |
| ...(options.sandboxGeneration ? { sandboxGeneration: options.sandboxGeneration } : {}), | |
| ...(options.producer ? { producer: options.producer } : {}), | |
| ...(options.leased | |
| ? { | |
| relay: { | |
| owner: "another-server", | |
| fence: 1, | |
| expiresAt: new Date(now.getTime() + 60_000), | |
| }, | |
| } | |
| : {}), | |
| }, | |
| }); | |
| return { conversationId, workspaceId }; | |
| } | |