From b4104e29b5ff70e88abf719fd430e60b0841db96 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sun, 9 Aug 2026 09:03:20 -0700 Subject: [PATCH] fix(sessions): stop active runs before archiving (#120892) * fix(sessions): stop active runs before archiving Archive now fences the exact session, stops and drains active work before commit, preserves main/global-main/unknown protections and Delete behavior, and keeps patchMany per-target ordering. * test(gateway): provide archive lifecycle context * test(ui): expect active sessions to remain archivable * refactor(gateway): keep archive drains internal * fix(gateway): keep abort lifecycle contract acyclic * fix(sessions): fence sharing across archive drains * test(gateway): type archive lifecycle responses * fix(sessions): reclaim cloud workers before archive * fix(sessions): align archive request integration * test(gateway): omit default deferred type * docs(sessions): clarify cloud archive retries --- docs/cli/sessions.md | 10 +- docs/gateway/cloud-workers.md | 2 + docs/gateway/protocol.md | 2 +- docs/web/control-ui.md | 4 +- .../tools/sessions-tool.self-archive.test.ts | 8 +- src/agents/tools/sessions-tool.ts | 8 +- src/commands/sessions-lifecycle.test.ts | 2 +- src/commands/sessions-lifecycle.ts | 3 +- src/gateway/chat-abort-lifecycle-internal.ts | 67 ++ src/gateway/chat-abort.ts | 3 + .../server-methods.authorization.test.ts | 74 ++ .../chat-abort-authorization.ts | 18 +- .../server-methods/chat-abort-runtime.ts | 78 +- .../sessions-archive-lifecycle.ts | 213 ++++ ...ions-mutations.archive-attribution.test.ts | 2 + .../sessions-mutations.perf.test.ts | 2 + .../server-methods/sessions-patch-archive.ts | 266 +++++ .../server-methods/sessions-patch-engine.ts | 483 ++++---- src/gateway/server-methods/sessions-shared.ts | 7 +- .../server-methods/sessions-sharing.ts | 31 +- src/gateway/server-runtime-subscriptions.ts | 6 + .../server.sessions.archive-lifecycle.test.ts | 1055 +++++++++++++++++ ....sessions.archive-worker-placement.test.ts | 446 +++++++ .../server.sessions.compaction.test.ts | 15 +- .../server.sessions.delete-lifecycle.test.ts | 20 - .../server.sessions.list-changed.test.ts | 3 + .../server.sessions.reset-cleanup.test.ts | 2 +- src/gateway/session-sharing.test.ts | 1 + src/gateway/session-sharing.ts | 6 +- src/gateway/test-helpers.mocks.ts | 1 + .../inference-control-internal.ts | 28 + .../worker-environments/inference.test.ts | 59 + src/gateway/worker-environments/inference.ts | 82 +- src/gateway/worker-environments/service.ts | 13 +- .../session-placement-lifecycle.ts | 60 +- src/sessions/session-lifecycle-admission.ts | 2 - src/shared/session-archive-timeout.ts | 3 + ui/src/components/session-menu.test.ts | 20 +- ui/src/components/session-menu.ts | 8 +- .../session-organizer-batch-mutations.test.ts | 20 +- .../session-organizer-batch-mutations.ts | 13 +- ui/src/components/sidebar-menus-render.ts | 3 + .../session-management.archive.e2e.test.ts | 36 +- .../e2e/session-management.groups.e2e.test.ts | 7 +- ui/src/lib/sessions/list-options.test.ts | 5 +- ui/src/lib/sessions/session-key.test.ts | 29 + ui/src/lib/sessions/session-key.ts | 29 +- ui/src/lib/sessions/session-requests.ts | 9 +- ui/src/pages/sessions/sessions-page.test.ts | 24 + ui/src/pages/sessions/sessions-page.ts | 15 +- 50 files changed, 2918 insertions(+), 385 deletions(-) create mode 100644 src/gateway/chat-abort-lifecycle-internal.ts create mode 100644 src/gateway/server-methods/sessions-archive-lifecycle.ts create mode 100644 src/gateway/server-methods/sessions-patch-archive.ts create mode 100644 src/gateway/server.sessions.archive-lifecycle.test.ts create mode 100644 src/gateway/server.sessions.archive-worker-placement.test.ts create mode 100644 src/gateway/worker-environments/inference-control-internal.ts create mode 100644 src/shared/session-archive-timeout.ts diff --git a/docs/cli/sessions.md b/docs/cli/sessions.md index 8b298f37d242..b0412cd6033c 100644 --- a/docs/cli/sessions.md +++ b/docs/cli/sessions.md @@ -93,9 +93,13 @@ openclaw sessions archive "agent:main:scratch-1" --json Archive uses the same `sessions.patch` lifecycle operation as the Control UI. It keeps the transcript, marks the session archived, and removes the session -from the default active list. Already archived sessions are successful no-ops. -Use `--dry-run` to validate every key and preview the result without changing -session state. +from the default active list. For a cloud-worker session with an active +placement, the Gateway first stops the worker, reconciles its workspace, and +reclaims the environment. If the placement is still transitioning or failed +without proof that its environment is gone, the session remains unarchived; +wait for the placement to settle, then retry. Agent main sessions remain +protected. Already archived sessions are successful no-ops. Use `--dry-run` to +validate every key and preview the result without changing session state. ## Delete sessions diff --git a/docs/gateway/cloud-workers.md b/docs/gateway/cloud-workers.md index c318624798d3..15e2535adc08 100644 --- a/docs/gateway/cloud-workers.md +++ b/docs/gateway/cloud-workers.md @@ -170,6 +170,8 @@ While a fenced result is still reconciling, a new turn waits up to 15 seconds fo When the work is complete and no turn is running, open the session menu and choose **Stop cloud worker…**. The Gateway performs one final workspace reconciliation before it destroys the environment. A placement already in `draining` or `reconciling` is finishing teardown; wait for its badge to become `reclaimed` before deleting the session. +Archiving a non-main cloud-worker session with an active placement also performs this safe stop and reclaim before the Gateway records it as archived. If the placement is still transitioning or failed without proof that its environment is gone, the session remains unarchived; wait for the placement to settle, then retry. Restoring the session retains the reclaimed placement metadata so the next turn can dispatch a fresh worker with the same workspace profile. + For a broken or runaway attached worker, an operator can call `environments.destroy` with `{ "force": true }` as a last resort. Forced teardown durably marks the placement failed and abandons any unreconciled remote result before destroying the environment. The equivalent administrative RPC is: diff --git a/docs/gateway/protocol.md b/docs/gateway/protocol.md index 03aa77ea62ed..5dfd7151ee43 100644 --- a/docs/gateway/protocol.md +++ b/docs/gateway/protocol.md @@ -643,7 +643,7 @@ methods. Treat this as feature discovery, not a full enumeration of - `sessions.send` sends a message into an existing session. - `sessions.steer` is the interrupt-and-steer variant for an active session. - `sessions.abort` aborts active work for a session. Pass `key` plus optional `runId`, or `runId` alone for active runs the gateway can resolve to a session. Supplying `runId` keeps cancellation scoped to that run. Set `clearQueued: true` on a key-only non-global request to also discard followup and lane queues owned by that session. Existing callers that omit `clearQueued` preserve those queues. The literal `global` key keeps the existing agent-qualified `chat.abort` ownership rules and does not perform non-global followup or lane cleanup. - - `sessions.patch` updates session metadata/overrides and reports the resolved canonical model plus effective `agentRuntime`. Session organization fields and the per-session `model` override require `operator.write`; thinking, fast, verbose, trace, reasoning, and other privileged overrides require `operator.admin`. Only an admin model selection can persist as the configured agent default. Spawn lineage (`spawnedBy`, `spawnedWorkspaceDir`, `spawnedCwd`, `spawnDepth`, `subagentRole`, `subagentControlScope`) is no longer publicly patchable; those facts are written once by trusted creation paths, and requests that still send them are rejected. + - `sessions.patch` updates session metadata/overrides and reports the resolved canonical model plus effective `agentRuntime`. Session organization fields and the per-session `model` override require `operator.write`; thinking, fast, verbose, trace, reasoning, and other privileged overrides require `operator.admin`. Only an admin model selection can persist as the configured agent default. With `archived: true`, the Gateway protects agent main sessions (including `global` when global scope is configured) and the `unknown` sentinel; for every other real session it first fences new admission, cancels exact-session active, pending, queued, reply, embedded, and worker work, and waits for admission and runtime terminal-persistence drains before committing `archivedAt`. A cancellation, drain, or persistence failure returns retryable `UNAVAILABLE` and leaves the session unarchived. `sessions.patchMany` prepares archive targets in input order inside the same batch lifecycle fence and returns ordered per-target outcomes. Spawn lineage (`spawnedBy`, `spawnedWorkspaceDir`, `spawnedCwd`, `spawnDepth`, `subagentRole`, `subagentControlScope`) is no longer publicly patchable; those facts are written once by trusted creation paths, and requests that still send them are rejected. - `sessions.reset`, `sessions.delete`, and `sessions.compact` perform session maintenance. - `sessions.get` returns the full stored session row. - Chat execution still uses `chat.history`, `chat.send`, `chat.abort`, and `chat.inject`. `chat.history` is display-normalized for UI clients: inline directive tags are stripped from visible text, plain-text tool-call XML payloads (`...`, `...`, `...`, `...`, and truncated tool-call blocks) and leaked ASCII/full-width model control tokens are stripped, pure silent-token assistant rows (exact `NO_REPLY` / `no_reply`) are omitted, and oversized rows can be replaced with placeholders. diff --git a/docs/web/control-ui.md b/docs/web/control-ui.md index bc11ad8c67ce..b60a178a2200 100644 --- a/docs/web/control-ui.md +++ b/docs/web/control-ui.md @@ -264,7 +264,7 @@ select it to open the owning Approvals page. - Channels: built-in plus bundled/external plugin channels status, QR login, and per-channel config (`channels.status`, `web.login.*`, `config.patch`). - Channel probe refreshes keep the previous snapshot visible while slow provider checks finish, and label partial snapshots when a probe or audit exceeds its UI budget. - - Threads (a workspace page at `/sessions`, with a **Worktrees** tab alongside it): list configured-agent sessions by default, pin frequent sessions, rename them, archive or restore inactive sessions, fall back from stale unconfigured agent session keys, and apply per-session model/thinking/fast/verbose/trace/reasoning overrides (`sessions.list`, `sessions.patch`). A three-way **Active / Archived / All** filter controls both this page and the sidebar; All dims archived rows and labels them explicitly. Archived sessions keep their transcripts, are never auto-pruned, and remain shelved until explicitly unarchived or deleted. Rows show an unread dot for active sessions with activity since they were last read, with mark-unread/mark-read actions (`sessions.patch { unread }`), and a Fork action that branches the transcript into a new session (`sessions.create { parentSessionKey, fork: true }`). Overview tiles above the table summarize the loaded roster (session count, live runs, unread sessions, total tokens, and archived count when available), each row carries a kind glyph with a live-run dot, status renders as a plain dot plus label, and the Tokens column shows a context-window usage meter when the session reports token and context sizes. Row management actions live in a per-row menu (kebab button or right-click) mirroring the sidebar's session menu, and the row drawer carries the agent runtime and run duration alongside the other session details. + - Threads (a workspace page at `/sessions`, with a **Worktrees** tab alongside it): list configured-agent sessions by default, pin frequent sessions, rename them, archive or restore sessions, fall back from stale unconfigured agent session keys, and apply per-session model/thinking/fast/verbose/trace/reasoning overrides (`sessions.list`, `sessions.patch`). A three-way **Active / Archived / All** filter controls both this page and the sidebar; All dims archived rows and labels them explicitly. Archived sessions keep their transcripts, are never auto-pruned, and remain shelved until explicitly unarchived or deleted. Rows show an unread dot for active sessions with activity since they were last read, with mark-unread/mark-read actions (`sessions.patch { unread }`), and a Fork action that branches the transcript into a new session (`sessions.create { parentSessionKey, fork: true }`). Overview tiles above the table summarize the loaded roster (session count, live runs, unread sessions, total tokens, and archived count when available), each row carries a kind glyph with a live-run dot, status renders as a plain dot plus label, and the Tokens column shows a context-window usage meter when the session reports token and context sizes. Row management actions live in a per-row menu (kebab button or right-click) mirroring the sidebar's session menu, and the row drawer carries the agent runtime and run duration alongside the other session details. - Native Claude and Codex sidebar catalogs stream one host at a time, then reconcile after node connectivity changes, on page focus, and at most every 30 seconds while visible. Catalog changes trigger a faster follow-up pass, so sessions created in the native tools appear without reloading the Control UI. Claude Desktop rows also retain their local custom-group label when present; OpenClaw reads that mapping from Desktop's local store and never writes it. - Session grouping: a Group by control organizes the sessions table into sections by custom groups, channel, kind, agent, or date. Custom groups persist per session via `sessions.patch` (`category`), so sessions started from message channels (Discord, Telegram, WhatsApp, ...) can be categorized too; assign groups by dragging rows onto a section, or with the per-row group selector, and create groups with the New group action. - Memory (a tab on the Agents page, scoped to the selected agent): dreaming status, enable/disable toggle, and Dream Diary reader (`doctor.memory.status`, `doctor.memory.dreamDiary`, `config.patch`). When the `memory-wiki` plugin is enabled, the Diary view adds **Imported Insights** and **Memory Wiki** sub-tabs that browse imported source chats and the compiled wiki — clustered synthesis, entity, and concept pages plus annotated sources and reports, with claims, open questions, contradictions, and inline page previews (`wiki.importInsights`, `wiki.overview`, `wiki.get`). @@ -463,7 +463,7 @@ Capability toggles stay disabled until the Gateway, session, and runtime config - `chat.inject` appends an assistant note to the session transcript and broadcasts a `chat` event for UI-only updates (no agent run, no channel delivery). - The sidebar lists every loaded active session by agent section and pinned/channel/work/custom/Chats buckets with a single New Session action that opens the draft dialog. Opening a visible row moves only the highlight. Sessions can be dropped onto Pinned to pin them, or onto a custom group or Chats to move them; custom groups are collapsible and drag-reorderable, group names and order sync through the gateway, and collapsed state stays in the browser. A new dashboard session asynchronously gets a concise generated title from its first non-command message; explicit names and authenticated sender identity remain separate, so account names are never used as generated titles. When New Session creates a worktree without an explicit worktree name, OpenClaw also uses the session label or generated title for its branch name, falling back to a readable crustacean-themed name. Set `agents.defaults.utilityModel` (or `agents.entries.*.utilityModel`) to route this separate model call to a lower-cost model; if that distinct model fails, title generation retries once with the primary model. Expanding another agent section browses that agent's sessions without leaving the open chat. - Thread search lives in the command palette (⌘K, or the search button in the top-left control cluster): typing a query follows a bounded number of matching pages across agents, filters internal child/cron rows, and lists visible matches next to navigation commands. The Threads page keeps the exhaustive searchable list with filters. - - Each sidebar row keeps direct pin access plus a full context menu for unread state, rename, fork, grouping, archive, and delete. Multi-selected rows (Cmd/Ctrl-click, Shift-click for ranges) get a batch menu covering unread state, grouping, archive, and delete; batch Archive reports per-session failures while archiving eligible rows, whereas batch Delete stays disabled unless every selected session is archivable. An active run and an agent's main session cannot be archived. Archiving or deleting the currently selected session switches Chat back to that agent's main session. + - Each sidebar row keeps direct pin access plus a full context menu for unread state, rename, fork, grouping, archive, and delete. Multi-selected rows (Cmd/Ctrl-click, Shift-click for ranges) get a batch menu covering unread state, grouping, archive, and delete; batch Archive reports per-session failures while archiving eligible rows, whereas batch Delete keeps its separate idle-or-already-archived eligibility. Archive stays disabled for agent main sessions (including `global` in global scope) and the `unknown` sentinel. For any other session, including one with active work, the Gateway stops and fully drains that session's work before archiving it. The selected archived session stays open with an archived notice and **Unarchive** action; deleting the selected session switches Chat back to that agent's main session. - In the macOS app, the OpenClaw mark uses the otherwise-empty native titlebar strip next to the window controls instead of consuming a sidebar row. - On desktop widths, chat controls stay on one compact row and collapse while scrolling down the transcript; scrolling up, returning to the top, or reaching the bottom restores the controls. - The session header shows a small facepile beside the workspace chip when other people are viewing the same session; it lists up to four viewer avatars with an overflow count and disappears when you are alone. diff --git a/src/agents/tools/sessions-tool.self-archive.test.ts b/src/agents/tools/sessions-tool.self-archive.test.ts index aa1167c08128..d4234f2a9222 100644 --- a/src/agents/tools/sessions-tool.self-archive.test.ts +++ b/src/agents/tools/sessions-tool.self-archive.test.ts @@ -243,7 +243,7 @@ describe("sessions tool self-archive", () => { identities: [sessionKey, sessionId], assertAllowed: () => {}, }); - throw new Error("Cannot archive a session with an active run."); + throw Object.assign(new Error("Session did not finish stopping."), { retryable: true }); } return { ok: true }; }); @@ -308,7 +308,7 @@ describe("sessions tool self-archive", () => { }); competingAdmission.release(); competingTurnFinished = true; - throw new Error("Cannot archive a session with an active run."); + throw Object.assign(new Error("Session did not finish stopping."), { retryable: true }); } return { ok: true }; }); @@ -409,7 +409,9 @@ describe("sessions tool self-archive", () => { const callGateway = vi.fn(async () => { attempts += 1; if (attempts <= 10) { - throw new Error("Cannot archive a session with an active run."); + throw Object.assign(new Error("Session did not finish stopping."), { + retryable: true, + }); } return { ok: true }; }); diff --git a/src/agents/tools/sessions-tool.ts b/src/agents/tools/sessions-tool.ts index b5ea8a93f5cd..784a50a09696 100644 --- a/src/agents/tools/sessions-tool.ts +++ b/src/agents/tools/sessions-tool.ts @@ -18,7 +18,6 @@ import { isIncognitoSessionKey, resolveAgentIdFromSessionKey } from "../../routi import { getCurrentSessionWorkAdmissionRelease, getSessionWorkAdmissionRelease, - SESSION_ARCHIVE_ACTIVE_RUN_ERROR, } from "../../sessions/session-lifecycle-admission.js"; import { resolveDefaultAgentId } from "../agent-scope-config.js"; import { stringEnum } from "../schema/typebox.js"; @@ -436,7 +435,6 @@ export function createSessionsTool(opts: SessionsToolOptions = {}): AnyAgentTool // admitted owner, or retry a transient gateway disconnect, // instead of losing an archive that was already scheduled. const message = formatErrorMessage(error); - const activeRun = message.includes(SESSION_ARCHIVE_ACTIVE_RUN_ERROR); const retryableGatewayFailure = error instanceof GatewayTransportError || isTransientNetworkError(error) || @@ -444,12 +442,10 @@ export function createSessionsTool(opts: SessionsToolOptions = {}): AnyAgentTool error !== null && "retryable" in error && error.retryable === true); - if (!activeRun && !retryableGatewayFailure) { + if (!retryableGatewayFailure) { throw error; } - if (!activeRun) { - log.warn(`retrying deferred self-archive for ${key}: ${message}`); - } + log.warn(`retrying deferred self-archive for ${key}: ${message}`); const retryAfterRelease = getSessionWorkAdmissionRelease({ scope: storePath, identities: archiveIdentities, diff --git a/src/commands/sessions-lifecycle.test.ts b/src/commands/sessions-lifecycle.test.ts index 38a6603da1df..c82b36bed99f 100644 --- a/src/commands/sessions-lifecycle.test.ts +++ b/src/commands/sessions-lifecycle.test.ts @@ -90,7 +90,7 @@ describe("sessions lifecycle commands", () => { expectedSessionId: "session-1", archived: true, }, - { defaultTimeoutMs: 30_000 }, + { defaultTimeoutMs: 10 * 60_000 }, ); expect(runtime.writeJson).toHaveBeenCalledWith( { diff --git a/src/commands/sessions-lifecycle.ts b/src/commands/sessions-lifecycle.ts index 51eda0f8b825..d4b973ffa439 100644 --- a/src/commands/sessions-lifecycle.ts +++ b/src/commands/sessions-lifecycle.ts @@ -3,6 +3,7 @@ import { formatCliCommand } from "../cli/command-format.js"; import { callGatewayFromCliWithTransport } from "../cli/gateway-rpc.js"; import { formatErrorMessage } from "../infra/errors.js"; import { type RuntimeEnv, writeRuntimeJson } from "../runtime.js"; +import { SESSION_ARCHIVE_REQUEST_TIMEOUT_MS } from "../shared/session-archive-timeout.js"; import { createClackPrompter } from "../wizard/clack-prompter.js"; type SessionsLifecycleCliOptions = { @@ -266,7 +267,7 @@ async function runSessionsLifecycleCommand( ...(session.sessionId ? { expectedSessionId: session.sessionId } : {}), archived: true, }, - { defaultTimeoutMs: 30_000 }, + { defaultTimeoutMs: SESSION_ARCHIVE_REQUEST_TIMEOUT_MS }, )) as SessionsPatchResult; if (response?.ok !== true || response.entry?.archivedAt === undefined) { throw new Error("Gateway did not confirm that the session was archived."); diff --git a/src/gateway/chat-abort-lifecycle-internal.ts b/src/gateway/chat-abort-lifecycle-internal.ts new file mode 100644 index 000000000000..014473a2e710 --- /dev/null +++ b/src/gateway/chat-abort-lifecycle-internal.ts @@ -0,0 +1,67 @@ +const terminalPersistenceErrorByEntry = new WeakMap(); +const removalWaitersByEntry = new WeakMap void>>(); + +export function markChatAbortTerminalPersistenceError(entry: object, error: unknown): void { + if (error === undefined) { + terminalPersistenceErrorByEntry.delete(entry); + return; + } + terminalPersistenceErrorByEntry.set(entry, error); +} + +export function notifyChatAbortControllerRemoved(entry: object): void { + const waiters = removalWaitersByEntry.get(entry); + removalWaitersByEntry.delete(entry); + for (const resolve of waiters ?? []) { + resolve(); + } +} + +/** Waits for captured run registrations and their terminal persistence owner to leave. */ +export async function waitForChatAbortControllerRemoval(params: { + entries: ReadonlyMap; + targets: ReadonlyArray<{ runId: string; entry: TEntry }>; + timeoutMs: number; +}): Promise { + const registeredWaiters: Array<{ entry: TEntry; resolve: () => void }> = []; + const removals = params.targets.flatMap(({ runId, entry }) => { + if (params.entries.get(runId) !== entry) { + return []; + } + return [ + new Promise((resolve) => { + const waiters = removalWaitersByEntry.get(entry) ?? new Set<() => void>(); + waiters.add(resolve); + removalWaitersByEntry.set(entry, waiters); + registeredWaiters.push({ entry, resolve }); + }), + ]; + }); + if (removals.length === 0) { + return true; + } + let timer: ReturnType | undefined; + try { + const removed = await Promise.race([ + Promise.all(removals).then(() => true), + new Promise((resolve) => { + timer = setTimeout(() => resolve(false), Math.max(0, params.timeoutMs)); + timer.unref?.(); + }), + ]); + return ( + removed && params.targets.every(({ entry }) => !terminalPersistenceErrorByEntry.has(entry)) + ); + } finally { + if (timer) { + clearTimeout(timer); + } + for (const { entry, resolve } of registeredWaiters) { + const waiters = removalWaitersByEntry.get(entry); + waiters?.delete(resolve); + if (waiters?.size === 0) { + removalWaitersByEntry.delete(entry); + } + } + } +} diff --git a/src/gateway/chat-abort.ts b/src/gateway/chat-abort.ts index 0a9035e29ab2..513906494062 100644 --- a/src/gateway/chat-abort.ts +++ b/src/gateway/chat-abort.ts @@ -18,6 +18,7 @@ import { } from "../infra/agent-events.js"; import { jsonUtf8Bytes } from "../infra/json-utf8-bytes.js"; import { parseAgentSessionKey } from "../routing/session-key.js"; +import { notifyChatAbortControllerRemoved } from "./chat-abort-lifecycle-internal.js"; import { projectLiveAssistantBufferedText } from "./live-chat-projector.js"; import { createChatAbortMarker, @@ -564,6 +565,8 @@ export function removeChatAbortControllerEntry( entry.onRemoved?.(); } catch { // Removal owns state cleanup even if a caller-provided release hook fails. + } finally { + notifyChatAbortControllerRemoved(entry); } return true; } diff --git a/src/gateway/server-methods.authorization.test.ts b/src/gateway/server-methods.authorization.test.ts index db586d09731a..8d5b79513925 100644 --- a/src/gateway/server-methods.authorization.test.ts +++ b/src/gateway/server-methods.authorization.test.ts @@ -362,6 +362,8 @@ describe("sessions.patchMany orchestration", () => { broadcastToConnIds: vi.fn(), getSessionEventSubscriberConnIds: () => new Set(), chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + dedupe: new Map(), ...overrides, }) as never; @@ -729,6 +731,78 @@ describe("sessions.patchMany orchestration", () => { }); }); + it("isolates archive preparation authorization per target and continues in input order", async () => { + await withOpenClawTestState({ scenario: "minimal" }, async () => { + for (let index = 0; index < 3; index += 1) { + await upsertSessionEntry( + { agentId: "main", sessionKey: `agent:main:archive-auth-${index}` }, + { sessionId: `session-archive-auth-${index}`, updatedAt: 1 }, + ); + } + const respond = vi.fn(); + const assertCurrent = vi.fn(() => { + throw new Error("outer all-target guard must not be used"); + }); + const assertTargetCurrent = vi.fn(({ sessionKey }: { sessionKey: string }) => { + if (sessionKey.endsWith("-1")) { + throw new SessionMutationAuthorizationChangedError({ + code: "INVALID_REQUEST", + message: "archive authorization changed; retry the request", + }); + } + }); + + await sessionMutationHandlers["sessions.patchMany"]!({ + params: { + targets: [0, 1, 2].map((index) => ({ + key: `agent:main:archive-auth-${index}`, + })), + patch: { archived: true }, + }, + respond, + context: context(), + client: { connect: { scopes: ["operator.write"] } }, + sessionMutationAuthorization: { assertCurrent, assertTargetCurrent }, + } as never); + + expect(assertCurrent).not.toHaveBeenCalled(); + expect(assertTargetCurrent.mock.calls.map(([target]) => target.sessionKey)).toEqual([ + "agent:main:archive-auth-0", + "agent:main:archive-auth-1", + "agent:main:archive-auth-2", + "agent:main:archive-auth-0", + "agent:main:archive-auth-2", + ]); + expect(respond).toHaveBeenCalledWith( + true, + { + outcomes: [ + { ok: true, key: "agent:main:archive-auth-0" }, + { + ok: false, + key: "agent:main:archive-auth-1", + error: { + code: "INVALID_REQUEST", + message: "archive authorization changed; retry the request", + }, + }, + { ok: true, key: "agent:main:archive-auth-2" }, + ], + }, + undefined, + ); + expect( + loadSessionEntry({ agentId: "main", sessionKey: "agent:main:archive-auth-0" }), + ).toHaveProperty("archivedAt"); + expect( + loadSessionEntry({ agentId: "main", sessionKey: "agent:main:archive-auth-1" }), + ).not.toHaveProperty("archivedAt"); + expect( + loadSessionEntry({ agentId: "main", sessionKey: "agent:main:archive-auth-2" }), + ).toHaveProperty("archivedAt"); + }); + }); + it("converts an unexpected target exception into an ordered isolated failure", async () => { await withOpenClawTestState({ scenario: "minimal" }, async () => { for (let index = 0; index < 3; index += 1) { diff --git a/src/gateway/server-methods/chat-abort-authorization.ts b/src/gateway/server-methods/chat-abort-authorization.ts index 467abae673e4..f18cfd54baae 100644 --- a/src/gateway/server-methods/chat-abort-authorization.ts +++ b/src/gateway/server-methods/chat-abort-authorization.ts @@ -267,6 +267,7 @@ export function resolveAuthorizedPreRegisteredRunsForSessionKeys(params: { requester: ChatAbortRequester; keyPrefix: string; preserveSideRuns?: boolean; + includeProtectedRuns?: boolean; excludeRunIds?: ReadonlySet; }) { const sessionKeys = new Set( @@ -318,8 +319,9 @@ export function resolveAuthorizedPreRegisteredRunsForSessionKeys(params: { } const requesterCanAbort = canRequesterAbortPreRegisteredRun(run.payload, params.requester); const isProtected = - run.payload.controlUiVisible === false || - (params.preserveSideRuns && normalizeUnknownText(run.payload.turnKind) === "btw"); + params.includeProtectedRuns !== true && + (run.payload.controlUiVisible === false || + (params.preserveSideRuns && normalizeUnknownText(run.payload.turnKind) === "btw")); if (isProtected) { // Broad lifecycle cleanup still needs ownership, while ordinary chat.abort // must keep treating hidden or preserved work as a non-match. @@ -351,6 +353,7 @@ export function resolveAuthorizedRunsForSessionKeys(params: { defaultAgentId: string; requester: ChatAbortRequester; preserveSideRuns?: boolean; + includeProtectedRuns?: boolean; excludeRunIds?: ReadonlySet; }) { const sessionKeys = new Set( @@ -364,7 +367,11 @@ export function resolveAuthorizedRunsForSessionKeys(params: { ), ); const agentId = normalizeOptionalText(params.agentId)?.toLowerCase(); - const authorizedRuns: Array<{ runId: string; sessionKey: string }> = []; + const authorizedRuns: Array<{ + runId: string; + sessionKey: string; + entry: ChatAbortControllerEntry; + }> = []; const matchedRunIds: string[] = []; let hasUnauthorizedRuns = false; let hasUnauthorizedProtectedRuns = false; @@ -388,7 +395,8 @@ export function resolveAuthorizedRunsForSessionKeys(params: { matchedRunIds.push(runId); const requesterCanAbort = canRequesterAbortChatRun(active, params.requester); const isProtected = - active.controlUiVisible === false || (params.preserveSideRuns && active.turnKind === "btw"); + params.includeProtectedRuns !== true && + (active.controlUiVisible === false || (params.preserveSideRuns && active.turnKind === "btw")); if (isProtected) { // Broad lifecycle cleanup still needs ownership, while ordinary chat.abort // must keep treating hidden or preserved work as a non-match. @@ -399,7 +407,7 @@ export function resolveAuthorizedRunsForSessionKeys(params: { continue; } if (requesterCanAbort) { - authorizedRuns.push({ runId, sessionKey: active.sessionKey }); + authorizedRuns.push({ runId, sessionKey: active.sessionKey, entry: active }); } else { hasUnauthorizedRuns = true; } diff --git a/src/gateway/server-methods/chat-abort-runtime.ts b/src/gateway/server-methods/chat-abort-runtime.ts index f0c2fa4421e9..819fb7bf184b 100644 --- a/src/gateway/server-methods/chat-abort-runtime.ts +++ b/src/gateway/server-methods/chat-abort-runtime.ts @@ -25,6 +25,8 @@ import type { GatewayRequestContext } from "./types.js"; type AbortOrigin = "rpc" | "stop-command"; +const SESSION_LIFECYCLE_ABORT_REQUESTER: ChatAbortRequester = { isAdmin: true }; + type AbortedPartialSnapshot = { runId: string; sessionId: string; @@ -34,24 +36,20 @@ type AbortedPartialSnapshot = { }; function collectSessionAbortPartials(params: { - chatAbortControllers: Map; chatRunState: GatewayRequestContext["chatRunState"]; - runIds: ReadonlySet; + runs: ReadonlyArray<{ runId: string; entry: ChatAbortControllerEntry }>; abortOrigin: AbortOrigin; }): AbortedPartialSnapshot[] { const out: AbortedPartialSnapshot[] = []; - for (const [runId, active] of params.chatAbortControllers) { - if (!params.runIds.has(runId)) { - continue; - } + for (const { runId, entry } of params.runs) { const text = params.chatRunState.resolveBuffer(runId).text; if (!text || !text.trim()) { continue; } out.push({ runId, - sessionId: active.sessionId, - agentId: active.agentId, + sessionId: entry.sessionId, + agentId: entry.agentId, text, abortOrigin: params.abortOrigin, }); @@ -134,6 +132,56 @@ function resolveAuthorizedQueuedTurnsForSession(params: { }; } +type SessionAbortOwnerParams = { + context: GatewayRequestContext; + sessionKeys: string[]; + sessionId?: string; + agentId?: string; + defaultAgentId: string; +}; + +/** Authoritative active, pending, or queued Gateway owner for an exact session. */ +export function hasGatewaySessionAbortOwner(params: SessionAbortOwnerParams): boolean { + const active = resolveAuthorizedRunsForSessionKeys({ + chatAbortControllers: params.context.chatAbortControllers, + sessionKeys: params.sessionKeys, + sessionIds: [params.sessionId], + agentId: params.agentId, + defaultAgentId: params.defaultAgentId, + requester: SESSION_LIFECYCLE_ABORT_REQUESTER, + includeProtectedRuns: true, + }); + if (active.authorizedRuns.length > 0) { + return true; + } + const queued = resolveAuthorizedQueuedTurnsForSession({ + context: params.context, + sessionKeys: params.sessionKeys, + sessionId: params.sessionId, + agentId: params.agentId, + defaultAgentId: params.defaultAgentId, + requester: SESSION_LIFECYCLE_ABORT_REQUESTER, + }); + if (queued.authorized.length > 0) { + return true; + } + for (const keyPrefix of ["agent:", PENDING_CHAT_SEND_DEDUPE_PREFIX]) { + const pending = resolveAuthorizedPreRegisteredRunsForSessionKeys({ + context: params.context, + sessionKeys: params.sessionKeys, + agentId: params.agentId, + defaultAgentId: params.defaultAgentId, + requester: SESSION_LIFECYCLE_ABORT_REQUESTER, + keyPrefix, + includeProtectedRuns: true, + }); + if (pending.authorizedRuns.length > 0) { + return true; + } + } + return false; +} + export function cancelWorkerInferenceForSession(params: { context: GatewayRequestContext; sessionId?: string; @@ -164,7 +212,13 @@ export async function abortChatRunsForSessionKeyWithPartials(params: { stopReason?: string; requester: ChatAbortRequester; preserveSideRuns?: boolean; + /** Exact lifecycle owners may include hidden and side runs for this one session. */ + includeProtectedRuns?: boolean; excludeRunIds?: ReadonlySet; + /** Captures exact registrations before cancellation can remove them. */ + onControllerTargets?: ( + targets: Array<{ runId: string; entry: ChatAbortControllerEntry }>, + ) => void; /** Internal session-wide cleanup after exact resolution and all matching owner checks. */ onAuthorizedAfterQueuedAbort?: () => boolean; }): Promise<{ aborted: boolean; runIds: string[]; unauthorized: boolean }> { @@ -191,6 +245,7 @@ export async function abortChatRunsForSessionKeyWithPartials(params: { defaultAgentId: params.defaultAgentId, requester: params.requester, preserveSideRuns: params.preserveSideRuns, + includeProtectedRuns: params.includeProtectedRuns, excludeRunIds: params.excludeRunIds, }); const { @@ -206,6 +261,7 @@ export async function abortChatRunsForSessionKeyWithPartials(params: { requester: params.requester, keyPrefix: "agent:", preserveSideRuns: params.preserveSideRuns, + includeProtectedRuns: params.includeProtectedRuns, excludeRunIds: params.excludeRunIds, }); const { @@ -221,6 +277,7 @@ export async function abortChatRunsForSessionKeyWithPartials(params: { requester: params.requester, keyPrefix: PENDING_CHAT_SEND_DEDUPE_PREFIX, preserveSideRuns: params.preserveSideRuns, + includeProtectedRuns: params.includeProtectedRuns, excludeRunIds: params.excludeRunIds, }); const hasAuthorizedGatewayRuns = @@ -264,6 +321,7 @@ export async function abortChatRunsForSessionKeyWithPartials(params: { // Keep ordinary chat.abort's admin worker behavior; only the injected broad // lifecycle path must preserve hidden or explicitly preserved Gateway runs. const canCancelWorkerSession = !params.onAuthorizedAfterQueuedAbort || !hasProtectedLifecycleRuns; + params.onControllerTargets?.(authorizedRuns); if (!hasAuthorizedGatewayRuns) { // The injected lifecycle callback must not turn a persisted session id into // a bypass around a matching connection or protected run owner. @@ -288,11 +346,9 @@ export async function abortChatRunsForSessionKeyWithPartials(params: { unauthorized: false, }; } - const authorizedRunIdSet = new Set(authorizedRuns.map((run) => run.runId)); const snapshots = collectSessionAbortPartials({ - chatAbortControllers: params.context.chatAbortControllers, chatRunState: params.context.chatRunState, - runIds: authorizedRunIdSet, + runs: authorizedRuns, abortOrigin: params.abortOrigin, }); // Abort queued owners before any active-work signal can promote a successor. diff --git a/src/gateway/server-methods/sessions-archive-lifecycle.ts b/src/gateway/server-methods/sessions-archive-lifecycle.ts new file mode 100644 index 000000000000..63e633206e2c --- /dev/null +++ b/src/gateway/server-methods/sessions-archive-lifecycle.ts @@ -0,0 +1,213 @@ +import { resolveEmbeddedSessionLane } from "../../agents/embedded-agent-runner/lanes.js"; +// Archive-owned cancellation and authoritative lifecycle drains. +import { + abortEmbeddedAgentRun, + isEmbeddedAgentRunInProgress, + waitForEmbeddedAgentRunEnd, +} from "../../agents/embedded-agent-runner/runs.js"; +import { clearSessionQueues } from "../../auto-reply/reply/queue/cleanup.js"; +import { hasPendingFollowupQueueWork } from "../../auto-reply/reply/queue/state.js"; +import { + abortReplyRunBySessionId, + isReplyRunActiveForSessionId, + replyRunRegistry, + waitForReplyRunEndBySessionId, +} from "../../auto-reply/reply/reply-run-registry.js"; +import { withTimeout } from "../../infra/fs-safe.js"; +import { getCommandLaneSnapshot } from "../../process/command-queue.js"; +import { + interruptSessionWorkAdmissions, + isCompetingSessionWorkAdmissionActive, + SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS, +} from "../../sessions/session-lifecycle-admission.js"; +import { waitForChatAbortControllerRemoval } from "../chat-abort-lifecycle-internal.js"; +import { + beginWorkerInferenceSessionDrain, + type WorkerInferenceSessionDrain, +} from "../worker-environments/inference-control-internal.js"; +import { asWorkerInferenceControl } from "../worker-environments/inference-control.js"; +import type { WorkerSessionPlacementStore } from "../worker-environments/placement-store.js"; +import { prepareSessionWorkerPlacementForArchive } from "../worker-environments/session-placement-lifecycle.js"; +import { + abortChatRunsForSessionKeyWithPartials, + createChatAbortOps, + hasGatewaySessionAbortOwner, +} from "./chat-abort-runtime.js"; +import type { GatewayRequestContext } from "./types.js"; + +type ArchivePlacementService = NonNullable & + Partial>; + +type ArchiveInferenceDrainService = { + beginInferenceSessionDrain(sessionId: string): WorkerInferenceSessionDrain; +}; + +function asArchiveInferenceDrainService(value: unknown): ArchiveInferenceDrainService | undefined { + if (typeof value !== "object" || value === null) { + return undefined; + } + return typeof (value as { beginInferenceSessionDrain?: unknown }).beginInferenceSessionDrain === + "function" + ? (value as ArchiveInferenceDrainService) + : undefined; +} + +type SessionArchiveLifecycleParams = { + context: GatewayRequestContext; + storePath: string; + sessionKeys: string[]; + sessionId?: string; + agentId: string; + sessionKey: string; + defaultAgentId: string; + lifecycleIdentities: string[]; +}; + +export type SessionArchiveLifecycleDrain = { + release(): void; + hasAuthoritativeWork(): boolean; +}; + +function hasAuthoritativeSessionWork( + params: SessionArchiveLifecycleParams, + workerDrain: WorkerInferenceSessionDrain | undefined, + workIdentities: string[], +): boolean { + const sessionId = params.sessionId; + return ( + isCompetingSessionWorkAdmissionActive(params.storePath, params.lifecycleIdentities) || + params.sessionKeys.some((key) => replyRunRegistry.isActive(key)) || + Boolean(sessionId && isReplyRunActiveForSessionId(sessionId)) || + Boolean(sessionId && isEmbeddedAgentRunInProgress(sessionId)) || + hasPendingFollowupQueueWork(workIdentities) || + workIdentities.some( + (key) => getCommandLaneSnapshot(resolveEmbeddedSessionLane(key)).queuedCount > 0, + ) || + hasGatewaySessionAbortOwner({ + context: params.context, + sessionKeys: params.sessionKeys, + sessionId, + agentId: params.agentId, + defaultAgentId: params.defaultAgentId, + }) || + Boolean( + sessionId && + params.context.workerSessionPlacementService?.getMany([sessionId]).get(sessionId)?.turnClaim, + ) || + workerDrain?.hasWork() === true + ); +} + +/** Fence is already active when this starts; retain the returned runtime fence through commit. */ +export async function prepareSessionArchiveLifecycle( + params: SessionArchiveLifecycleParams, +): Promise { + // Reject transient/live-failed placement before archive cancellation has side effects. + await prepareSessionWorkerPlacementForArchive({ ...params, reclaimActive: false }); + const timeoutMs = SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS; + const workIdentities = Array.from( + new Set([...params.sessionKeys, ...(params.sessionId ? [params.sessionId] : [])]), + ); + const workerService = params.context.workerEnvironmentService; + const workerControl = asWorkerInferenceControl(workerService); + let workerDrain: WorkerInferenceSessionDrain | undefined; + if (params.sessionId) { + // Lightweight contexts may expose the drain directly without widening the public service. + workerDrain = + beginWorkerInferenceSessionDrain(workerService, params.sessionId) ?? + asArchiveInferenceDrainService(workerService)?.beginInferenceSessionDrain(params.sessionId); + if (!workerDrain && workerControl?.hasInferenceForSession(params.sessionId) === true) { + throw new Error("Worker inference drain is unavailable"); + } + } + + try { + let controllerDrain = Promise.resolve(true); + const abortResult = await abortChatRunsForSessionKeyWithPartials({ + context: params.context, + ops: createChatAbortOps(params.context), + sessionKey: params.sessionKeys[0]!, + sessionKeyAliases: params.sessionKeys.slice(1), + sessionId: params.sessionId, + agentId: params.agentId, + defaultAgentId: params.defaultAgentId, + abortOrigin: "rpc", + stopReason: "archive", + requester: { isAdmin: true }, + includeProtectedRuns: true, + onControllerTargets: (targets) => { + controllerDrain = waitForChatAbortControllerRemoval({ + entries: params.context.chatAbortControllers, + targets, + timeoutMs, + }); + }, + onAuthorizedAfterQueuedAbort: () => { + const cleared = clearSessionQueues(workIdentities); + let aborted = cleared.followupCleared > 0 || cleared.laneCleared > 0; + for (const key of params.sessionKeys) { + aborted = replyRunRegistry.abort(key) || aborted; + } + if (params.sessionId) { + aborted = abortReplyRunBySessionId(params.sessionId) || aborted; + aborted = abortEmbeddedAgentRun(params.sessionId) || aborted; + } + return aborted; + }, + }); + if (abortResult.unauthorized) { + throw new Error("Archive cancellation lost session ownership"); + } + + const admittedWork = interruptSessionWorkAdmissions({ + scope: params.storePath, + identities: params.lifecycleIdentities, + timeoutMs, + }); + const replyWork = Promise.all([ + ...params.sessionKeys.map((key) => replyRunRegistry.waitForIdle(key, timeoutMs)), + ...(params.sessionId ? [waitForReplyRunEndBySessionId(params.sessionId, timeoutMs)] : []), + ]).then((results) => results.every(Boolean)); + const embeddedWork = params.sessionId + ? waitForEmbeddedAgentRunEnd(params.sessionId, timeoutMs) + : Promise.resolve(true); + const placementService = params.context.workerSessionPlacementService as + | ArchivePlacementService + | undefined; + const placement = params.sessionId + ? placementService?.getMany([params.sessionId]).get(params.sessionId) + : undefined; + const placementWork = placement?.turnClaim + ? placementService?.waitForTurnClaimRelease + ? placementService + .waitForTurnClaimRelease(params.sessionId!, { timeoutMs }) + .then(() => true) + : Promise.resolve(false) + : Promise.resolve(true); + const workerWork = workerDrain + ? withTimeout(workerDrain.drained, timeoutMs, "worker inference archive drain").then( + () => true, + ) + : Promise.resolve(true); + const drains = await Promise.all([ + controllerDrain, + admittedWork, + replyWork, + embeddedWork, + placementWork, + workerWork, + ]); + if (!drains.every(Boolean)) { + throw new Error("Session work did not fully drain before archive"); + } + // Fresh exact placement must be reclaimed before archivedAt can commit. + await prepareSessionWorkerPlacementForArchive({ ...params, reclaimActive: true }); + return { + release: () => workerDrain?.release(), + hasAuthoritativeWork: () => hasAuthoritativeSessionWork(params, workerDrain, workIdentities), + }; + } catch (error) { + workerDrain?.release(); + throw error; + } +} diff --git a/src/gateway/server-methods/sessions-mutations.archive-attribution.test.ts b/src/gateway/server-methods/sessions-mutations.archive-attribution.test.ts index 619950b0c8b2..3c6e51ca8dc5 100644 --- a/src/gateway/server-methods/sessions-mutations.archive-attribution.test.ts +++ b/src/gateway/server-methods/sessions-mutations.archive-attribution.test.ts @@ -59,6 +59,8 @@ function context(): GatewayRequestContext { broadcastToConnIds: vi.fn(), getSessionEventSubscriberConnIds: () => new Set(), chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + dedupe: new Map(), } as unknown as GatewayRequestContext; } diff --git a/src/gateway/server-methods/sessions-mutations.perf.test.ts b/src/gateway/server-methods/sessions-mutations.perf.test.ts index 2a429e2f60ca..15a10b2d0ea3 100644 --- a/src/gateway/server-methods/sessions-mutations.perf.test.ts +++ b/src/gateway/server-methods/sessions-mutations.perf.test.ts @@ -195,6 +195,8 @@ test("sessions.patchMany archives 30 human sessions without transcript hydration broadcastToConnIds: vi.fn(), getSessionEventSubscriberConnIds: () => new Set(), chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + dedupe: new Map(), cron: { list: cronList, updateWithPrecondition: cronUpdate, diff --git a/src/gateway/server-methods/sessions-patch-archive.ts b/src/gateway/server-methods/sessions-patch-archive.ts new file mode 100644 index 000000000000..9f60f4a48caf --- /dev/null +++ b/src/gateway/server-methods/sessions-patch-archive.ts @@ -0,0 +1,266 @@ +import { err, ok, type Result } from "@openclaw/normalization-core/result"; +import { + ErrorCodes, + errorShape, + type ErrorShape, + type SessionCreatedActor, + type SessionsPatchParams, +} from "../../../packages/gateway-protocol/src/index.js"; +import { resolveDefaultAgentId } from "../../agents/agent-scope.js"; +import type { ModelCatalogEntry } from "../../agents/model-catalog.js"; +import type { SessionEntry } from "../../config/sessions.js"; +import type { OpenClawConfig } from "../../config/types.openclaw.js"; +import { formatErrorMessage } from "../../infra/errors.js"; +import { resolveMissingAgentHarnessSessionError } from "../../sessions/agent-harness-session-key.js"; +import { resolvePluginSessionOwnershipError } from "../session-plugin-ownership.js"; +import { + resolveCanonicalGatewaySessionStoreKey, + resolveGatewaySessionStoreTargetWithStore, +} from "../session-utils.js"; +import { projectSessionsPatchEntry } from "../sessions-patch.js"; +import { + prepareSessionArchiveLifecycle, + type SessionArchiveLifecycleDrain, +} from "./sessions-archive-lifecycle.js"; +import { + isAgentMainSessionKey, + resolveSessionWorkerPlacementPatchError, + sessionLog, +} from "./sessions-shared.js"; +import type { GatewayRequestContext } from "./types.js"; + +export type SessionPatchArchivePreparation = { + canonicalKey: string; + drain: SessionArchiveLifecycleDrain; + entry?: SessionEntry; +}; + +type SessionPatchArchiveTarget = { + archiveActor: SessionCreatedActor | undefined; + canonicalKey: string; + fullPatch: SessionsPatchParams; + initialEntry?: SessionEntry; + initialStoreKeys: string[]; + key: string; + lifecycleIdentities: Array; + requestedAgentId?: string; + storePath: string; +}; + +function archiveChangedError(key: string): ErrorShape { + return errorShape(ErrorCodes.INVALID_REQUEST, `Session ${key} changed before patch. Retry.`); +} + +function archiveUnavailableError(key: string, message: "active" | "stopping"): ErrorShape { + return errorShape( + ErrorCodes.UNAVAILABLE, + message === "active" + ? `Session ${key} is still active; retry the archive.` + : `Session ${key} did not finish stopping; retry the archive.`, + { retryable: true }, + ); +} + +function protectedArchiveError(cfg: OpenClawConfig, canonicalKey: string): ErrorShape | undefined { + if (canonicalKey === "unknown") { + return errorShape(ErrorCodes.INVALID_REQUEST, "Cannot archive the unknown session sentinel."); + } + if (canonicalKey === "global" || isAgentMainSessionKey(cfg, canonicalKey)) { + return errorShape(ErrorCodes.INVALID_REQUEST, "Cannot archive an agent's main session."); + } + return undefined; +} + +function archiveTargetChanged(params: { + baselineEntry: SessionEntry | undefined; + currentEntry: SessionEntry | undefined; + patch: SessionsPatchParams; +}): boolean { + const { baselineEntry, currentEntry, patch } = params; + const expectedSessionChanged = + (patch.expectedSessionId !== undefined && + currentEntry?.sessionId !== patch.expectedSessionId) || + (patch.expectedLifecycleRevision !== undefined && + currentEntry?.lifecycleRevision !== patch.expectedLifecycleRevision); + const generationChanged = + baselineEntry !== undefined && + currentEntry !== undefined && + (currentEntry.sessionId !== baselineEntry.sessionId || + currentEntry.lifecycleRevision !== baselineEntry.lifecycleRevision); + return ( + expectedSessionChanged || + (baselineEntry !== undefined && currentEntry === undefined) || + (baselineEntry === undefined && currentEntry !== undefined) || + generationChanged + ); +} + +export async function prepareSessionPatchArchive(params: { + commitGuard: () => ErrorShape | undefined; + cfg: OpenClawConfig; + context: GatewayRequestContext; + loadGatewayModelCatalog: () => Promise; + pluginOwnerId?: string; + target: SessionPatchArchiveTarget; +}): Promise> { + const { cfg, target } = params; + const freshResolved = resolveGatewaySessionStoreTargetWithStore({ + cfg, + key: target.key, + ...(target.requestedAgentId ? { agentId: target.requestedAgentId } : {}), + exactRead: true, + }); + if (freshResolved.storePath !== target.storePath) { + return err(archiveChangedError(target.key)); + } + const fresh = resolveCanonicalGatewaySessionStoreKey({ + cfg, + key: target.key, + store: freshResolved.store, + agentId: target.requestedAgentId, + }); + const freshCanonicalKey = fresh.target.canonicalKey ?? target.key; + const ownershipError = resolvePluginSessionOwnershipError({ + action: "patch", + entry: fresh.entry, + key: freshCanonicalKey, + pluginOwnerId: params.pluginOwnerId, + }); + if (ownershipError) { + return err(ownershipError); + } + if ( + freshCanonicalKey !== target.canonicalKey || + archiveTargetChanged({ + currentEntry: fresh.entry, + baselineEntry: target.initialEntry, + patch: target.fullPatch, + }) + ) { + return err(archiveChangedError(target.key)); + } + const missingHarnessSessionError = resolveMissingAgentHarnessSessionError( + freshCanonicalKey, + fresh.entry, + ); + if (missingHarnessSessionError) { + return err(errorShape(ErrorCodes.INVALID_REQUEST, missingHarnessSessionError)); + } + const protectedError = protectedArchiveError(cfg, freshCanonicalKey); + if (protectedError) { + return err(protectedError); + } + const placementError = resolveSessionWorkerPlacementPatchError({ + agentId: freshResolved.agentId, + cfg, + context: params.context, + entry: fresh.entry, + key: target.key, + patch: target.fullPatch, + sessionKey: freshCanonicalKey, + validateModelRuntime: false, + }); + if (placementError) { + return err(errorShape(ErrorCodes.INVALID_REQUEST, placementError)); + } + const freshCandidateKeys = new Set(fresh.target.storeKeys); + const preview = await projectSessionsPatchEntry({ + cfg, + existingEntry: fresh.entry, + isLabelInUse: (label) => + Object.entries(freshResolved.store).some( + ([sessionKey, entry]) => !freshCandidateKeys.has(sessionKey) && entry.label === label, + ), + storeKey: fresh.primaryKey, + agentId: target.requestedAgentId, + patch: target.fullPatch, + archivedBy: target.archiveActor, + loadGatewayModelCatalog: params.loadGatewayModelCatalog, + }); + if (!preview.ok) { + return err(preview.error); + } + const previewPlacementError = resolveSessionWorkerPlacementPatchError({ + agentId: freshResolved.agentId, + cfg, + context: params.context, + entry: preview.entry, + key: target.key, + patch: target.fullPatch, + sessionKey: freshCanonicalKey, + validateModelRuntime: true, + }); + if (previewPlacementError) { + return err(errorShape(ErrorCodes.INVALID_REQUEST, previewPlacementError)); + } + const authorizationError = params.commitGuard(); + if (authorizationError) { + return err(authorizationError); + } + + try { + const drain = await prepareSessionArchiveLifecycle({ + context: params.context, + storePath: target.storePath, + sessionKeys: Array.from( + new Set([ + target.key, + target.canonicalKey, + ...target.initialStoreKeys, + freshCanonicalKey, + ...fresh.target.storeKeys, + ]), + ), + sessionId: fresh.entry?.sessionId, + sessionKey: freshCanonicalKey, + agentId: freshResolved.agentId, + defaultAgentId: resolveDefaultAgentId(cfg), + lifecycleIdentities: target.lifecycleIdentities.filter((identity): identity is string => + Boolean(identity), + ), + }); + return ok({ + canonicalKey: freshCanonicalKey, + drain, + ...(fresh.entry ? { entry: fresh.entry } : {}), + }); + } catch (error) { + sessionLog.warn( + `sessions.patch: archive drain failed for ${target.canonicalKey}: ${formatErrorMessage(error)}`, + ); + return err(archiveUnavailableError(target.key, "stopping")); + } +} + +export function validateSessionPatchArchiveProjection(params: { + cfg: OpenClawConfig; + existingEntry: SessionEntry | undefined; + fullPatch: SessionsPatchParams; + key: string; + pluginOwnerId?: string; + preparation: SessionPatchArchivePreparation; + primaryKey: string; +}): ErrorShape | undefined { + if (params.preparation.drain.hasAuthoritativeWork()) { + return archiveUnavailableError(params.key, "active"); + } + if ( + params.primaryKey !== params.preparation.canonicalKey || + archiveTargetChanged({ + currentEntry: params.existingEntry, + baselineEntry: params.preparation.entry, + patch: params.fullPatch, + }) + ) { + return archiveChangedError(params.key); + } + return ( + protectedArchiveError(params.cfg, params.primaryKey) ?? + resolvePluginSessionOwnershipError({ + action: "patch", + entry: params.existingEntry, + key: params.primaryKey, + pluginOwnerId: params.pluginOwnerId, + }) + ); +} diff --git a/src/gateway/server-methods/sessions-patch-engine.ts b/src/gateway/server-methods/sessions-patch-engine.ts index 3d1b9792e3fd..310fc5e5565b 100644 --- a/src/gateway/server-methods/sessions-patch-engine.ts +++ b/src/gateway/server-methods/sessions-patch-engine.ts @@ -6,8 +6,6 @@ import { type SessionsPatchManyTarget, type SessionsPatchParams, } from "../../../packages/gateway-protocol/src/index.js"; -import { resolveDefaultAgentId } from "../../agents/agent-scope.js"; -import { replyRunRegistry } from "../../auto-reply/reply/reply-run-registry.js"; import type { SessionEntry } from "../../config/sessions.js"; import { isInternalSessionEffectsKey } from "../../config/sessions/internal-session-key.js"; import { @@ -18,12 +16,7 @@ import { SessionLabelOwnerIndex } from "../../config/sessions/session-entry-sele import { disableCronJobsBoundToSessions } from "../../cron/job-session-bindings.js"; import { formatErrorMessage } from "../../infra/errors.js"; import { resolveMissingAgentHarnessSessionError } from "../../sessions/agent-harness-session-key.js"; -import { - isSessionLifecycleMutationActive, - isSessionWorkAdmissionActive, - runExclusiveSessionLifecycleMutation, - SESSION_ARCHIVE_ACTIVE_RUN_ERROR, -} from "../../sessions/session-lifecycle-admission.js"; +import { runExclusiveSessionLifecycleMutation } from "../../sessions/session-lifecycle-admission.js"; import { ADMIN_SCOPE } from "../operator-scopes.js"; import { ensureSessionGroupRegistered } from "../session-groups.js"; import { triggerSessionPatchHook } from "../session-patch-hooks.js"; @@ -39,15 +32,15 @@ import { } from "../session-utils.js"; import { projectSessionsPatchEntry } from "../sessions-patch.js"; import { gatewayClientSessionCreator } from "./gateway-client-identity.js"; -import { hasVisibleActiveSessionRun } from "./session-active-runs.js"; import { appendSessionAudit } from "./session-audit.js"; import { emitSessionsChanged } from "./session-change-event.js"; -import { persistSessionPatchModelSelection } from "./sessions-patch-model-selection.js"; import { - isAgentMainSessionKey, - resolveSessionWorkerPlacementPatchError, - sessionLog, -} from "./sessions-shared.js"; + prepareSessionPatchArchive, + type SessionPatchArchivePreparation, + validateSessionPatchArchiveProjection, +} from "./sessions-patch-archive.js"; +import { persistSessionPatchModelSelection } from "./sessions-patch-model-selection.js"; +import { resolveSessionWorkerPlacementPatchError, sessionLog } from "./sessions-shared.js"; import type { GatewayClient, GatewayRequestContext, @@ -64,10 +57,13 @@ type MutationTarget = PatchTargetIdentity & { }; type PreparedPatchTarget = { + archivePreparation?: SessionPatchArchivePreparation; + archiveActor: ReturnType; canonicalKey: string; fullPatch: SessionsPatchParams; index: number; initialEntry?: SessionEntry; + initialStoreKeys: string[]; key: string; lifecycleIdentities: Array; requestedAgentId?: string; @@ -140,7 +136,7 @@ async function executeSessionPatchMutations(params: { ? params.client.connect.scopes : []; const callerCanManageCron = params.client === null || callerScopes.includes(ADMIN_SCOPE); - const defaultAgentId = resolveDefaultAgentId(cfg); + const pluginOwnerId = params.client?.internal?.pluginRuntimeOwnerId; const targetDiscoveryCache = new Map(); const preflightTargets = params.targets.map((input) => { const key = input.key.trim(); @@ -235,22 +231,16 @@ async function executeSessionPatchMutations(params: { }; continue; } - const lifecycleIdentities = [canonicalKey, key, initialEntry?.sessionId]; - if ( - fullPatch.archived === true && - isSessionLifecycleMutationActive(resolved.storePath, lifecycleIdentities) - ) { - outcomes[index] = { - ok: false, - error: errorShape(ErrorCodes.INVALID_REQUEST, SESSION_ARCHIVE_ACTIVE_RUN_ERROR), - }; - continue; - } + const lifecycleIdentities = Array.from( + new Set([key, canonicalKey, ...candidateKeys, initialEntry?.sessionId]), + ); const preparedTarget: PreparedPatchTarget = { + archiveActor, canonicalKey, fullPatch, index, ...(initialEntry ? { initialEntry } : {}), + initialStoreKeys: [...candidateKeys], key, lifecycleIdentities, ...(requestedAgentId ? { requestedAgentId } : {}), @@ -272,233 +262,266 @@ async function executeSessionPatchMutations(params: { }; if (prepared.length > 0) { - await runExclusiveSessionLifecycleMutation({ - targets: prepared.map((target) => ({ - scope: target.storePath, - identities: target.lifecycleIdentities, - })), - run: async () => { - const groups = new Map(); - for (const target of prepared) { - const groupKey = `${target.storePath}\0${target.targetAgentId}`; - const group = groups.get(groupKey); - if (group) { - group.push(target); - } else { - groups.set(groupKey, [target]); + try { + await runExclusiveSessionLifecycleMutation({ + targets: prepared.map((target) => ({ + scope: target.storePath, + identities: target.lifecycleIdentities, + })), + prepare: async () => { + await Promise.all( + prepared + .filter((target) => target.fullPatch.archived === true) + .map(async (target) => { + try { + const result = await prepareSessionPatchArchive({ + cfg, + commitGuard: params.targets[target.index]!.commitGuard, + context: params.context, + loadGatewayModelCatalog: () => loadModelCatalog(target.targetAgentId), + ...(pluginOwnerId ? { pluginOwnerId } : {}), + target, + }); + if (result.ok) { + target.archivePreparation = result.value; + } else { + outcomes[target.index] = result; + } + } catch (error) { + outcomes[target.index] = { + ok: false, + error: unexpectedPatchError(target.key, error), + }; + } + }), + ); + }, + run: async () => { + const groups = new Map(); + for (const target of prepared) { + if (target.fullPatch.archived === true && !target.archivePreparation) { + continue; + } + const groupKey = `${target.storePath}\0${target.targetAgentId}`; + const group = groups.get(groupKey); + if (group) { + group.push(target); + } else { + groups.set(groupKey, [target]); + } } - } - await Promise.all( - [...groups.values()].map(async (group) => { - const first = group[0]!; - try { - const groupOutcomes = await applySqliteSessionEntryCanonicalReplacements({ - agentId: first.targetAgentId, - storePath: first.storePath, - skipMaintenance: true, - update: async (entries) => { - const workingStore = Object.fromEntries( - entries.flatMap(({ entry, sessionKey }) => - isInternalSessionEffectsKey(sessionKey) ? [] : [[sessionKey, entry] as const], - ), - ); - const labelOwners = new SessionLabelOwnerIndex(workingStore); - const replacements: SessionEntryCanonicalReplacement[] = []; - const projectedOutcomes: MutationOutcome[] = []; - for (const target of group) { - try { - // Preflight facts can stale behind the writer queue; resolve this snapshot - // again so a new legacy alias is rejected rather than promoted or deleted. - const { - entry: existingEntry, - primaryKey, - target: currentTarget, - } = resolveCanonicalGatewaySessionStoreKey({ - cfg, - key: target.key, - store: workingStore, - ...(target.requestedAgentId ? { agentId: target.requestedAgentId } : {}), - }); - const candidateKeys = currentTarget.storeKeys; - const ownershipError = pluginOwnershipError({ - client: params.client, - entry: existingEntry, - key: primaryKey, - }); - if (ownershipError) { - projectedOutcomes.push({ ok: false, error: ownershipError }); - continue; - } - const expectedSessionChanged = - (target.fullPatch.expectedSessionId !== undefined && - existingEntry?.sessionId !== target.fullPatch.expectedSessionId) || - (target.fullPatch.expectedLifecycleRevision !== undefined && - existingEntry?.lifecycleRevision !== - target.fullPatch.expectedLifecycleRevision); - const lifecycleEntryRemoved = - target.initialEntry !== undefined && existingEntry === undefined; - const archiveTargetChanged = - target.fullPatch.archived === true && - (target.initialEntry === undefined - ? existingEntry !== undefined - : existingEntry !== undefined && - (existingEntry.sessionId !== target.initialEntry.sessionId || - existingEntry.lifecycleRevision !== - target.initialEntry.lifecycleRevision)); - if (expectedSessionChanged || lifecycleEntryRemoved || archiveTargetChanged) { - projectedOutcomes.push({ - ok: false, - error: errorShape( - ErrorCodes.INVALID_REQUEST, - `Session ${target.key} changed before patch. Retry.`, - ), + await Promise.all( + [...groups.values()].map(async (group) => { + const first = group[0]!; + try { + const groupOutcomes = await applySqliteSessionEntryCanonicalReplacements({ + agentId: first.targetAgentId, + storePath: first.storePath, + skipMaintenance: true, + update: async (entries) => { + const workingStore = Object.fromEntries( + entries.flatMap(({ entry, sessionKey }) => + isInternalSessionEffectsKey(sessionKey) + ? [] + : [[sessionKey, entry] as const], + ), + ); + const labelOwners = new SessionLabelOwnerIndex(workingStore); + const replacements: SessionEntryCanonicalReplacement[] = []; + const projectedOutcomes: MutationOutcome[] = []; + for (const target of group) { + try { + // Preflight facts can stale behind the writer queue; resolve this snapshot + // again so a new legacy alias is rejected rather than promoted or deleted. + const { + entry: existingEntry, + primaryKey, + target: currentTarget, + } = resolveCanonicalGatewaySessionStoreKey({ + cfg, + key: target.key, + store: workingStore, + ...(target.requestedAgentId ? { agentId: target.requestedAgentId } : {}), }); - continue; - } - if (target.fullPatch.archived === true) { - if (primaryKey === "global" || isAgentMainSessionKey(cfg, primaryKey)) { - projectedOutcomes.push({ - ok: false, - error: errorShape( - ErrorCodes.INVALID_REQUEST, - "Cannot archive an agent's main session.", - ), - }); + const candidateKeys = currentTarget.storeKeys; + const ownershipError = pluginOwnershipError({ + client: params.client, + entry: existingEntry, + key: primaryKey, + }); + if (ownershipError) { + projectedOutcomes.push({ ok: false, error: ownershipError }); continue; } - const activeIdentities = [primaryKey, target.key, existingEntry?.sessionId]; + const expectedSessionChanged = + (target.fullPatch.expectedSessionId !== undefined && + existingEntry?.sessionId !== target.fullPatch.expectedSessionId) || + (target.fullPatch.expectedLifecycleRevision !== undefined && + existingEntry?.lifecycleRevision !== + target.fullPatch.expectedLifecycleRevision); + const lifecycleEntryRemoved = + target.initialEntry !== undefined && existingEntry === undefined; + const archiveTargetChanged = + target.fullPatch.archived === true && + (target.initialEntry === undefined + ? existingEntry !== undefined + : existingEntry !== undefined && + (existingEntry.sessionId !== target.initialEntry.sessionId || + existingEntry.lifecycleRevision !== + target.initialEntry.lifecycleRevision)); if ( - isSessionWorkAdmissionActive(target.storePath, activeIdentities) || - replyRunRegistry.isActive(primaryKey) || - replyRunRegistry.isActive(target.key) || - hasVisibleActiveSessionRun({ - context: params.context, - requestedKey: target.key, - canonicalKey: primaryKey, - sessionId: existingEntry?.sessionId, - defaultAgentId, - }) + expectedSessionChanged || + lifecycleEntryRemoved || + archiveTargetChanged ) { projectedOutcomes.push({ ok: false, error: errorShape( ErrorCodes.INVALID_REQUEST, - SESSION_ARCHIVE_ACTIVE_RUN_ERROR, + `Session ${target.key} changed before patch. Retry.`, ), }); continue; } - } - const wasArchivedBeforePatch = existingEntry?.archivedAt !== undefined; - const projected = await projectSessionsPatchEntry({ - cfg, - existingEntry, - isLabelInUse: (label) => labelOwners.isLabelInUse(label, candidateKeys), - storeKey: primaryKey, - agentId: target.requestedAgentId, - patch: target.fullPatch, - archivedBy: archiveActor, - loadGatewayModelCatalog: () => loadModelCatalog(target.targetAgentId), - }); - if (!projected.ok) { - projectedOutcomes.push(projected); - continue; - } - const placementPatchError = resolveSessionWorkerPlacementPatchError({ - agentId: target.targetAgentId, - cfg, - context: params.context, - entry: projected.entry, - key: target.key, - patch: target.fullPatch, - sessionKey: primaryKey, - validateModelRuntime: true, - }); - if (placementPatchError) { + if (target.fullPatch.archived === true) { + const archiveError = validateSessionPatchArchiveProjection({ + cfg, + existingEntry, + fullPatch: target.fullPatch, + key: target.key, + ...(pluginOwnerId ? { pluginOwnerId } : {}), + preparation: target.archivePreparation!, + primaryKey, + }); + if (archiveError) { + projectedOutcomes.push({ + ok: false, + error: archiveError, + }); + continue; + } + } + const wasArchivedBeforePatch = existingEntry?.archivedAt !== undefined; + const projected = await projectSessionsPatchEntry({ + cfg, + existingEntry, + isLabelInUse: (label) => labelOwners.isLabelInUse(label, candidateKeys), + storeKey: primaryKey, + agentId: target.requestedAgentId, + patch: target.fullPatch, + archivedBy: archiveActor, + loadGatewayModelCatalog: () => loadModelCatalog(target.targetAgentId), + }); + if (!projected.ok) { + projectedOutcomes.push(projected); + continue; + } + const placementPatchError = resolveSessionWorkerPlacementPatchError({ + agentId: target.targetAgentId, + cfg, + context: params.context, + entry: projected.entry, + key: target.key, + patch: target.fullPatch, + sessionKey: primaryKey, + validateModelRuntime: true, + }); + if (placementPatchError) { + projectedOutcomes.push({ + ok: false, + error: errorShape(ErrorCodes.INVALID_REQUEST, placementPatchError), + }); + continue; + } + const authorizationFailure = params.targets[target.index]!.commitGuard(); + if (authorizationFailure) { + projectedOutcomes.push({ ok: false, error: authorizationFailure }); + continue; + } + const previousSessionKeys = candidateKeys.filter( + (sessionKey) => sessionKey !== primaryKey && workingStore[sessionKey], + ); + replacements.push({ + entry: projected.entry, + previousSessionKeys, + sessionKey: primaryKey, + }); + const cloned = labelOwners.replaceEntry( + candidateKeys, + primaryKey, + projected.entry, + ); + projectedOutcomes.push({ + ok: true, + archiveStateChanged: + typeof target.fullPatch.archived === "boolean" && + wasArchivedBeforePatch !== (cloned.archivedAt !== undefined), + entry: cloned, + }); + } catch (error) { projectedOutcomes.push({ ok: false, - error: errorShape(ErrorCodes.INVALID_REQUEST, placementPatchError), + error: unexpectedPatchError(target.key, error), }); - continue; } - const authorizationFailure = params.targets[target.index]!.commitGuard(); - if (authorizationFailure) { - projectedOutcomes.push({ ok: false, error: authorizationFailure }); - continue; - } - const previousSessionKeys = candidateKeys.filter( - (sessionKey) => sessionKey !== primaryKey && workingStore[sessionKey], - ); - replacements.push({ - entry: projected.entry, - previousSessionKeys, - sessionKey: primaryKey, - }); - const cloned = labelOwners.replaceEntry( - candidateKeys, - primaryKey, - projected.entry, - ); - projectedOutcomes.push({ - ok: true, - archiveStateChanged: - typeof target.fullPatch.archived === "boolean" && - wasArchivedBeforePatch !== (cloned.archivedAt !== undefined), - entry: cloned, - }); - } catch (error) { - projectedOutcomes.push({ - ok: false, - error: unexpectedPatchError(target.key, error), - }); } - } - return { replacements, result: projectedOutcomes }; - }, - }); - for (const [groupIndex, target] of group.entries()) { - outcomes[target.index] = groupOutcomes[groupIndex]!; + return { replacements, result: projectedOutcomes }; + }, + }); + for (const [groupIndex, target] of group.entries()) { + outcomes[target.index] = groupOutcomes[groupIndex]!; + } + } catch (error) { + for (const target of group) { + outcomes[target.index] = { + ok: false, + error: unexpectedPatchError(target.key, error), + }; + } } - } catch (error) { - for (const target of group) { - outcomes[target.index] = { - ok: false, - error: unexpectedPatchError(target.key, error), - }; - } - } - }), - ); + }), + ); - for (const target of prepared) { - const outcome = outcomes[target.index]; - if (!outcome?.ok || !archiveActor) { - continue; - } - if (!outcome.archiveStateChanged) { - continue; - } - const action = outcome.entry.archivedAt === undefined ? "unarchived" : "archived"; - try { - await appendSessionAudit({ - cfg, - target: { - agentId: target.targetAgentId, - entry: outcome.entry, - sessionKey: target.canonicalKey, - storePath: target.storePath, - }, - text: `${action} by ${archiveActor.label ?? archiveActor.id}`, - now: Date.now(), - }); - } catch (error) { - sessionLog.warn( - `sessions.patch: ${action} audit note failed for ${target.canonicalKey}; archive kept: ${formatErrorMessage(error)}`, - ); + for (const target of prepared) { + const outcome = outcomes[target.index]; + if (!outcome?.ok || !archiveActor) { + continue; + } + if (!outcome.archiveStateChanged) { + continue; + } + const action = outcome.entry.archivedAt === undefined ? "unarchived" : "archived"; + try { + await appendSessionAudit({ + cfg, + target: { + agentId: target.targetAgentId, + entry: outcome.entry, + sessionKey: target.canonicalKey, + storePath: target.storePath, + }, + text: `${action} by ${archiveActor.label ?? archiveActor.id}`, + now: Date.now(), + }); + } catch (error) { + sessionLog.warn( + `sessions.patch: ${action} audit note failed for ${target.canonicalKey}; archive kept: ${formatErrorMessage(error)}`, + ); + } } + }, + }); + } finally { + for (const target of prepared) { + try { + target.archivePreparation?.drain.release(); + } catch (error) { + sessionLog.warn( + `sessions.patch: archive drain release failed for ${target.canonicalKey}: ${formatErrorMessage(error)}`, + ); } - }, - }); + } + } } let patched = false; diff --git a/src/gateway/server-methods/sessions-shared.ts b/src/gateway/server-methods/sessions-shared.ts index 2abdeb860d79..609e56d9e39d 100644 --- a/src/gateway/server-methods/sessions-shared.ts +++ b/src/gateway/server-methods/sessions-shared.ts @@ -26,6 +26,7 @@ import { isWorkerPlacementSessionRuntimeSupported, resolveWorkerPlacementSessionRuntime, } from "../worker-environments/placement-session-runtime.js"; +import { isWorkerPlacementSafeForArchive } from "../worker-environments/session-placement-lifecycle.js"; export { resolveSessionWorkerPlacementMutationError, retireSessionWorkerPlacementBeforeMutation, @@ -60,8 +61,10 @@ export function resolveSessionWorkerPlacementPatchError(params: { if (!placement || placement.state === "local") { return undefined; } - if (params.patch.archived !== undefined) { - return `Session ${params.key} cannot change archive state while cloud worker placement is ${placement.state}.`; + if (params.patch.archived === false) { + if (!isWorkerPlacementSafeForArchive(params.context, placement)) { + return `Session ${params.key} cannot change archive state while cloud worker placement is ${placement.state}.`; + } } if (!params.validateModelRuntime || params.patch.model === undefined || !params.entry) { return undefined; diff --git a/src/gateway/server-methods/sessions-sharing.ts b/src/gateway/server-methods/sessions-sharing.ts index fb9c1a3afab1..bdbc6ed2d787 100644 --- a/src/gateway/server-methods/sessions-sharing.ts +++ b/src/gateway/server-methods/sessions-sharing.ts @@ -17,7 +17,7 @@ import { removeSessionMember, } from "../../config/sessions.js"; import { patchSessionEntry } from "../../config/sessions/session-accessor.js"; -import { runQueuedStoreWrite, type StoreWriterQueue } from "../../shared/store-writer-queue.js"; +import { runExclusiveSessionLifecycleMutation } from "../../sessions/session-lifecycle-admission.js"; import { listProfiles } from "../../state/user-profiles.js"; import { allowedSessionVisibilities, @@ -34,17 +34,16 @@ import { emitSessionsChanged } from "./session-change-event.js"; import type { GatewayClient, GatewayRequestContext, GatewayRequestHandlers } from "./types.js"; import { assertValidParams } from "./validation.js"; -const sharingMutationQueues = new Map(); - function runExclusiveSharingMutation( target: NonNullable>, run: () => Promise, ): Promise { - return runQueuedStoreWrite({ - queues: sharingMutationQueues, - storePath: `${target.storePath}\0${target.canonicalKey}`, - label: "session-sharing-mutation", - fn: run, + // Sharing and lifecycle mutations share one exact-row fence so authorization + // cannot change between archive's stop and commit boundaries. + return runExclusiveSessionLifecycleMutation({ + scope: target.storePath, + identities: [target.canonicalKey, target.storeKey, ...target.storeKeys, target.entry.sessionId], + run, }); } @@ -91,9 +90,9 @@ function requireManageableTarget(params: { return { target, role }; } -// Manager authorization runs before the exclusive queue, so a session can be +// Manager authorization runs before the lifecycle fence, so a session can be // reset or recreated under the same key while a mutation waits. Requiring the -// same session instance and a still-valid manager role inside the queue keeps +// same session instance and a still-valid manager role inside the fence keeps // a stale owner from mutating the replacement session's sharing state. function requireCurrentManagedTarget(params: { cfg: ReturnType; @@ -213,10 +212,9 @@ export const sessionSharingHandlers: GatewayRequestHandlers = { sessionKey: current.canonicalKey, storePath: current.storePath, }; - // The entry-store write queue is separate from the sharing queue, so a - // reset/recreate can replace the row between the check above and this - // write. Re-check the instance inside the atomic patch and no-op if it - // changed, so a stale owner cannot stamp the replacement's visibility. + // The lifecycle fence excludes canonical reset/recreate. Keep the exact + // session-id check at the storage boundary so an out-of-band row + // replacement still cannot inherit this visibility change. let sessionChanged = false; await patchSessionEntry(scope, (entry) => { if (entry.sessionId !== current.entry.sessionId) { @@ -239,9 +237,8 @@ export const sessionSharingHandlers: GatewayRequestHandlers = { now, }); } catch (error) { - // Roll back only if this is still the same instance we patched; a - // concurrent reset could otherwise stamp the old restricted value onto - // a fresh (shared-default) replacement. + // Roll back only the exact instance and value we patched; an unexpected + // storage-owner replacement must not inherit the old visibility. await patchSessionEntry(scope, (entry) => entry.sessionId === current.entry.sessionId && resolveSessionVisibility(entry) === visibility diff --git a/src/gateway/server-runtime-subscriptions.ts b/src/gateway/server-runtime-subscriptions.ts index e351ed1f93a3..c7359d8d841b 100644 --- a/src/gateway/server-runtime-subscriptions.ts +++ b/src/gateway/server-runtime-subscriptions.ts @@ -15,6 +15,7 @@ import { onInternalSessionTranscriptUpdate } from "../sessions/transcript-events import { createLazyPromise, createLazyPromiseLoader } from "../shared/lazy-runtime.js"; import { isTerminalTaskStatus } from "../tasks/task-executor-policy.js"; import type { TaskRegistryObserverEvent } from "../tasks/task-registry.store.js"; +import { markChatAbortTerminalPersistenceError } from "./chat-abort-lifecycle-internal.js"; import { type ChatAbortControllerEntry, removeChatAbortControllerEntry, @@ -142,6 +143,7 @@ export function startGatewayEventSubscriptions(params: { entry.projectSessionActive = false; entry.projectSessionTerminalPending = false; entry.projectSessionTerminalPersisted = false; + markChatAbortTerminalPersistenceError(entry, undefined); queueMicrotask(() => { const current = params.chatAbortControllers.get(candidateRunId); if ( @@ -168,6 +170,7 @@ export function startGatewayEventSubscriptions(params: { entry.projectSessionTerminalPending = false; entry.projectSessionTerminalPersisted = true; entry.projectSessionTerminalPersistence = undefined; + markChatAbortTerminalPersistenceError(entry, undefined); } } }, @@ -184,6 +187,9 @@ export function startGatewayEventSubscriptions(params: { if (entry) { entry.projectSessionTerminalPending = false; entry.projectSessionTerminalPersistence = persistence; + void persistence.catch((error: unknown) => { + markChatAbortTerminalPersistenceError(entry, error); + }); if (entry.registrationCleanupRequested === true) { void persistence .catch(() => undefined) diff --git a/src/gateway/server.sessions.archive-lifecycle.test.ts b/src/gateway/server.sessions.archive-lifecycle.test.ts new file mode 100644 index 000000000000..6ff165b3884d --- /dev/null +++ b/src/gateway/server.sessions.archive-lifecycle.test.ts @@ -0,0 +1,1055 @@ +// Archive lifecycle tests protect fence-before-cancel, terminal drains, and sentinels. +import { afterEach, expect, test, vi } from "vitest"; +import { SessionManager } from "../agents/sessions/session-manager.js"; +import { loadSessionEntry, upsertSessionEntry } from "../config/sessions/session-accessor.js"; +import { onAgentEvent } from "../infra/agent-events.js"; +import { + beginSessionWorkAdmission, + isSessionLifecycleMutationActive, + runExclusiveSessionLifecycleMutation, +} from "../sessions/session-lifecycle-admission.js"; +import { createDeferred } from "../shared/deferred.js"; +import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; +import { markChatAbortTerminalPersistenceError } from "./chat-abort-lifecycle-internal.js"; +import { registerChatAbortController, removeChatAbortControllerEntry } from "./chat-abort.js"; +import { createChatRunState } from "./server-chat-state.js"; +import type { GatewayClient, GatewayRequestContext, RespondFn } from "./server-methods/types.js"; +import { + resolveSessionMutationAuthorization, + resolveSessionSharingTarget, +} from "./session-sharing.js"; +import { embeddedRunMock, writeSessionStore } from "./test-helpers.js"; +import { + directSessionReq, + expectNoSessionQueueCleanup, + getGatewayConfigModule, + getSessionsHandlers, + sessionStoreEntry, + setupGatewaySessionsHandlerTestHarness, +} from "./test/server-sessions.test-helpers.js"; +import { registerWorkerInferenceSessionDrain } from "./worker-environments/inference-control-internal.js"; +import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-record.js"; + +const sessionAuditGate = vi.hoisted(() => ({ + entered: vi.fn(), + wait: undefined as Promise | undefined, +})); + +vi.mock("./server-methods/session-audit.js", async () => { + const actual = await vi.importActual( + "./server-methods/session-audit.js", + ); + return { + ...actual, + appendSessionAudit: async (...args: Parameters) => { + if (sessionAuditGate.wait) { + sessionAuditGate.entered(); + await sessionAuditGate.wait; + } + await actual.appendSessionAudit(...args); + }, + }; +}); + +const { + createConfiguredGlobalAgentSessionStore, + createSessionStoreDir, + resetConfiguredGlobalAgentSessionStore, +} = setupGatewaySessionsHandlerTestHarness(); + +afterEach(() => { + closeOpenClawStateDatabaseForTest(); +}); + +function activeRunContext(params: { + runId: string; + sessionId: string; + sessionKey: string; + persistence: ReturnType>; + ownerConnId?: string; +}) { + const chatAbortControllers = new Map(); + const registration = registerChatAbortController({ + chatAbortControllers, + runId: params.runId, + sessionId: params.sessionId, + sessionKey: params.sessionKey, + timeoutMs: 60_000, + ownerConnId: params.ownerConnId, + }); + if (!registration.entry) { + throw new Error("expected active run registration"); + } + const entry = registration.entry; + const unsubscribe = onAgentEvent((event) => { + if ( + event.runId !== params.runId || + event.stream !== "lifecycle" || + event.data.phase !== "end" + ) { + return; + } + entry.projectSessionTerminalPending = false; + entry.projectSessionTerminalPersistence = params.persistence.promise; + void params.persistence.promise.then( + () => { + entry.projectSessionTerminalPersistence = undefined; + entry.projectSessionTerminalPersisted = true; + removeChatAbortControllerEntry(chatAbortControllers, params.runId, entry); + }, + (error: unknown) => { + markChatAbortTerminalPersistenceError(entry, error); + removeChatAbortControllerEntry(chatAbortControllers, params.runId, entry); + }, + ); + }); + const chatRunState = createChatRunState(); + return { + context: { + agentRunSeq: new Map([[params.runId, 0]]), + broadcast: vi.fn(), + cancelRunBoundApprovals: vi.fn(), + chatAbortControllers, + chatRunState, + logGateway: { warn: vi.fn() }, + nodeSendToSession: vi.fn(), + removeChatRun: vi.fn(() => ({ + sessionKey: params.sessionKey, + clientRunId: params.runId, + })), + }, + controller: registration.controller, + unsubscribe, + }; +} + +function identifiedClient(profileId: string): GatewayClient { + return { + connId: `${profileId}-connection`, + authenticatedUserId: `${profileId}@example.com`, + authenticatedUserProfile: { + profileId, + displayName: profileId, + hasAvatar: false, + updatedAt: 1, + }, + connect: { + minProtocol: 1, + maxProtocol: 1, + client: { + id: "openclaw-control-ui", + version: "test", + platform: "test", + mode: "webchat", + }, + role: "operator", + scopes: ["operator.read", "operator.write"], + }, + }; +} + +function workerPlacement(params: { + sessionId: string; + sessionKey: string; + state: WorkerSessionPlacementRecord["state"]; + agentId?: string; + environmentId?: string | null; +}): WorkerSessionPlacementRecord { + return { + sessionId: params.sessionId, + sessionKey: params.sessionKey, + agentId: params.agentId ?? "main", + state: params.state, + generation: 2, + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 2, + stateChangedAtMs: 2, + environmentId: + params.environmentId !== undefined + ? params.environmentId + : params.state === "local" || params.state === "requested" + ? null + : "worker-environment", + activeOwnerEpoch: ["active", "draining", "reconciling", "reclaimed", "failed"].includes( + params.state, + ) + ? 1 + : null, + workspaceBaseManifestRef: + params.state === "local" || + params.state === "requested" || + params.state === "provisioning" || + params.state === "syncing" + ? null + : "manifest-ref", + remoteWorkspaceDir: + params.state === "local" || + params.state === "requested" || + params.state === "provisioning" || + params.state === "syncing" + ? null + : "/workspace", + workerBundleHash: + params.state === "local" || params.state === "requested" || params.state === "provisioning" + ? null + : "bundle-hash", + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: params.state === "failed" ? "worker recovery stopped" : null, + } as WorkerSessionPlacementRecord; +} + +function placementReader(current: () => WorkerSessionPlacementRecord | undefined) { + return { + getMany(sessionIds: readonly string[]) { + const placement = current(); + return new Map( + placement && sessionIds.includes(placement.sessionId) + ? [[placement.sessionId, placement]] + : [], + ); + }, + }; +} + +async function archiveLifecycleRequestContext( + overrides: Record, +): Promise { + const { getRuntimeConfig } = await getGatewayConfigModule(); + const loadGatewayModelCatalog = async () => []; + return { + broadcast: vi.fn(), + broadcastToConnIds: vi.fn(), + chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + dedupe: new Map(), + getSessionEventSubscriberConnIds: () => new Set(), + getRuntimeConfig, + loadGatewayModelCatalog, + readPreparedGatewayModelCatalog: loadGatewayModelCatalog, + ...overrides, + } as unknown as GatewayRequestContext; +} + +type LifecycleHandlerResponse = { + ok: boolean; + payload?: unknown; + error?: Parameters[2]; +}; + +async function invokeArchiveHandler(params: { + authorization: NonNullable< + ReturnType["authorization"] + >; + client: GatewayClient; + context: GatewayRequestContext; + sessionKey: string; +}): Promise { + const handlers = await getSessionsHandlers(); + let response: LifecycleHandlerResponse | undefined; + const respond: RespondFn = (ok, payload, error) => { + response = { ok, payload, error }; + }; + await handlers["sessions.patch"]?.({ + req: {} as never, + params: { key: params.sessionKey, archived: true }, + client: params.client, + context: params.context, + isWebchatConnect: () => false, + sessionMutationAuthorization: params.authorization, + respond, + } as never); + if (!response) { + throw new Error("sessions.patch did not respond"); + } + return response; +} + +async function invokeVisibilityHandler(params: { + client: GatewayClient; + context: GatewayRequestContext; + sessionKey: string; + visibility: "draft" | "shared"; +}): Promise { + const handlers = await getSessionsHandlers(); + let response: LifecycleHandlerResponse | undefined; + const respond: RespondFn = (ok, payload, error) => { + response = { ok, payload, error }; + }; + await handlers["session.visibility.set"]?.({ + params: { sessionKey: params.sessionKey, visibility: params.visibility }, + client: params.client, + context: params.context, + respond, + } as never); + if (!response) { + throw new Error("session.visibility.set did not respond"); + } + return response; +} + +test("sessions.patch cancels active work and commits only after admission and terminal persistence drain", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-active"; + const sessionId = "session-archive-active"; + const runId = "run-archive-active"; + await writeSessionStore({ + entries: { [sessionKey]: sessionStoreEntry(sessionId) }, + }); + let interrupted = false; + const admission = await beginSessionWorkAdmission({ + scope: storePath, + identities: [sessionKey, sessionId], + assertAllowed: () => {}, + onInterrupt: () => { + interrupted = true; + }, + }); + const persistence = createDeferred(); + const active = activeRunContext({ + runId, + sessionId, + sessionKey, + persistence, + ownerConnId: "different-connection", + }); + try { + const archive = directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: active.context, + client: { connId: "archive-writer", connect: { scopes: ["operator.write"] } } as never, + }, + ); + await vi.waitFor(() => { + expect(interrupted).toBe(true); + expect(active.controller.signal.aborted).toBe(true); + }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + + let replacementAdmitted = false; + const replacement = beginSessionWorkAdmission({ + scope: storePath, + identities: [sessionKey, sessionId], + assertAllowed: () => { + replacementAdmitted = true; + if (loadSessionEntry({ storePath, sessionKey })?.archivedAt !== undefined) { + throw new Error("archived"); + } + }, + }).then( + (lease) => lease, + (error: unknown) => error, + ); + await Promise.resolve(); + expect(replacementAdmitted).toBe(false); + + admission.release(); + await Promise.resolve(); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + persistence.resolve(); + + const archived = await archive; + expect(archived.ok).toBe(true); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toEqual(expect.any(Number)); + expect(await replacement).toBeInstanceOf(Error); + } finally { + admission.release(); + active.unsubscribe(); + } +}); + +test("sharing revocation fences archive before cancellation and forces fresh authorization", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-sharing-revocation"; + const sessionId = "session-archive-sharing-revocation"; + const runId = "run-archive-sharing-revocation"; + const owner = identifiedClient("archive-owner"); + const viewer = identifiedClient("archive-viewer"); + await writeSessionStore({ + entries: { + [sessionKey]: sessionStoreEntry(sessionId, { + createdActor: { type: "human", id: "archive-owner" }, + visibility: "shared", + }), + }, + }); + let interrupted = false; + const admission = await beginSessionWorkAdmission({ + scope: storePath, + identities: [sessionKey, sessionId], + assertAllowed: () => {}, + onInterrupt: () => { + interrupted = true; + }, + }); + const persistence = createDeferred(); + const active = activeRunContext({ runId, sessionId, sessionKey, persistence }); + const requestContext = await archiveLifecycleRequestContext(active.context); + const placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const reclaim = vi.fn(); + requestContext.workerSessionPlacementService = placementReader(() => placement); + requestContext.workerPlacementDispatchService = { dispatch: vi.fn(), reclaim }; + const authorized = resolveSessionMutationAuthorization({ + client: viewer, + method: "sessions.patch", + requestParams: { key: sessionKey, archived: true }, + context: requestContext, + }); + expect(authorized.error).toBeNull(); + if (!authorized.authorization) { + throw new Error("expected captured archive authorization"); + } + const sharingTarget = resolveSessionSharingTarget({ + cfg: requestContext.getRuntimeConfig(), + sessionKey, + }); + if (!sharingTarget) { + throw new Error("expected resolved sharing target"); + } + + const releaseAudit = createDeferred(); + sessionAuditGate.entered.mockClear(); + sessionAuditGate.wait = releaseAudit.promise; + let sharing: Promise | undefined; + let archive: Promise | undefined; + + try { + let sharingSettled = false; + sharing = invokeVisibilityHandler({ + client: owner, + context: requestContext, + sessionKey, + visibility: "draft", + }).finally(() => { + sharingSettled = true; + }); + await vi.waitFor(() => { + expect(sessionAuditGate.entered).toHaveBeenCalledOnce(); + expect( + isSessionLifecycleMutationActive(sharingTarget.storePath, [sessionKey, sessionId]), + ).toBe(true); + }); + expect(sharingSettled).toBe(false); + expect(loadSessionEntry({ storePath, sessionKey })?.visibility).toBe("draft"); + + let archiveSettled = false; + archive = invokeArchiveHandler({ + authorization: authorized.authorization, + client: viewer, + context: requestContext, + sessionKey, + }).finally(() => { + archiveSettled = true; + }); + await Promise.resolve(); + expect(archiveSettled).toBe(false); + expect(interrupted).toBe(false); + expect(active.controller.signal.aborted).toBe(false); + expectNoSessionQueueCleanup(); + expect(reclaim).not.toHaveBeenCalled(); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + + releaseAudit.resolve(); + expect(await sharing).toMatchObject({ ok: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.visibility).toBe("draft"); + + expect(await archive).toMatchObject({ + ok: false, + error: { details: { code: "SESSION_PARTICIPATION_REQUIRED" } }, + }); + expect(interrupted).toBe(false); + expect(active.controller.signal.aborted).toBe(false); + expectNoSessionQueueCleanup(); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + } finally { + sessionAuditGate.wait = undefined; + releaseAudit.resolve(); + admission.release(); + await Promise.allSettled([...(sharing ? [sharing] : []), ...(archive ? [archive] : [])]); + active.unsubscribe(); + } +}); + +test("archive retains the lifecycle fence until drain and commit before sharing proceeds", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-before-sharing"; + const sessionId = "session-archive-before-sharing"; + const runId = "run-archive-before-sharing"; + const owner = identifiedClient("archive-owner"); + await writeSessionStore({ + entries: { + [sessionKey]: sessionStoreEntry(sessionId, { + createdActor: { type: "human", id: "archive-owner" }, + visibility: "shared", + }), + }, + }); + const admission = await beginSessionWorkAdmission({ + scope: storePath, + identities: [sessionKey, sessionId], + assertAllowed: () => {}, + }); + const persistence = createDeferred(); + const active = activeRunContext({ runId, sessionId, sessionKey, persistence }); + const requestContext = await archiveLifecycleRequestContext(active.context); + let placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const reclaimGate = createDeferred(); + const reclaim = vi.fn(async () => { + await reclaimGate.promise; + placement = workerPlacement({ sessionId, sessionKey, state: "reclaimed" }); + return placement as Extract; + }); + requestContext.workerSessionPlacementService = placementReader(() => placement); + requestContext.workerPlacementDispatchService = { dispatch: vi.fn(), reclaim }; + const authorized = resolveSessionMutationAuthorization({ + client: owner, + method: "sessions.patch", + requestParams: { key: sessionKey, archived: true }, + context: requestContext, + }); + expect(authorized.error).toBeNull(); + if (!authorized.authorization) { + throw new Error("expected captured archive authorization"); + } + let archive: Promise | undefined; + let sharing: Promise | undefined; + + try { + archive = invokeArchiveHandler({ + authorization: authorized.authorization, + client: owner, + context: requestContext, + sessionKey, + }); + await vi.waitFor(() => expect(active.controller.signal.aborted).toBe(true)); + + let sharingSettled = false; + sharing = invokeVisibilityHandler({ + client: owner, + context: requestContext, + sessionKey, + visibility: "draft", + }).finally(() => { + sharingSettled = true; + }); + await Promise.resolve(); + expect(sharingSettled).toBe(false); + expect(loadSessionEntry({ storePath, sessionKey })?.visibility).toBe("shared"); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + + admission.release(); + persistence.resolve(); + await vi.waitFor(() => expect(reclaim).toHaveBeenCalledOnce()); + expect(sharingSettled).toBe(false); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + reclaimGate.resolve(); + expect(await archive).toMatchObject({ ok: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toEqual(expect.any(Number)); + expect(await sharing).toMatchObject({ ok: true }); + expect(loadSessionEntry({ storePath, sessionKey })).toMatchObject({ + archivedAt: expect.any(Number), + visibility: "draft", + }); + } finally { + admission.release(); + persistence.resolve(); + reclaimGate.resolve(); + await Promise.allSettled([...(archive ? [archive] : []), ...(sharing ? [sharing] : [])]); + active.unsubscribe(); + } +}); + +test("alias archive lets the canonical cloud reclaim barrier reenter without deadlock", async () => { + const { storePath } = await createSessionStoreDir(); + const aliasKey = "aaa-archive-cloud-alias"; + const sessionKey = `agent:main:${aliasKey}`; + const sessionId = "session-archive-cloud-alias"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + let placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const reclaimEntered = createDeferred(); + const allowNestedReclaim = createDeferred(); + const contenderRelease = createDeferred(); + const reclaim = vi.fn(async () => { + reclaimEntered.resolve(); + await allowNestedReclaim.promise; + await runExclusiveSessionLifecycleMutation({ + scope: storePath, + identities: [aliasKey, sessionKey, sessionId], + run: async () => {}, + }); + placement = workerPlacement({ sessionId, sessionKey, state: "reclaimed" }); + return placement as Extract; + }); + const archive = directSessionReq( + "sessions.patch", + { key: aliasKey, archived: true }, + { + context: { + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + await reclaimEntered.promise; + const contender = runExclusiveSessionLifecycleMutation({ + scope: storePath, + identities: [aliasKey], + run: async () => await contenderRelease.promise, + }); + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + allowNestedReclaim.resolve(); + let timer: ReturnType | undefined; + + try { + const result = await Promise.race([ + archive, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error("nested alias reclaim deadlocked")), 2_000); + }), + ]); + expect(result.ok).toBe(true); + expect(reclaim).toHaveBeenCalledOnce(); + expect(placement.state).toBe("reclaimed"); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toEqual(expect.any(Number)); + } finally { + if (timer) { + clearTimeout(timer); + } + contenderRelease.resolve(); + allowNestedReclaim.resolve(); + await Promise.allSettled([archive, contender]); + } +}); + +test("sessions.patch returns retryable UNAVAILABLE when runtime drain does not settle", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-stuck"; + const sessionId = "session-archive-stuck"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + embeddedRunMock.activeIds.add(sessionId); + embeddedRunMock.waitResults.set(sessionId, false); + + const archived = await directSessionReq("sessions.patch", { key: sessionKey, archived: true }); + + expect(archived.ok).toBe(false); + expect(archived.error).toMatchObject({ code: "UNAVAILABLE", retryable: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); +}); + +test("sessions.patch rechecks authoritative worker work before projection and releases the drain", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-worker-recheck"; + const sessionId = "session-archive-worker-recheck"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const release = vi.fn(); + const workerEnvironmentService = { + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }; + registerWorkerInferenceSessionDrain(workerEnvironmentService, () => ({ + drained: Promise.resolve(), + hasWork: () => true, + release, + })); + + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + workerEnvironmentService, + }, + }, + ); + + expect(archived.ok).toBe(false); + expect(archived.error).toMatchObject({ code: "UNAVAILABLE", retryable: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + expect(release).toHaveBeenCalledOnce(); +}); + +test("sessions.patch fails closed when active worker inference has no archive drain", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-worker-drain-unavailable"; + const sessionId = "session-archive-worker-drain-unavailable"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + workerEnvironmentService: { + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => true), + resolveInferenceSessionForRunId: vi.fn(), + }, + }, + }, + ); + + expect(archived.ok).toBe(false); + expect(archived.error).toMatchObject({ code: "UNAVAILABLE", retryable: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); +}); + +test("sessions.patch retains the archive drain through the ordered audit append", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-drain-audit"; + const sessionId = "session-archive-drain-audit"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const release = vi.fn(); + const append = vi.spyOn(SessionManager, "appendMessageToTranscript"); + try { + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + client: { + authenticatedUserId: "archive-reviewer@example.com", + authenticatedUserProfile: { + profileId: "archive-reviewer", + displayName: "Archive Reviewer", + hasAvatar: false, + updatedAt: 1, + }, + connect: { scopes: ["operator.write"] }, + } as never, + context: { + workerEnvironmentService: { + beginInferenceSessionDrain: vi.fn(() => ({ + drained: Promise.resolve(), + hasWork: () => false, + release, + })), + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + }, + }, + ); + + expect(archived.ok).toBe(true); + expect(append).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(append.mock.invocationCallOrder[0]).toBeLessThan(release.mock.invocationCallOrder[0]!); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toEqual(expect.any(Number)); + } finally { + append.mockRestore(); + } +}); + +test("sessions.patch returns UNAVAILABLE when terminal persistence fails", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-persistence-failure"; + const sessionId = "session-archive-persistence-failure"; + const runId = "run-archive-persistence-failure"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const persistence = createDeferred(); + const active = activeRunContext({ runId, sessionId, sessionKey, persistence }); + try { + const archive = directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: active.context, + }, + ); + await vi.waitFor(() => expect(active.controller.signal.aborted).toBe(true)); + persistence.reject(new Error("disk full")); + + const archived = await archive; + expect(archived.ok).toBe(false); + expect(archived.error).toMatchObject({ code: "UNAVAILABLE", retryable: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + } finally { + active.unsubscribe(); + } +}); + +test("sessions.patch rejects main and global archives before cancellation side effects", async () => { + const { storePath } = await createSessionStoreDir(); + await writeSessionStore({ entries: { main: sessionStoreEntry("session-main") } }); + embeddedRunMock.activeIds.add("session-main"); + + const main = await directSessionReq("sessions.patch", { key: "main", archived: true }); + expect(main.ok).toBe(false); + expect(main.error?.message).toContain("main session"); + expect(embeddedRunMock.abortCalls).toEqual([]); + expectNoSessionQueueCleanup(); + expect(loadSessionEntry({ storePath, sessionKey: "main" })?.archivedAt).toBeUndefined(); + + const globalFixture = await createConfiguredGlobalAgentSessionStore(); + try { + embeddedRunMock.activeIds.add("sess-main-global"); + const global = await directSessionReq("sessions.patch", { + key: "global", + agentId: "main", + archived: true, + }); + expect(global.ok).toBe(false); + expect(global.error?.message).toContain("main session"); + expect(embeddedRunMock.abortCalls).toEqual([]); + expectNoSessionQueueCleanup(); + } finally { + await resetConfiguredGlobalAgentSessionStore(globalFixture); + } +}); + +test("sessions.patch rejects unknown without materializing a session entry", async () => { + const { storePath } = await createSessionStoreDir(); + await writeSessionStore({ entries: {} }); + + const archived = await directSessionReq("sessions.patch", { key: "unknown", archived: true }); + + expect(archived.ok).toBe(false); + expect(archived.error?.message).toContain("unknown session sentinel"); + expect(loadSessionEntry({ storePath, sessionKey: "unknown" })).toBeUndefined(); + expectNoSessionQueueCleanup(); +}); + +test("sessions.patchMany independently archives active and idle sessions in target order", async () => { + const { storePath } = await createSessionStoreDir(); + const activeKey = "agent:main:archive-batch-active"; + const idleKey = "agent:main:archive-batch-idle"; + const activeSessionId = "session-batch-active"; + await writeSessionStore({ + entries: { + [activeKey]: sessionStoreEntry(activeSessionId), + [idleKey]: sessionStoreEntry("session-batch-idle"), + }, + }); + embeddedRunMock.activeIds.add(activeSessionId); + embeddedRunMock.waitResults.set(activeSessionId, true); + + const result = await directSessionReq<{ outcomes: Array<{ key: string; ok: boolean }> }>( + "sessions.patchMany", + { + targets: [{ key: activeKey }, { key: idleKey }], + patch: { archived: true }, + }, + ); + + expect(result.ok).toBe(true); + expect(result.payload?.outcomes).toEqual([ + { key: activeKey, ok: true }, + { key: idleKey, ok: true }, + ]); + expect(loadSessionEntry({ storePath, sessionKey: activeKey })?.archivedAt).toEqual( + expect.any(Number), + ); + expect(loadSessionEntry({ storePath, sessionKey: idleKey })?.archivedAt).toEqual( + expect.any(Number), + ); +}); + +test("sessions.patchMany prepares independent archive drains concurrently and releases in target order", async () => { + const { storePath } = await createSessionStoreDir(); + const firstKey = "agent:main:archive-batch-concurrent-first"; + const secondKey = "agent:main:archive-batch-concurrent-second"; + const firstSessionId = "session-batch-concurrent-first"; + const secondSessionId = "session-batch-concurrent-second"; + await writeSessionStore({ + entries: { + [firstKey]: sessionStoreEntry(firstSessionId), + [secondKey]: sessionStoreEntry(secondSessionId), + }, + }); + const firstDrained = createDeferred(); + const firstRelease = vi.fn(); + const secondRelease = vi.fn(); + const beginInferenceSessionDrain = vi.fn((sessionId: string) => ({ + drained: sessionId === firstSessionId ? firstDrained.promise : Promise.resolve(), + hasWork: () => false, + release: sessionId === firstSessionId ? firstRelease : secondRelease, + })); + + const archive = directSessionReq<{ outcomes: Array<{ key: string; ok: boolean }> }>( + "sessions.patchMany", + { + targets: [{ key: firstKey }, { key: secondKey }], + patch: { archived: true }, + }, + { + context: { + workerEnvironmentService: { + beginInferenceSessionDrain, + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + }, + }, + ); + + await vi.waitFor(() => expect(beginInferenceSessionDrain).toHaveBeenCalledTimes(2)); + expect(beginInferenceSessionDrain.mock.calls.map(([sessionId]) => sessionId)).toEqual([ + firstSessionId, + secondSessionId, + ]); + expect(firstRelease).not.toHaveBeenCalled(); + expect(secondRelease).not.toHaveBeenCalled(); + firstDrained.resolve(); + + const result = await archive; + expect(result.payload?.outcomes).toEqual([ + { key: firstKey, ok: true }, + { key: secondKey, ok: true }, + ]); + expect(firstRelease).toHaveBeenCalledOnce(); + expect(secondRelease).toHaveBeenCalledOnce(); + expect(firstRelease.mock.invocationCallOrder[0]).toBeLessThan( + secondRelease.mock.invocationCallOrder[0]!, + ); + expect(loadSessionEntry({ storePath, sessionKey: firstKey })?.archivedAt).toEqual( + expect.any(Number), + ); + expect(loadSessionEntry({ storePath, sessionKey: secondKey })?.archivedAt).toEqual( + expect.any(Number), + ); +}); + +test("sessions.patchMany attempts every archive drain release without masking success", async () => { + const { storePath } = await createSessionStoreDir(); + const firstKey = "agent:main:archive-release-throws-first"; + const secondKey = "agent:main:archive-release-after-throw"; + const firstSessionId = "session-archive-release-throws-first"; + const secondSessionId = "session-archive-release-after-throw"; + await writeSessionStore({ + entries: { + [firstKey]: sessionStoreEntry(firstSessionId), + [secondKey]: sessionStoreEntry(secondSessionId), + }, + }); + const firstRelease = vi.fn(() => { + throw new Error("release failed"); + }); + const secondRelease = vi.fn(); + + const result = await directSessionReq<{ outcomes: Array<{ key: string; ok: boolean }> }>( + "sessions.patchMany", + { + targets: [{ key: firstKey }, { key: secondKey }], + patch: { archived: true }, + }, + { + context: { + workerEnvironmentService: { + beginInferenceSessionDrain: vi.fn((sessionId: string) => ({ + drained: Promise.resolve(), + hasWork: () => false, + release: sessionId === firstSessionId ? firstRelease : secondRelease, + })), + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + }, + }, + ); + + expect(result.ok).toBe(true); + expect(result.payload?.outcomes).toEqual([ + { key: firstKey, ok: true }, + { key: secondKey, ok: true }, + ]); + expect(firstRelease).toHaveBeenCalledOnce(); + expect(secondRelease).toHaveBeenCalledOnce(); + expect(loadSessionEntry({ storePath, sessionKey: firstKey })?.archivedAt).toEqual( + expect.any(Number), + ); + expect(loadSessionEntry({ storePath, sessionKey: secondKey })?.archivedAt).toEqual( + expect.any(Number), + ); +}); + +test("sessions.patchMany isolates a failed archive drain and continues later targets", async () => { + const { storePath } = await createSessionStoreDir(); + const stuckKey = "agent:main:archive-batch-stuck"; + const idleKey = "agent:main:archive-batch-after-stuck"; + const stuckSessionId = "session-batch-stuck"; + await writeSessionStore({ + entries: { + [stuckKey]: sessionStoreEntry(stuckSessionId), + [idleKey]: sessionStoreEntry("session-batch-after-stuck"), + }, + }); + embeddedRunMock.activeIds.add(stuckSessionId); + embeddedRunMock.waitResults.set(stuckSessionId, false); + + const result = await directSessionReq<{ + outcomes: Array<{ error?: { code: string; retryable?: boolean }; key: string; ok: boolean }>; + }>("sessions.patchMany", { + targets: [{ key: stuckKey }, { key: idleKey }], + patch: { archived: true }, + }); + + expect(result.ok).toBe(true); + expect(result.payload?.outcomes).toEqual([ + { + key: stuckKey, + ok: false, + error: expect.objectContaining({ code: "UNAVAILABLE", retryable: true }), + }, + { key: idleKey, ok: true }, + ]); + expect(loadSessionEntry({ storePath, sessionKey: stuckKey })?.archivedAt).toBeUndefined(); + expect(loadSessionEntry({ storePath, sessionKey: idleKey })?.archivedAt).toEqual( + expect.any(Number), + ); +}); + +test("sessions.patch rejects a generation replaced after the exact preparation read", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-generation-race"; + const sessionId = "session-archive-generation-race"; + const runId = "run-archive-generation-race"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const persistence = createDeferred(); + const active = activeRunContext({ runId, sessionId, sessionKey, persistence }); + let placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const dispatch = vi.fn(); + const reclaim = vi.fn(async () => { + placement = workerPlacement({ sessionId, sessionKey, state: "reclaimed" }); + return placement as Extract; + }); + try { + const archive = directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + ...active.context, + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch, reclaim }, + }, + }, + ); + await vi.waitFor(() => expect(active.controller.signal.aborted).toBe(true)); + await upsertSessionEntry( + { storePath, sessionKey }, + { sessionId: "session-archive-generation-replacement", updatedAt: 2 }, + ); + persistence.resolve(); + + const archived = await archive; + expect(archived.ok).toBe(false); + expect(archived.error).toMatchObject({ code: "INVALID_REQUEST" }); + expect(loadSessionEntry({ storePath, sessionKey })).toMatchObject({ + sessionId: "session-archive-generation-replacement", + }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + expect(reclaim).toHaveBeenCalledOnce(); + expect(placement.state).toBe("reclaimed"); + expect(dispatch).not.toHaveBeenCalled(); + } finally { + active.unsubscribe(); + } +}); diff --git a/src/gateway/server.sessions.archive-worker-placement.test.ts b/src/gateway/server.sessions.archive-worker-placement.test.ts new file mode 100644 index 000000000000..df6799e93c81 --- /dev/null +++ b/src/gateway/server.sessions.archive-worker-placement.test.ts @@ -0,0 +1,446 @@ +import { afterEach, expect, test, vi } from "vitest"; +import { loadSessionEntry } from "../config/sessions/session-accessor.js"; +import { createDeferred } from "../shared/deferred.js"; +import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; +import { embeddedRunMock, writeSessionStore } from "./test-helpers.js"; +import { + directSessionReq, + expectNoSessionQueueCleanup, + sessionStoreEntry, + setupGatewaySessionsHandlerTestHarness, +} from "./test/server-sessions.test-helpers.js"; +import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-record.js"; + +const { createSessionStoreDir } = setupGatewaySessionsHandlerTestHarness(); + +afterEach(() => { + closeOpenClawStateDatabaseForTest(); +}); + +function workerPlacement(params: { + sessionId: string; + sessionKey: string; + state: WorkerSessionPlacementRecord["state"]; + agentId?: string; + environmentId?: string | null; +}): WorkerSessionPlacementRecord { + return { + sessionId: params.sessionId, + sessionKey: params.sessionKey, + agentId: params.agentId ?? "main", + state: params.state, + generation: 2, + turnClaim: null, + createdAtMs: 1, + updatedAtMs: 2, + stateChangedAtMs: 2, + environmentId: + params.environmentId !== undefined + ? params.environmentId + : params.state === "local" || params.state === "requested" + ? null + : "worker-environment", + activeOwnerEpoch: ["active", "draining", "reconciling", "reclaimed", "failed"].includes( + params.state, + ) + ? 1 + : null, + workspaceBaseManifestRef: + params.state === "local" || + params.state === "requested" || + params.state === "provisioning" || + params.state === "syncing" + ? null + : "manifest-ref", + remoteWorkspaceDir: + params.state === "local" || + params.state === "requested" || + params.state === "provisioning" || + params.state === "syncing" + ? null + : "/workspace", + workerBundleHash: + params.state === "local" || params.state === "requested" || params.state === "provisioning" + ? null + : "bundle-hash", + lastTranscriptAckCursor: null, + lastLiveEventAckCursor: null, + recoveryError: params.state === "failed" ? "worker recovery stopped" : null, + } as WorkerSessionPlacementRecord; +} + +function placementReader(current: () => WorkerSessionPlacementRecord | undefined) { + return { + getMany(sessionIds: readonly string[]) { + const placement = current(); + return new Map( + placement && sessionIds.includes(placement.sessionId) + ? [[placement.sessionId, placement]] + : [], + ); + }, + }; +} + +test("sessions.patch reclaims the exact active cloud placement before archive metadata commits", async () => { + const { storePath } = await createSessionStoreDir(); + const requestedKey = "archive-cloud-active"; + const sessionKey = `agent:main:${requestedKey}`; + const sessionId = "session-archive-cloud-active"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + let placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const reclaimGate = createDeferred(); + const reclaim = vi.fn(async () => { + await reclaimGate.promise; + placement = workerPlacement({ sessionId, sessionKey, state: "reclaimed" }); + return placement as Extract; + }); + + const archive = directSessionReq( + "sessions.patch", + { key: requestedKey, archived: true }, + { + context: { + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + + await vi.waitFor(() => expect(reclaim).toHaveBeenCalledOnce()); + expect(reclaim).toHaveBeenCalledWith({ sessionId, sessionKey, agentId: "main" }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + reclaimGate.resolve(); + + await expect(archive).resolves.toMatchObject({ ok: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toEqual(expect.any(Number)); +}); + +test.each(["rejected", "unavailable"] as const)( + "sessions.patch leaves active placement unarchived and releases its drain when reclaim is %s", + async (failure) => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = `agent:main:archive-cloud-${failure}`; + const sessionId = `session-archive-cloud-${failure}`; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const release = vi.fn(); + const reclaim = vi.fn(async () => { + throw new Error("provider reclaim rejected"); + }); + const workerPlacementDispatchService = + failure === "rejected" ? { dispatch: vi.fn(), reclaim } : { dispatch: vi.fn() }; + + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + workerEnvironmentService: { + beginInferenceSessionDrain: () => ({ + drained: Promise.resolve(), + hasWork: () => false, + release, + }), + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService, + }, + }, + ); + + expect(archived).toMatchObject({ + ok: false, + error: { code: "UNAVAILABLE", retryable: true }, + }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); + expect(release).toHaveBeenCalledOnce(); + expect(reclaim).toHaveBeenCalledTimes(failure === "rejected" ? 1 : 0); + }, +); + +test("sessions.patch rejects a mismatched reclaimed identity without archiving", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-cloud-identity"; + const sessionId = "session-archive-cloud-identity"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const reclaim = vi.fn(async () => + workerPlacement({ + sessionId, + sessionKey: "agent:main:wrong-session", + state: "reclaimed", + }), + ); + + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + + expect(archived).toMatchObject({ + ok: false, + error: { code: "UNAVAILABLE", retryable: true }, + }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); +}); + +test("sessions.patch rejects a placement identity changed during the runtime drain", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:archive-cloud-fresh-placement"; + const sessionId = "session-archive-cloud-fresh-placement"; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + let placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + const drainGate = createDeferred(); + const drainStarted = vi.fn(); + const release = vi.fn(); + const reclaim = vi.fn(); + + const archive = directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + workerEnvironmentService: { + beginInferenceSessionDrain: () => { + drainStarted(); + return { drained: drainGate.promise, hasWork: () => false, release }; + }, + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + + await vi.waitFor(() => expect(drainStarted).toHaveBeenCalledOnce()); + placement = workerPlacement({ + sessionId, + sessionKey: "agent:main:replacement-placement", + state: "active", + }); + drainGate.resolve(); + + await expect(archive).resolves.toMatchObject({ + ok: false, + error: { code: "UNAVAILABLE", retryable: true }, + }); + expect(reclaim).not.toHaveBeenCalled(); + expect(release).toHaveBeenCalledOnce(); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); +}); + +test.each([ + { name: "requested", state: "requested" as const }, + { name: "provisioning", state: "provisioning" as const }, + { name: "syncing", state: "syncing" as const }, + { name: "starting", state: "starting" as const }, + { name: "draining", state: "draining" as const }, + { name: "reconciling", state: "reconciling" as const }, + { name: "failed with a live environment", state: "failed" as const, live: true }, + { name: "failed with an unknown environment", state: "failed" as const }, +])("sessions.patch rejects $name before cancellation or reclaim", async (testCase) => { + const { storePath } = await createSessionStoreDir(); + const caseId = testCase.name.replaceAll(" ", "-"); + const sessionKey = `agent:main:archive-cloud-${caseId}`; + const sessionId = `session-archive-cloud-${caseId}`; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placement = workerPlacement({ sessionId, sessionKey, state: testCase.state }); + const reclaim = vi.fn(); + embeddedRunMock.activeIds.add(sessionId); + + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + ...(testCase.live + ? { workerEnvironmentService: { get: () => ({ state: "attached", leaseId: "lease" }) } } + : {}), + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + + expect(archived).toMatchObject({ + ok: false, + error: { code: "UNAVAILABLE", retryable: true }, + }); + expect(reclaim).not.toHaveBeenCalled(); + expect(embeddedRunMock.abortCalls).toEqual([]); + expectNoSessionQueueCleanup(); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); +}); + +test.each([ + { name: "local", state: "local" as const }, + { name: "reclaimed", state: "reclaimed" as const }, + { name: "failed after its environment is gone", state: "failed" as const, gone: true }, +])("sessions.patch archives $name placement without reclaim", async (testCase) => { + const { storePath } = await createSessionStoreDir(); + const caseId = testCase.name.replaceAll(" ", "-"); + const sessionKey = `agent:main:archive-cloud-${caseId}`; + const sessionId = `session-archive-cloud-${caseId}`; + await writeSessionStore({ entries: { [sessionKey]: sessionStoreEntry(sessionId) } }); + const placement = workerPlacement({ sessionId, sessionKey, state: testCase.state }); + const reclaim = vi.fn(); + + const archived = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: true }, + { + context: { + ...(testCase.gone + ? { + workerEnvironmentService: { + get: () => ({ state: "destroyed" }), + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + } + : {}), + workerSessionPlacementService: placementReader(() => placement), + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + + expect(archived).toMatchObject({ ok: true }); + expect(reclaim).not.toHaveBeenCalled(); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toEqual(expect.any(Number)); +}); + +test.each([ + { name: "reclaimed", state: "reclaimed" as const }, + { name: "failed after its environment is gone", state: "failed" as const, gone: true }, +])("sessions.patch restores $name placement", async (testCase) => { + const { storePath } = await createSessionStoreDir(); + const caseId = testCase.name.replaceAll(" ", "-"); + const sessionKey = `agent:main:restore-cloud-${caseId}`; + const sessionId = `session-restore-cloud-${caseId}`; + await writeSessionStore({ + entries: { [sessionKey]: sessionStoreEntry(sessionId, { archivedAt: 1 }) }, + }); + const placement = workerPlacement({ sessionId, sessionKey, state: testCase.state }); + + const restored = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: false }, + { + context: { + ...(testCase.gone + ? { + workerEnvironmentService: { + get: () => ({ state: "destroyed" }), + cancelInferenceForSession: vi.fn(() => []), + hasInferenceForSession: vi.fn(() => false), + resolveInferenceSessionForRunId: vi.fn(), + }, + } + : {}), + workerSessionPlacementService: placementReader(() => placement), + }, + }, + ); + + expect(restored).toMatchObject({ ok: true }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBeUndefined(); +}); + +test("sessions.patch keeps restore blocked for an active cloud placement", async () => { + const { storePath } = await createSessionStoreDir(); + const sessionKey = "agent:main:restore-cloud-active"; + const sessionId = "session-restore-cloud-active"; + await writeSessionStore({ + entries: { [sessionKey]: sessionStoreEntry(sessionId, { archivedAt: 1 }) }, + }); + const placement = workerPlacement({ sessionId, sessionKey, state: "active" }); + + const restored = await directSessionReq( + "sessions.patch", + { key: sessionKey, archived: false }, + { context: { workerSessionPlacementService: placementReader(() => placement) } }, + ); + + expect(restored).toMatchObject({ ok: false, error: { code: "INVALID_REQUEST" } }); + expect(loadSessionEntry({ storePath, sessionKey })?.archivedAt).toBe(1); +}); + +test("sessions.patchMany isolates a reclaim failure and archives a later target in input order", async () => { + const { storePath } = await createSessionStoreDir(); + const failedKey = "agent:main:archive-batch-reclaim-failed"; + const laterKey = "agent:main:archive-batch-reclaim-later"; + const failedSessionId = "session-batch-reclaim-failed"; + const laterSessionId = "session-batch-reclaim-later"; + await writeSessionStore({ + entries: { + [failedKey]: sessionStoreEntry(failedSessionId), + [laterKey]: sessionStoreEntry(laterSessionId), + }, + }); + const placements = new Map([ + [ + failedSessionId, + workerPlacement({ sessionId: failedSessionId, sessionKey: failedKey, state: "active" }), + ], + [ + laterSessionId, + workerPlacement({ sessionId: laterSessionId, sessionKey: laterKey, state: "local" }), + ], + ]); + const reclaim = vi.fn(async () => { + throw new Error("reclaim failed"); + }); + + const result = await directSessionReq<{ + outcomes: Array<{ error?: { code: string; retryable?: boolean }; key: string; ok: boolean }>; + }>( + "sessions.patchMany", + { + targets: [{ key: failedKey }, { key: laterKey }], + patch: { archived: true }, + }, + { + context: { + workerSessionPlacementService: { + getMany: (sessionIds: readonly string[]) => + new Map( + sessionIds.flatMap((sessionId) => { + const placement = placements.get(sessionId); + return placement ? [[sessionId, placement] as const] : []; + }), + ), + }, + workerPlacementDispatchService: { dispatch: vi.fn(), reclaim }, + }, + }, + ); + + expect(result.payload?.outcomes).toEqual([ + { + key: failedKey, + ok: false, + error: expect.objectContaining({ code: "UNAVAILABLE", retryable: true }), + }, + { key: laterKey, ok: true }, + ]); + expect(reclaim).toHaveBeenCalledOnce(); + expect(loadSessionEntry({ storePath, sessionKey: failedKey })?.archivedAt).toBeUndefined(); + expect(loadSessionEntry({ storePath, sessionKey: laterKey })?.archivedAt).toEqual( + expect.any(Number), + ); +}); diff --git a/src/gateway/server.sessions.compaction.test.ts b/src/gateway/server.sessions.compaction.test.ts index f474739e2805..6f11858a8a09 100644 --- a/src/gateway/server.sessions.compaction.test.ts +++ b/src/gateway/server.sessions.compaction.test.ts @@ -1596,7 +1596,7 @@ test("sessions.compact refuses real compaction while a worker inference owns the expectNoSessionQueueCleanup(); }); -test("sessions.patch rejects archive while terminal compaction owns the session", async () => { +test("sessions.patch waits for terminal compaction before archiving the session", async () => { const { storePath } = await createSessionStoreDir(); const sessionKey = "agent:main:dashboard:compact-race"; await seedSessionEntry({ @@ -1627,9 +1627,15 @@ test("sessions.patch rejects archive while terminal compaction owns the session" await vi.waitFor(() => { expect(embeddedRunMock.compactEmbeddedAgentSession).toHaveBeenCalledTimes(1); }); - const archived = await rpcReq(ws, "sessions.patch", { key: sessionKey, archived: true }); - expect(archived.ok).toBe(false); - expect(archived.error?.message).toContain("active run"); + let archiveSettled = false; + const archiveResult = rpcReq(ws, "sessions.patch", { key: sessionKey, archived: true }).then( + (result) => { + archiveSettled = true; + return result; + }, + ); + await Promise.resolve(); + expect(archiveSettled).toBe(false); compaction.resolve({ ok: true, @@ -1642,6 +1648,7 @@ test("sessions.patch rejects archive while terminal compaction owns the session" }, }); expect((await compactResult).ok).toBe(true); + expect((await archiveResult).ok).toBe(true); ws.close(); }); diff --git a/src/gateway/server.sessions.delete-lifecycle.test.ts b/src/gateway/server.sessions.delete-lifecycle.test.ts index f77a4cd080c6..d2fc4f85bfd9 100644 --- a/src/gateway/server.sessions.delete-lifecycle.test.ts +++ b/src/gateway/server.sessions.delete-lifecycle.test.ts @@ -679,26 +679,6 @@ test("sessions.delete keeps lifecycle admission blocked through session unbindin } }); -test("sessions.patch rejects archiving active runs", async () => { - await createSessionStoreDir(); - await writeSessionStore({ - entries: { - "discord:group:dev": sessionStoreEntry("sess-active"), - }, - }); - embeddedRunMock.activeIds.add("sess-active"); - - const archived = await directSessionReq("sessions.patch", { - key: "discord:group:dev", - archived: true, - }); - - expect(archived.ok).toBe(false); - expect(archived.error).toMatchObject({ - message: "Cannot archive a session with an active run.", - }); -}); - test("sessions.delete limits plugin-runtime cleanup to sessions owned by that plugin", async () => { const { dir, storePath } = await createSessionStoreDir(); await writeSingleLineSession(dir, "sess-owned", "owned"); diff --git a/src/gateway/server.sessions.list-changed.test.ts b/src/gateway/server.sessions.list-changed.test.ts index a7f44d8e0093..157fae88660f 100644 --- a/src/gateway/server.sessions.list-changed.test.ts +++ b/src/gateway/server.sessions.list-changed.test.ts @@ -180,6 +180,9 @@ async function invokeSessionMutation({ respond, context: { broadcastToConnIds, + chatAbortControllers: new Map(), + chatQueuedTurns: new Map(), + dedupe: new Map(), getSessionEventSubscriberConnIds: () => subscribedConnIds, loadGatewayModelCatalog: async () => ({ providers: [] }), getRuntimeConfig, diff --git a/src/gateway/server.sessions.reset-cleanup.test.ts b/src/gateway/server.sessions.reset-cleanup.test.ts index 0ca44ebc2a54..f1bd521efc79 100644 --- a/src/gateway/server.sessions.reset-cleanup.test.ts +++ b/src/gateway/server.sessions.reset-cleanup.test.ts @@ -619,7 +619,7 @@ test("sessions.reset rejects a concurrent archive during lifecycle rotation", as expect(reset.ok).toBe(true); expect(archived).toMatchObject({ ok: false, - error: { message: "Cannot archive a session with an active run." }, + error: { message: `Session ${sessionKey} changed before patch. Retry.` }, }); const entry = loadSessionEntry({ storePath, sessionKey }); expect(entry?.archivedAt).toBeUndefined(); diff --git a/src/gateway/session-sharing.test.ts b/src/gateway/session-sharing.test.ts index f0124e8b0e13..4694f258c597 100644 --- a/src/gateway/session-sharing.test.ts +++ b/src/gateway/session-sharing.test.ts @@ -82,6 +82,7 @@ function target(createdActor?: { type: "human"; id: string; label?: string }): S ...(createdActor ? { createdActor } : {}), }, storeKey: "agent:main:main", + storeKeys: ["agent:main:main"], storePath: "/tmp/sessions.json", }; } diff --git a/src/gateway/session-sharing.ts b/src/gateway/session-sharing.ts index 8599f0e028eb..98bd9a77203c 100644 --- a/src/gateway/session-sharing.ts +++ b/src/gateway/session-sharing.ts @@ -36,13 +36,12 @@ import { resolveGatewaySessionStoreTargetWithStore, } from "./session-utils.js"; -const ADMIN_SCOPE = "operator.admin"; - type SessionSharingTarget = { agentId: string; canonicalKey: string; entry: SessionEntry; storeKey: string; + storeKeys: string[]; storePath: string; }; @@ -80,7 +79,7 @@ export function resolveSessionVisibility( export function isGatewayAdmin(client: Pick | null): boolean { // Internal/plugin-runtime runs reach authorization with a client that has no // connect handshake; treat a connect-less client as a non-admin, never a crash. - return client?.connect?.scopes?.includes(ADMIN_SCOPE) === true; + return client?.connect?.scopes?.includes("operator.admin") === true; } export function allowedSessionVisibilities(cfg: OpenClawConfig): SessionVisibility[] { @@ -124,6 +123,7 @@ export function resolveSessionSharingTarget(params: { canonicalKey: target.canonicalKey, entry: match.entry, storeKey: match.key, + storeKeys: target.storeKeys, storePath: target.storePath, } : null; diff --git a/src/gateway/test-helpers.mocks.ts b/src/gateway/test-helpers.mocks.ts index 9a75bd201ce7..a3f6b52a483e 100644 --- a/src/gateway/test-helpers.mocks.ts +++ b/src/gateway/test-helpers.mocks.ts @@ -38,6 +38,7 @@ function createEmbeddedRunMockExports() { embeddedRunMock.waitCalls.push(sessionId); const ended = embeddedRunMock.waitResults.get(sessionId) ?? true; if (ended) { + embeddedRunMock.activeIds.delete(sessionId); embeddedRunMock.endWaiters.get(sessionId)?.(true); } else if (embeddedRunMock.resolveEndBeforeTimeoutIds.delete(sessionId)) { embeddedRunMock.endWaiters.get(sessionId)?.(true); diff --git a/src/gateway/worker-environments/inference-control-internal.ts b/src/gateway/worker-environments/inference-control-internal.ts new file mode 100644 index 000000000000..60b4d0e74a99 --- /dev/null +++ b/src/gateway/worker-environments/inference-control-internal.ts @@ -0,0 +1,28 @@ +export type WorkerInferenceSessionDrain = { + drained: Promise; + hasWork(): boolean; + release(): void; +}; + +type BeginWorkerInferenceSessionDrain = (sessionId: string) => WorkerInferenceSessionDrain; + +// Archive lifecycle needs a stronger control without widening the inferred public service shape. +// The weak registration follows the concrete service instance's lifetime. +const sessionDrainByService = new WeakMap(); + +export function registerWorkerInferenceSessionDrain( + service: object, + beginDrain: BeginWorkerInferenceSessionDrain, +): void { + sessionDrainByService.set(service, beginDrain); +} + +export function beginWorkerInferenceSessionDrain( + service: unknown, + sessionId: string, +): WorkerInferenceSessionDrain | undefined { + if (typeof service !== "object" || service === null) { + return undefined; + } + return sessionDrainByService.get(service)?.(sessionId); +} diff --git a/src/gateway/worker-environments/inference.test.ts b/src/gateway/worker-environments/inference.test.ts index 55da5ef9aba5..053e520ed97b 100644 --- a/src/gateway/worker-environments/inference.test.ts +++ b/src/gateway/worker-environments/inference.test.ts @@ -6,6 +6,7 @@ import type { } from "../../../packages/gateway-protocol/src/schema/worker-inference.js"; import { createDeferred } from "../../../test/helpers/promise.js"; import type { WorkerConnectionIdentity } from "./connection-identity.js"; +import type { WorkerInferenceSessionDrain } from "./inference-control-internal.js"; import type { WorkerInferenceStore } from "./inference-store.js"; import { createWorkerInferenceManager, @@ -20,6 +21,17 @@ function waitForFast( return vi.waitFor(callback, { interval: 1, ...options }); } +function beginSessionDrain( + manager: ReturnType, + sessionId: string, +): WorkerInferenceSessionDrain { + return ( + manager as typeof manager & { + beginSessionDrain(sessionId: string): WorkerInferenceSessionDrain; + } + ).beginSessionDrain(sessionId); +} + const REQUEST: WorkerInferenceStartParams = { runEpoch: 3, sessionId: "s", @@ -185,6 +197,53 @@ describe("worker inference manager", () => { await instance.stop(); }); + it("blocks replacement inference until an exact session drain settles", async () => { + const pending = createDeferred(); + const execute = vi.fn(async () => await pending.promise); + const instance = makeManager(execute); + accept(instance); + await waitForFast(() => expect(execute).toHaveBeenCalledOnce()); + + const drain = beginSessionDrain(instance, REQUEST.sessionId); + expect(drain.hasWork()).toBe(true); + expect( + instance.start({ + identity: IDENTITY, + request: { ...REQUEST, runId: "replacement", turnId: "replacement" }, + sink: createSink().sink, + }), + ).toEqual({ ok: false, reason: "cancelled" }); + + pending.resolve(ERROR); + await drain.drained; + expect(drain.hasWork()).toBe(false); + drain.release(); + expect( + instance.start({ + identity: IDENTITY, + request: { ...REQUEST, runId: "replacement", turnId: "replacement" }, + sink: createSink().sink, + }), + ).toMatchObject({ ok: true }); + await instance.stop(); + }); + + it("rejects an inference drain when terminal persistence fails", async () => { + const store = createMemoryStore(); + vi.spyOn(store, "complete").mockImplementation(() => { + throw new Error("write failed"); + }); + const pending = createDeferred(); + const instance = makeManager(async () => await pending.promise, store); + accept(instance); + + const drain = beginSessionDrain(instance, REQUEST.sessionId); + pending.resolve(ERROR); + await expect(drain.drained).rejects.toThrow("terminal persistence failed"); + drain.release(); + await instance.stop(); + }); + it("settles cumulatively oversized output while preserving the abort reason", async () => { let signal: AbortSignal | undefined; const store = createMemoryStore(); diff --git a/src/gateway/worker-environments/inference.ts b/src/gateway/worker-environments/inference.ts index c0a99747c413..a98f9abeb635 100644 --- a/src/gateway/worker-environments/inference.ts +++ b/src/gateway/worker-environments/inference.ts @@ -20,6 +20,7 @@ import { withTimeout } from "../../infra/fs-safe.js"; import { boundedJsonUtf8Bytes } from "../../infra/json-utf8-bytes.js"; import { runWithGatewayIndependentRootWorkContinuation } from "../../process/gateway-work-admission.js"; import type { WorkerConnectionIdentity } from "./connection-identity.js"; +import type { WorkerInferenceSessionDrain } from "./inference-control-internal.js"; import { createWorkerInferenceStore, type WorkerInferenceStore, @@ -228,6 +229,7 @@ export function createWorkerInferenceManager(options: { const now = options.now ?? Date.now; const active = new Map(); const operations = new Map, string>(); + const drainingSessionIds = new Set(); let stopping = false; store.recoverPending(terminalError("provider-error")); @@ -421,7 +423,7 @@ export function createWorkerInferenceManager(options: { sink: WorkerInferenceSink; revalidate?: RevalidateInference; }): WorkerInferenceStartApplicationResult => { - if (stopping) { + if (stopping || drainingSessionIds.has(params.request.sessionId)) { return { ok: false, reason: "cancelled" }; } const identityError = matchesIdentity(params.identity, params.request); @@ -611,12 +613,16 @@ export function createWorkerInferenceManager(options: { const cancelWhere = ( predicate: (entry: ActiveInference) => boolean, reason: WorkerInferenceErrorReason, - ): void => { + onCancel?: (entry: ActiveInference) => void, + ): boolean => { + let terminalPersistenceFailed = false; for (const entry of active.values()) { if (predicate(entry)) { - settleAbort(entry, reason); + onCancel?.(entry); + terminalPersistenceFailed = !settleAbort(entry, reason) || terminalPersistenceFailed; } } + return terminalPersistenceFailed; }; const cancelEnvironment = ( @@ -628,29 +634,70 @@ export function createWorkerInferenceManager(options: { const cancelSession = (sessionId: string, runId?: string): string[] => { const cancelledRunIds = new Set(); - for (const entry of active.values()) { - if ( - entry.request.sessionId === sessionId && - (runId === undefined || entry.request.runId === runId) - ) { - cancelledRunIds.add(entry.request.runId); - } - } cancelWhere( (entry) => entry.request.sessionId === sessionId && (runId === undefined || entry.request.runId === runId), "cancelled", + (entry) => cancelledRunIds.add(entry.request.runId), ); return [...cancelledRunIds].toSorted(); }; - const hasSession = (sessionId: string, runId?: string): boolean => - [...active.values()].some( - (entry) => + const hasSession = (sessionId: string, runId?: string): boolean => { + for (const entry of active.values()) { + if ( entry.request.sessionId === sessionId && - (runId === undefined || entry.request.runId === runId), + (runId === undefined || entry.request.runId === runId) + ) { + return true; + } + } + return false; + }; + + const hasSessionOperation = (sessionId: string): boolean => { + for (const operationSessionId of operations.values()) { + if (operationSessionId === sessionId) { + return true; + } + } + return false; + }; + + const beginSessionDrain = (sessionId: string): WorkerInferenceSessionDrain => { + if (drainingSessionIds.has(sessionId)) { + throw new Error(`Worker inference drain already owns session ${sessionId}`); + } + // Block first so cancellation cannot race a replacement provider operation. + drainingSessionIds.add(sessionId); + const terminalPersistenceFailed = cancelWhere( + (entry) => entry.request.sessionId === sessionId, + "cancelled", ); + const providerOperations: Promise[] = []; + for (const [operation, operationSessionId] of operations) { + if (operationSessionId === sessionId) { + providerOperations.push(operation); + } + } + let released = false; + return { + drained: Promise.allSettled(providerOperations).then(() => { + if (terminalPersistenceFailed) { + throw new Error(`Worker inference terminal persistence failed for session ${sessionId}`); + } + }), + hasWork: () => hasSession(sessionId) || hasSessionOperation(sessionId), + release: () => { + if (released) { + return; + } + released = true; + drainingSessionIds.delete(sessionId); + }, + }; + }; const resolveSessionIdForRunId = (runId: string): string | undefined => { const sessionIds = new Set(); @@ -672,7 +719,7 @@ export function createWorkerInferenceManager(options: { ).catch(() => undefined); }; - return { + const manager = { start, cancel, cancelEnvironment, @@ -681,4 +728,7 @@ export function createWorkerInferenceManager(options: { resolveSessionIdForRunId, stop, }; + // Archive-only control stays non-enumerable so the manager's inferred contract remains stable. + Object.defineProperty(manager, "beginSessionDrain", { value: beginSessionDrain }); + return manager; } diff --git a/src/gateway/worker-environments/service.ts b/src/gateway/worker-environments/service.ts index 3f3a5a160bcd..2391ecdc50d8 100644 --- a/src/gateway/worker-environments/service.ts +++ b/src/gateway/worker-environments/service.ts @@ -49,6 +49,10 @@ import { type WorkerCredentialBinding, type WorkerCredentialDeliveryClaim, } from "./credential.js"; +import { + registerWorkerInferenceSessionDrain, + type WorkerInferenceSessionDrain, +} from "./inference-control-internal.js"; import type { WorkerInferenceStore } from "./inference-store.js"; import { createWorkerInferenceManager, @@ -234,6 +238,9 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService now, ...(options.inferenceStore ? { store: options.inferenceStore } : {}), }); + const inferenceWithDrain = inference as typeof inference & { + beginSessionDrain(sessionId: string): WorkerInferenceSessionDrain; + }; let reconcileInFlight: Promise | undefined; let interval: ReturnType | undefined; let unsubscribeSessionIdentityMutation: (() => void) | undefined; @@ -1506,7 +1513,7 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService }); }; - return { + const service = { list: () => store.list().map(project), get: (environmentId: string) => { const record = store.get(environmentId); @@ -1637,6 +1644,10 @@ export function createWorkerEnvironmentService(options: WorkerEnvironmentService start, stop, }; + registerWorkerInferenceSessionDrain(service, (sessionId) => + inferenceWithDrain.beginSessionDrain(sessionId), + ); + return service; } export type WorkerEnvironmentService = ReturnType; diff --git a/src/gateway/worker-environments/session-placement-lifecycle.ts b/src/gateway/worker-environments/session-placement-lifecycle.ts index a2d8c10e625c..1ae0c31b1dfe 100644 --- a/src/gateway/worker-environments/session-placement-lifecycle.ts +++ b/src/gateway/worker-environments/session-placement-lifecycle.ts @@ -3,10 +3,14 @@ import type { WorkerSessionPlacementRetirement, WorkerSessionPlacementStore, } from "./placement-store.js"; -import type { WorkerEnvironmentServiceContract } from "./service-contract.js"; +import type { + WorkerEnvironmentServiceContract, + WorkerPlacementDispatchContract, +} from "./service-contract.js"; export type SessionWorkerPlacementContext = { workerEnvironmentService?: Pick; + workerPlacementDispatchService?: Pick; workerSessionPlacementService?: Pick & Partial>; }; @@ -60,6 +64,19 @@ export function isFailedWorkerPlacementEnvironmentGone(params: { } } +export function isWorkerPlacementSafeForArchive( + context: SessionWorkerPlacementContext, + placement: Placement, +): boolean { + if (placement.state === "failed") { + return isFailedWorkerPlacementEnvironmentGone({ + environmentService: context.workerEnvironmentService, + placement, + }); + } + return placement.state === "local" || placement.state === "reclaimed"; +} + function retirementGuard(placement: RetirablePlacement): SessionWorkerPlacementMutationGuard { return { status: "retirement-required", @@ -127,3 +144,44 @@ export function resolveSessionWorkerPlacementMutationError( const guard = resolveSessionWorkerPlacementMutationGuard(params); return guard.status === "blocked" ? guard.error : undefined; } + +export async function prepareSessionWorkerPlacementForArchive(params: { + agentId: string; + context: SessionWorkerPlacementContext; + reclaimActive: boolean; + sessionId?: string; + sessionKey: string; +}): Promise { + const { agentId, context, sessionId, sessionKey } = params; + if (!sessionId) { + return; + } + const request = { agentId, sessionId, sessionKey }; + const placement = context.workerSessionPlacementService?.getMany([sessionId]).get(sessionId); + if (!placement) { + return; + } + const matches = (candidate: Placement) => + candidate.sessionId === sessionId && + candidate.sessionKey === sessionKey && + candidate.agentId === agentId; + if (!matches(placement)) { + throw new Error(`Session ${sessionKey} cloud worker placement identity changed.`); + } + if (isWorkerPlacementSafeForArchive(context, placement)) { + return; + } + if (placement.state !== "active") { + throw new Error(`Session ${sessionKey} cannot archive from placement ${placement.state}.`); + } + if (!params.reclaimActive) { + return; + } + if (!context.workerPlacementDispatchService?.reclaim) { + throw new Error(`Session ${sessionKey} cloud worker reclaim is unavailable.`); + } + const reclaimed: Placement = await context.workerPlacementDispatchService.reclaim(request); + if (reclaimed.state !== "reclaimed" || !matches(reclaimed)) { + throw new Error(`Session ${sessionKey} cloud worker reclaim identity changed.`); + } +} diff --git a/src/sessions/session-lifecycle-admission.ts b/src/sessions/session-lifecycle-admission.ts index a16ebb38df8b..f2276c0cf30d 100644 --- a/src/sessions/session-lifecycle-admission.ts +++ b/src/sessions/session-lifecycle-admission.ts @@ -22,8 +22,6 @@ export { } from "./session-work-admission-handoff.js"; export const SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS = 15_000; -/** Stable gateway error for an archive rejected by an admitted or projected run. */ -export const SESSION_ARCHIVE_ACTIVE_RUN_ERROR = "Cannot archive a session with an active run."; type SessionWorkAdmission = HandoffSessionWorkAdmission & { interrupt?: () => void; released: Promise; diff --git a/src/shared/session-archive-timeout.ts b/src/shared/session-archive-timeout.ts new file mode 100644 index 000000000000..b7b09acd8dc4 --- /dev/null +++ b/src/shared/session-archive-timeout.ts @@ -0,0 +1,3 @@ +/** Cloud workspace reconciliation may take several minutes before archive can commit. */ +export const SESSION_ARCHIVE_REQUEST_TIMEOUT_MS = 10 * 60_000; +export const SESSION_ARCHIVE_REQUEST_OPTIONS = { timeoutMs: SESSION_ARCHIVE_REQUEST_TIMEOUT_MS }; diff --git a/ui/src/components/session-menu.test.ts b/ui/src/components/session-menu.test.ts index 50735d783f99..e212363349c1 100644 --- a/ui/src/components/session-menu.test.ts +++ b/ui/src/components/session-menu.test.ts @@ -36,6 +36,7 @@ async function mountMenu( work?: SessionMenuWork | null; workboard?: { captured: boolean; busy: boolean } | null; archiveAllowed?: boolean; + deleteAllowed?: boolean; cloudWorkerStopAllowed?: boolean; selectionCount?: number; lastActive?: string; @@ -68,6 +69,8 @@ async function mountMenu( .actionDisabledReasons=${options.actionDisabledReasons ?? {}} .forkDisabled=${false} .archiveAllowed=${options.archiveAllowed ?? true} + .deleteAllowed=${options.deleteAllowed ?? + (session.archived || (options.archiveAllowed ?? true))} .cloudWorkerStopAllowed=${options.cloudWorkerStopAllowed ?? false} .groups=${options.groups ?? []} .canOpenChat=${options.canOpenChat ?? true} @@ -227,13 +230,24 @@ describe("session menu", () => { expect(menuItem(menu, "Pin session").disabled).toBe(true); }); - it("disables archive and delete when an active session cannot be archived", async () => { - const menu = await mountMenu({ archiveAllowed: false }); + it("enables archive while preserving disabled delete for an active session", async () => { + const menu = await mountMenu({ archiveAllowed: true, deleteAllowed: false }); - expect(menuItem(menu, "Archive session").disabled).toBe(true); + expect(menuItem(menu, "Archive session").disabled).toBe(false); expect(menuItem(menu, "Delete…").disabled).toBe(true); }); + it("keeps batch archive enabled while independently guarding delete", async () => { + const menu = await mountMenu({ + selectionCount: 2, + archiveAllowed: false, + deleteAllowed: false, + }); + + expect(menuItem(menu, "Archive 2").disabled).toBe(false); + expect(menuItem(menu, "Delete 2…").disabled).toBe(true); + }); + it("closes before dispatching Pin", async () => { const calls: string[] = []; const menu = await mountMenu({ diff --git a/ui/src/components/session-menu.ts b/ui/src/components/session-menu.ts index c8804a39db98..bdbf96ff2d10 100644 --- a/ui/src/components/session-menu.ts +++ b/ui/src/components/session-menu.ts @@ -78,9 +78,8 @@ class SessionMenu extends OpenClawLightDomElement { Record > = {}; @property({ attribute: false }) forkDisabled = false; - // Single-session Archive and Delete use this local eligibility guard. - // Batch Archive delegates mixed rows to patchMany's authoritative ordered outcomes. @property({ attribute: false }) archiveAllowed = false; + @property({ attribute: false }) deleteAllowed = false; @property({ attribute: false }) cloudWorkerStopAllowed = false; @property({ attribute: false }) groups: readonly string[] = []; @property({ attribute: false }) canOpenChat = false; @@ -614,10 +613,7 @@ class SessionMenu extends OpenClawLightDomElement { variant="danger" data-shortcut="d" aria-keyshortcuts="D" - ?disabled=${this.actionDisabled( - "delete", - !(session.archived || this.archiveAllowed), - )} + ?disabled=${this.actionDisabled("delete", !this.deleteAllowed)} title=${this.actionTitle("delete")} >