Spaces:
Paused
Paused
| import { beforeEach, describe, expect, it, vi } from "vitest"; | |
| import { ObjectId } from "mongodb"; | |
| import type { Workspace } from "$lib/types/Workspace"; | |
| const mocks = vi.hoisted(() => { | |
| class Rpc { | |
| static instances: Rpc[] = []; | |
| commands: Array<Record<string, unknown>> = []; | |
| lastSequence = 0; | |
| constructor() { | |
| Rpc.instances.push(this); | |
| } | |
| async connect() {} | |
| subscribe() { | |
| return () => undefined; | |
| } | |
| subscribeReplayGap() { | |
| return () => undefined; | |
| } | |
| async command(command: Record<string, unknown>) { | |
| this.commands.push(command); | |
| if (command.type === "get_state") { | |
| return { type: "response", success: true, data: { isStreaming: false } }; | |
| } | |
| return { type: "response", success: true, data: {} }; | |
| } | |
| close() {} | |
| } | |
| return { Rpc }; | |
| }); | |
| vi.mock("./PiRpcClient", () => ({ PiRpcClient: mocks.Rpc })); | |
| vi.mock("./SandboxClient", () => ({ | |
| SandboxClient: { | |
| connect: () => ({ | |
| proxyHeaders: {}, | |
| proxyUrlFor: () => "ws://bridge.test/ws", | |
| }), | |
| }, | |
| })); | |
| import { AgentSessionManager } from "./AgentSessionManager"; | |
| beforeEach(() => { | |
| mocks.Rpc.instances.length = 0; | |
| }); | |
| describe("AgentSessionManager dispatch fencing", () => { | |
| it("never sends prompt or abort when a new-turn signal was already aborted", async () => { | |
| const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" }); | |
| const controller = new AbortController(); | |
| controller.abort(); | |
| await expect( | |
| manager | |
| .runBatches("do work", { | |
| signal: controller.signal, | |
| commandId: "prompt-a", | |
| }) | |
| .next() | |
| ).rejects.toMatchObject({ name: "AbortError" }); | |
| expect(mocks.Rpc.instances[0]?.commands).toEqual([]); | |
| }); | |
| it("runs the last ownership assertion before prompt and sends nothing when it fails", async () => { | |
| const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" }); | |
| const assertion = vi.fn(async () => { | |
| throw new Error("lease lost"); | |
| }); | |
| await expect( | |
| manager | |
| .runBatches("do work", { | |
| commandId: "prompt-a", | |
| beforePromptDispatch: assertion, | |
| }) | |
| .next() | |
| ).rejects.toThrow("lease lost"); | |
| expect(assertion).toHaveBeenCalledOnce(); | |
| expect(mocks.Rpc.instances[0]?.commands).toEqual([]); | |
| }); | |
| it("fences a resume abort and carries the expected prompt id", async () => { | |
| const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" }); | |
| const controller = new AbortController(); | |
| controller.abort(); | |
| const assertion = vi.fn(async () => undefined); | |
| await manager | |
| .resumeBatches({ | |
| signal: controller.signal, | |
| expectedPromptId: "prompt-a", | |
| beforeAbortDispatch: assertion, | |
| }) | |
| .next(); | |
| await vi.waitFor(() => expect(assertion).toHaveBeenCalledOnce()); | |
| expect(mocks.Rpc.instances[0]?.commands).toContainEqual({ | |
| type: "abort", | |
| _chatUiExpectedPromptId: "prompt-a", | |
| }); | |
| }); | |
| it("detaching a stale relay never aborts pi", async () => { | |
| const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" }); | |
| const detached = new AbortController(); | |
| detached.abort(); | |
| await expect( | |
| manager.resumeBatches({ detachSignal: detached.signal, expectedPromptId: "prompt-a" }).next() | |
| ).rejects.toMatchObject({ name: "AgentSessionDetachedError" }); | |
| expect(mocks.Rpc.instances[0]?.commands).toEqual([]); | |
| }); | |
| }); | |
| function workspace(): Workspace { | |
| const now = new Date(); | |
| return { | |
| _id: new ObjectId(), | |
| conversationId: new ObjectId(), | |
| sessionId: "owner", | |
| repo: { kind: "blank" }, | |
| flavor: "cpu-basic", | |
| model: "test/model", | |
| thinkingLevel: "medium", | |
| policy: "ask", | |
| sandbox: { | |
| status: "ready", | |
| flavor: "cpu-basic", | |
| id: "sandbox", | |
| generation: "incarnation", | |
| namespace: "owner", | |
| baseUrl: "https://sandbox.hf.jobs", | |
| nonce: "nonce", | |
| image: "node:22", | |
| createdAt: now, | |
| }, | |
| spend: { accruedUsd: 0 }, | |
| createdAt: now, | |
| updatedAt: now, | |
| }; | |
| } | |