From 4c4eeaba2b03f9825276c76b6d9972df8693bdf8 Mon Sep 17 00:00:00 2001 From: Chris Farhood Date: Wed, 13 May 2026 18:52:28 -0400 Subject: [PATCH] test: cover agentId threading on plugin lease RPCs and call sites Adds focused tests for every code path the agentId addition touches: - environment-runtime.test.ts (4 new tests): - plugin-driver acquireLease forwards agentId in RPC payload when present - plugin-driver acquireLease omits agentId from RPC payload when null - sandbox-provider acquireLease forwards agentId when present - sandbox-provider resumeLease forwards agentId when reuseLease=true matches - seedEnvironment helper now exposes the seeded agentId - environment-run-orchestrator.test.ts (2 new tests): - acquireForRun threads agentId through to runtime.acquireRunLease - logActivity records the same agentId on environment.lease_acquired - new vi.hoisted mocks for environmentService.getById + ensureLocalEnvironment - agent-test-environment-routes.test.ts (1 new assertion): - ad-hoc operator test-environment probe calls acquireRunLease with agentId: null and heartbeatRunId: null (no agent context) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../agent-test-environment-routes.test.ts | 9 + .../environment-run-orchestrator.test.ts | 78 ++++- .../src/__tests__/environment-runtime.test.ts | 295 ++++++++++++++++++ 3 files changed, 380 insertions(+), 2 deletions(-) diff --git a/server/src/__tests__/agent-test-environment-routes.test.ts b/server/src/__tests__/agent-test-environment-routes.test.ts index 3063fb70..d951a5ee 100644 --- a/server/src/__tests__/agent-test-environment-routes.test.ts +++ b/server/src/__tests__/agent-test-environment-routes.test.ts @@ -196,6 +196,15 @@ describe("agent test-environment route", () => { }), status: "failed", }); + // Ad-hoc operator probes have no agent context — the route must pass + // agentId: null so plugin-backed providers don't accidentally scope the + // probe lease against some leftover agentId from the heartbeat path. + expect(mockEnvironmentRuntime.acquireRunLease).toHaveBeenCalledWith( + expect.objectContaining({ + agentId: null, + heartbeatRunId: null, + }), + ); }); it("returns a diagnostic result instead of probing the host when the requested environment is missing", async () => { diff --git a/server/src/__tests__/environment-run-orchestrator.test.ts b/server/src/__tests__/environment-run-orchestrator.test.ts index c4dbb58b..93a91ff1 100644 --- a/server/src/__tests__/environment-run-orchestrator.test.ts +++ b/server/src/__tests__/environment-run-orchestrator.test.ts @@ -10,6 +10,8 @@ const mockBuildWorkspaceRealizationRequest = vi.hoisted(() => vi.fn()); const mockUpdateLeaseMetadata = vi.hoisted(() => vi.fn()); const mockUpdateExecutionWorkspace = vi.hoisted(() => vi.fn()); const mockLogActivity = vi.hoisted(() => vi.fn()); +const mockEnvironmentsEnsureLocal = vi.hoisted(() => vi.fn()); +const mockEnvironmentsGetById = vi.hoisted(() => vi.fn()); vi.mock("../services/environment-execution-target.js", () => ({ resolveEnvironmentExecutionTarget: mockResolveEnvironmentExecutionTarget, @@ -26,8 +28,8 @@ vi.mock("../services/workspace-realization.js", () => ({ vi.mock("../services/environments.js", () => ({ environmentService: vi.fn(() => ({ - ensureLocalEnvironment: vi.fn(), - getById: vi.fn(), + ensureLocalEnvironment: mockEnvironmentsEnsureLocal, + getById: mockEnvironmentsGetById, acquireLease: vi.fn(), releaseLease: vi.fn(), updateLeaseMetadata: mockUpdateLeaseMetadata, @@ -548,3 +550,75 @@ describe("environmentRunOrchestrator — realizeForRun", () => { expect(mockResolveEnvironmentExecutionTarget).not.toHaveBeenCalled(); }); }); + +describe("environmentRunOrchestrator — acquireForRun threads agentId", () => { + const mockDb = {} as any; + + beforeEach(() => { + vi.clearAllMocks(); + // selectedEnvironmentId !== defaultEnvironmentId in our inputs so the + // resolver goes through getById rather than ensureLocalEnvironment. + mockEnvironmentsGetById.mockResolvedValue(makeEnvironment("sandbox")); + mockResolveEnvironmentExecutionTarget.mockResolvedValue({ + kind: "local", + environmentId: "env-1", + leaseId: "lease-1", + }); + mockAdapterExecutionTargetToRemoteSpec.mockReturnValue(null); + }); + + function makeAcquireInput(overrides: { agentId?: string } = {}) { + return { + companyId: "company-1", + // distinct from defaultEnvironmentId so resolveEnvironment hits getById + selectedEnvironmentId: "env-1", + defaultEnvironmentId: "env-default", + adapterType: "claude_local", + issueId: null as string | null, + heartbeatRunId: "run-1", + agentId: overrides.agentId ?? "agent-uuid-abc", + persistedExecutionWorkspace: null, + }; + } + + it("passes agentId from acquireForRun's input through to runtime.acquireRunLease", async () => { + const runtime = makeMockRuntime({ + acquireRunLease: vi.fn().mockResolvedValue({ + lease: makeLease(), + leaseContext: { executionWorkspaceId: null, executionWorkspaceMode: null }, + }), + }); + const orchestrator = environmentRunOrchestrator(mockDb, { environmentRuntime: runtime }); + + await orchestrator.acquireForRun(makeAcquireInput({ agentId: "agent-uuid-abc" })); + + expect(runtime.acquireRunLease).toHaveBeenCalledWith( + expect.objectContaining({ + agentId: "agent-uuid-abc", + heartbeatRunId: "run-1", + companyId: "company-1", + }), + ); + }); + + it("logs the lease-acquired activity with the same agentId it threads to the runtime", async () => { + const runtime = makeMockRuntime({ + acquireRunLease: vi.fn().mockResolvedValue({ + lease: makeLease(), + leaseContext: { executionWorkspaceId: null, executionWorkspaceMode: null }, + }), + }); + const orchestrator = environmentRunOrchestrator(mockDb, { environmentRuntime: runtime }); + + await orchestrator.acquireForRun(makeAcquireInput({ agentId: "agent-uuid-xyz" })); + + expect(mockLogActivity).toHaveBeenCalledWith( + mockDb, + expect.objectContaining({ + action: "environment.lease_acquired", + agentId: "agent-uuid-xyz", + actorId: "agent-uuid-xyz", + }), + ); + }); +}); diff --git a/server/src/__tests__/environment-runtime.test.ts b/server/src/__tests__/environment-runtime.test.ts index c850da81..3b8bf356 100644 --- a/server/src/__tests__/environment-runtime.test.ts +++ b/server/src/__tests__/environment-runtime.test.ts @@ -216,6 +216,7 @@ describeEmbeddedPostgres("environmentRuntimeService", () => { return { companyId, + agentId, environment: { id: environmentId, companyId, @@ -1394,4 +1395,298 @@ describeEmbeddedPostgres("environmentRuntimeService", () => { expect(sshRelease).not.toHaveBeenCalled(); expect(acquired.lease.metadata?.driver).toBe("local"); }); + + // ------------------------------------------------------------------------- + // agentId is threaded through plugin RPC params (see protocol.ts — + // PluginEnvironmentAcquireLeaseParams.agentId and + // PluginEnvironmentResumeLeaseParams.agentId). Plugin-backed sandbox + // providers can use this to scope lease state (subdirs, PVCs, etc.) per + // agent without callbacks or DB lookups. The runtime must forward it when + // present and omit it when null/undefined so older plugin SDKs that don't + // declare the field aren't surprised. + // ------------------------------------------------------------------------- + + it("plugin-driver acquireLease: forwards agentId in the RPC payload when present", async () => { + const pluginId = randomUUID(); + const workerManager = { + isRunning: vi.fn(() => true), + call: vi.fn(async (_pluginId: string, method: string) => { + if (method === "environmentAcquireLease") { + return { providerLeaseId: "plugin-lease-agent", metadata: { remoteCwd: "/workspace" } }; + } + return undefined; + }), + } as unknown as PluginWorkerManager; + const runtimeWithPlugin = environmentRuntimeService(db, { pluginWorkerManager: workerManager }); + const { companyId, agentId, environment, runId } = await seedEnvironment({ + driver: "plugin", + name: "Plugin agentId fwd", + config: { + pluginKey: "acme.environments", + driverKey: "fake-plugin", + driverConfig: { template: "base" }, + }, + }); + await db.insert(plugins).values({ + id: pluginId, + pluginKey: "acme.environments", + packageName: "@acme/paperclip-environments", + version: "1.0.0", + apiVersion: 1, + categories: ["automation"], + manifestJson: { + id: "acme.environments", + apiVersion: 1, + version: "1.0.0", + displayName: "Acme", + description: "Test", + author: "Acme", + categories: ["automation"], + capabilities: ["environment.drivers.register"], + entrypoints: { worker: "dist/worker.js" }, + environmentDrivers: [{ driverKey: "fake-plugin", displayName: "Fake", configSchema: { type: "object" } }], + }, + status: "ready", + installOrder: 1, + updatedAt: new Date(), + } as any); + + await runtimeWithPlugin.acquireRunLease({ + companyId, + environment, + issueId: null, + agentId, + heartbeatRunId: runId, + persistedExecutionWorkspace: null, + }); + + expect(workerManager.call).toHaveBeenCalledWith( + pluginId, + "environmentAcquireLease", + expect.objectContaining({ agentId }), + ); + }); + + it("plugin-driver acquireLease: omits agentId from RPC payload when null", async () => { + const pluginId = randomUUID(); + const workerManager = { + isRunning: vi.fn(() => true), + call: vi.fn(async (_pluginId: string, method: string) => { + if (method === "environmentAcquireLease") { + return { providerLeaseId: "plugin-lease-no-agent", metadata: { remoteCwd: "/workspace" } }; + } + return undefined; + }), + } as unknown as PluginWorkerManager; + const runtimeWithPlugin = environmentRuntimeService(db, { pluginWorkerManager: workerManager }); + const { companyId, environment, runId } = await seedEnvironment({ + driver: "plugin", + name: "Plugin agentId null", + config: { + pluginKey: "acme.environments", + driverKey: "fake-plugin", + driverConfig: { template: "base" }, + }, + }); + await db.insert(plugins).values({ + id: pluginId, + pluginKey: "acme.environments", + packageName: "@acme/paperclip-environments", + version: "1.0.0", + apiVersion: 1, + categories: ["automation"], + manifestJson: { + id: "acme.environments", + apiVersion: 1, + version: "1.0.0", + displayName: "Acme", + description: "Test", + author: "Acme", + categories: ["automation"], + capabilities: ["environment.drivers.register"], + entrypoints: { worker: "dist/worker.js" }, + environmentDrivers: [{ driverKey: "fake-plugin", displayName: "Fake", configSchema: { type: "object" } }], + }, + status: "ready", + installOrder: 1, + updatedAt: new Date(), + } as any); + + await runtimeWithPlugin.acquireRunLease({ + companyId, + environment, + issueId: null, + agentId: null, + heartbeatRunId: runId, + persistedExecutionWorkspace: null, + }); + + const payload = (workerManager.call as ReturnType).mock.calls.find( + ([, method]) => method === "environmentAcquireLease", + )?.[2] as Record; + expect(payload).toBeDefined(); + expect(payload.agentId).toBeUndefined(); + expect("agentId" in payload).toBe(false); + }); + + it("sandbox-provider acquireLease: forwards agentId when present", async () => { + const pluginId = randomUUID(); + const workerManager = { + isRunning: vi.fn((id: string) => id === pluginId), + call: vi.fn(async (_pluginId: string, method: string) => { + if (method === "environmentAcquireLease") { + return { providerLeaseId: "sandbox-agent-1", metadata: { reuseLease: false } }; + } + throw new Error(`Unexpected plugin method: ${method}`); + }), + } as unknown as PluginWorkerManager; + const runtimeWithPlugin = environmentRuntimeService(db, { pluginWorkerManager: workerManager }); + const { companyId, agentId, environment, runId } = await seedEnvironment({ + driver: "sandbox", + name: "Sandbox agentId fwd", + config: { + provider: "fake-plugin", + image: "fake:test", + timeoutMs: 30_000, + reuseLease: false, + }, + }); + await db.insert(plugins).values({ + id: pluginId, + pluginKey: "acme.sandbox", + packageName: "@acme/paperclip-sandbox", + version: "1.0.0", + apiVersion: 1, + categories: ["automation"], + manifestJson: { + id: "acme.sandbox", + apiVersion: 1, + version: "1.0.0", + displayName: "Acme Sandbox", + description: "Test", + author: "Acme", + categories: ["automation"], + capabilities: ["environment.drivers.register"], + entrypoints: { worker: "dist/worker.js" }, + environmentDrivers: [{ + driverKey: "fake-plugin", + kind: "sandbox_provider", + displayName: "Fake", + configSchema: { type: "object" }, + }], + }, + status: "ready", + installOrder: 1, + updatedAt: new Date(), + } as any); + + await runtimeWithPlugin.acquireRunLease({ + companyId, + environment, + issueId: null, + agentId, + heartbeatRunId: runId, + persistedExecutionWorkspace: null, + }); + + expect(workerManager.call).toHaveBeenCalledWith( + pluginId, + "environmentAcquireLease", + expect.objectContaining({ agentId }), + expect.any(Number), + ); + }); + + it("sandbox-provider resumeLease: forwards agentId when present", async () => { + const pluginId = randomUUID(); + const calls: { method: string; params: Record }[] = []; + const workerManager = { + isRunning: vi.fn((id: string) => id === pluginId), + call: vi.fn(async (_pluginId: string, method: string, params: Record) => { + calls.push({ method, params }); + if (method === "environmentAcquireLease") { + return { providerLeaseId: "sandbox-resume-1", metadata: { reuseLease: true } }; + } + if (method === "environmentResumeLease") { + return { providerLeaseId: "sandbox-resume-1", metadata: { reuseLease: true } }; + } + throw new Error(`Unexpected plugin method: ${method}`); + }), + } as unknown as PluginWorkerManager; + const runtimeWithPlugin = environmentRuntimeService(db, { pluginWorkerManager: workerManager }); + const { companyId, agentId, environment, runId } = await seedEnvironment({ + driver: "sandbox", + name: "Sandbox agentId resume", + config: { + provider: "fake-plugin", + image: "fake:test", + timeoutMs: 30_000, + reuseLease: true, + }, + }); + await db.insert(plugins).values({ + id: pluginId, + pluginKey: "acme.sandbox", + packageName: "@acme/paperclip-sandbox", + version: "1.0.0", + apiVersion: 1, + categories: ["automation"], + manifestJson: { + id: "acme.sandbox", + apiVersion: 1, + version: "1.0.0", + displayName: "Acme Sandbox", + description: "Test", + author: "Acme", + categories: ["automation"], + capabilities: ["environment.drivers.register"], + entrypoints: { worker: "dist/worker.js" }, + environmentDrivers: [{ + driverKey: "fake-plugin", + kind: "sandbox_provider", + displayName: "Fake", + configSchema: { type: "object" }, + }], + }, + status: "ready", + installOrder: 1, + updatedAt: new Date(), + } as any); + + // First acquire seeds a reusable lease row in DB + await runtimeWithPlugin.acquireRunLease({ + companyId, + environment, + issueId: null, + agentId, + heartbeatRunId: runId, + persistedExecutionWorkspace: null, + }); + + // Second acquire on the same environment + reuseLease=true exercises the + // resume path (host's matcher finds the reusable lease, plugin's + // resumeLease is invoked). + const newRunId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: newRunId, + companyId, + agentId, + invocationSource: "manual", + status: "running", + createdAt: new Date(), + updatedAt: new Date(), + } as any); + await runtimeWithPlugin.acquireRunLease({ + companyId, + environment, + issueId: null, + agentId, + heartbeatRunId: newRunId, + persistedExecutionWorkspace: null, + }); + + const resumeCall = calls.find((c) => c.method === "environmentResumeLease"); + expect(resumeCall).toBeDefined(); + expect(resumeCall?.params.agentId).toBe(agentId); + }); });