diff --git a/docs/plan/runners.md b/docs/plan/runners.md index 32ed708b44f0..409b0a6f14f0 100644 --- a/docs/plan/runners.md +++ b/docs/plan/runners.md @@ -190,8 +190,11 @@ stated honestly (revision 1 undersold this): silent superseded-pairing pruning share one client-invalidation, credential, environment, and placement reconciliation flow. Explicit RPCs wait for the credential fence before success returns; periodic reconciliation retries - failed provider or placement cleanup. Unreferenced - terminal environment rows retain + failed provider or placement cleanup. Placement startup, identity-mutation, + and periodic reconciliation also compare each durable placement with the + canonical session entry: confirmed absence force-fences live environments + and exact-CAS retires safe terminal rows, while unreadable session evidence + retains the placement. Unreferenced terminal environment rows retain seven days of operator diagnostics, then prune in bounded post-reconcile batches; any surviving placement keeps its environment provenance. Device-side GC of per-session workspace dirs and superseded bundles is a diff --git a/src/config/sessions/session-accessor.readonly.test.ts b/src/config/sessions/session-accessor.readonly.test.ts index 4c859264fe3d..b3769083821a 100644 --- a/src/config/sessions/session-accessor.readonly.test.ts +++ b/src/config/sessions/session-accessor.readonly.test.ts @@ -1,6 +1,10 @@ import fs from "node:fs"; import { afterEach, describe, expect, it } from "vitest"; -import { cleanupTempDirs, makeTempDir } from "../../../test/helpers/temp-dir.js"; +import { + cleanupTempDirs, + makeTempDir, + useAutoCleanupTempDirTracker, +} from "../../../test/helpers/temp-dir.js"; import { closeOpenClawAgentDatabasesForTest, isOpenClawAgentDatabaseOpen, @@ -15,11 +19,13 @@ import { hasSessionEntriesByStatusReadOnly, listSessionEntriesCore, listSessionEntriesReadOnly, + readSessionIdentityEvidence, resolveTranscriptSessionKeyBySessionId, upsertSessionEntryCore, } from "./session-accessor.js"; const tempDirs: string[] = []; +const autoTempDirs = useAutoCleanupTempDirTracker(afterEach); function countRegisteredAgentDatabases(env: NodeJS.ProcessEnv): number { const row = openOpenClawStateDatabase({ env }) @@ -138,6 +144,75 @@ describe("session accessor readonly listing", () => { expect(countRegisteredAgentDatabases(env)).toBe(0); }); + it("probes session identity by exact key and indexed current session id", async () => { + const stateDir = autoTempDirs.make("openclaw-session-readonly-evidence-"); + const env = { OPENCLAW_STATE_DIR: stateDir }; + const agentId = "worker-1"; + const sessionKey = "agent:worker-1:moved"; + const sessionId = "session-1"; + await upsertSessionEntryCore({ agentId, env, sessionKey }, { sessionId, updatedAt: 1 }); + const storePath = resolveOpenClawAgentSqlitePath({ agentId, env }); + closeOpenClawAgentDatabasesForTest(); + + expect(readSessionIdentityEvidence({ agentId, sessionId, sessionKey, storePath })).toEqual({ + status: "current", + sessionKey, + }); + expect( + readSessionIdentityEvidence({ + agentId, + sessionId, + sessionKey: "agent:worker-1:old-key", + storePath, + }), + ).toEqual({ status: "current", sessionKey }); + expect( + readSessionIdentityEvidence({ + agentId, + sessionId: "missing-session", + sessionKey, + storePath, + }), + ).toEqual({ status: "absent" }); + }); + + it("reports migration-invalid session evidence as unknown", async () => { + const stateDir = autoTempDirs.make("openclaw-session-readonly-evidence-invalid-"); + const env = { OPENCLAW_STATE_DIR: stateDir }; + const agentId = "worker-1"; + const sessionKey = "agent:worker-1:main"; + const sessionId = "session-1"; + await upsertSessionEntryCore({ agentId, env, sessionKey }, { sessionId, updatedAt: 1 }); + const database = openOpenClawAgentDatabase({ agentId, env }); + database.db.exec("PRAGMA user_version = 999;"); + const storePath = database.path; + closeOpenClawAgentDatabasesForTest(); + + expect(readSessionIdentityEvidence({ agentId, sessionId, sessionKey, storePath })).toEqual({ + status: "unknown", + reason: "read-failed", + }); + }); + + it("uses the current-session-id index for fallback identity probes", async () => { + const stateDir = autoTempDirs.make("openclaw-session-readonly-evidence-index-"); + const env = { OPENCLAW_STATE_DIR: stateDir }; + const agentId = "worker-1"; + const database = openOpenClawAgentDatabase({ agentId, env }); + const detail = database.db + .prepare( + "EXPLAIN QUERY PLAN SELECT session_key FROM session_nodes WHERE current_session_id = ? LIMIT 2", + ) + .all("session-1") + .map((row) => { + const rowDetail = (row as { detail?: unknown }).detail; + return typeof rowDetail === "string" ? rowDetail : ""; + }) + .join(" "); + + expect(detail).toContain("idx_agent_session_nodes_current_session_id"); + }); + it("does not register a populated database during readonly health-style listing", async () => { const stateDir = makeTempDir(tempDirs, "openclaw-session-readonly-registry-"); const env = { OPENCLAW_STATE_DIR: stateDir }; diff --git a/src/config/sessions/session-accessor.sqlite-entry-availability.ts b/src/config/sessions/session-accessor.sqlite-entry-availability.ts index 0232174b17d7..5d84e2bdbcd7 100644 --- a/src/config/sessions/session-accessor.sqlite-entry-availability.ts +++ b/src/config/sessions/session-accessor.sqlite-entry-availability.ts @@ -1,4 +1,7 @@ -import { executeSqliteQueryTakeFirstSync } from "../../infra/kysely-sync.js"; +import { + executeSqliteQuerySync, + executeSqliteQueryTakeFirstSync, +} from "../../infra/kysely-sync.js"; import { withOpenClawAgentDatabaseReadOnly } from "../../state/openclaw-agent-db-readonly.js"; import type { ExactSessionEntry, SessionAccessScope } from "./session-accessor.sqlite-contract.js"; import { readExactSessionEntryRowValidated } from "./session-accessor.sqlite-entry-store.js"; @@ -10,6 +13,14 @@ import { } from "./session-accessor.sqlite-scope.js"; import type { SessionEntry } from "./types.js"; +export type SessionIdentityEvidenceResult = + | { status: "current"; sessionKey: string } + | { status: "absent" } + | { + status: "unknown"; + reason: "ambiguous" | "read-failed" | "row-invalid" | "schema-missing" | "table-missing"; + }; + type ExactSessionEntryReadOnlyResult = | { found: true; value: ExactSessionEntry | undefined } | { @@ -70,3 +81,62 @@ export function loadExactSessionEntryReadOnlyResult( }, }; } + +/** Indexed exact-key/session-id probe that preserves unreadable state as unknown. */ +export function readSessionIdentityEvidence(params: { + agentId: string; + sessionId: string; + sessionKey: string; + storePath: string; +}): SessionIdentityEvidenceResult { + const resolved = resolveSqliteScope({ + agentId: params.agentId, + sessionKey: params.sessionKey, + storePath: params.storePath, + }); + let result: + | { found: true; value: SessionIdentityEvidenceResult } + | { found: false; reason: "database-missing" | "schema-missing" | "table-missing" }; + try { + result = withOpenClawAgentDatabaseReadOnly((database): SessionIdentityEvidenceResult => { + const exact = readExactSessionEntryRowValidated(database, resolved.sessionKey)?.entry; + if (exact?.sessionId === params.sessionId) { + return { status: "current", sessionKey: resolved.sessionKey }; + } + const rows = executeSqliteQuerySync( + database.db, + getSessionKysely(database.db) + .selectFrom("session_nodes") + .select(["session_key", "entry_valid"]) + .where("current_session_id", "=", params.sessionId) + .limit(2), + ).rows; + if (rows.length === 0) { + return { status: "absent" }; + } + if (rows.length !== 1) { + return { status: "unknown", reason: "ambiguous" }; + } + const row = rows[0]; + if (row?.entry_valid === -1) { + return { status: "absent" }; + } + const sessionKey = row?.session_key; + if (!sessionKey || row.entry_valid !== 1) { + return { status: "unknown", reason: "row-invalid" }; + } + const entry = readExactSessionEntryRowValidated(database, sessionKey)?.entry; + return entry?.sessionId === params.sessionId + ? { status: "current", sessionKey } + : { status: "unknown", reason: "row-invalid" }; + }, toDatabaseOptions(resolved)); + } catch { + return { status: "unknown", reason: "read-failed" }; + } + if (result.found) { + return result.value; + } + return result.reason === "database-missing" + ? { status: "absent" } + : { status: "unknown", reason: result.reason }; +} diff --git a/src/config/sessions/session-accessor.ts b/src/config/sessions/session-accessor.ts index 27fc2d757f3b..895886699390 100644 --- a/src/config/sessions/session-accessor.ts +++ b/src/config/sessions/session-accessor.ts @@ -156,6 +156,10 @@ export { updateResolvedSessionEntry, upsertSessionEntryCore, } from "./session-accessor.entry.js"; +export { + readSessionIdentityEvidence, + type SessionIdentityEvidenceResult, +} from "./session-accessor.sqlite-entry-availability.js"; export { createSessionEntryWithTranscript, forkSessionEntryFromParentTarget, diff --git a/src/gateway/server-worker-placement-session-evidence.test.ts b/src/gateway/server-worker-placement-session-evidence.test.ts new file mode 100644 index 000000000000..743af18c5e61 --- /dev/null +++ b/src/gateway/server-worker-placement-session-evidence.test.ts @@ -0,0 +1,61 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { useAutoCleanupTempDirTracker } from "../../test/helpers/temp-dir.js"; +import { upsertSessionEntryCore } from "../config/sessions/session-accessor.js"; +import { + closeOpenClawAgentDatabasesForTest, + openOpenClawAgentDatabase, +} from "../state/openclaw-agent-db.js"; +import { withEnvAsync } from "../test-utils/env.js"; +import { resolveWorkerPlacementSessionEvidence } from "./server-worker-placement-session-evidence.js"; +import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-record.js"; + +const tempDirs = useAutoCleanupTempDirTracker(afterEach); + +afterEach(() => { + closeOpenClawAgentDatabasesForTest(); +}); + +function localPlacement( + sessionId: string, + sessionKey: string, +): Extract { + return { + sessionId, + sessionKey, + agentId: "main", + state: "local", + generation: 1, + turnClaim: null, + environmentId: null, + activeOwnerEpoch: null, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + terminalReason: null, + terminalAtMs: null, + createdAtMs: 1, + updatedAtMs: 1, + stateChangedAtMs: 1, + }; +} + +describe("worker placement session evidence", () => { + it("keeps a placement when its session database is migration-invalid", async () => { + const stateDir = tempDirs.make("openclaw-placement-session-evidence-"); + await withEnvAsync({ OPENCLAW_STATE_DIR: stateDir }, async () => { + const sessionId = "session-1"; + const sessionKey = "agent:main:main"; + await upsertSessionEntryCore({ agentId: "main", sessionKey }, { sessionId, updatedAt: 1 }); + const database = openOpenClawAgentDatabase({ agentId: "main" }); + database.db.exec("PRAGMA user_version = 999;"); + closeOpenClawAgentDatabasesForTest(); + + await expect( + resolveWorkerPlacementSessionEvidence(localPlacement(sessionId, sessionKey)), + ).resolves.toBe("unknown"); + }); + }); +}); diff --git a/src/gateway/server-worker-placement-session-evidence.ts b/src/gateway/server-worker-placement-session-evidence.ts new file mode 100644 index 000000000000..0eabd16eee1d --- /dev/null +++ b/src/gateway/server-worker-placement-session-evidence.ts @@ -0,0 +1,33 @@ +import { getRuntimeConfig } from "../config/config.js"; +import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; +import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-record.js"; + +const loadPlacementSessionEvidenceRuntime = createLazyRuntimeModule(async () => { + const [sessionUtils, sessionAccessor] = await Promise.all([ + import("./session-utils.js"), + import("../config/sessions/session-accessor.js"), + ]); + return { + readSessionIdentityEvidence: sessionAccessor.readSessionIdentityEvidence, + resolveGatewaySessionStoreTarget: sessionUtils.resolveGatewaySessionStoreTarget, + }; +}); + +/** Resolves authoritative session existence without treating unreadable state as absence. */ +export async function resolveWorkerPlacementSessionEvidence( + placement: WorkerSessionPlacementRecord, +): Promise<"current" | "absent" | "unknown"> { + const runtime = await loadPlacementSessionEvidenceRuntime(); + const target = runtime.resolveGatewaySessionStoreTarget({ + cfg: getRuntimeConfig(), + key: placement.sessionKey, + agentId: placement.agentId, + }); + const evidence = runtime.readSessionIdentityEvidence({ + agentId: target.agentId, + sessionId: placement.sessionId, + sessionKey: target.canonicalKey, + storePath: target.storePath, + }); + return evidence.status; +} diff --git a/src/gateway/server-worker-placement-startup.test.ts b/src/gateway/server-worker-placement-startup.test.ts index ffe184cf1e7d..82ca7ef3d9b9 100644 --- a/src/gateway/server-worker-placement-startup.test.ts +++ b/src/gateway/server-worker-placement-startup.test.ts @@ -4,6 +4,7 @@ import { createDeferredCore } from "../shared/deferred.js"; const runtimeFactoryMocks = vi.hoisted(() => ({ createDispatch: vi.fn(), createDiskSpace: vi.fn(), + resolveSessionEvidence: vi.fn(), })); vi.mock("./worker-environments/placement-dispatch.js", async (importOriginal) => { @@ -15,6 +16,10 @@ vi.mock("./worker-environments/placement-dispatch.js", async (importOriginal) => }; }); +vi.mock("./server-worker-placement-session-evidence.js", () => ({ + resolveWorkerPlacementSessionEvidence: runtimeFactoryMocks.resolveSessionEvidence, +})); + vi.mock("./worker-environments/placement-disk-space.js", async (importOriginal) => { const actual = await importOriginal(); @@ -24,152 +29,7 @@ vi.mock("./worker-environments/placement-disk-space.js", async (importOriginal) }; }); -import { - coordinateWorkerPlacementDispatch, - createGatewayWorkerPlacementRuntime, - type GatewayWorkerPlacementRuntime, -} from "./server-worker-placement-startup.js"; -import type { WorkerPlacementDispatchRequest } from "./worker-environments/service-contract.js"; - -type DispatchService = GatewayWorkerPlacementRuntime["dispatchService"]; - -const REQUEST: WorkerPlacementDispatchRequest = { - sessionId: "session-1", - sessionKey: "agent:main:session-1", - agentId: "main", - profileId: "test", -}; - -describe("worker placement dispatch coordinator", () => { - it("forwards the optional internal transition observer", async () => { - const observer = vi.fn(); - const dispatch = vi.fn().mockResolvedValue({ state: "active" }); - const service = { - dispatch, - forceDestroyEnvironment: vi.fn(), - reclaim: vi.fn(), - reconcile: vi.fn(), - reconcileActive: vi.fn(), - } as unknown as DispatchService; - - await coordinateWorkerPlacementDispatch(service).dispatch(REQUEST, observer); - - expect(dispatch).toHaveBeenCalledWith(REQUEST, observer); - }); - - it("coalesces an identical dispatch and rejects a conflicting in-flight request", async () => { - const dispatchStarted = createDeferredCore(); - const releaseDispatch = createDeferredCore(); - const active = { state: "active" }; - const dispatch = vi.fn(async () => { - dispatchStarted.resolve(); - await releaseDispatch.promise; - return active; - }); - const service = { - dispatch, - forceDestroyEnvironment: vi.fn(), - reclaim: vi.fn(), - reconcile: vi.fn(), - reconcileActive: vi.fn(), - } as unknown as DispatchService; - const coordinated = coordinateWorkerPlacementDispatch(service); - - const first = coordinated.dispatch(REQUEST); - await dispatchStarted.promise; - await expect( - coordinated.dispatch({ ...REQUEST, profileId: "another-profile" }), - ).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`); - await expect( - coordinated.dispatch({ - ...REQUEST, - inheritedProfile: { - providerId: "fake", - profileSnapshot: { settings: { region: "parent" } }, - }, - }), - ).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`); - const retry = coordinated.dispatch(REQUEST); - releaseDispatch.resolve(); - - const [firstResult, retryResult] = await Promise.all([first, retry]); - expect(retryResult).toBe(firstResult); - expect(dispatch).toHaveBeenCalledOnce(); - - await coordinated.dispatch({ ...REQUEST, profileId: "another-profile" }); - expect(dispatch).toHaveBeenCalledTimes(2); - }); - - it("joins a retry before a queued reconciliation after dispatch failure", async () => { - const dispatchStarted = createDeferredCore(); - const releaseDispatch = createDeferredCore(); - const dispatchError = new Error("provision failed"); - const dispatch = vi.fn(async () => { - dispatchStarted.resolve(); - await releaseDispatch.promise; - throw dispatchError; - }); - const reconcileActive = vi.fn(); - const service = { - dispatch, - forceDestroyEnvironment: vi.fn(), - reclaim: vi.fn(), - reconcile: vi.fn(), - reconcileActive, - } as unknown as DispatchService; - const coordinated = coordinateWorkerPlacementDispatch(service); - - const first = coordinated.dispatch(REQUEST); - await dispatchStarted.promise; - const reconciliation = coordinated.reconcileActive(); - const retry = coordinated.dispatch(REQUEST); - const outcomes = Promise.allSettled([first, retry]); - releaseDispatch.resolve(); - - expect(await outcomes).toEqual([ - { status: "rejected", reason: dispatchError }, - { status: "rejected", reason: dispatchError }, - ]); - await reconciliation; - expect(dispatch).toHaveBeenCalledOnce(); - expect(reconcileActive).toHaveBeenCalledOnce(); - - await expect(coordinated.dispatch({ ...REQUEST, profileId: "another-profile" })).rejects.toBe( - dispatchError, - ); - expect(dispatch).toHaveBeenCalledTimes(2); - }); - - it("coalesces full sweeps but runs a fresh targeted pass with its environment id", async () => { - const fullSweepStarted = createDeferredCore(); - const releaseFullSweep = createDeferredCore(); - const reconcileActive = vi.fn(async (environmentId?: string) => { - if (environmentId === undefined) { - fullSweepStarted.resolve(); - await releaseFullSweep.promise; - } - }); - const service = { - dispatch: vi.fn(), - forceDestroyEnvironment: vi.fn(), - reclaim: vi.fn(), - reconcile: vi.fn(), - reconcileActive, - } as unknown as DispatchService; - const coordinated = coordinateWorkerPlacementDispatch(service); - - const firstFullSweep = coordinated.reconcileActive(); - const secondFullSweep = coordinated.reconcileActive(); - await fullSweepStarted.promise; - const targetedSweep = coordinated.reconcileActive("worker-target"); - - expect(reconcileActive).toHaveBeenCalledTimes(1); - releaseFullSweep.resolve(); - await Promise.all([firstFullSweep, secondFullSweep, targetedSweep]); - - expect(reconcileActive.mock.calls).toEqual([[], ["worker-target"]]); - }); -}); +import { createGatewayWorkerPlacementRuntime } from "./server-worker-placement-startup.js"; describe("worker placement startup health lifetime", () => { it("samples disk on schedule while reconciliation is stuck and drains both on stop", async () => { @@ -205,6 +65,9 @@ describe("worker placement startup health lifetime", () => { const warn = vi.fn(); const runtime = createGatewayWorkerPlacementRuntime({ placements: { + get: () => undefined, + list: () => [], + retireSessionPlacement: vi.fn(), pruneOrphanedWorkspaceReconciliations: () => [], listWorkspaceReconciliationOwners: () => [], } as never, @@ -233,7 +96,7 @@ describe("worker placement startup health lifetime", () => { releaseScheduledHealth.reject(healthError); await Promise.resolve(); expect(stopSettled).toBe(false); - expect(environments.stop).toHaveBeenCalledOnce(); + expect(environments.stop).not.toHaveBeenCalled(); releaseReconcile.resolve(); await stopping; @@ -244,4 +107,87 @@ describe("worker placement startup health lifetime", () => { vi.useRealTimers(); } }); + + it("drains deferred startup session evidence before stopping environments", async () => { + const evidence = createDeferredCore<"current">(); + runtimeFactoryMocks.resolveSessionEvidence.mockImplementation(async () => evidence.promise); + runtimeFactoryMocks.createDiskSpace.mockReturnValue({ + read: vi.fn(), + version: vi.fn(() => 0), + sweep: vi.fn().mockResolvedValue(undefined), + }); + runtimeFactoryMocks.createDispatch.mockReturnValue({ + dispatch: vi.fn(), + forceDestroyEnvironment: vi.fn(), + reclaim: vi.fn(), + reconcile: vi.fn().mockResolvedValue(undefined), + reconcileActive: vi.fn().mockResolvedValue(undefined), + }); + const placement = { + sessionId: "session-startup", + sessionKey: "agent:main:startup", + agentId: "main", + state: "local", + generation: 1, + turnClaim: null, + environmentId: null, + activeOwnerEpoch: null, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + terminalReason: null, + terminalAtMs: null, + createdAtMs: 1, + updatedAtMs: 1, + stateChangedAtMs: 1, + } as const; + const environments = { + start: vi.fn(), + stop: vi.fn().mockResolvedValue(undefined), + }; + const runtime = createGatewayWorkerPlacementRuntime({ + placements: { + get: () => placement, + list: () => [placement], + retireSessionPlacement: vi.fn(), + pruneOrphanedWorkspaceReconciliations: () => [], + listWorkspaceReconciliationOwners: () => [], + } as never, + environments: environments as never, + admitNewPlacements: true, + revokeSessionAuthority: vi.fn(), + warn: vi.fn(), + }); + let closeStarted = false; + let sidecar: { stop: () => Promise } | undefined; + const starting = runtime.startRuntime({ + isClosePreludeStarted: () => closeStarted, + registerSidecar: (registered) => { + sidecar = registered; + }, + }); + await vi.waitFor(() => expect(runtimeFactoryMocks.resolveSessionEvidence).toHaveBeenCalled()); + closeStarted = true; + const stopping = sidecar?.stop(); + const repeatedStop = sidecar?.stop(); + if (!stopping || !repeatedStop) { + throw new Error("startup did not register its placement sidecar"); + } + let repeatedStopSettled = false; + void repeatedStop.then(() => { + repeatedStopSettled = true; + }); + + await Promise.resolve(); + expect(repeatedStop).toBe(stopping); + expect(repeatedStopSettled).toBe(false); + expect(environments.stop).not.toHaveBeenCalled(); + evidence.resolve("current"); + await expect(starting).resolves.toBeNull(); + await Promise.all([stopping, repeatedStop]); + expect(environments.stop).toHaveBeenCalledOnce(); + }); }); diff --git a/src/gateway/server-worker-placement-startup.ts b/src/gateway/server-worker-placement-startup.ts index 9f4a1fb5f33d..a3480861dec7 100644 --- a/src/gateway/server-worker-placement-startup.ts +++ b/src/gateway/server-worker-placement-startup.ts @@ -1,4 +1,3 @@ -import { isDeepStrictEqual } from "node:util"; import { installSessionPlacementAdmissionProvider } from "../agents/session-placement-admission.js"; import { clearSessionQueues } from "../auto-reply/reply/queue/cleanup.js"; import { getRuntimeConfig } from "../config/config.js"; @@ -10,16 +9,16 @@ import { runExclusiveSessionLifecycleMutation, SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS, } from "../sessions/session-lifecycle-admission.js"; +import { onSessionIdentityMutation } from "../sessions/session-lifecycle-events.js"; import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; +import { resolveWorkerPlacementSessionEvidence } from "./server-worker-placement-session-evidence.js"; import { createWorkerPlacementDiskSpaceMonitor } from "./worker-environments/placement-disk-space.js"; -import { - createWorkerPlacementDispatchService, - type WorkerPlacementDispatchService, -} from "./worker-environments/placement-dispatch.js"; +import { coordinateWorkerPlacementDispatch } from "./worker-environments/placement-dispatch-coordinator.js"; +import { createWorkerPlacementDispatchService } from "./worker-environments/placement-dispatch.js"; import { FORCED_WORKER_ABANDONMENT_ERROR } from "./worker-environments/placement-force-abandon.js"; +import { createPlacementSessionRetirement } from "./worker-environments/placement-session-retirement.js"; import type { WorkerSessionPlacementStore } from "./worker-environments/placement-store.js"; import { createReclaimedPlacementRedispatch } from "./worker-environments/reclaimed-placement-redispatch.js"; -import type { WorkerPlacementDispatchRequest } from "./worker-environments/service-contract.js"; import type { WorkerEnvironmentService } from "./worker-environments/service.js"; import { createWorkerSessionTurnPlacementProvider } from "./worker-environments/worker-turn-launcher.js"; import { createWorkerWorkspaceOperationCoordinator } from "./worker-environments/workspace-operation-coordinator.js"; @@ -110,126 +109,6 @@ function resolveWorkerPlacementSessionTarget(params: { return { config: params.config, target, entry, worktree }; } -/** Serializes reconciliation sweeps against in-flight dispatches so a sweep never - * observes a placement mid-transition. Dispatches wait out any pending sweep. */ -export function coordinateWorkerPlacementDispatch( - service: WorkerPlacementDispatchService, -): WorkerPlacementDispatchService { - let activeDispatchCount = 0; - let reconciliation: Promise | undefined; - const dispatchIdleWaiters = new Set<() => void>(); - const waitForDispatchIdle = (): Promise => { - if (activeDispatchCount === 0) { - return Promise.resolve(); - } - return new Promise((resolve) => { - dispatchIdleWaiters.add(resolve); - }); - }; - const runReconciliation = (operation: () => Promise): Promise => { - if (reconciliation) { - return reconciliation; - } - const current = (async () => { - await waitForDispatchIdle(); - await operation(); - })(); - reconciliation = current; - const clearCurrent = () => { - if (reconciliation === current) { - reconciliation = undefined; - } - }; - void current.then(clearCurrent, clearCurrent); - return current; - }; - const runExclusivePlacementOperation = (operation: () => Promise): Promise => { - const current = (async () => { - const pendingReconciliation = reconciliation; - if (pendingReconciliation) { - await pendingReconciliation.catch(() => undefined); - } - await waitForDispatchIdle(); - return await operation(); - })(); - const barrier = current.then( - () => undefined, - () => undefined, - ); - reconciliation = barrier; - return current.finally(() => { - if (reconciliation === barrier) { - reconciliation = undefined; - } - }); - }; - const runPlacementOperation = async (operation: () => Promise): Promise => { - for (;;) { - const pendingReconciliation = reconciliation; - if (!pendingReconciliation) { - break; - } - await pendingReconciliation.catch(() => undefined); - } - activeDispatchCount += 1; - try { - return await operation(); - } finally { - activeDispatchCount -= 1; - if (activeDispatchCount === 0) { - const waiters = [...dispatchIdleWaiters]; - dispatchIdleWaiters.clear(); - for (const resolve of waiters) { - resolve(); - } - } - } - }; - const dispatchInFlight = new Map< - string, - { - request: WorkerPlacementDispatchRequest; - operation: ReturnType; - } - >(); - return { - dispatch: async (request, onTransition) => { - const inFlight = dispatchInFlight.get(request.sessionId); - if (inFlight) { - if ( - inFlight.request.sessionKey !== request.sessionKey || - inFlight.request.agentId !== request.agentId || - inFlight.request.profileId !== request.profileId || - inFlight.request.deviceId !== request.deviceId || - !isDeepStrictEqual(inFlight.request.inheritedProfile, request.inheritedProfile) - ) { - throw new Error(`Session ${request.sessionKey} is already dispatching another request`); - } - return await inFlight.operation; - } - const operation = runPlacementOperation(() => service.dispatch(request, onTransition)); - dispatchInFlight.set(request.sessionId, { request, operation }); - try { - return await operation; - } finally { - if (dispatchInFlight.get(request.sessionId)?.operation === operation) { - dispatchInFlight.delete(request.sessionId); - } - } - }, - forceDestroyEnvironment: (environmentId, onCleanupError) => - runExclusivePlacementOperation(() => - service.forceDestroyEnvironment(environmentId, onCleanupError), - ), - reclaim: async (request) => await runPlacementOperation(() => service.reclaim(request)), - reconcile: () => runReconciliation(service.reconcile), - reconcileActive: (environmentId) => - environmentId === undefined - ? runReconciliation(() => service.reconcileActive()) - : runExclusivePlacementOperation(() => service.reconcileActive(environmentId)), - }; -} - type WorkerPlacementSidecar = { stop: () => Promise }; export type GatewayWorkerPlacementRuntimeParams = { @@ -521,6 +400,13 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme workspaceOperations, }), ); + const sessionRetirement = createPlacementSessionRetirement({ + placements: params.placements, + environments: params.environments, + forceDestroyEnvironment: dispatchService.forceDestroyEnvironment, + resolveSessionEvidence: resolveWorkerPlacementSessionEvidence, + warn: params.warn, + }); const admissionProvider = createWorkerSessionTurnPlacementProvider({ environments: params.environments, placements: params.placements, @@ -608,7 +494,10 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme } return trackOperation( placementReconcile, - dispatchService.reconcileActive(), + (async () => { + await sessionRetirement.reconcile(); + await dispatchService.reconcileActive(); + })(), "Worker placement reconcile sweep failed", ); }; @@ -626,26 +515,38 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme // Session-lifetime sampling covers idle placements independently of provider health. void sweepDiskSpace(); }; - const sidecar: WorkerPlacementSidecar = { - stop: async () => { - if (stopped) { + const uninstallSessionIdentityMutation = onSessionIdentityMutation((mutation) => { + const previousSessionId = mutation.previous.sessionId; + const currentSessionId = "current" in mutation ? mutation.current.sessionId : undefined; + if (previousSessionId && previousSessionId !== currentSessionId) { + const pending = placementReconcile.current; + if (!pending) { + void reconcileActivePlacements(); return; } + void pending.then(reconcileActivePlacements, reconcileActivePlacements); + } + }); + let stopPromise: Promise | undefined; + const sidecar: WorkerPlacementSidecar = { + stop: () => { + if (stopPromise) { + return stopPromise; + } stopped = true; clearInterval(placementReconcileInterval); placementReconcileInterval = undefined; + uninstallSessionIdentityMutation(); uninstallPlacementAdmission(); - const environmentStop = params.environments.stop(); - const stopResults = await Promise.allSettled([ - ...[placementReconcile.current, diskSpaceSweep.current].filter( - (operation): operation is Promise => operation !== undefined, - ), - environmentStop, - ]); - const environmentStopResult = stopResults.at(-1); - if (environmentStopResult?.status === "rejected") { - throw environmentStopResult.reason; - } + stopPromise = (async () => { + await Promise.allSettled( + [placementReconcile.current, diskSpaceSweep.current].filter( + (operation): operation is Promise => operation !== undefined, + ), + ); + await params.environments.stop(); + })(); + return stopPromise; }, }; // Close must see the drain handle before reconciliation can yield. @@ -665,7 +566,10 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme await sidecar.stop(); return null; } - const startupReconcile = dispatchService.reconcile(); + const startupReconcile = (async () => { + await dispatchService.reconcile(); + await sessionRetirement.reconcile(); + })(); placementReconcile.current = startupReconcile; try { try { diff --git a/src/gateway/server.worker-placement-startup-context.test.ts b/src/gateway/server.worker-placement-startup-context.test.ts index d9232c16b2e0..e2b5ec97da6d 100644 --- a/src/gateway/server.worker-placement-startup-context.test.ts +++ b/src/gateway/server.worker-placement-startup-context.test.ts @@ -1,4 +1,5 @@ -import { afterEach, expect, test } from "vitest"; +import { afterEach, expect, test, vi } from "vitest"; +import { deleteSessionEntryLifecycle } from "../config/sessions/session-accessor.js"; import { getFallbackGatewayContext } from "./server-plugin-fallback-context.js"; import { startGatewayServerHarness, type GatewayServerHarness } from "./server.e2e-ws-harness.js"; import { loadSessionEntry } from "./session-utils.js"; @@ -104,6 +105,37 @@ test( lifecycleRevision: resetLifecycleRevision, }); expect(placements.get(resetSessionId)).toBeUndefined(); + + const createdForExternalDelete = await rpcReq<{ key?: string; sessionId?: string }>( + ws, + "sessions.create", + { agentId: "main", key: "startup-placement-external-delete" }, + ); + const externalSessionId = createdForExternalDelete.payload?.sessionId; + const externalSessionKey = createdForExternalDelete.payload?.key; + if (!externalSessionId || !externalSessionKey) { + throw new Error("external-delete session creation did not return placement identity"); + } + const externalClaim = placements.claimTurn({ + sessionId: externalSessionId, + sessionKey: externalSessionKey, + agentId: "main", + owner: { kind: "local" }, + claimId: "startup-placement-external-delete-claim", + runId: "startup-placement-external-delete-run", + }); + placements.releaseTurn(externalClaim); + const externalTarget = loadSessionEntry(externalSessionKey); + await deleteSessionEntryLifecycle({ + archiveTranscript: false, + storePath: externalTarget.storePath, + target: { + canonicalKey: externalTarget.canonicalKey, + storeKeys: externalTarget.storeKeys, + }, + }); + await vi.waitFor(() => expect(placements.get(externalSessionId)).toBeUndefined()); + expect(getFallbackGatewayContext()?.workerEnvironmentService).toBeDefined(); ws.close(); }, diff --git a/src/gateway/worker-environments/placement-dispatch-coordinator.test.ts b/src/gateway/worker-environments/placement-dispatch-coordinator.test.ts new file mode 100644 index 000000000000..28e08034823e --- /dev/null +++ b/src/gateway/worker-environments/placement-dispatch-coordinator.test.ts @@ -0,0 +1,145 @@ +import { describe, expect, it, vi } from "vitest"; +import { createDeferredCore } from "../../shared/deferred.js"; +import { coordinateWorkerPlacementDispatch } from "./placement-dispatch-coordinator.js"; +import type { WorkerPlacementDispatchService } from "./placement-dispatch.js"; +import type { WorkerPlacementDispatchRequest } from "./service-contract.js"; + +type DispatchService = WorkerPlacementDispatchService; + +const REQUEST: WorkerPlacementDispatchRequest = { + sessionId: "session-1", + sessionKey: "agent:main:session-1", + agentId: "main", + profileId: "test", +}; + +describe("worker placement dispatch coordinator", () => { + it("forwards the optional internal transition observer", async () => { + const observer = vi.fn(); + const dispatch = vi.fn().mockResolvedValue({ state: "active" }); + const service = { + dispatch, + forceDestroyEnvironment: vi.fn(), + reclaim: vi.fn(), + reconcile: vi.fn(), + reconcileActive: vi.fn(), + } as unknown as DispatchService; + + await coordinateWorkerPlacementDispatch(service).dispatch(REQUEST, observer); + + expect(dispatch).toHaveBeenCalledWith(REQUEST, observer); + }); + + it("coalesces an identical dispatch and rejects a conflicting in-flight request", async () => { + const dispatchStarted = createDeferredCore(); + const releaseDispatch = createDeferredCore(); + const active = { state: "active" }; + const dispatch = vi.fn(async () => { + dispatchStarted.resolve(); + await releaseDispatch.promise; + return active; + }); + const service = { + dispatch, + forceDestroyEnvironment: vi.fn(), + reclaim: vi.fn(), + reconcile: vi.fn(), + reconcileActive: vi.fn(), + } as unknown as DispatchService; + const coordinated = coordinateWorkerPlacementDispatch(service); + + const first = coordinated.dispatch(REQUEST); + await dispatchStarted.promise; + await expect( + coordinated.dispatch({ ...REQUEST, profileId: "another-profile" }), + ).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`); + await expect( + coordinated.dispatch({ + ...REQUEST, + inheritedProfile: { + providerId: "fake", + profileSnapshot: { settings: { region: "parent" } }, + }, + }), + ).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`); + const retry = coordinated.dispatch(REQUEST); + releaseDispatch.resolve(); + + const [firstResult, retryResult] = await Promise.all([first, retry]); + expect(retryResult).toBe(firstResult); + expect(dispatch).toHaveBeenCalledOnce(); + + await coordinated.dispatch({ ...REQUEST, profileId: "another-profile" }); + expect(dispatch).toHaveBeenCalledTimes(2); + }); + + it("joins a retry before a queued reconciliation after dispatch failure", async () => { + const dispatchStarted = createDeferredCore(); + const releaseDispatch = createDeferredCore(); + const dispatchError = new Error("provision failed"); + const dispatch = vi.fn(async () => { + dispatchStarted.resolve(); + await releaseDispatch.promise; + throw dispatchError; + }); + const reconcileActive = vi.fn(); + const service = { + dispatch, + forceDestroyEnvironment: vi.fn(), + reclaim: vi.fn(), + reconcile: vi.fn(), + reconcileActive, + } as unknown as DispatchService; + const coordinated = coordinateWorkerPlacementDispatch(service); + + const first = coordinated.dispatch(REQUEST); + await dispatchStarted.promise; + const reconciliation = coordinated.reconcileActive(); + const retry = coordinated.dispatch(REQUEST); + const outcomes = Promise.allSettled([first, retry]); + releaseDispatch.resolve(); + + expect(await outcomes).toEqual([ + { status: "rejected", reason: dispatchError }, + { status: "rejected", reason: dispatchError }, + ]); + await reconciliation; + expect(dispatch).toHaveBeenCalledOnce(); + expect(reconcileActive).toHaveBeenCalledOnce(); + + await expect(coordinated.dispatch({ ...REQUEST, profileId: "another-profile" })).rejects.toBe( + dispatchError, + ); + expect(dispatch).toHaveBeenCalledTimes(2); + }); + + it("coalesces full sweeps but runs a fresh targeted pass with its environment id", async () => { + const fullSweepStarted = createDeferredCore(); + const releaseFullSweep = createDeferredCore(); + const reconcileActive = vi.fn(async (environmentId?: string) => { + if (environmentId === undefined) { + fullSweepStarted.resolve(); + await releaseFullSweep.promise; + } + }); + const service = { + dispatch: vi.fn(), + forceDestroyEnvironment: vi.fn(), + reclaim: vi.fn(), + reconcile: vi.fn(), + reconcileActive, + } as unknown as DispatchService; + const coordinated = coordinateWorkerPlacementDispatch(service); + + const firstFullSweep = coordinated.reconcileActive(); + const secondFullSweep = coordinated.reconcileActive(); + await fullSweepStarted.promise; + const targetedSweep = coordinated.reconcileActive("worker-target"); + + expect(reconcileActive).toHaveBeenCalledTimes(1); + releaseFullSweep.resolve(); + await Promise.all([firstFullSweep, secondFullSweep, targetedSweep]); + + expect(reconcileActive.mock.calls).toEqual([[], ["worker-target"]]); + }); +}); diff --git a/src/gateway/worker-environments/placement-dispatch-coordinator.ts b/src/gateway/worker-environments/placement-dispatch-coordinator.ts new file mode 100644 index 000000000000..029cf511f860 --- /dev/null +++ b/src/gateway/worker-environments/placement-dispatch-coordinator.ts @@ -0,0 +1,122 @@ +import { isDeepStrictEqual } from "node:util"; +import type { WorkerPlacementDispatchService } from "./placement-dispatch.js"; +import type { WorkerPlacementDispatchRequest } from "./service-contract.js"; + +/** Serializes reconciliation sweeps against dispatches and deduplicates exact requests. */ +export function coordinateWorkerPlacementDispatch( + service: WorkerPlacementDispatchService, +): WorkerPlacementDispatchService { + let activeDispatchCount = 0; + let reconciliation: Promise | undefined; + const dispatchIdleWaiters = new Set<() => void>(); + const waitForDispatchIdle = (): Promise => { + if (activeDispatchCount === 0) { + return Promise.resolve(); + } + return new Promise((resolve) => { + dispatchIdleWaiters.add(resolve); + }); + }; + const runReconciliation = (operation: () => Promise): Promise => { + if (reconciliation) { + return reconciliation; + } + const current = (async () => { + await waitForDispatchIdle(); + await operation(); + })(); + reconciliation = current; + const clearCurrent = () => { + if (reconciliation === current) { + reconciliation = undefined; + } + }; + void current.then(clearCurrent, clearCurrent); + return current; + }; + const runExclusivePlacementOperation = (operation: () => Promise): Promise => { + const current = (async () => { + const pendingReconciliation = reconciliation; + if (pendingReconciliation) { + await pendingReconciliation.catch(() => undefined); + } + await waitForDispatchIdle(); + return await operation(); + })(); + const barrier = current.then( + () => undefined, + () => undefined, + ); + reconciliation = barrier; + return current.finally(() => { + if (reconciliation === barrier) { + reconciliation = undefined; + } + }); + }; + const runPlacementOperation = async (operation: () => Promise): Promise => { + for (;;) { + const pendingReconciliation = reconciliation; + if (!pendingReconciliation) { + break; + } + await pendingReconciliation.catch(() => undefined); + } + activeDispatchCount += 1; + try { + return await operation(); + } finally { + activeDispatchCount -= 1; + if (activeDispatchCount === 0) { + const waiters = [...dispatchIdleWaiters]; + dispatchIdleWaiters.clear(); + for (const resolve of waiters) { + resolve(); + } + } + } + }; + const dispatchInFlight = new Map< + string, + { + request: WorkerPlacementDispatchRequest; + operation: ReturnType; + } + >(); + return { + dispatch: async (request, onTransition) => { + const inFlight = dispatchInFlight.get(request.sessionId); + if (inFlight) { + if ( + inFlight.request.sessionKey !== request.sessionKey || + inFlight.request.agentId !== request.agentId || + inFlight.request.profileId !== request.profileId || + inFlight.request.deviceId !== request.deviceId || + !isDeepStrictEqual(inFlight.request.inheritedProfile, request.inheritedProfile) + ) { + throw new Error(`Session ${request.sessionKey} is already dispatching another request`); + } + return await inFlight.operation; + } + const operation = runPlacementOperation(() => service.dispatch(request, onTransition)); + dispatchInFlight.set(request.sessionId, { request, operation }); + try { + return await operation; + } finally { + if (dispatchInFlight.get(request.sessionId)?.operation === operation) { + dispatchInFlight.delete(request.sessionId); + } + } + }, + forceDestroyEnvironment: (environmentId, onCleanupError) => + runExclusivePlacementOperation(() => + service.forceDestroyEnvironment(environmentId, onCleanupError), + ), + reclaim: async (request) => await runPlacementOperation(() => service.reclaim(request)), + reconcile: () => runReconciliation(service.reconcile), + reconcileActive: (environmentId) => + environmentId === undefined + ? runReconciliation(() => service.reconcileActive()) + : runExclusivePlacementOperation(() => service.reconcileActive(environmentId)), + }; +} diff --git a/src/gateway/worker-environments/placement-session-retirement.test.ts b/src/gateway/worker-environments/placement-session-retirement.test.ts new file mode 100644 index 000000000000..152d90554d95 --- /dev/null +++ b/src/gateway/worker-environments/placement-session-retirement.test.ts @@ -0,0 +1,182 @@ +import { describe, expect, it, vi } from "vitest"; +import type { WorkerSessionPlacementRecord } from "./placement-record.js"; +import { createPlacementSessionRetirement } from "./placement-session-retirement.js"; +import type { WorkerSessionPlacementRetirement } from "./placement-store.js"; + +function localPlacement( + sessionId: string, +): Extract { + return { + sessionId, + sessionKey: `agent:main:${sessionId}`, + agentId: "main", + state: "local", + generation: 1, + turnClaim: null, + environmentId: null, + activeOwnerEpoch: null, + workspaceBaseManifestRef: null, + remoteWorkspaceDir: null, + workerBundleHash: null, + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: null, + terminalReason: null, + terminalAtMs: null, + createdAtMs: 1, + updatedAtMs: 1, + stateChangedAtMs: 1, + }; +} + +function activePlacement( + sessionId: string, +): Extract { + return { + ...localPlacement(sessionId), + state: "active", + generation: 2, + turnClaim: null, + environmentId: `environment:${sessionId}`, + activeOwnerEpoch: 3, + workspaceBaseManifestRef: "manifest", + remoteWorkspaceDir: "/workspace", + workerBundleHash: "a".repeat(64), + }; +} + +function failedPlacement( + placement: Extract, +): Extract { + return { + ...placement, + state: "failed", + generation: placement.generation + 1, + turnClaim: null, + recoveryError: "forced teardown", + terminalReason: "failed", + terminalAtMs: 2, + }; +} + +function createHarness(records: WorkerSessionPlacementRecord[]) { + const placements = new Map(records.map((record) => [record.sessionId, record])); + const environments = new Map< + string, + { environmentId: string; state: "attached" | "destroyed"; leaseId: string | null } + >( + records.flatMap((record) => + record.environmentId + ? [ + [ + record.environmentId, + { + environmentId: record.environmentId, + state: "attached" as const, + leaseId: `lease:${record.environmentId}`, + }, + ] as const, + ] + : [], + ), + ); + const retired: WorkerSessionPlacementRetirement[] = []; + const forceDestroyEnvironment = vi.fn(async (environmentId: string) => { + environments.set(environmentId, { + environmentId, + state: "destroyed", + leaseId: null, + }); + for (const record of placements.values()) { + if (record.environmentId === environmentId && record.state === "active") { + placements.set(record.sessionId, failedPlacement(record)); + } + } + }); + const retirement = createPlacementSessionRetirement({ + placements: { + get: (sessionId) => placements.get(sessionId), + list: () => [...placements.values()], + retireSessionPlacement: (input) => { + const current = placements.get(input.sessionId); + if ( + current?.state !== input.expectedState || + current.generation !== input.expectedGeneration || + current.turnClaim + ) { + throw new Error("placement changed"); + } + placements.delete(input.sessionId); + retired.push(input); + }, + }, + environments: { + get: (environmentId) => environments.get(environmentId) as never, + }, + forceDestroyEnvironment, + resolveSessionEvidence: async () => "absent", + warn: vi.fn(), + }); + return { forceDestroyEnvironment, placements, retired, retirement }; +} + +describe("placement session retirement", () => { + it("retires an exact local placement after its session disappears", async () => { + const harness = createHarness([localPlacement("session-local")]); + + await harness.retirement.reconcile(); + + expect(harness.retired).toEqual([ + { + sessionId: "session-local", + expectedState: "local", + expectedGeneration: 1, + }, + ]); + expect(harness.forceDestroyEnvironment).not.toHaveBeenCalled(); + }); + + it("fences a live environment before retiring its failed placement", async () => { + const harness = createHarness([activePlacement("session-active")]); + + await harness.retirement.reconcile(); + + expect(harness.forceDestroyEnvironment).toHaveBeenCalledWith( + "environment:session-active", + expect.any(Function), + ); + expect(harness.retired).toEqual([ + { + sessionId: "session-active", + expectedState: "failed", + expectedGeneration: 3, + }, + ]); + }); + + it("retains current and unknown session evidence", async () => { + const current = localPlacement("session-current"); + const unknown = localPlacement("session-unknown"); + const harness = createHarness([current, unknown]); + const retirement = createPlacementSessionRetirement({ + placements: { + get: (sessionId) => harness.placements.get(sessionId), + list: () => [...harness.placements.values()], + retireSessionPlacement: () => { + throw new Error("must not retire"); + }, + }, + environments: { get: () => undefined }, + forceDestroyEnvironment: async () => { + throw new Error("must not destroy"); + }, + resolveSessionEvidence: async (placement) => + placement.sessionId === current.sessionId ? "current" : "unknown", + warn: vi.fn(), + }); + + await retirement.reconcile(); + + expect(harness.placements.size).toBe(2); + }); +}); diff --git a/src/gateway/worker-environments/placement-session-retirement.ts b/src/gateway/worker-environments/placement-session-retirement.ts new file mode 100644 index 000000000000..f71440852005 --- /dev/null +++ b/src/gateway/worker-environments/placement-session-retirement.ts @@ -0,0 +1,110 @@ +import type { WorkerSessionPlacementRecord } from "./placement-record.js"; +import type { WorkerSessionPlacementStore } from "./placement-store.js"; +import type { WorkerEnvironmentService } from "./service.js"; +import { isFailedWorkerPlacementEnvironmentGone } from "./session-placement-lifecycle.js"; + +type PlacementSessionEvidence = "current" | "absent" | "unknown"; + +type PlacementSessionRetirementDeps = { + placements: Pick; + environments: Pick; + forceDestroyEnvironment: ( + environmentId: string, + onCleanupError?: (error: unknown) => void, + ) => Promise; + resolveSessionEvidence: ( + placement: WorkerSessionPlacementRecord, + ) => Promise; + warn: (message: string) => void; +}; + +export function createPlacementSessionRetirement(deps: PlacementSessionRetirementDeps) { + const retireCurrent = (placement: WorkerSessionPlacementRecord): boolean => { + if (placement.turnClaim) { + return false; + } + const retirement = + placement.state === "local" || placement.state === "reclaimed" + ? { + sessionId: placement.sessionId, + expectedState: placement.state, + expectedGeneration: placement.generation, + } + : placement.state === "failed" && + isFailedWorkerPlacementEnvironmentGone({ + environmentService: deps.environments, + placement, + }) + ? { + sessionId: placement.sessionId, + expectedState: placement.state, + expectedGeneration: placement.generation, + } + : undefined; + if (!retirement) { + return false; + } + deps.placements.retireSessionPlacement(retirement); + return true; + }; + + const reconcilePlacement = async (placement: WorkerSessionPlacementRecord): Promise => { + const evidence = await deps.resolveSessionEvidence(placement); + if (evidence !== "absent") { + return; + } + + let current = deps.placements.get(placement.sessionId); + if (!current) { + return; + } + try { + if (retireCurrent(current)) { + return; + } + } catch { + return; + } + + const environmentId = current.environmentId; + if (!environmentId) { + return; + } + try { + await deps.forceDestroyEnvironment(environmentId, (error) => { + deps.warn( + `Worker placement orphan cleanup deferred for ${current?.sessionId ?? placement.sessionId}: ${String(error)}`, + ); + }); + } catch (error) { + deps.warn( + `Worker placement orphan teardown failed for ${current.sessionId}: ${String(error)}`, + ); + return; + } + + current = deps.placements.get(placement.sessionId); + if (!current) { + return; + } + try { + retireCurrent(current); + } catch { + // A concurrent placement transition owns the next reconciliation pass. + } + }; + + const reconcile = async (): Promise => { + for (const placement of deps.placements.list()) { + try { + await reconcilePlacement(placement); + } catch (error) { + deps.warn( + `Worker placement session evidence check failed for ${placement.sessionId}: ${String(error)}`, + ); + } + } + }; + + return { reconcile }; +}