From 88d6a2c8d635be601f37dc4ed90636d26ef950db Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sun, 2 Aug 2026 16:08:07 -0700 Subject: [PATCH] perf(gateway): reduce sessions.list read amplification under streaming load (#118207) * perf(gateway): reduce sessions list read amplification * test(gateway): stabilize title batch perf proof * fix(gateway): satisfy sessions list CI contracts * perf(gateway): reuse warm batch title fields * fix(gateway): satisfy strict title cache types * test(gateway): use transcript write scope in cache test * fix(gateway): fence session list cache on active runs * refactor(infra): extract agent run registry * fix(infra): keep run-registry context types module-local --- packages/sdk/src/index.e2e.test.ts | 3 +- .../agent-command.live-model-switch.test.ts | 6 +- src/agents/agent-command.ts | 2 +- src/agents/command/acp-execution.ts | 6 +- src/agents/command/session-preparation.ts | 6 +- .../run.overflow-compaction.test.ts | 3 +- .../run/lane-controller.ts | 3 +- .../run/session-prompt-state.ts | 2 +- .../main-session-restart-recovery-marking.ts | 6 +- .../main-session-restart-recovery.test.ts | 2 +- src/agents/subagent-registry-read.ts | 2 +- src/agents/subagent-registry-sweeper.ts | 2 +- .../subagent-registry.archive.e2e.test.ts | 4 +- src/agents/subagent-registry.test.ts | 4 +- src/agents/tools/cron-tool.pacing.test.ts | 2 +- src/agents/tools/cron-tool.ts | 2 +- .../tools/media-generate-background-shared.ts | 2 +- .../tools/video-generate-background.test.ts | 3 +- .../agent-runner-execution-lifecycle.test.ts | 8 +- .../agent-runner-execution.test-support.ts | 10 + .../reply/agent-runner-execution.ts | 3 +- src/auto-reply/reply/agent-runner-memory.ts | 3 +- .../reply/agent-runner.media-paths.test.ts | 9 + src/auto-reply/reply/followup-runner.test.ts | 2 +- src/auto-reply/reply/followup-runner.ts | 2 +- src/commands/agent.acp.test.ts | 4 + ...ssion-accessor.sqlite-data-version.test.ts | 48 +- .../session-accessor.sqlite-entry-cache.ts | 77 ++- .../session-accessor.sqlite-entry-store.ts | 2 +- .../session-accessor.sqlite-title-probes.ts | 208 ++++++++ src/config/sessions/session-accessor.ts | 4 + .../isolated-agent/run-execution.runtime.ts | 2 +- .../run.message-tool-policy.test.ts | 34 +- src/cron/isolated-agent/run.ts | 10 +- src/gateway/chat-abort.test.ts | 7 +- src/gateway/server-chat.agent-events.test.ts | 8 +- src/gateway/server-chat.ts | 3 +- src/gateway/server-maintenance.ts | 2 +- .../agent-run-admission-phase.ts | 2 +- .../server-methods/agent-run-dispatch.ts | 2 +- .../agent.sessions-and-models.test-utils.ts | 2 + .../server-methods/agent.test-harness.ts | 7 + .../server-methods/chat-send-admission.ts | 7 +- .../server-methods/chat-send-attachments.ts | 2 +- .../chat-send-dispatch-errors.ts | 2 +- .../server-methods/chat-send-handler.ts | 6 +- .../chat.directive-tags.test.ts | 2 +- .../session-active-runs.test.ts | 2 +- .../server-methods/session-active-runs.ts | 2 +- .../session-change-event.test.ts | 125 +++++ .../server-methods/session-change-event.ts | 110 +++- .../server-methods/sessions-list-cache.ts | 124 +++++ .../sessions-read-cache.test.ts | 172 ++++++- src/gateway/server-methods/sessions-read.ts | 183 ++++--- src/gateway/server-runtime-subscriptions.ts | 7 +- src/gateway/server-session-key.test.ts | 3 +- src/gateway/server-session-key.ts | 2 +- src/gateway/server-startup-finish.ts | 13 +- ...erver.agent.gateway-server-agent-b.test.ts | 3 +- .../server.chat.gateway-server-chat.test.ts | 3 +- .../server.node-chat-subscriptions.test.ts | 3 +- ...essions.list-store-materialization.test.ts | 7 +- src/gateway/session-message-events.test.ts | 7 +- src/gateway/session-observer.ts | 2 +- .../session-sharing-snapshot-cache.test.ts | 73 +++ src/gateway/session-sharing-snapshot-cache.ts | 84 +++- .../session-transcript-readers.test.ts | 88 +++- .../session-transcript-title-reader.ts | 88 +++- src/gateway/session-utils-list.ts | 46 +- src/gateway/session-utils.perf.test.ts | 55 +- src/gateway/session-utils.search.test.ts | 3 +- src/gateway/session-utils.subagent.test.ts | 3 +- src/gateway/sessions-resolve.test.ts | 3 + src/gateway/sessions-resolve.ts | 5 +- .../worker-environments/live-events.test.ts | 12 +- .../worker-environments/live-events.ts | 6 +- src/infra/agent-events.test.ts | 78 ++- src/infra/agent-events.ts | 453 ++--------------- src/infra/agent-run-registry.ts | 474 ++++++++++++++++++ src/tasks/task-registry.maintenance.ts | 2 +- src/tasks/task-registry.test.ts | 7 +- src/worker/worker.fault-injection.test.ts | 5 +- test/scripts/audit-seams.test.ts | 2 +- 83 files changed, 2056 insertions(+), 737 deletions(-) create mode 100644 src/config/sessions/session-accessor.sqlite-title-probes.ts create mode 100644 src/gateway/server-methods/session-change-event.test.ts create mode 100644 src/gateway/server-methods/sessions-list-cache.ts create mode 100644 src/gateway/session-sharing-snapshot-cache.test.ts create mode 100644 src/infra/agent-run-registry.ts diff --git a/packages/sdk/src/index.e2e.test.ts b/packages/sdk/src/index.e2e.test.ts index a08d29be1787..9c40999f8d05 100644 --- a/packages/sdk/src/index.e2e.test.ts +++ b/packages/sdk/src/index.e2e.test.ts @@ -4,7 +4,8 @@ import net from "node:net"; import { afterEach, describe, expect, it } from "vitest"; import { WebSocketServer, type WebSocket } from "ws"; import { installGatewayTestHooks, startServer } from "../../../src/gateway/test-helpers.js"; -import { emitAgentEvent, registerAgentRunContext } from "../../../src/infra/agent-events.js"; +import { emitAgentEvent } from "../../../src/infra/agent-events.js"; +import { registerAgentRunContext } from "../../../src/infra/agent-run-registry.js"; import { rawDataToString } from "../../../src/infra/ws.js"; import { withTimeout } from "../../../src/utils/with-timeout.js"; import { GatewayClientTransport, OpenClaw } from "./index.js"; diff --git a/src/agents/agent-command.live-model-switch.test.ts b/src/agents/agent-command.live-model-switch.test.ts index f626ede7205d..8b64c109e84a 100644 --- a/src/agents/agent-command.live-model-switch.test.ts +++ b/src/agents/agent-command.live-model-switch.test.ts @@ -355,15 +355,17 @@ vi.mock("../infra/agent-events.js", () => ({ assertAgentRunLifecycleGenerationCurrent: (...args: unknown[]) => state.assertLifecycleCurrentMock(...args), captureAgentRunLifecycleGeneration: () => "test-generation", - clearAgentRunContext: (...args: unknown[]) => state.clearAgentRunContextMock(...args), emitAgentEvent: (...args: unknown[]) => state.emitAgentEventMock(...args), getAgentEventLifecycleGeneration: () => "test-generation", isAgentEventLifecycleGenerationCurrent: (generation: string) => generation === "test-generation", onAgentEvent: vi.fn(), registerAgentEventLifecycleRotationHandler: vi.fn(), - registerAgentRunContext: (...args: unknown[]) => state.registerAgentRunContextMock(...args), withAgentRunLifecycleGeneration: (_generation: string, run: () => unknown) => run(), })); +vi.mock("../infra/agent-run-registry.js", () => ({ + clearAgentRunContext: (...args: unknown[]) => state.clearAgentRunContextMock(...args), + registerAgentRunContext: (...args: unknown[]) => state.registerAgentRunContextMock(...args), +})); vi.mock("../infra/outbound/session-context.js", () => ({ buildOutboundSessionContext: () => ({}), diff --git a/src/agents/agent-command.ts b/src/agents/agent-command.ts index abdd514e43ee..20fe5547f96b 100644 --- a/src/agents/agent-command.ts +++ b/src/agents/agent-command.ts @@ -11,9 +11,9 @@ import { withLocalGatewayRequestScope } from "../gateway/local-request-context.j import { assertAgentRunLifecycleGenerationCurrent, captureAgentRunLifecycleGeneration, - clearAgentRunContext, withAgentRunLifecycleGeneration, } from "../infra/agent-events.js"; +import { clearAgentRunContext } from "../infra/agent-run-registry.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; import { isSubagentSessionKey } from "../routing/session-key.js"; import { defaultRuntime, type RuntimeEnv } from "../runtime.js"; diff --git a/src/agents/command/acp-execution.ts b/src/agents/command/acp-execution.ts index 540667dce1d5..8cf79314a29a 100644 --- a/src/agents/command/acp-execution.ts +++ b/src/agents/command/acp-execution.ts @@ -2,10 +2,8 @@ import { resolveInlineAgentImageAttachments } from "../../auto-reply/reply/agent import type { CliDeps } from "../../cli/deps.types.js"; import type { SessionEntry } from "../../config/sessions/types.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { - assertAgentRunLifecycleGenerationCurrent, - registerAgentRunContext, -} from "../../infra/agent-events.js"; +import { assertAgentRunLifecycleGenerationCurrent } from "../../infra/agent-events.js"; +import { registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { formatErrorMessage } from "../../infra/errors.js"; import { createSubsystemLogger } from "../../logging/subsystem.js"; import { normalizeAgentId, resolveAgentIdFromSessionKey } from "../../routing/session-key.js"; diff --git a/src/agents/command/session-preparation.ts b/src/agents/command/session-preparation.ts index aef9e4917b02..72b68e00a56a 100644 --- a/src/agents/command/session-preparation.ts +++ b/src/agents/command/session-preparation.ts @@ -1,10 +1,8 @@ import type { ThinkLevel, VerboseLevel } from "../../auto-reply/thinking.js"; import type { SessionEntry } from "../../config/sessions/types.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { - assertAgentRunLifecycleGenerationCurrent, - registerAgentRunContext, -} from "../../infra/agent-events.js"; +import { assertAgentRunLifecycleGenerationCurrent } from "../../infra/agent-events.js"; +import { registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { applyVerboseOverride } from "../../sessions/level-overrides.js"; import { recordSessionHumanDirectMessage } from "../../sessions/session-state-events.js"; import { resolveEffectiveAgentSkillFilter } from "../../skills/discovery/agent-filter.js"; diff --git a/src/agents/embedded-agent-runner/run.overflow-compaction.test.ts b/src/agents/embedded-agent-runner/run.overflow-compaction.test.ts index 64c68eea7edc..78e26292fee9 100644 --- a/src/agents/embedded-agent-runner/run.overflow-compaction.test.ts +++ b/src/agents/embedded-agent-runner/run.overflow-compaction.test.ts @@ -11,13 +11,12 @@ import { rollbackAgentHarnessSessionEntryLifecycle, } from "../../config/sessions/session-accessor.js"; import { - claimAgentRunContext, getAgentEventLifecycleGeneration, - getAgentRunContext, resetAgentEventsForTest, rotateAgentEventLifecycleGeneration, withAgentRunLifecycleGeneration, } from "../../infra/agent-events.js"; +import { claimAgentRunContext, getAgentRunContext } from "../../infra/agent-run-registry.js"; import { AGENT_HARNESS_SESSION_KEY_RESERVED_MESSAGE } from "../../sessions/agent-harness-session-key.js"; import type { AgentHarness } from "../harness/types.js"; import type { AgentInternalEvent } from "../internal-events.js"; diff --git a/src/agents/embedded-agent-runner/run/lane-controller.ts b/src/agents/embedded-agent-runner/run/lane-controller.ts index 97364de0c39e..cbf4da0d618c 100644 --- a/src/agents/embedded-agent-runner/run/lane-controller.ts +++ b/src/agents/embedded-agent-runner/run/lane-controller.ts @@ -1,10 +1,9 @@ import { assertAgentRunLifecycleGenerationCurrent, - claimAgentRunContext, getAgentEventLifecycleGeneration, - getAgentRunContext, withAgentRunLifecycleGeneration, } from "../../../infra/agent-events.js"; +import { claimAgentRunContext, getAgentRunContext } from "../../../infra/agent-run-registry.js"; import { enqueueCommandInLane, getCommandLaneSnapshot } from "../../../process/command-queue.js"; import type { CommandQueueEnqueueOptions } from "../../../process/command-queue.types.js"; import { withSessionPlacementTurnAdmission } from "../../session-placement-admission.js"; diff --git a/src/agents/embedded-agent-runner/run/session-prompt-state.ts b/src/agents/embedded-agent-runner/run/session-prompt-state.ts index 8695058f51b6..2fb31326a260 100644 --- a/src/agents/embedded-agent-runner/run/session-prompt-state.ts +++ b/src/agents/embedded-agent-runner/run/session-prompt-state.ts @@ -1,5 +1,5 @@ import type { ContextEngineSessionTarget } from "../../../context-engine/types.js"; -import { registerAgentRunContext } from "../../../infra/agent-events.js"; +import { registerAgentRunContext } from "../../../infra/agent-run-registry.js"; import { formatErrorMessage } from "../../../infra/errors.js"; import { resolveAgentRunSessionTarget } from "../../run-session-target.js"; import { log } from "../logger.js"; diff --git a/src/agents/main-session-restart-recovery-marking.ts b/src/agents/main-session-restart-recovery-marking.ts index 50188ee62355..0a0df0612efb 100644 --- a/src/agents/main-session-restart-recovery-marking.ts +++ b/src/agents/main-session-restart-recovery-marking.ts @@ -9,10 +9,8 @@ import { import { applySessionEntryReplacements } from "../config/sessions/session-accessor.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; import { resolveGatewaySessionStoreTarget } from "../gateway/session-utils.js"; -import { - getAgentEventLifecycleGeneration, - listAgentRunsForSession, -} from "../infra/agent-events.js"; +import { getAgentEventLifecycleGeneration } from "../infra/agent-events.js"; +import { listAgentRunsForSession } from "../infra/agent-run-registry.js"; import { listActiveEmbeddedRunSessionIds, listActiveEmbeddedRunSessionKeys, diff --git a/src/agents/main-session-restart-recovery.test.ts b/src/agents/main-session-restart-recovery.test.ts index 16820774fcd1..a96bfed48ce6 100644 --- a/src/agents/main-session-restart-recovery.test.ts +++ b/src/agents/main-session-restart-recovery.test.ts @@ -21,10 +21,10 @@ import { callGateway } from "../gateway/call.js"; import type { GatewayRecoveryRuntime } from "../gateway/server-instance-runtime.types.js"; import { getAgentEventLifecycleGeneration, - registerAgentRunContext, resetAgentEventsForTest, rotateAgentEventLifecycleGeneration, } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { initializeGlobalHookRunner, resetGlobalHookRunner, diff --git a/src/agents/subagent-registry-read.ts b/src/agents/subagent-registry-read.ts index 34cf7d575958..76bce59b283d 100644 --- a/src/agents/subagent-registry-read.ts +++ b/src/agents/subagent-registry-read.ts @@ -3,7 +3,7 @@ * * Combines persisted snapshots with in-memory live runs for UI, announce, control, and recovery paths. */ -import { getAgentRunContext } from "../infra/agent-events.js"; +import { getAgentRunContext } from "../infra/agent-run-registry.js"; import { getSubagentRunsForChildSession, subagentRuns } from "./subagent-registry-memory.js"; import { buildLatestSubagentRunReadIndexFromRuns, diff --git a/src/agents/subagent-registry-sweeper.ts b/src/agents/subagent-registry-sweeper.ts index 1aeda039c58a..79f5537825c6 100644 --- a/src/agents/subagent-registry-sweeper.ts +++ b/src/agents/subagent-registry-sweeper.ts @@ -1,5 +1,5 @@ import type { callGateway } from "../gateway/call.js"; -import { getAgentRunContext } from "../infra/agent-events.js"; +import { getAgentRunContext } from "../infra/agent-run-registry.js"; import { isFastTestRuntimeEnv } from "../infra/env.js"; import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js"; import { emitSessionLifecycleEvent } from "../sessions/session-lifecycle-events.js"; diff --git a/src/agents/subagent-registry.archive.e2e.test.ts b/src/agents/subagent-registry.archive.e2e.test.ts index 38f9a9c6f3b2..f05d13943a61 100644 --- a/src/agents/subagent-registry.archive.e2e.test.ts +++ b/src/agents/subagent-registry.archive.e2e.test.ts @@ -58,11 +58,13 @@ vi.mock("../tasks/task-status-access.js", () => ({ vi.mock("../infra/agent-events.js", () => ({ getAgentEventLifecycleGeneration: () => "test-generation", - getAgentRunContext: vi.fn(() => undefined), isAgentEventLifecycleGenerationCurrent: (generation: string) => generation === "test-generation", onAgentEvent: vi.fn((_handler: unknown) => noop), registerAgentEventLifecycleRotationHandler: vi.fn(), })); +vi.mock("../infra/agent-run-registry.js", () => ({ + getAgentRunContext: vi.fn(() => undefined), +})); vi.mock("../config/config.js", async () => { const actual = await vi.importActual("../config/config.js"); diff --git a/src/agents/subagent-registry.test.ts b/src/agents/subagent-registry.test.ts index 15f3597a077f..ebb09ee808b2 100644 --- a/src/agents/subagent-registry.test.ts +++ b/src/agents/subagent-registry.test.ts @@ -209,11 +209,13 @@ vi.mock("../gateway/call.js", () => ({ vi.mock("../infra/agent-events.js", () => ({ getAgentEventLifecycleGeneration: () => "test-generation", - getAgentRunContext: mocks.getAgentRunContext, isAgentEventLifecycleGenerationCurrent: (generation: string) => generation === "test-generation", onAgentEvent: mocks.onAgentEvent, registerAgentEventLifecycleRotationHandler: vi.fn(), })); +vi.mock("../infra/agent-run-registry.js", () => ({ + getAgentRunContext: mocks.getAgentRunContext, +})); vi.mock("../config/config.js", () => { return { diff --git a/src/agents/tools/cron-tool.pacing.test.ts b/src/agents/tools/cron-tool.pacing.test.ts index 43e55253d89a..2ed26de64246 100644 --- a/src/agents/tools/cron-tool.pacing.test.ts +++ b/src/agents/tools/cron-tool.pacing.test.ts @@ -3,7 +3,7 @@ import { claimAgentRunContext, clearAgentRunContext, consumeCronNextCheckProposal, -} from "../../infra/agent-events.js"; +} from "../../infra/agent-run-registry.js"; import { createCronTool } from "./cron-tool.js"; const RUN_ID = "paced-run"; diff --git a/src/agents/tools/cron-tool.ts b/src/agents/tools/cron-tool.ts index c2e2362b14fa..191d96d102d3 100644 --- a/src/agents/tools/cron-tool.ts +++ b/src/agents/tools/cron-tool.ts @@ -12,7 +12,7 @@ import { normalizeCronJobCreate, normalizeCronJobPatch } from "../../cron/normal import type { CronDelivery } from "../../cron/types.js"; import { normalizeHttpWebhookUrl } from "../../cron/webhook-url.js"; import { GatewayClientRequestError } from "../../gateway/client.js"; -import { recordCronNextCheckProposal } from "../../infra/agent-events.js"; +import { recordCronNextCheckProposal } from "../../infra/agent-run-registry.js"; import { normalizeAgentId } from "../../routing/session-key.js"; import { parseAgentSessionKey } from "../../sessions/session-key-utils.js"; import { isRecord } from "../../utils.js"; diff --git a/src/agents/tools/media-generate-background-shared.ts b/src/agents/tools/media-generate-background-shared.ts index 93d9a647825d..072f7bf03972 100644 --- a/src/agents/tools/media-generate-background-shared.ts +++ b/src/agents/tools/media-generate-background-shared.ts @@ -8,7 +8,7 @@ import { getCliSessionBinding } from "../../config/sessions/cli-session-binding. import { loadSessionEntryReadOnly } from "../../config/sessions/session-accessor.js"; import { runWithoutOwnedSessionTranscriptWrites } from "../../config/sessions/transcript-write-context.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { clearAgentRunContext, registerAgentRunContext } from "../../infra/agent-events.js"; +import { clearAgentRunContext, registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { formatErrorMessage } from "../../infra/errors.js"; import { createSubsystemLogger } from "../../logging/subsystem.js"; import { parseCronRunScopeSuffix } from "../../sessions/session-key-utils.js"; diff --git a/src/agents/tools/video-generate-background.test.ts b/src/agents/tools/video-generate-background.test.ts index d1558a0cf886..c7ed5a16ec4d 100644 --- a/src/agents/tools/video-generate-background.test.ts +++ b/src/agents/tools/video-generate-background.test.ts @@ -1,7 +1,8 @@ // Video generation background tests cover detached task lifecycle, keepalive // progress and completion delivery through the durable requester-agent handoff. import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { getAgentRunContext, resetAgentEventsForTest } from "../../infra/agent-events.js"; +import { resetAgentEventsForTest } from "../../infra/agent-events.js"; +import { getAgentRunContext } from "../../infra/agent-run-registry.js"; import { VIDEO_GENERATION_TASK_KIND } from "../video-generation-task-status.js"; import { announceDeliveryMocks, diff --git a/src/auto-reply/reply/agent-runner-execution-lifecycle.test.ts b/src/auto-reply/reply/agent-runner-execution-lifecycle.test.ts index 56cb7cdb7f10..0efc644f68bb 100644 --- a/src/auto-reply/reply/agent-runner-execution-lifecycle.test.ts +++ b/src/auto-reply/reply/agent-runner-execution-lifecycle.test.ts @@ -586,8 +586,8 @@ describe("executeAgentTurn: run lifecycle and ownership", () => { }); it("registers run ownership before asynchronous image preflight", async () => { - const agentEvents = await import("../../infra/agent-events.js"); - const registerAgentRunContext = vi.mocked(agentEvents.registerAgentRunContext); + const agentRunRegistry = await import("../../infra/agent-run-registry.js"); + const registerAgentRunContext = vi.mocked(agentRunRegistry.registerAgentRunContext); let resolveImages: (() => void) | undefined; state.resolveCurrentTurnImagesMock.mockImplementationOnce( () => @@ -617,8 +617,8 @@ describe("executeAgentTurn: run lifecycle and ownership", () => { }); it("clears run ownership when image preflight fails", async () => { - const agentEvents = await import("../../infra/agent-events.js"); - const clearAgentRunContext = vi.mocked(agentEvents.clearAgentRunContext); + const agentRunRegistry = await import("../../infra/agent-run-registry.js"); + const clearAgentRunContext = vi.mocked(agentRunRegistry.clearAgentRunContext); state.resolveCurrentTurnImagesMock.mockRejectedValueOnce(new Error("invalid image metadata")); const executeAgentTurn = await getExecuteAgentTurnForTest(); diff --git a/src/auto-reply/reply/agent-runner-execution.test-support.ts b/src/auto-reply/reply/agent-runner-execution.test-support.ts index f973ad748090..185367572a25 100644 --- a/src/auto-reply/reply/agent-runner-execution.test-support.ts +++ b/src/auto-reply/reply/agent-runner-execution.test-support.ts @@ -176,6 +176,16 @@ vi.mock("../../infra/agent-events.js", async () => { registerAgentRunContext: vi.fn(), }; }); +vi.mock("../../infra/agent-run-registry.js", async () => { + const actual = await vi.importActual( + "../../infra/agent-run-registry.js", + ); + return { + ...actual, + clearAgentRunContext: vi.fn(), + registerAgentRunContext: vi.fn(), + }; +}); vi.mock("../../runtime.js", () => ({ defaultRuntime: { diff --git a/src/auto-reply/reply/agent-runner-execution.ts b/src/auto-reply/reply/agent-runner-execution.ts index 1de85494c1b3..08d87626540c 100644 --- a/src/auto-reply/reply/agent-runner-execution.ts +++ b/src/auto-reply/reply/agent-runner-execution.ts @@ -23,10 +23,9 @@ import type { SessionEntry } from "../../config/sessions.js"; import { logVerbose } from "../../globals.js"; import { captureAgentRunLifecycleGeneration, - clearAgentRunContext, - registerAgentRunContext, withAgentRunLifecycleGeneration, } from "../../infra/agent-events.js"; +import { clearAgentRunContext, registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { emitAgentRunStatusEvent } from "../../infra/agent-run-status-events.js"; import { isDiagnosticsEnabled } from "../../infra/diagnostic-events.js"; import { formatErrorMessage } from "../../infra/errors.js"; diff --git a/src/auto-reply/reply/agent-runner-memory.ts b/src/auto-reply/reply/agent-runner-memory.ts index f7d3f4b873e8..4566946ce72e 100644 --- a/src/auto-reply/reply/agent-runner-memory.ts +++ b/src/auto-reply/reply/agent-runner-memory.ts @@ -49,7 +49,8 @@ import type { OpenClawConfig } from "../../config/types.openclaw.js"; import { readSessionMessagesAsync } from "../../gateway/session-transcript-readers.js"; import { logVerbose } from "../../globals.js"; import { isAbortError } from "../../infra/abort-signal.js"; -import { emitAgentEvent, registerAgentRunContext } from "../../infra/agent-events.js"; +import { emitAgentEvent } from "../../infra/agent-events.js"; +import { registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { formatErrorMessage } from "../../infra/errors.js"; import { resolveMemoryFlushPlan } from "../../plugins/memory-state.js"; import { CommandLane } from "../../process/lanes.js"; diff --git a/src/auto-reply/reply/agent-runner.media-paths.test.ts b/src/auto-reply/reply/agent-runner.media-paths.test.ts index 040b6787da9c..ee51e3619231 100644 --- a/src/auto-reply/reply/agent-runner.media-paths.test.ts +++ b/src/auto-reply/reply/agent-runner.media-paths.test.ts @@ -98,6 +98,15 @@ vi.mock("../../infra/agent-events.js", async () => { registerAgentRunContext: vi.fn(), }; }); +vi.mock("../../infra/agent-run-registry.js", async () => { + const actual = await vi.importActual( + "../../infra/agent-run-registry.js", + ); + return { + ...actual, + registerAgentRunContext: vi.fn(), + }; +}); vi.mock("../../agents/embedded-agent.js", () => ({ abortEmbeddedAgentRun: abortEmbeddedAgentRunMock, diff --git a/src/auto-reply/reply/followup-runner.test.ts b/src/auto-reply/reply/followup-runner.test.ts index 4caef68eefff..79f502858942 100644 --- a/src/auto-reply/reply/followup-runner.test.ts +++ b/src/auto-reply/reply/followup-runner.test.ts @@ -15,7 +15,7 @@ const state = vi.hoisted(() => ({ clearRunContext: vi.fn(), })); -vi.mock("../../infra/agent-events.js", () => ({ +vi.mock("../../infra/agent-run-registry.js", () => ({ clearAgentRunContext: (...args: unknown[]) => state.clearRunContext(...args), })); diff --git a/src/auto-reply/reply/followup-runner.ts b/src/auto-reply/reply/followup-runner.ts index 4734ec4049bc..27b0a9fd9959 100644 --- a/src/auto-reply/reply/followup-runner.ts +++ b/src/auto-reply/reply/followup-runner.ts @@ -1,6 +1,6 @@ /** Composes queued admission, canonical execution, accounting, and delivery. */ import { hasCompletedSourceReplyDeliveryEvidence } from "../../agents/embedded-agent-runner/delivery-evidence.js"; -import { clearAgentRunContext } from "../../infra/agent-events.js"; +import { clearAgentRunContext } from "../../infra/agent-run-registry.js"; import { formatErrorMessage } from "../../infra/errors.js"; import { defaultRuntime } from "../../runtime.js"; import { accountFollowupTurn } from "./agent-runner-result-accounting.js"; diff --git a/src/commands/agent.acp.test.ts b/src/commands/agent.acp.test.ts index 53cfb9a20878..a685552a44eb 100644 --- a/src/commands/agent.acp.test.ts +++ b/src/commands/agent.acp.test.ts @@ -52,6 +52,10 @@ const attemptExecutionMocks = vi.hoisted(() => ({ })); vi.mock("../infra/agent-events.js", () => agentEventMocks); +vi.mock("../infra/agent-run-registry.js", () => ({ + clearAgentRunContext: agentEventMocks.clearAgentRunContext, + registerAgentRunContext: agentEventMocks.registerAgentRunContext, +})); vi.mock("../agents/command/delivery.runtime.js", () => ({ deliverAgentCommandResult: vi.fn( diff --git a/src/config/sessions/session-accessor.sqlite-data-version.test.ts b/src/config/sessions/session-accessor.sqlite-data-version.test.ts index 640a420f9dde..b204d240d163 100644 --- a/src/config/sessions/session-accessor.sqlite-data-version.test.ts +++ b/src/config/sessions/session-accessor.sqlite-data-version.test.ts @@ -455,7 +455,7 @@ describe("SQLite session entry cache", () => { expect(listProjectionCalls).toHaveBeenCalledTimes(2); }); - it("fully reloads after a tracked same-process upsert", async () => { + it("patches only the tracked row after a same-process upsert", async () => { const scope = createSessionScope("write-through"); const siblingScope = { ...scope, sessionKey: "agent:main:write-through-sibling" }; await upsertSessionEntry(scope, { @@ -468,17 +468,53 @@ describe("SQLite session entry cache", () => { sessionId: "write-through-sibling", updatedAt: 1, }); - listSessionEntries({ ...scope, clone: false, projection: "list" }); + const before = listSessionEntries({ ...scope, clone: false, projection: "list" }); + const siblingBefore = before.find((row) => row.sessionKey === siblingScope.sessionKey)?.entry; + parseSessionEntryCalls.mockClear(); + listProjectionCalls.mockClear(); await upsertSessionEntry(scope, { label: "projection-probe-after", updatedAt: 2 }); parseSessionEntryCalls.mockClear(); listProjectionCalls.mockClear(); + const after = listSessionEntries({ ...scope, clone: false, projection: "list" }); - expect(listSessionEntries({ ...scope, clone: false, projection: "list" })[0]?.entry.label).toBe( + expect(after.find((row) => row.sessionKey === scope.sessionKey)?.entry.label).toBe( "projection-probe-after", ); - expect(parseSessionEntryCalls).toHaveBeenCalledTimes(2); - expect(listProjectionCalls).toHaveBeenCalledTimes(2); + expect(after.find((row) => row.sessionKey === siblingScope.sessionKey)?.entry).toBe( + siblingBefore, + ); + expect(parseSessionEntryCalls).not.toHaveBeenCalled(); + expect(listProjectionCalls).toHaveBeenCalledOnce(); + }); + + it("adds a tracked upsert to a warm snapshot without reparsing siblings", async () => { + const scope = createSessionScope("write-through-insert"); + await upsertSessionEntry(scope, { + label: "projection-probe-existing", + sessionId: "write-through-existing", + updatedAt: 1, + }); + const existing = listSessionEntries({ ...scope, clone: false, projection: "list" })[0]?.entry; + const insertedScope = { ...scope, sessionKey: "agent:main:write-through-inserted" }; + + parseSessionEntryCalls.mockClear(); + listProjectionCalls.mockClear(); + await upsertSessionEntry(insertedScope, { + label: "projection-probe-inserted", + sessionId: "write-through-inserted", + updatedAt: 2, + }); + parseSessionEntryCalls.mockClear(); + listProjectionCalls.mockClear(); + const after = listSessionEntries({ ...scope, clone: false, projection: "list" }); + + expect(after.map((row) => row.sessionKey)).toEqual( + [scope.sessionKey, insertedScope.sessionKey].toSorted(), + ); + expect(after.find((row) => row.sessionKey === scope.sessionKey)?.entry).toBe(existing); + expect(parseSessionEntryCalls).not.toHaveBeenCalled(); + expect(listProjectionCalls).toHaveBeenCalledOnce(); }); it("does not let a tracked write mask an earlier raw connection write", async () => { @@ -508,7 +544,7 @@ describe("SQLite session entry cache", () => { label: "tracked-after", sessionId: "tracked", }); - expect(parseSessionEntryCalls).toHaveBeenCalledTimes(2); + expect(parseSessionEntryCalls).toHaveBeenCalledOnce(); }); it("invalidates cached keys when transcript creation inserts a placeholder node", async () => { diff --git a/src/config/sessions/session-accessor.sqlite-entry-cache.ts b/src/config/sessions/session-accessor.sqlite-entry-cache.ts index 13ee489c6398..95471347688b 100644 --- a/src/config/sessions/session-accessor.sqlite-entry-cache.ts +++ b/src/config/sessions/session-accessor.sqlite-entry-cache.ts @@ -37,7 +37,9 @@ const MAX_INCREMENTAL_ENTRY_READ_KEYS = 500; // One parsed snapshot per opened agent database bounds memory to the process's database set. // Weak connection ownership lets closed read-only and evicted database handles release their // snapshots. The connection-local validity token plus tracked-write invalidation keeps live -// snapshots current; without both, every read would re-query and re-parse every entry_json document. +// snapshots current; narrow tracked upserts patch one authoritative row after commit, while +// structural/unknown writes invalidate. Without both, every read would re-query and re-parse +// every entry_json document. const sessionEntryCaches = new WeakMap(); function readDataVersion(database: DatabaseSync): number { @@ -203,7 +205,8 @@ export function readSqliteSessionEntryCache( if (cached && cached.validityToken.dataVersion === validityToken.dataVersion) { // updated_at is entry-controlled, not a rowversion. Other connections can rewrite entry_json // without advancing it, so data_version changes must fully reload or same-ms rewrites go stale. - // Accessor-owned writes on this connection hard-invalidate; unrelated local writes can safely diff. + // Tracked single-row upserts patch their row but retain this old token; unrelated local writes + // and any other same-connection changes are still discovered by this incremental diff. const revalidated = incrementallyRevalidateSessionEntrySnapshot( database, cached, @@ -242,6 +245,74 @@ function invalidateTrackedCache(database: OpenClawAgentDatabase): void { invalidate(); } -export function publishSqliteSessionEntryCacheInvalidation(database: OpenClawAgentDatabase): void { +function publishTrackedCacheUpdate(database: OpenClawAgentDatabase, publish: () => void): void { + if (deferOpenClawAgentPostCommitPublication(database, publish)) { + return; + } + if (database.db.isTransaction) { + throw new Error( + "SQLite session entry writes must use runOpenClawAgentWriteTransaction for cache publication", + ); + } + publish(); +} + +function publishSqliteSessionEntryCacheUpsert( + database: OpenClawAgentDatabase, + row: { + current_session_id: string; + entry_json: string; + session_key: string; + updated_at: number; + }, +): void { + const entry = parseSqliteSessionEntryJson({ + current_session_id: row.current_session_id, + entry_json: row.entry_json, + updated_at: row.updated_at, + }); + if (!entry) { + invalidateTrackedCache(database); + return; + } + publishTrackedCacheUpdate(database, () => { + const cached = sessionEntryCaches.get(database.db); + if (!cached) { + return; + } + const entries = new Map(cached.entries); + entries.set(row.session_key, entry); + const listProjections = new Map(cached.listProjections); + listProjections.delete(row.session_key); + const updatedAtByKey = new Map(cached.updatedAtByKey); + const knownKey = updatedAtByKey.has(row.session_key); + updatedAtByKey.set(row.session_key, row.updated_at); + // Patch only the authoritative row but retain the old validity token. The next read + // must still reconcile any other local total_changes, while a changed data_version + // forces a full reload; advancing either here could mask an earlier unknown write. + sessionEntryCaches.set(database.db, { + entries, + keys: knownKey ? cached.keys : [...cached.keys, row.session_key].toSorted(), + listEntries: createLazyListProjections(entries, listProjections), + listProjections, + updatedAtByKey, + validityToken: cached.validityToken, + }); + }); +} + +export function publishSqliteSessionEntryCacheInvalidation( + database: OpenClawAgentDatabase, + row?: { + current_session_id: string; + entry_json: string; + session_key: string; + updated_at: number; + }, +): void { + if (row) { + publishSqliteSessionEntryCacheUpsert(database, row); + return; + } invalidateTrackedCache(database); } diff --git a/src/config/sessions/session-accessor.sqlite-entry-store.ts b/src/config/sessions/session-accessor.sqlite-entry-store.ts index 40e0527f7ed3..90abc9c1358c 100644 --- a/src/config/sessions/session-accessor.sqlite-entry-store.ts +++ b/src/config/sessions/session-accessor.sqlite-entry-store.ts @@ -740,7 +740,7 @@ export function writeSessionEntry( updatedAt, }); } - publishSqliteSessionEntryCacheInvalidation(database); + publishSqliteSessionEntryCacheInvalidation(database, sessionNode); } /** Resolves the parent fork decision using SQLite transcript rows when totals are stale. */ diff --git a/src/config/sessions/session-accessor.sqlite-title-probes.ts b/src/config/sessions/session-accessor.sqlite-title-probes.ts new file mode 100644 index 000000000000..e66f5abb3b42 --- /dev/null +++ b/src/config/sessions/session-accessor.sqlite-title-probes.ts @@ -0,0 +1,208 @@ +import { sql } from "kysely"; +import { executeSqliteQuerySync, getNodeSqliteKysely } from "../../infra/kysely-sync.js"; +import { runSqliteDeferredTransactionSync } from "../../infra/sqlite-transaction.js"; +import type { DB as OpenClawAgentKyselyDatabase } from "../../state/openclaw-agent-db.generated.js"; +import { + openOpenClawAgentDatabase, + type OpenClawAgentDatabase, +} from "../../state/openclaw-agent-db.js"; +import type { + SessionTranscriptReadScope, + TranscriptEvent, +} from "./session-accessor.sqlite-contract.js"; +import { + resolveSqliteTranscriptReadScope, + toDatabaseOptions, +} from "./session-accessor.sqlite-scope.js"; + +type TitleProbeDatabase = Pick< + OpenClawAgentKyselyDatabase, + | "session_transcript_active_events" + | "session_transcript_index_state" + | "session_windows" + | "transcript_events" + | "transcript_rewrite_watermarks" +>; + +export type SessionTranscriptTitleProbe = { + generation: string | null; + head: Array<{ event: TranscriptEvent; seq: number }>; + maxSeq: number | null; + tail: Array<{ event: TranscriptEvent; seq: number }>; + totalMessages: number; +}; + +const SESSION_TITLE_PROBE_MESSAGES = 20; +const SESSION_TITLE_PROBE_QUERY_CHUNK_SIZE = 400; + +function getTitleProbeKysely(database: OpenClawAgentDatabase) { + return getNodeSqliteKysely(database.db); +} + +function parseEventType(eventJson: string | null): string | undefined { + if (!eventJson) { + return undefined; + } + try { + const event = JSON.parse(eventJson) as { type?: unknown }; + return typeof event.type === "string" ? event.type : undefined; + } catch { + return undefined; + } +} + +function sqliteTranscriptBoundaryEventType() { + return /* kysely-allow-raw: boundary type lives inside canonical transcript JSON. */ sql`json_extract(boundary_event.event_json, '$.type')`; +} + +function readTitleProbeChunk( + database: OpenClawAgentDatabase, + sessionIds: readonly string[], +): Map { + const db = getTitleProbeKysely(database); + const rows = runSqliteDeferredTransactionSync( + database.db, + () => + executeSqliteQuerySync( + database.db, + db + .selectFrom("session_windows as window") + .leftJoin( + "session_transcript_index_state as state", + "state.session_id", + "window.session_id", + ) + .leftJoin( + "transcript_rewrite_watermarks as rewrite", + "rewrite.session_id", + "window.session_id", + ) + .leftJoin("session_transcript_active_events as active", (join) => + join + .onRef("active.session_id", "=", "window.session_id") + .on("active.message_position", "is not", null) + .on((eb) => + eb.or([ + eb("active.message_position", "<", SESSION_TITLE_PROBE_MESSAGES), + eb( + "active.message_position", + ">=", + eb("state.active_message_count", "-", SESSION_TITLE_PROBE_MESSAGES), + ), + ]), + ), + ) + .leftJoin("transcript_events as event", (join) => + join + .onRef("event.session_id", "=", "active.session_id") + .onRef("event.seq", "=", "active.event_seq"), + ) + .select((eb) => [ + "window.session_id", + "state.active_message_count", + "state.indexed_seq", + "state.needs_rebuild", + "rewrite.generation", + "active.message_position", + "event.event_json", + eb + .selectFrom("transcript_events as latest") + .select("latest.seq") + .whereRef("latest.session_id", "=", "window.session_id") + .orderBy("latest.seq", "desc") + .limit(1) + .as("latest_seq"), + eb + .selectFrom("session_transcript_active_events as boundary") + .innerJoin("transcript_events as boundary_event", (join) => + join + .onRef("boundary_event.session_id", "=", "boundary.session_id") + .onRef("boundary_event.seq", "=", "boundary.event_seq"), + ) + .select("boundary_event.event_json") + .whereRef("boundary.session_id", "=", "window.session_id") + .where("boundary.message_position", "is", null) + // excluding latest-reset sessions prevents pre-reset text from leaking. + .where(sqliteTranscriptBoundaryEventType(), "in", ["reset", "compaction"]) + .orderBy("boundary.active_position", "desc") + .limit(1) + .as("latest_boundary_json"), + ]) + .where("window.session_id", "in", sessionIds) + .orderBy("window.session_id", "asc") + .orderBy("active.message_position", "asc"), + ).rows, + { databaseLabel: database.path, operationLabel: "sessions.list.title-probes" }, + ); + const probes = new Map(); + for (const row of rows) { + const emptyTranscript = row.latest_seq === null; + const projectionCurrent = row.needs_rebuild === 0 && row.indexed_seq === row.latest_seq; + if ( + (!emptyTranscript && !projectionCurrent) || + parseEventType(row.latest_boundary_json) === "reset" + ) { + continue; + } + const totalMessages = row.active_message_count ?? 0; + const probe = probes.get(row.session_id) ?? { + generation: row.generation ?? null, + head: [], + maxSeq: row.latest_seq ?? null, + tail: [], + totalMessages, + }; + if (row.event_json !== null && row.message_position !== null) { + const event = { + event: JSON.parse(row.event_json) as TranscriptEvent, + seq: row.message_position + 1, + }; + if (row.message_position < SESSION_TITLE_PROBE_MESSAGES) { + probe.head.push(event); + } + if (row.message_position >= totalMessages - SESSION_TITLE_PROBE_MESSAGES) { + probe.tail.push(event); + } + } + probes.set(row.session_id, probe); + } + return probes; +} + +/** Reads bounded title probes in one statement per opened store (chunked for SQLite limits). */ +export function readSessionTranscriptTitleProbeBatch( + scopes: readonly SessionTranscriptReadScope[], +): Array { + const results: Array = Array.from({ + length: scopes.length, + }); + const groups = new Map< + string, + { database: OpenClawAgentDatabase; items: Array<{ index: number; sessionId: string }> } + >(); + for (const [index, scope] of scopes.entries()) { + const resolved = resolveSqliteTranscriptReadScope(scope); + const database = openOpenClawAgentDatabase(toDatabaseOptions(resolved)); + const group = groups.get(database.path) ?? { database, items: [] }; + group.items.push({ index, sessionId: resolved.sessionId }); + groups.set(database.path, group); + } + for (const group of groups.values()) { + const sessionIds = [...new Set(group.items.map((item) => item.sessionId))]; + const probes = new Map(); + for ( + let offset = 0; + offset < sessionIds.length; + offset += SESSION_TITLE_PROBE_QUERY_CHUNK_SIZE + ) { + const chunk = sessionIds.slice(offset, offset + SESSION_TITLE_PROBE_QUERY_CHUNK_SIZE); + for (const [sessionId, probe] of readTitleProbeChunk(group.database, chunk)) { + probes.set(sessionId, probe); + } + } + for (const item of group.items) { + results[item.index] = probes.get(item.sessionId); + } + } + return results; +} diff --git a/src/config/sessions/session-accessor.ts b/src/config/sessions/session-accessor.ts index 987c61b05848..77689cdf11d8 100644 --- a/src/config/sessions/session-accessor.ts +++ b/src/config/sessions/session-accessor.ts @@ -233,6 +233,10 @@ export { SessionTranscriptProjectionUnavailableError, waitForSessionTranscriptProjection, } from "./session-accessor.sqlite-active-events.js"; +export { + readSessionTranscriptTitleProbeBatch, + type SessionTranscriptTitleProbe, +} from "./session-accessor.sqlite-title-probes.js"; export type { SessionTranscriptBoundedMessageTailPage, SessionTranscriptMessageAnchorPage, diff --git a/src/cron/isolated-agent/run-execution.runtime.ts b/src/cron/isolated-agent/run-execution.runtime.ts index 6197b679a1de..bf3c69caf985 100644 --- a/src/cron/isolated-agent/run-execution.runtime.ts +++ b/src/cron/isolated-agent/run-execution.runtime.ts @@ -15,7 +15,7 @@ export { } from "../../agents/embedded-agent-runner/result-fallback-classifier.js"; export { isCliProvider } from "../../agents/model-selection-cli.js"; export { normalizeVerboseLevel } from "../../auto-reply/thinking.shared.js"; -export { registerAgentRunContext } from "../../infra/agent-events.js"; +export { registerAgentRunContext } from "../../infra/agent-run-registry.js"; export { logWarn } from "../../logger.js"; import { createLazyImportLoader } from "../../shared/lazy-promise.js"; diff --git a/src/cron/isolated-agent/run.message-tool-policy.test.ts b/src/cron/isolated-agent/run.message-tool-policy.test.ts index e1a1b027b088..54b268d1efc8 100644 --- a/src/cron/isolated-agent/run.message-tool-policy.test.ts +++ b/src/cron/isolated-agent/run.message-tool-policy.test.ts @@ -960,7 +960,7 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { }); resolveCronSessionMock.mockReturnValue(cronSession); const { getAgentRunContext, registerAgentRunContext } = - await import("../../infra/agent-events.js"); + await import("../../infra/agent-run-registry.js"); registerAgentRunContext("test-session-id", { sessionKey: "agent:default:cron:message-tool-policy", verboseLevel: "off", @@ -974,12 +974,9 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { it("does not let old cron cleanup clear a newer same-id run context", async () => { mockRunCronFallbackPassthrough(); - const { - claimAgentRunContext, - clearAgentRunContext, - getAgentRunContext, - rotateAgentEventLifecycleGeneration, - } = await import("../../infra/agent-events.js"); + const { claimAgentRunContext, clearAgentRunContext, getAgentRunContext } = + await import("../../infra/agent-run-registry.js"); + const { rotateAgentEventLifecycleGeneration } = await import("../../infra/agent-events.js"); let newerLifecycleGeneration = ""; runEmbeddedAgentMock.mockImplementationOnce( async (runParams: { @@ -1020,8 +1017,8 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { return { status: "available" }; }); }); - const { getAgentRunContext, rotateAgentEventLifecycleGeneration } = - await import("../../infra/agent-events.js"); + const { getAgentRunContext } = await import("../../infra/agent-run-registry.js"); + const { rotateAgentEventLifecycleGeneration } = await import("../../infra/agent-events.js"); const runPromise = runCronIsolatedAgentTurn(makeParams()); await preflightStarted; @@ -1047,7 +1044,7 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { }); resolveCronSessionMock.mockReturnValue(cronSession); const { clearAgentRunContext, getAgentRunContext, registerAgentRunContext } = - await import("../../infra/agent-events.js"); + await import("../../infra/agent-run-registry.js"); registerAgentRunContext("test-session-id", { sessionKey: "agent:default:cron:message-tool-policy", verboseLevel: "off", @@ -1070,7 +1067,7 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { it("releases a shared cron run context created by this invocation", async () => { mockRunCronFallbackPassthrough(); runEmbeddedAgentMock.mockRejectedValueOnce(new Error("runner failed")); - const { getAgentRunContext } = await import("../../infra/agent-events.js"); + const { getAgentRunContext } = await import("../../infra/agent-run-registry.js"); const currentSessionJob = makeMessageToolPolicyJob() as unknown as Record; currentSessionJob.sessionTarget = "current"; @@ -1092,8 +1089,9 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { process.env.OPENCLAW_TEST_FAST = "1"; mockRunCronFallbackPassthrough(); resolveCronSessionMock.mockImplementation(() => makeCronSession()); - const { claimAgentRunContext, getAgentEventLifecycleGeneration, getAgentRunContext } = - await import("../../infra/agent-events.js"); + const { claimAgentRunContext, getAgentRunContext } = + await import("../../infra/agent-run-registry.js"); + const { getAgentEventLifecycleGeneration } = await import("../../infra/agent-events.js"); let invocationCount = 0; let releaseFirst = () => {}; let releaseSecond = () => {}; @@ -1154,12 +1152,10 @@ describe("runCronIsolatedAgentTurn message tool policy", () => { it("releases a stale shared cron context replaced by this invocation", async () => { mockRunCronFallbackPassthrough(); runEmbeddedAgentMock.mockRejectedValueOnce(new Error("runner failed")); - const { - claimAgentRunContext, - getAgentEventLifecycleGeneration, - getAgentRunContext, - rotateAgentEventLifecycleGeneration, - } = await import("../../infra/agent-events.js"); + const { claimAgentRunContext, getAgentRunContext } = + await import("../../infra/agent-run-registry.js"); + const { getAgentEventLifecycleGeneration, rotateAgentEventLifecycleGeneration } = + await import("../../infra/agent-events.js"); claimAgentRunContext("test-session-id", { sessionKey: "agent:default:cron:message-tool-policy", sessionId: "test-session-id", diff --git a/src/cron/isolated-agent/run.ts b/src/cron/isolated-agent/run.ts index deb8c8c83e8b..b8a8fe13f6c2 100644 --- a/src/cron/isolated-agent/run.ts +++ b/src/cron/isolated-agent/run.ts @@ -5,13 +5,15 @@ import type { CliDeps } from "../../cli/outbound-send-deps.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; import { assertAgentRunLifecycleGenerationCurrent, - claimAgentRunContext, - consumeCronNextCheckProposal, getAgentEventLifecycleGeneration, - getAgentRunContext, - releaseAgentRunContext, withAgentRunLifecycleGeneration, } from "../../infra/agent-events.js"; +import { + claimAgentRunContext, + consumeCronNextCheckProposal, + getAgentRunContext, + releaseAgentRunContext, +} from "../../infra/agent-run-registry.js"; import { isDiagnosticsEnabled } from "../../infra/diagnostic-events.js"; import { isFastTestRuntimeEnv } from "../../infra/env.js"; import { createDiagnosticMessageLifecycle } from "../../logging/message-lifecycle.js"; diff --git a/src/gateway/chat-abort.test.ts b/src/gateway/chat-abort.test.ts index 9137e1037871..8f6d2431b193 100644 --- a/src/gateway/chat-abort.test.ts +++ b/src/gateway/chat-abort.test.ts @@ -2,11 +2,8 @@ // abort fanout, history snapshots, and cleanup of buffered streaming state. import { afterEach, describe, expect, it, vi } from "vitest"; import { isAgentRunRestartAbortReason } from "../agents/run-termination.js"; -import { - clearAgentRunContext, - onAgentEvent, - registerAgentRunContext, -} from "../infra/agent-events.js"; +import { onAgentEvent } from "../infra/agent-events.js"; +import { clearAgentRunContext, registerAgentRunContext } from "../infra/agent-run-registry.js"; import { abortChatRunById, abortChatRunsForProvider, diff --git a/src/gateway/server-chat.agent-events.test.ts b/src/gateway/server-chat.agent-events.test.ts index b73cb449c619..5cb27b16d4d6 100644 --- a/src/gateway/server-chat.agent-events.test.ts +++ b/src/gateway/server-chat.agent-events.test.ts @@ -9,14 +9,16 @@ import { } from "../agents/internal-runtime-context.js"; import { formatChannelProgressDraftLine } from "../channels/streaming.js"; import { - claimAgentRunContext, emitAgentEvent as emitRuntimeAgentEvent, emitAgentEventForOwner, onAgentRuntimeEvent, - registerAgentRunContext, - releaseAgentRunContext, resetAgentEventsForTest, } from "../infra/agent-events.js"; +import { + claimAgentRunContext, + registerAgentRunContext, + releaseAgentRunContext, +} from "../infra/agent-run-registry.js"; import { subscribePluginSessionsChanged } from "../plugins/gateway-events.js"; const persistGatewaySessionLifecycleEventMock = vi.fn(); diff --git a/src/gateway/server-chat.ts b/src/gateway/server-chat.ts index 6feef24f4583..9d3166ed8e7e 100644 --- a/src/gateway/server-chat.ts +++ b/src/gateway/server-chat.ts @@ -20,9 +20,8 @@ import { type AgentEventPayload, type AgentEventRuntimePayload, getAgentEventLifecycleGeneration, - getAgentRunContext, - getAgentRunContextOwnerStatus, } from "../infra/agent-events.js"; +import { getAgentRunContext, getAgentRunContextOwnerStatus } from "../infra/agent-run-registry.js"; import { formatErrorMessage } from "../infra/errors.js"; import { resolveHeartbeatVisibility } from "../infra/heartbeat-visibility.js"; import { logError } from "../logger.js"; diff --git a/src/gateway/server-maintenance.ts b/src/gateway/server-maintenance.ts index d36eaf04a839..d4bc828e5470 100644 --- a/src/gateway/server-maintenance.ts +++ b/src/gateway/server-maintenance.ts @@ -8,7 +8,7 @@ import { WORKTREE_GC_INTERVAL_MS, } from "../agents/worktrees/service.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; -import { sweepStaleRunContexts } from "../infra/agent-events.js"; +import { sweepStaleRunContexts } from "../infra/agent-run-registry.js"; import { pruneOrphanedDeliveryQueueMedia } from "../infra/outbound/delivery-queue-media-spool.js"; import { cleanOldMedia } from "../media/store.js"; import { startSkillCuratorMaintenance } from "../skills/workshop/curator.js"; diff --git a/src/gateway/server-methods/agent-run-admission-phase.ts b/src/gateway/server-methods/agent-run-admission-phase.ts index 1ce268b5d70e..9c6ada0332a1 100644 --- a/src/gateway/server-methods/agent-run-admission-phase.ts +++ b/src/gateway/server-methods/agent-run-admission-phase.ts @@ -19,7 +19,7 @@ import { resolveEffectiveAgentRuntime } from "../../agents/thinking-runtime.js"; import { resolveAgentTimeoutMs } from "../../agents/timeout.js"; import type { SessionEntry } from "../../config/sessions.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { claimAgentRunContext } from "../../infra/agent-events.js"; +import { claimAgentRunContext } from "../../infra/agent-run-registry.js"; import type { InputProvenance } from "../../sessions/input-provenance.js"; import type { SessionWorkAdmissionLease } from "../../sessions/session-lifecycle-admission.js"; import { registerChatAbortController, resolveAgentRunExpiresAtMs } from "../chat-abort.js"; diff --git a/src/gateway/server-methods/agent-run-dispatch.ts b/src/gateway/server-methods/agent-run-dispatch.ts index d11437130564..ba749de9e0a0 100644 --- a/src/gateway/server-methods/agent-run-dispatch.ts +++ b/src/gateway/server-methods/agent-run-dispatch.ts @@ -8,7 +8,7 @@ import { isAgentRunRestartAbortReason } from "../../agents/run-termination.js"; import { normalizeAgentRunTimeoutPhase } from "../../agents/run-timeout-attribution.js"; import { agentCommandFromGatewayIngress } from "../../commands/agent.js"; import { isAbortError } from "../../infra/abort-signal.js"; -import { clearAgentRunContext } from "../../infra/agent-events.js"; +import { clearAgentRunContext } from "../../infra/agent-run-registry.js"; import { readErrorName } from "../../infra/errors.js"; import { defaultRuntime } from "../../runtime.js"; import { createRunningTaskRun } from "../../tasks/detached-task-runtime.js"; diff --git a/src/gateway/server-methods/agent.sessions-and-models.test-utils.ts b/src/gateway/server-methods/agent.sessions-and-models.test-utils.ts index fd5967a11db7..8a06c9323a83 100644 --- a/src/gateway/server-methods/agent.sessions-and-models.test-utils.ts +++ b/src/gateway/server-methods/agent.sessions-and-models.test-utils.ts @@ -1378,6 +1378,7 @@ describe("gateway agent handler", () => { storePath: "/tmp/sessions.json", }, ); + vi.advanceTimersByTime(100); expect(broadcastToConnIds.mock.calls.map((callValue) => callValue[1]?.reason)).toEqual([ "create", "send", @@ -1728,6 +1729,7 @@ describe("gateway agent handler", () => { storePath: "/tmp/sessions.json", }, ); + vi.advanceTimersByTime(100); expect(broadcastToConnIds.mock.calls.map((callLocal) => callLocal[1]?.reason)).toEqual([ "create", "send", diff --git a/src/gateway/server-methods/agent.test-harness.ts b/src/gateway/server-methods/agent.test-harness.ts index 91c905bc19a6..f1d23a569b8c 100644 --- a/src/gateway/server-methods/agent.test-harness.ts +++ b/src/gateway/server-methods/agent.test-harness.ts @@ -216,6 +216,13 @@ vi.mock("../../infra/agent-events.js", () => ({ registerAgentRunContext: mocks.registerAgentRunContext, onAgentEvent: vi.fn(), })); +vi.mock("../../infra/agent-run-registry.js", () => ({ + claimAgentRunContext: mocks.registerAgentRunContext, + clearAgentRunContext: mocks.clearAgentRunContext, + getAgentRunContext: vi.fn(() => undefined), + hasProjectedAgentRunForSession: vi.fn(() => false), + registerAgentRunContext: mocks.registerAgentRunContext, +})); vi.mock("../../agents/subagent-registry-read.js", () => ({ getLatestSubagentRunByChildSessionKey: mocks.getLatestSubagentRunByChildSessionKey, diff --git a/src/gateway/server-methods/chat-send-admission.ts b/src/gateway/server-methods/chat-send-admission.ts index 5903601a974f..79be87d828be 100644 --- a/src/gateway/server-methods/chat-send-admission.ts +++ b/src/gateway/server-methods/chat-send-admission.ts @@ -9,11 +9,8 @@ import { readSessionTranscriptActiveLeafEvents, resolveSessionTranscriptActiveLeafEntryId, } from "../../config/sessions/session-accessor.js"; -import { - claimAgentRunContext, - clearAgentRunContext, - getAgentEventLifecycleGeneration, -} from "../../infra/agent-events.js"; +import { getAgentEventLifecycleGeneration } from "../../infra/agent-events.js"; +import { claimAgentRunContext, clearAgentRunContext } from "../../infra/agent-run-registry.js"; import { beginSessionWorkAdmission } from "../../sessions/session-lifecycle-admission.js"; import { registerChatAbortController, resolveChatRunExpiresAtMs } from "../chat-abort.js"; import { PENDING_CHAT_SEND_DEDUPE_PREFIX, type DedupeEntry } from "../server-shared.js"; diff --git a/src/gateway/server-methods/chat-send-attachments.ts b/src/gateway/server-methods/chat-send-attachments.ts index ae3317e5d001..d2fbcc26875c 100644 --- a/src/gateway/server-methods/chat-send-attachments.ts +++ b/src/gateway/server-methods/chat-send-attachments.ts @@ -9,7 +9,7 @@ import { } from "../../auto-reply/reply/stage-sandbox-media.js"; import type { MsgContext, TemplateContext } from "../../auto-reply/templating.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { clearAgentRunContext } from "../../infra/agent-events.js"; +import { clearAgentRunContext } from "../../infra/agent-run-registry.js"; import { measureDiagnosticsTimelineSpan } from "../../infra/diagnostics-timeline.js"; import { formatErrorMessage } from "../../infra/errors.js"; import { parseInboundMediaUri } from "../../media/media-reference.js"; diff --git a/src/gateway/server-methods/chat-send-dispatch-errors.ts b/src/gateway/server-methods/chat-send-dispatch-errors.ts index 0f859e3aa735..b6b81675aebf 100644 --- a/src/gateway/server-methods/chat-send-dispatch-errors.ts +++ b/src/gateway/server-methods/chat-send-dispatch-errors.ts @@ -1,6 +1,6 @@ import { ErrorCodes, errorShape } from "../../../packages/gateway-protocol/src/index.js"; import { resolveDefaultAgentId } from "../../agents/agent-scope.js"; -import { clearAgentRunContext } from "../../infra/agent-events.js"; +import { clearAgentRunContext } from "../../infra/agent-run-registry.js"; import { retainGatewayRootWorkAdmissionContinuation } from "../../process/gateway-work-admission.js"; import type { UserTurnTranscriptRecorder } from "../../sessions/user-turn-transcript.js"; import { chatAbortMarkerTimestampMs } from "../server-chat-state.js"; diff --git a/src/gateway/server-methods/chat-send-handler.ts b/src/gateway/server-methods/chat-send-handler.ts index b1b8931fa8f2..85402987295c 100644 --- a/src/gateway/server-methods/chat-send-handler.ts +++ b/src/gateway/server-methods/chat-send-handler.ts @@ -9,10 +9,8 @@ import { resolveDefaultAgentId } from "../../agents/agent-scope.js"; import { resolveProviderIdForAuth } from "../../agents/provider-auth-aliases.js"; import { createAgentRunRestartAbortError } from "../../agents/run-termination.js"; import { dispatchInboundMessageWithProjectedDispatcher } from "../../auto-reply/dispatch.js"; -import { - clearAgentRunContext, - getAgentEventLifecycleGeneration, -} from "../../infra/agent-events.js"; +import { getAgentEventLifecycleGeneration } from "../../infra/agent-events.js"; +import { clearAgentRunContext } from "../../infra/agent-run-registry.js"; import { emitDiagnosticsTimelineEvent, measureDiagnosticsTimelineSpan, diff --git a/src/gateway/server-methods/chat.directive-tags.test.ts b/src/gateway/server-methods/chat.directive-tags.test.ts index 645480506654..8ec659ccc541 100644 --- a/src/gateway/server-methods/chat.directive-tags.test.ts +++ b/src/gateway/server-methods/chat.directive-tags.test.ts @@ -29,7 +29,7 @@ import { } from "../../config/sessions/session-accessor.js"; import { resolveMirroredTranscriptText } from "../../config/sessions/transcript-mirror.js"; import { withOwnedSessionTranscriptWrites } from "../../config/sessions/transcript-write-context.js"; -import { getAgentRunContext } from "../../infra/agent-events.js"; +import { getAgentRunContext } from "../../infra/agent-run-registry.js"; import { runExclusiveSessionLifecycleMutation } from "../../sessions/session-lifecycle-admission.js"; import { createDeferred } from "../../test-utils/deferred.js"; import { withEnvAsync } from "../../test-utils/env.js"; diff --git a/src/gateway/server-methods/session-active-runs.test.ts b/src/gateway/server-methods/session-active-runs.test.ts index d1f55771ab26..7c217f952fbf 100644 --- a/src/gateway/server-methods/session-active-runs.test.ts +++ b/src/gateway/server-methods/session-active-runs.test.ts @@ -12,7 +12,7 @@ import { buildProjectedAgentRunIndex, clearAgentRunContext, registerAgentRunContext, -} from "../../infra/agent-events.js"; +} from "../../infra/agent-run-registry.js"; import { collectTrackedActiveSessionRuns, hasTrackedActiveSessionRun, diff --git a/src/gateway/server-methods/session-active-runs.ts b/src/gateway/server-methods/session-active-runs.ts index cfb5907feacc..f0ac9dd9b5a9 100644 --- a/src/gateway/server-methods/session-active-runs.ts +++ b/src/gateway/server-methods/session-active-runs.ts @@ -2,7 +2,7 @@ import { isEmbeddedAgentRunInProgress } from "../../agents/embedded-agent-runner import { hasProjectedAgentRunForSession, type ProjectedAgentRunIndex, -} from "../../infra/agent-events.js"; +} from "../../infra/agent-run-registry.js"; import { normalizeAgentId } from "../../routing/session-key.js"; import type { GatewayRequestContext } from "./types.js"; diff --git a/src/gateway/server-methods/session-change-event.test.ts b/src/gateway/server-methods/session-change-event.test.ts new file mode 100644 index 000000000000..e01ba94d6af3 --- /dev/null +++ b/src/gateway/server-methods/session-change-event.test.ts @@ -0,0 +1,125 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { GatewayRequestContext } from "./types.js"; + +const mocks = vi.hoisted(() => ({ + invalidate: vi.fn(), + loadRow: vi.fn(), + rowLabel: "first", +})); + +vi.mock("../session-sharing.js", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, invalidateSessionSharingSnapshot: mocks.invalidate }; +}); + +vi.mock("../session-utils.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + loadGatewaySessionRow: mocks.loadRow.mockImplementation((key: string) => ({ + key, + label: mocks.rowLabel, + sessionId: `${key}-id`, + })), + }; +}); + +vi.mock("../session-event-payload.js", () => ({ + buildGatewaySessionEventFields: ({ + sessionRow, + }: { + sessionRow: { key: string; label: string }; + }) => ({ key: sessionRow.key, label: sessionRow.label }), +})); + +vi.mock("./session-active-runs.js", () => ({ + resolveVisibleActiveSessionRunState: () => ({ active: false, runIds: [] }), +})); + +const { emitSessionsChanged, flushPendingSessionsChangedEvents, readSessionsMutationVersion } = + await import("./session-change-event.js"); + +function createContext(receivers = new Set(["conn-1"])) { + return { + broadcastToConnIds: vi.fn(), + chatAbortControllers: new Map(), + getRuntimeConfig: () => ({}), + getSessionEventSubscriberConnIds: () => receivers, + } as unknown as GatewayRequestContext; +} + +beforeEach(() => { + vi.useFakeTimers(); + mocks.invalidate.mockClear(); + mocks.loadRow.mockClear(); + mocks.rowLabel = "first"; +}); + +afterEach(() => { + flushPendingSessionsChangedEvents(); + vi.useRealTimers(); +}); + +describe("sessions.changed coalescing", () => { + it("emits a leading row and one trailing row with the latest state", () => { + const context = createContext(); + const initialVersion = readSessionsMutationVersion(context); + + emitSessionsChanged(context, { reason: "create", sessionKey: "agent:main:chat" }); + mocks.rowLabel = "latest"; + emitSessionsChanged(context, { reason: "update", sessionKey: "agent:main:chat" }); + emitSessionsChanged(context, { reason: "send", sessionKey: "agent:main:chat" }); + + expect(context.broadcastToConnIds).toHaveBeenCalledOnce(); + expect(mocks.loadRow).toHaveBeenCalledOnce(); + vi.advanceTimersByTime(100); + + expect(context.broadcastToConnIds).toHaveBeenCalledTimes(2); + expect(mocks.loadRow).toHaveBeenCalledTimes(2); + expect(vi.mocked(context.broadcastToConnIds).mock.calls[1]?.[1]).toMatchObject({ + label: "latest", + reason: "send", + }); + expect(readSessionsMutationVersion(context)).toBe(initialVersion + 3); + expect(mocks.invalidate).toHaveBeenCalledTimes(3); + }); + + it("keeps different session keys independent", () => { + const context = createContext(); + + emitSessionsChanged(context, { reason: "update", sessionKey: "agent:main:first" }); + emitSessionsChanged(context, { reason: "update", sessionKey: "agent:main:second" }); + + expect(context.broadcastToConnIds).toHaveBeenCalledTimes(2); + expect(mocks.loadRow).toHaveBeenCalledTimes(2); + }); + + it("advances the mutation fence without loading rows when nobody receives events", () => { + const context = createContext(new Set()); + const initialVersion = readSessionsMutationVersion(context); + + emitSessionsChanged(context, { reason: "update", sessionKey: "agent:main:chat" }); + + expect(readSessionsMutationVersion(context)).toBe(initialVersion + 1); + expect(mocks.invalidate).toHaveBeenCalledOnce(); + expect(mocks.loadRow).not.toHaveBeenCalled(); + expect(context.broadcastToConnIds).not.toHaveBeenCalled(); + }); + + it("flushes the latest trailing row and clears its shutdown timer", () => { + const context = createContext(); + emitSessionsChanged(context, { reason: "create", sessionKey: "agent:main:chat" }); + mocks.rowLabel = "shutdown-latest"; + emitSessionsChanged(context, { reason: "send", sessionKey: "agent:main:chat" }); + + flushPendingSessionsChangedEvents(context); + expect(context.broadcastToConnIds).toHaveBeenCalledTimes(2); + expect(vi.mocked(context.broadcastToConnIds).mock.calls[1]?.[1]).toMatchObject({ + label: "shutdown-latest", + reason: "send", + }); + + vi.advanceTimersByTime(100); + expect(context.broadcastToConnIds).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/gateway/server-methods/session-change-event.ts b/src/gateway/server-methods/session-change-event.ts index 8cd73bf2adad..805cdc6879a9 100644 --- a/src/gateway/server-methods/session-change-event.ts +++ b/src/gateway/server-methods/session-change-event.ts @@ -14,17 +14,39 @@ type SessionChangedPayload = { compacted?: boolean; }; -export function emitSessionsChanged( - context: Pick< - GatewayRequestContext, - | "broadcastToConnIds" - | "chatAbortControllers" - | "getRuntimeConfig" - | "getSessionEventSubscriberConnIds" - >, +type SessionChangeContext = Pick< + GatewayRequestContext, + | "broadcastToConnIds" + | "chatAbortControllers" + | "getRuntimeConfig" + | "getSessionEventSubscriberConnIds" +>; + +type PendingSessionChange = { + context: SessionChangeContext; + dirty: boolean; + key: string; + payload: SessionChangedPayload; + timer: ReturnType | null; +}; + +const SESSIONS_CHANGED_DEBOUNCE_MS = 100; +const sessionsMutationVersions = new WeakMap(); +const pendingChangesByContext = new WeakMap>(); +const pendingSessionChanges = new Set(); + +export function readSessionsMutationVersion(context: object): number { + return sessionsMutationVersions.get(context) ?? 0; +} + +function sessionChangeKey(payload: SessionChangedPayload): string { + return `${payload.agentId ?? ""}\0${payload.sessionKey ?? ""}`; +} + +function broadcastSessionsChanged( + context: SessionChangeContext, payload: SessionChangedPayload, -) { - invalidateSessionSharingSnapshot(payload.sessionKey); +): void { const connIds = context.getSessionEventSubscriberConnIds(); if (!hasSessionChangeReceivers(connIds)) { return; @@ -79,3 +101,71 @@ export function emitSessionsChanged( }, ); } + +function finishPendingSessionChange(pending: PendingSessionChange): void { + if (pending.timer) { + clearTimeout(pending.timer); + pending.timer = null; + } + pendingSessionChanges.delete(pending); + const byKey = pendingChangesByContext.get(pending.context); + if (byKey?.get(pending.key) === pending) { + byKey.delete(pending.key); + } + if (pending.dirty) { + broadcastSessionsChanged(pending.context, pending.payload); + } +} + +/** Flush trailing notifications and release every debounce timer before gateway shutdown. */ +export function flushPendingSessionsChangedEvents(context?: object): void { + for (const pending of pendingSessionChanges) { + if (!context || pending.context === context) { + finishPendingSessionChange(pending); + } + } +} + +export function emitSessionsChanged(context: SessionChangeContext, payload: SessionChangedPayload) { + // This counter is the sessions.list projection fence: every mutation advances it + // synchronously, before event coalescing, so work started on an older value is never + // joined or cached by a request that begins after the mutation. + sessionsMutationVersions.set(context, readSessionsMutationVersion(context) + 1); + invalidateSessionSharingSnapshot(payload.sessionKey); + const connIds = context.getSessionEventSubscriberConnIds(); + if (!hasSessionChangeReceivers(connIds)) { + return; + } + const key = sessionChangeKey(payload); + const byKey = pendingChangesByContext.get(context) ?? new Map(); + pendingChangesByContext.set(context, byKey); + const pending = byKey.get(key); + if (pending) { + pending.payload = payload; + pending.dirty = true; + if (pending.timer) { + clearTimeout(pending.timer); + } + pending.timer = setTimeout( + () => finishPendingSessionChange(pending), + SESSIONS_CHANGED_DEBOUNCE_MS, + ); + pending.timer.unref?.(); + return; + } + + // Lead after a quiet period for responsive UI, then coalesce a burst into one trailing + // rebuild. The trailing row is loaded only when emitted, so it reflects the newest state. + const next: PendingSessionChange = { + context, + dirty: false, + key, + payload, + timer: null, + }; + next.timer = setTimeout(() => finishPendingSessionChange(next), SESSIONS_CHANGED_DEBOUNCE_MS); + next.timer.unref?.(); + byKey.set(key, next); + pendingSessionChanges.add(next); + broadcastSessionsChanged(context, payload); +} diff --git a/src/gateway/server-methods/sessions-list-cache.ts b/src/gateway/server-methods/sessions-list-cache.ts new file mode 100644 index 000000000000..efd108239a67 --- /dev/null +++ b/src/gateway/server-methods/sessions-list-cache.ts @@ -0,0 +1,124 @@ +import type { SessionsListParams } from "../../../packages/gateway-protocol/src/index.js"; +import type { OpenClawConfig } from "../../config/types.openclaw.js"; +import { readAgentRunIndexVersion } from "../../infra/agent-run-registry.js"; +import { isGatewayAdmin } from "../session-sharing.js"; +import { gatewayClientSessionCreator } from "./gateway-client-identity.js"; +import { readSessionsMutationVersion } from "./session-change-event.js"; +import type { GatewayClient, GatewayRequestContext, RespondFn } from "./types.js"; + +type SessionListFence = { + agentRunIndexVersion: number; + sessionsMutationVersion: number; +}; +type SessionListOperation = SessionListFence & { promise: Promise }; +type SessionListCompleted = SessionListFence & { result: unknown }; +type SessionListState = { + completed: Map; + config: OpenClawConfig; + inFlight: Map; +}; + +const SESSIONS_LIST_COMPLETED_CACHE_LIMIT = 64; +const sessionListsByContext = new WeakMap(); + +function readSessionListFence(context: GatewayRequestContext): SessionListFence { + return { + agentRunIndexVersion: readAgentRunIndexVersion(), + sessionsMutationVersion: readSessionsMutationVersion(context), + }; +} + +function matchesSessionListFence(value: SessionListFence, fence: SessionListFence): boolean { + return ( + value.agentRunIndexVersion === fence.agentRunIndexVersion && + value.sessionsMutationVersion === fence.sessionsMutationVersion + ); +} + +function sessionListVisibilityIdentity(client: GatewayClient | null): string { + if (isGatewayAdmin(client)) { + return "admin"; + } + const profileId = gatewayClientSessionCreator(client)?.id; + return profileId ? `profile:${profileId}` : "anonymous"; +} + +function sessionListWorkKey(params: SessionsListParams, client: GatewayClient | null): string { + return JSON.stringify([ + sessionListVisibilityIdentity(client), + Object.entries(params).toSorted(([left], [right]) => left.localeCompare(right)), + ]); +} + +function sessionListState( + context: GatewayRequestContext, + config: OpenClawConfig, +): SessionListState { + let state = sessionListsByContext.get(context); + if (!state || state.config !== config) { + state = { completed: new Map(), config, inFlight: new Map() }; + sessionListsByContext.set(context, state); + } + return state; +} + +function rememberCompletedSessionList( + state: SessionListState, + workKey: string, + completed: SessionListCompleted, +): void { + state.completed.delete(workKey); + state.completed.set(workKey, completed); + while (state.completed.size > SESSIONS_LIST_COMPLETED_CACHE_LIMIT) { + const oldest = state.completed.keys().next().value; + if (oldest === undefined) { + break; + } + state.completed.delete(oldest); + } +} + +export async function respondWithCachedSessionList(params: { + client: GatewayClient | null; + config: OpenClawConfig; + context: GatewayRequestContext; + request: SessionsListParams; + respond: RespondFn; + run: () => Promise; +}): Promise { + const workKey = sessionListWorkKey(params.request, params.client); + const state = sessionListState(params.context, params.config); + // Every input that can change a projected row must fence reuse. Store mutations and + // live-run transitions have separate owners, so their monotonic counters stay separate. + const fence = readSessionListFence(params.context); + const completed = state.completed.get(workKey); + if (completed && matchesSessionListFence(completed, fence)) { + params.respond(true, completed.result, undefined); + return; + } + const pending = state.inFlight.get(workKey); + if (pending && matchesSessionListFence(pending, fence)) { + params.respond(true, await pending.promise, undefined); + return; + } + + // A request may share only work begun at the same fence. A transition during projection + // leaves current callers intact but fences every later caller and cache write. + const promise = Promise.resolve() + .then(params.run) + .then((result) => { + if (matchesSessionListFence(readSessionListFence(params.context), fence)) { + rememberCompletedSessionList(state, workKey, { ...fence, result }); + } + return result; + }); + const operation = { ...fence, promise }; + state.inFlight.set(workKey, operation); + try { + params.respond(true, await promise, undefined); + } finally { + if (state.inFlight.get(workKey) === operation) { + state.inFlight.delete(workKey); + } + } +} diff --git a/src/gateway/server-methods/sessions-read-cache.test.ts b/src/gateway/server-methods/sessions-read-cache.test.ts index 4e1fa280204b..1803ee89065b 100644 --- a/src/gateway/server-methods/sessions-read-cache.test.ts +++ b/src/gateway/server-methods/sessions-read-cache.test.ts @@ -1,11 +1,18 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { SessionsListParams } from "../../../packages/gateway-protocol/src/index.js"; import { upsertSessionEntry } from "../../config/sessions/session-accessor.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; +import { resetAgentEventsForTest } from "../../infra/agent-events.js"; +import { clearAgentRunContext, registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { withOpenClawTestState } from "../../test-utils/openclaw-test-state.js"; import type { GatewayClient, GatewayRequestContext, RespondFn } from "./types.js"; -const loader = vi.hoisted(() => ({ calls: vi.fn(), failNext: false })); +const loader = vi.hoisted(() => ({ + calls: vi.fn(), + failNext: false, + rowCalls: vi.fn(), + rowGate: undefined as Promise | undefined, +})); vi.mock("../session-utils.js", async (importOriginal) => { const actual = await importOriginal(); @@ -21,10 +28,18 @@ vi.mock("../session-utils.js", async (importOriginal) => { } return actual.loadCombinedSessionStoreForGateway(...args); }, + listSessionsFromStoreAsync: async ( + ...args: Parameters + ) => { + loader.rowCalls(...args); + await loader.rowGate; + return await actual.listSessionsFromStoreAsync(...args); + }, }; }); const { sessionReadHandlers } = await import("./sessions-read.js"); +const { emitSessionsChanged } = await import("./session-change-event.js"); function identifiedClient(profileId: string): GatewayClient { return { @@ -48,6 +63,7 @@ function requestContext(config: OpenClawConfig): GatewayRequestContext { return { chatAbortControllers: new Map(), getRuntimeConfig: () => config, + getSessionEventSubscriberConnIds: () => new Set(), loadGatewayModelCatalog: async () => [], logGateway: { debug: vi.fn() }, } as unknown as GatewayRequestContext; @@ -67,7 +83,11 @@ async function listSessions(params: { } as never); expect(responses).toHaveLength(1); expect(responses[0]?.[0]).toBe(true); - return responses[0]?.[1] as { sessions: Array<{ key: string }> }; + return responses[0]?.[1] as { + count: number; + nextOffset: number | null; + sessions: Array<{ hasActiveRun?: boolean; key: string }>; + }; } async function seedSessions(): Promise { @@ -114,10 +134,17 @@ async function seedSessions(): Promise { return config; } +beforeEach(() => { + resetAgentEventsForTest(); +}); + afterEach(() => { + resetAgentEventsForTest(); vi.restoreAllMocks(); loader.calls.mockClear(); loader.failNext = false; + loader.rowCalls.mockClear(); + loader.rowGate = undefined; }); describe("sessions.list single-flight", () => { @@ -164,6 +191,61 @@ describe("sessions.list single-flight", () => { }); }); + it("reuses a completed result until the session mutation version advances", async () => { + await withOpenClawTestState({ scenario: "minimal" }, async () => { + const config = await seedSessions(); + const context = requestContext(config); + const client = identifiedClient("owner@example.com"); + const request = { archived: "all" as const, limit: 100 }; + + const first = await listSessions({ client, context, request }); + const cached = await listSessions({ client, context, request }); + expect(cached).toBe(first); + expect(loader.calls).toHaveBeenCalledTimes(1); + + emitSessionsChanged(context, { reason: "test", sessionKey: "agent:main:active" }); + await listSessions({ client, context, request }); + expect(loader.calls).toHaveBeenCalledTimes(2); + }); + }); + + it("rebuilds a completed result when a projected run ends without a store mutation", async () => { + await withOpenClawTestState({ scenario: "minimal" }, async () => { + const config = await seedSessions(); + const context = requestContext(config); + const client = identifiedClient("owner@example.com"); + const request = { agentId: "main", archived: "all" as const, limit: 100 }; + const runId = "sessions-list-cache-active-run"; + registerAgentRunContext(runId, { + agentId: "main", + projectSessionActive: true, + sessionId: "main-active", + sessionKey: "agent:main:active", + }); + + const active = await listSessions({ client, context, request }); + expect(active.sessions.find((session) => session.key === "agent:main:active")).toMatchObject({ + hasActiveRun: true, + }); + const activeCached = await listSessions({ client, context, request }); + expect(activeCached).toBe(active); + expect(loader.calls).toHaveBeenCalledTimes(1); + + clearAgentRunContext(runId); + const settled = await listSessions({ client, context, request }); + expect(settled.sessions.find((session) => session.key === "agent:main:active")).toMatchObject( + { + hasActiveRun: false, + }, + ); + expect(loader.calls).toHaveBeenCalledTimes(2); + + const settledCached = await listSessions({ client, context, request }); + expect(settledCached).toBe(settled); + expect(loader.calls).toHaveBeenCalledTimes(2); + }); + }); + it("does not share filtered results across client identities", async () => { await withOpenClawTestState({ scenario: "minimal" }, async () => { const config = await seedSessions(); @@ -188,6 +270,63 @@ describe("sessions.list single-flight", () => { }); }); + it("refills a page from the loaded store when a selected row becomes hidden", async () => { + await withOpenClawTestState({ scenario: "minimal" }, async () => { + const config = await seedSessions(); + for (const [name, updatedAt] of [ + ["third", 500], + ["second", 600], + ["first", 700], + ] as const) { + await upsertSessionEntry( + { agentId: "main", sessionKey: `agent:main:page-${name}` }, + { + sessionId: `page-${name}`, + updatedAt, + createdActor: { type: "human", id: "owner@example.com" }, + visibility: "shared", + }, + ); + } + const context = requestContext(config); + const client = identifiedClient("viewer@example.com"); + let releaseRows!: () => void; + loader.rowGate = new Promise((resolve) => { + releaseRows = resolve; + }); + + const firstPage = listSessions({ + client, + context, + request: { agentId: "main", archived: "all", limit: 1 }, + }); + await vi.waitFor(() => expect(loader.rowCalls).toHaveBeenCalledOnce()); + await upsertSessionEntry( + { agentId: "main", sessionKey: "agent:main:page-first" }, + { visibility: "draft", updatedAt: 800 }, + ); + emitSessionsChanged(context, { + reason: "sharing", + sessionKey: "agent:main:page-first", + }); + releaseRows(); + + const repaired = await firstPage; + expect(repaired.sessions.map((session) => session.key)).toEqual(["agent:main:page-second"]); + expect(repaired).toMatchObject({ count: 1, nextOffset: 1 }); + expect(loader.calls).toHaveBeenCalledTimes(1); + expect(loader.rowCalls).toHaveBeenCalledTimes(2); + + loader.rowGate = undefined; + const next = await listSessions({ + client, + context, + request: { agentId: "main", archived: "all", limit: 1, offset: 1 }, + }); + expect(next.sessions.map((session) => session.key)).toEqual(["agent:main:page-third"]); + }); + }); + it("rejects followers and retries after an underlying store failure", async () => { await withOpenClawTestState({ scenario: "minimal" }, async () => { const config = await seedSessions(); @@ -210,29 +349,32 @@ describe("sessions.list single-flight", () => { it("does not share work that started before an intervening session mutation", async () => { await withOpenClawTestState({ scenario: "minimal" }, async () => { const config = await seedSessions(); - let releaseCatalog!: () => void; - const catalog = new Promise<[]>((resolve) => { - releaseCatalog = () => resolve([]); + let releaseRows!: () => void; + loader.rowGate = new Promise((resolve) => { + releaseRows = resolve; }); - const loadGatewayModelCatalog = vi.fn(async () => await catalog); - const context = { - ...requestContext(config), - loadGatewayModelCatalog, - } as GatewayRequestContext; + const context = requestContext(config); const client = identifiedClient("owner@example.com"); const request = { archived: "all" as const, limit: 100 }; const beforeMutation = listSessions({ client, context, request }); - await vi.waitFor(() => expect(loadGatewayModelCatalog).toHaveBeenCalledTimes(1)); + await vi.waitFor(() => expect(loader.rowCalls).toHaveBeenCalledTimes(1)); await upsertSessionEntry( { agentId: "main", sessionKey: "agent:main:created-mid-list" }, { sessionId: "created-mid-list", updatedAt: 500, visibility: "shared" }, ); + emitSessionsChanged(context, { + reason: "test", + sessionKey: "agent:main:created-mid-list", + }); const afterMutation = listSessions({ client, context, request }); - await vi.waitFor(() => expect(loadGatewayModelCatalog).toHaveBeenCalledTimes(2)); - releaseCatalog(); + await vi.waitFor(() => expect(loader.rowCalls).toHaveBeenCalledTimes(2)); + releaseRows(); - const [, fresh] = await Promise.all([beforeMutation, afterMutation]); + const [stale, fresh] = await Promise.all([beforeMutation, afterMutation]); + expect(stale.sessions.map((session) => session.key)).not.toContain( + "agent:main:created-mid-list", + ); expect(fresh.sessions.map((session) => session.key)).toContain("agent:main:created-mid-list"); expect(loader.calls).toHaveBeenCalledTimes(2); }); diff --git a/src/gateway/server-methods/sessions-read.ts b/src/gateway/server-methods/sessions-read.ts index e25f4343e8c5..49fc60b086c4 100644 --- a/src/gateway/server-methods/sessions-read.ts +++ b/src/gateway/server-methods/sessions-read.ts @@ -24,8 +24,7 @@ import { } from "../../config/sessions.js"; import { listSessionEntriesReadOnly } from "../../config/sessions/session-accessor.js"; import { searchSessionTranscripts } from "../../config/sessions/session-transcript-search.js"; -import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { buildProjectedAgentRunIndex } from "../../infra/agent-events.js"; +import { buildProjectedAgentRunIndex } from "../../infra/agent-run-registry.js"; import { measureDiagnosticsTimelineSpan, measureDiagnosticsTimelineSpanSync, @@ -77,46 +76,15 @@ import { resolveVisibleActiveSessionRunState, } from "./session-active-runs.js"; import { emitSessionsChanged } from "./session-change-event.js"; +import { respondWithCachedSessionList } from "./sessions-list-cache.js"; import { filterSessionStoreToConfiguredAgents, loadSessionEntriesForTarget, requireSessionKey, } from "./sessions-shared.js"; -import type { GatewayClient, GatewayRequestContext, GatewayRequestHandlers } from "./types.js"; +import type { GatewayRequestHandlers } from "./types.js"; import { assertValidParams } from "./validation.js"; -const sessionListsByContext = new WeakMap< - GatewayRequestContext, - { config: OpenClawConfig; inFlight: Map> } ->(); - -function sessionListVisibilityIdentity(client: GatewayClient | null): string { - if (isGatewayAdmin(client)) { - return "admin"; - } - const profileId = gatewayClientSessionCreator(client)?.id; - return profileId ? `profile:${profileId}` : "anonymous"; -} - -function sessionListWorkKey(params: SessionsListParams, client: GatewayClient | null): string { - return JSON.stringify([ - sessionListVisibilityIdentity(client), - Object.entries(params).toSorted(([left], [right]) => left.localeCompare(right)), - ]); -} - -function sessionListInflightMap( - context: GatewayRequestContext, - config: OpenClawConfig, -): Map> { - let state = sessionListsByContext.get(context); - if (!state || state.config !== config) { - state = { config, inFlight: new Map() }; - sessionListsByContext.set(context, state); - } - return state.inFlight; -} - export const sessionReadHandlers: GatewayRequestHandlers = { "sessions.search": async ({ params, respond, context, client }) => { if (!assertValidParams(params, validateSessionsSearchParams, "sessions.search", respond)) { @@ -266,52 +234,72 @@ export const sessionReadHandlers: GatewayRequestHandlers = { const p = params as SessionsListParams; const cfg = context.getRuntimeConfig(); const configuredAgentsOnly = p.configuredAgentsOnly === true; - const workKey = sessionListWorkKey(p, client); - const inFlight = sessionListInflightMap(context, cfg); - const pending = inFlight.get(workKey); - if (pending) { - respond(true, await pending, undefined); - return; - } const run = () => measureDiagnosticsTimelineSpan( "gateway.sessions.list", async function listVisibleSessions( - remainingVisibilityRetries = 1, + options: { + allowFullReload?: boolean; + excludedKeys?: ReadonlySet; + loaded?: { + durableStorePath?: string; + listStore: Record; + modelCatalog: Awaited>; + storePath: string; + }; + rowRepairAttempted?: boolean; + } = {}, ): Promise>> { - const modelCatalog = await measureDiagnosticsTimelineSpan( - "gateway.sessions.list.model_catalog", - () => - loadOptionalServerMethodModelCatalog( - context, - "sessions.list", - p.agentId ? { loadParams: { agentId: p.agentId } } : undefined, - ), - { - config: cfg, - phase: "sessions.list", - }, - ); - const { durableStorePath, storePath, store } = measureDiagnosticsTimelineSpanSync( - "gateway.sessions.list.store_load", - () => - loadCombinedSessionStoreForGateway(cfg, { - agentId: p.agentId, - projection: "list", - }), - { - config: cfg, - phase: "sessions.list", - attributes: { - agentId: p.agentId ?? null, - configuredAgentsOnly, + let loaded = options.loaded; + if (!loaded) { + const modelCatalog = await measureDiagnosticsTimelineSpan( + "gateway.sessions.list.model_catalog", + () => + loadOptionalServerMethodModelCatalog( + context, + "sessions.list", + p.agentId ? { loadParams: { agentId: p.agentId } } : undefined, + ), + { + config: cfg, + phase: "sessions.list", }, - }, - ); - const entryFilter = createSessionListEntryFilter({ client }); - const listStore = configuredAgentsOnly - ? filterSessionStoreToConfiguredAgents(cfg, store) - : store; + ); + const { durableStorePath, storePath, store } = measureDiagnosticsTimelineSpanSync( + "gateway.sessions.list.store_load", + () => + loadCombinedSessionStoreForGateway(cfg, { + agentId: p.agentId, + projection: "list", + }), + { + config: cfg, + phase: "sessions.list", + attributes: { + agentId: p.agentId ?? null, + configuredAgentsOnly, + }, + }, + ); + loaded = { + durableStorePath, + listStore: configuredAgentsOnly + ? filterSessionStoreToConfiguredAgents(cfg, store) + : store, + modelCatalog, + storePath, + }; + } + if (!loaded) { + throw new Error("sessions.list store input was not loaded"); + } + const { durableStorePath, listStore, modelCatalog, storePath } = loaded; + const visibilityFilter = createSessionListEntryFilter({ client }); + const entryFilter = + visibilityFilter || options.excludedKeys?.size + ? (key: string, entry: SessionEntry) => + !options.excludedKeys?.has(key) && (visibilityFilter?.(key, entry) ?? true) + : undefined; const result = await measureDiagnosticsTimelineSpan( "gateway.sessions.list.rows", () => @@ -477,12 +465,29 @@ export const sessionReadHandlers: GatewayRequestHandlers = { (session.visibility !== "draft" || session.sharingRole === "owner"), ); if (visibleSessions.length !== sessions.length) { - if (remainingVisibilityRetries === 0) { - throw new Error("session visibility changed during list reconciliation"); + const visibleKeys = new Set(visibleSessions.map((session) => session.key)); + const excludedKeys = new Set(options.excludedKeys); + for (const session of sessions) { + if (!visibleKeys.has(session.key)) { + excludedKeys.add(session.key); + } } - // Rebuild the complete canonical page so totals, offsets, creator - // facets, and replacement rows describe the same visible snapshot. - return await listVisibleSessions(remainingVisibilityRetries - 1); + if (!options.rowRepairAttempted) { + // Excluding only freshly rejected rows refills this page from the already-loaded + // store, preserving cursor continuity without multiplying catalog/store work. + return await listVisibleSessions({ + ...options, + excludedKeys, + loaded, + rowRepairAttempted: true, + }); + } + if (options.allowFullReload !== false) { + // A second visibility drift means the loaded snapshot cannot restore a coherent + // page. One full reload is the last resort; repeated drift below fails closed. + return await listVisibleSessions({ allowFullReload: false }); + } + return { ...result, count: visibleSessions.length, sessions: visibleSessions }; } return { ...result, @@ -498,23 +503,7 @@ export const sessionReadHandlers: GatewayRequestHandlers = { }, }, ); - // The delayed computation is the shared promise, so every failure reaches all followers. - const operation = new Promise((done) => { - setImmediate(done); - }).then(() => { - // Only the pre-start socket burst may share. Once loading begins, an intervening session - // mutation must make the next request build a fresh projection instead of joining this one. - inFlight.delete(workKey); - return run(); - }); - inFlight.set(workKey, operation); - try { - respond(true, await operation, undefined); - } finally { - if (inFlight.get(workKey) === operation) { - inFlight.delete(workKey); - } - } + await respondWithCachedSessionList({ client, config: cfg, context, request: p, respond, run }); }, "sessions.cleanup": async ({ params, respond, context }) => { if (!assertValidParams(params, validateSessionsCleanupParams, "sessions.cleanup", respond)) { diff --git a/src/gateway/server-runtime-subscriptions.ts b/src/gateway/server-runtime-subscriptions.ts index 21eb8c7a228d..614844d3eeb7 100644 --- a/src/gateway/server-runtime-subscriptions.ts +++ b/src/gateway/server-runtime-subscriptions.ts @@ -4,11 +4,8 @@ import { isAuditLedgerEnabled, resolveAuditMessageMode } from "../audit/audit-co import { createAuditEventRecorder } from "../audit/audit-recorder.js"; import { onTrustedMessageAuditEvent } from "../audit/message-audit-events.js"; import { getRuntimeConfig } from "../config/io.js"; -import { - clearAgentRunContext, - onAgentAuditEvent, - onAgentRuntimeEvent, -} from "../infra/agent-events.js"; +import { onAgentAuditEvent, onAgentRuntimeEvent } from "../infra/agent-events.js"; +import { clearAgentRunContext } from "../infra/agent-run-registry.js"; import { onTrustedToolExecutionEvent } from "../infra/diagnostic-events.js"; import { onHeartbeatEvent } from "../infra/heartbeat-events.js"; import type { SubsystemLogger } from "../logging/subsystem.js"; diff --git a/src/gateway/server-session-key.test.ts b/src/gateway/server-session-key.test.ts index e39119785a7c..69d0f00b1968 100644 --- a/src/gateway/server-session-key.test.ts +++ b/src/gateway/server-session-key.test.ts @@ -3,7 +3,8 @@ */ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../config/config.js"; -import { registerAgentRunContext, resetAgentEventsForTest } from "../infra/agent-events.js"; +import { resetAgentEventsForTest } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; const hoisted = vi.hoisted(() => ({ loadConfigMock: vi.fn<() => OpenClawConfig>(), diff --git a/src/gateway/server-session-key.ts b/src/gateway/server-session-key.ts index 0ce7cf2f4128..f481a6815a67 100644 --- a/src/gateway/server-session-key.ts +++ b/src/gateway/server-session-key.ts @@ -8,7 +8,7 @@ import { resolveDefaultAgentId } from "../agents/agent-scope.js"; import { getRuntimeConfig } from "../config/io.js"; import type { SessionEntry } from "../config/sessions.js"; import type { OpenClawConfig } from "../config/types.js"; -import { getAgentRunContext } from "../infra/agent-events.js"; +import { getAgentRunContext } from "../infra/agent-run-registry.js"; import { normalizeAgentId, parseAgentSessionKey, diff --git a/src/gateway/server-startup-finish.ts b/src/gateway/server-startup-finish.ts index 73eab196d0df..a1d01fa986f7 100644 --- a/src/gateway/server-startup-finish.ts +++ b/src/gateway/server-startup-finish.ts @@ -285,6 +285,14 @@ export async function finishGatewayStartup(params: { broadcastVoiceWakeRoutingChanged, }); }); + const sessionChangeSidecar = { + stop: async () => { + const { flushPendingSessionsChangedEvents } = + await import("./server-methods/session-change-event.js"); + flushPendingSessionsChangedEvents(gatewayRequestContext); + }, + }; + runtimeState.gatewayLifetimeSidecars.push(sessionChangeSidecar); pluginGatewayContext.current = gatewayRequestContext; const { createGatewayInstanceRuntime } = await import("./server-instance-runtime.js"); const gatewayInstanceRuntimeLocal = createGatewayInstanceRuntime({ @@ -469,9 +477,10 @@ export async function finishGatewayStartup(params: { } }, onGatewayLifetimeSidecars: (gatewayLifetimeSidecars) => { - runtimeState.gatewayLifetimeSidecars = gatewayLifetimeSidecars; + const lifetimeSidecars = [sessionChangeSidecar, ...gatewayLifetimeSidecars]; + runtimeState.gatewayLifetimeSidecars = lifetimeSidecars; stopPostReadySidecarsAfterCloseStarted({ - postReadySidecars: gatewayLifetimeSidecars, + postReadySidecars: lifetimeSidecars, closeStarted: lifecycle.closePreludeStarted, }); if (lifecycle.closePreludeStarted) { diff --git a/src/gateway/server.agent.gateway-server-agent-b.test.ts b/src/gateway/server.agent.gateway-server-agent-b.test.ts index 342c9bfacd68..a056ac3376e2 100644 --- a/src/gateway/server.agent.gateway-server-agent-b.test.ts +++ b/src/gateway/server.agent.gateway-server-agent-b.test.ts @@ -7,7 +7,8 @@ import { cleanupTempDirs, makeTempDir } from "../../test/helpers/temp-dir.js"; import { AcpRuntimeError } from "../acp/runtime/errors.js"; import type { ChannelPlugin } from "../channels/plugins/types.public.js"; import { loadSessionEntry } from "../config/sessions/session-accessor.js"; -import { emitAgentEvent, registerAgentRunContext } from "../infra/agent-events.js"; +import { emitAgentEvent } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { createChannelTestPluginBase, createDirectOutboundTestAdapter, diff --git a/src/gateway/server.chat.gateway-server-chat.test.ts b/src/gateway/server.chat.gateway-server-chat.test.ts index 152e2d723ee5..6bd95d739e55 100644 --- a/src/gateway/server.chat.gateway-server-chat.test.ts +++ b/src/gateway/server.chat.gateway-server-chat.test.ts @@ -7,7 +7,8 @@ import { beforeEach, describe, expect, test, vi } from "vitest"; import { WebSocket } from "ws"; import { loadSessionEntry } from "../config/sessions/session-accessor.js"; import { replaceSqliteTranscriptEvents } from "../config/sessions/session-accessor.sqlite.js"; -import { emitAgentEvent, registerAgentRunContext } from "../infra/agent-events.js"; +import { emitAgentEvent } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { getActiveGatewayRootWorkCount, isGatewaySubordinateWorkAdmissionClosed, diff --git a/src/gateway/server.node-chat-subscriptions.test.ts b/src/gateway/server.node-chat-subscriptions.test.ts index bd28ffdcbba2..b89d15d13e31 100644 --- a/src/gateway/server.node-chat-subscriptions.test.ts +++ b/src/gateway/server.node-chat-subscriptions.test.ts @@ -1,6 +1,7 @@ // Real gateway WebSocket coverage for canonical node chat subscriptions and reconnects. import { describe, expect, test, vi } from "vitest"; -import { emitAgentEvent, registerAgentRunContext } from "../infra/agent-events.js"; +import { emitAgentEvent } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { approveNodePairing, requestNodePairing } from "../infra/node-pairing.js"; import { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES } from "../utils/message-channel.js"; import { pairDeviceIdentity } from "./device-authz.test-helpers.js"; diff --git a/src/gateway/server.sessions.list-store-materialization.test.ts b/src/gateway/server.sessions.list-store-materialization.test.ts index b358bf964bf2..88fb7112d902 100644 --- a/src/gateway/server.sessions.list-store-materialization.test.ts +++ b/src/gateway/server.sessions.list-store-materialization.test.ts @@ -114,6 +114,7 @@ test("startup prewarm fills session snapshot and title caches before the first l sessionKey, storePath, }); + const titleBatchSpy = vi.spyOn(sessionAccessor, "readSessionTranscriptTitleProbeBatch"); const titlePageSpy = vi.spyOn(sessionAccessor, "readSessionTranscriptMessageEventPage"); let sidecar: ReturnType | undefined; vi.useFakeTimers(); @@ -144,7 +145,9 @@ test("startup prewarm fills session snapshot and title caches before the first l await vi.advanceTimersToNextTimerAsync(); await sessionPrewarm; sidecar.stop(); - expect(titlePageSpy).toHaveBeenCalled(); + expect(titleBatchSpy).toHaveBeenCalled(); + expect(titlePageSpy).not.toHaveBeenCalled(); + titleBatchSpy.mockClear(); titlePageSpy.mockClear(); vi.useRealTimers(); const cachedEntries = sessionAccessor.listSessionEntriesReadOnly({ @@ -160,6 +163,7 @@ test("startup prewarm fills session snapshot and title caches before the first l }); expect(result.ok).toBe(true); + expect(titleBatchSpy).not.toHaveBeenCalled(); expect(titlePageSpy).not.toHaveBeenCalled(); const afterListEntries = sessionAccessor.listSessionEntriesReadOnly({ agentId: "main", @@ -171,6 +175,7 @@ test("startup prewarm fills session snapshot and title caches before the first l } finally { sidecar?.stop(); vi.useRealTimers(); + titleBatchSpy.mockRestore(); titlePageSpy.mockRestore(); } }); diff --git a/src/gateway/session-message-events.test.ts b/src/gateway/session-message-events.test.ts index 98b6686dec9e..3631eaf688ee 100644 --- a/src/gateway/session-message-events.test.ts +++ b/src/gateway/session-message-events.test.ts @@ -22,11 +22,8 @@ import { } from "../config/sessions/session-accessor.js"; import { appendAssistantMessageToSessionTranscript } from "../config/sessions/transcript.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; -import { - claimAgentRunContext, - clearAgentRunContext, - emitAgentEvent, -} from "../infra/agent-events.js"; +import { emitAgentEvent } from "../infra/agent-events.js"; +import { claimAgentRunContext, clearAgentRunContext } from "../infra/agent-run-registry.js"; import { rawDataToString } from "../infra/ws.js"; import { emitSessionLifecycleEvent } from "../sessions/session-lifecycle-events.js"; import * as transcriptEvents from "../sessions/transcript-events.js"; diff --git a/src/gateway/session-observer.ts b/src/gateway/session-observer.ts index 5097b77f47c0..80eddd08fcc0 100644 --- a/src/gateway/session-observer.ts +++ b/src/gateway/session-observer.ts @@ -8,7 +8,7 @@ import { terminalHealthFor, } from "../agents/session-activity-notes.js"; import { resolveUtilityModelRefForAgent } from "../agents/utility-model.js"; -import { getAgentRunContext } from "../infra/agent-events.js"; +import { getAgentRunContext } from "../infra/agent-run-registry.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; import { createSessionObserverAudience } from "./session-observer-audience.js"; import { createSessionObserverCompletion } from "./session-observer-completion.js"; diff --git a/src/gateway/session-sharing-snapshot-cache.test.ts b/src/gateway/session-sharing-snapshot-cache.test.ts new file mode 100644 index 000000000000..10a9883146ba --- /dev/null +++ b/src/gateway/session-sharing-snapshot-cache.test.ts @@ -0,0 +1,73 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + invalidateSessionSharingSnapshot, + loadCachedSessionSharingSnapshot, + type SessionSharingSnapshot, +} from "./session-sharing-snapshot-cache.js"; + +const snapshot: SessionSharingSnapshot = { incognito: false, visibility: "shared" }; + +function loadSnapshot(params: { alias: string; canonical: string; resolve: () => void }) { + return loadCachedSessionSharingSnapshot({ + sessionKey: params.alias, + resolve: () => { + params.resolve(); + return { canonicalKey: params.canonical, snapshot }; + }, + }); +} + +afterEach(() => { + invalidateSessionSharingSnapshot(); +}); + +describe("session sharing snapshot reverse indexes", () => { + it("invalidates a canonical entry and every alias targeting it", () => { + const first = vi.fn(); + const second = vi.fn(); + loadSnapshot({ alias: "alias-one", canonical: "canonical", resolve: first }); + loadSnapshot({ alias: "alias-two", canonical: "canonical", resolve: second }); + expect(first).toHaveBeenCalledOnce(); + expect(second).toHaveBeenCalledOnce(); + + loadSnapshot({ alias: "alias-two", canonical: "canonical", resolve: second }); + expect(second).toHaveBeenCalledOnce(); + invalidateSessionSharingSnapshot("canonical"); + loadSnapshot({ alias: "alias-two", canonical: "canonical", resolve: second }); + expect(second).toHaveBeenCalledTimes(2); + }); + + it("invalidating one alias removes its canonical entry and sibling aliases", () => { + const first = vi.fn(); + const second = vi.fn(); + loadSnapshot({ alias: "alias-one", canonical: "canonical", resolve: first }); + loadSnapshot({ alias: "alias-two", canonical: "canonical", resolve: second }); + + invalidateSessionSharingSnapshot("alias-one"); + loadSnapshot({ alias: "alias-two", canonical: "canonical", resolve: second }); + expect(second).toHaveBeenCalledTimes(2); + }); + + it("keeps reverse indexes consistent across canonical eviction and alias reuse", () => { + const oldAlias = vi.fn(); + loadSnapshot({ alias: "reused-alias", canonical: "old-canonical", resolve: oldAlias }); + for (let index = 0; index < 2_048; index += 1) { + loadSnapshot({ + alias: `alias-${index}`, + canonical: `canonical-${index}`, + resolve: vi.fn(), + }); + } + + const newAlias = vi.fn(); + loadSnapshot({ alias: "reused-alias", canonical: "new-canonical", resolve: newAlias }); + expect(newAlias).toHaveBeenCalledOnce(); + invalidateSessionSharingSnapshot("old-canonical"); + loadSnapshot({ alias: "reused-alias", canonical: "new-canonical", resolve: newAlias }); + expect(newAlias).toHaveBeenCalledOnce(); + + invalidateSessionSharingSnapshot("reused-alias"); + loadSnapshot({ alias: "reused-alias", canonical: "new-canonical", resolve: newAlias }); + expect(newAlias).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/gateway/session-sharing-snapshot-cache.ts b/src/gateway/session-sharing-snapshot-cache.ts index ec1801e5e7f3..50c37d1c9ac1 100644 --- a/src/gateway/session-sharing-snapshot-cache.ts +++ b/src/gateway/session-sharing-snapshot-cache.ts @@ -10,65 +10,103 @@ export type SessionSharingSnapshot = { const snapshotCache = new Map(); const snapshotAliases = new Map(); +const snapshotKeysBySessionKey = new Map>(); +const aliasKeysBySessionKey = new Map>(); +const aliasKeysByCanonicalKey = new Map>(); function snapshotKey(sessionKey: string, agentId?: string): string { return `${agentId ?? ""}\0${sessionKey}`; } +function logicalSessionKey(key: string): string { + return key.slice(key.lastIndexOf("\0") + 1); +} + +function addReverseIndex(index: Map>, key: string, value: string): void { + const values = index.get(key) ?? new Set(); + values.add(value); + index.set(key, values); +} + +function removeReverseIndex(index: Map>, key: string, value: string): void { + const values = index.get(key); + values?.delete(value); + if (values?.size === 0) { + index.delete(key); + } +} + +function removeSnapshotAlias(alias: string): void { + const canonical = snapshotAliases.get(alias); + if (!canonical) { + return; + } + snapshotAliases.delete(alias); + removeReverseIndex(aliasKeysBySessionKey, logicalSessionKey(alias), alias); + removeReverseIndex(aliasKeysByCanonicalKey, canonical, alias); +} + +function removeSnapshot(key: string): void { + if (!snapshotCache.delete(key)) { + return; + } + removeReverseIndex(snapshotKeysBySessionKey, logicalSessionKey(key), key); + for (const alias of aliasKeysByCanonicalKey.get(key) ?? []) { + removeSnapshotAlias(alias); + } +} + function rememberSnapshot(key: string, snapshot: SessionSharingSnapshot): void { - snapshotCache.delete(key); + const known = snapshotCache.delete(key); snapshotCache.set(key, snapshot); + if (!known) { + addReverseIndex(snapshotKeysBySessionKey, logicalSessionKey(key), key); + } if (snapshotCache.size <= SNAPSHOT_CACHE_LIMIT) { return; } const oldest = snapshotCache.keys().next().value; if (oldest) { - snapshotCache.delete(oldest); - for (const [alias, canonical] of snapshotAliases) { - if (canonical === oldest) { - snapshotAliases.delete(alias); - } - } + removeSnapshot(oldest); } } function rememberSnapshotAlias(alias: string, canonical: string): void { - snapshotAliases.delete(alias); + removeSnapshotAlias(alias); snapshotAliases.set(alias, canonical); + addReverseIndex(aliasKeysBySessionKey, logicalSessionKey(alias), alias); + addReverseIndex(aliasKeysByCanonicalKey, canonical, alias); if (snapshotAliases.size <= SNAPSHOT_CACHE_LIMIT * 2) { return; } const oldest = snapshotAliases.keys().next().value; if (oldest) { - snapshotAliases.delete(oldest); + removeSnapshotAlias(oldest); } } export function invalidateSessionSharingSnapshot(sessionKey?: string): void { if (sessionKey) { - const matchingCanonicalKeys = new Set(); - for (const key of snapshotCache.keys()) { - if (key.endsWith(`\0${sessionKey}`)) { - matchingCanonicalKeys.add(key); - } - } - for (const [alias, canonical] of snapshotAliases) { - if (alias.endsWith(`\0${sessionKey}`) || canonical.endsWith(`\0${sessionKey}`)) { + // Mutation events carry the logical key without an agent id. These reverse indexes + // preserve that cross-agent invalidation contract while keeping work proportional to + // aliases/canonical entries for the affected session instead of the whole bounded cache. + const matchingCanonicalKeys = new Set(snapshotKeysBySessionKey.get(sessionKey)); + for (const alias of aliasKeysBySessionKey.get(sessionKey) ?? []) { + const canonical = snapshotAliases.get(alias); + if (canonical) { matchingCanonicalKeys.add(canonical); } } for (const key of matchingCanonicalKeys) { - snapshotCache.delete(key); - } - for (const [alias, canonical] of snapshotAliases) { - if (matchingCanonicalKeys.has(canonical)) { - snapshotAliases.delete(alias); - } + removeSnapshot(key); } return; } snapshotCache.clear(); snapshotAliases.clear(); + snapshotKeysBySessionKey.clear(); + aliasKeysBySessionKey.clear(); + aliasKeysByCanonicalKey.clear(); } export function loadCachedSessionSharingSnapshot(params: { diff --git a/src/gateway/session-transcript-readers.test.ts b/src/gateway/session-transcript-readers.test.ts index ae3d621399cc..2db0345e2292 100644 --- a/src/gateway/session-transcript-readers.test.ts +++ b/src/gateway/session-transcript-readers.test.ts @@ -24,7 +24,10 @@ import { readLatestSessionUsageFromTranscriptAsync, type SessionTranscriptReadScope, } from "./session-transcript-readers.js"; -import { readSessionTitleFieldsFromTranscript } from "./session-transcript-title-reader.js"; +import { + readSessionTitleFieldsFromTranscript, + readSessionTitleFieldsFromTranscriptBatch, +} from "./session-transcript-title-reader.js"; vi.mock("../config/sessions/session-accessor.js", async (importOriginal) => { const actual = await importOriginal(); @@ -32,6 +35,7 @@ vi.mock("../config/sessions/session-accessor.js", async (importOriginal) => { ...actual, readSessionTranscriptMessageEventPage: vi.fn(actual.readSessionTranscriptMessageEventPage), readSessionTranscriptMessageEvents: vi.fn(actual.readSessionTranscriptMessageEvents), + readSessionTranscriptTitleProbeBatch: vi.fn(actual.readSessionTranscriptTitleProbeBatch), }; }); @@ -400,6 +404,46 @@ describe("session transcript reader facade", () => { expect(sessionAccessor.readSessionTranscriptMessageEvents).not.toHaveBeenCalled(); }); + test("falls back to the canonical visible window for reset transcripts", async () => { + const sessionId = "reader-title-reset-window"; + const scope = await writeTranscript(sessionId, [ + { type: "session", version: 3, id: sessionId }, + { + type: "message", + id: "old", + parentId: null, + message: { role: "user", content: "hidden old prompt" }, + }, + { + type: "message", + id: "kept-user", + parentId: "old", + message: { role: "user", content: "kept prompt" }, + }, + { + type: "message", + id: "kept-assistant", + parentId: "kept-user", + message: { role: "assistant", content: "kept answer" }, + }, + { + type: "reset", + id: "reset-boundary", + parentId: "kept-assistant", + firstKeptEntryId: "kept-user", + }, + { + type: "message", + id: "post-reset", + parentId: "reset-boundary", + message: { role: "assistant", content: "newest answer" }, + }, + ]); + expect(readSessionTitleFieldsFromTranscriptBatch([scope])).toEqual([ + { firstUserMessage: "kept prompt", lastMessagePreview: "newest answer" }, + ]); + }); + test("bounds title probe reads independently of transcript length", async () => { const probeReadCount = async (sessionId: string, messageCount: number) => { const scope = await writeSqliteMessages( @@ -443,6 +487,48 @@ describe("session transcript reader facade", () => { expect(sessionAccessor.readSessionTranscriptMessageEventPage).not.toHaveBeenCalled(); }); + test("skips batch title probes while every cached transcript watermark is unchanged", async () => { + const scope = await writeSqliteMessages("reader-title-batch-cache-warm", [ + { role: "user", content: "cached batch prompt" }, + { role: "assistant", content: "cached batch reply" }, + ]); + expect(readSessionTitleFieldsFromTranscriptBatch([scope])).toEqual([ + { firstUserMessage: "cached batch prompt", lastMessagePreview: "cached batch reply" }, + ]); + vi.clearAllMocks(); + + expect(readSessionTitleFieldsFromTranscriptBatch([scope])).toEqual([ + { firstUserMessage: "cached batch prompt", lastMessagePreview: "cached batch reply" }, + ]); + expect(sessionAccessor.readSessionTranscriptTitleProbeBatch).not.toHaveBeenCalled(); + expect(sessionAccessor.readSessionTranscriptMessageEventPage).not.toHaveBeenCalled(); + }); + + test("reprobes cached batch title fields after an append advances max seq", async () => { + const sessionId = "reader-title-batch-cache-append"; + const scope = await writeSqliteMessages(sessionId, [ + { role: "user", content: "batch append prompt" }, + { role: "assistant", content: "first batch reply" }, + ]); + expect(readSessionTitleFieldsFromTranscriptBatch([scope])[0]?.lastMessagePreview).toBe( + "first batch reply", + ); + await persistSessionTranscriptTurn( + { agentId: "main", sessionId, sessionKey: `agent:main:${sessionId}`, storePath }, + { + messages: [{ message: { role: "assistant", content: "appended batch reply" } }], + touchSessionEntry: false, + }, + ); + vi.clearAllMocks(); + + expect(readSessionTitleFieldsFromTranscriptBatch([scope])[0]?.lastMessagePreview).toBe( + "appended batch reply", + ); + expect(sessionAccessor.readSessionTranscriptTitleProbeBatch).toHaveBeenCalledOnce(); + expect(sessionAccessor.readSessionTranscriptMessageEventPage).not.toHaveBeenCalled(); + }); + test("invalidates cached SQLite title fields after an append advances max seq", async () => { const sessionId = "reader-title-cache-append"; const scope = await writeSqliteMessages(sessionId, [ diff --git a/src/gateway/session-transcript-title-reader.ts b/src/gateway/session-transcript-title-reader.ts index b86f7fa29457..f34143bc1c39 100644 --- a/src/gateway/session-transcript-title-reader.ts +++ b/src/gateway/session-transcript-title-reader.ts @@ -2,6 +2,7 @@ // cache so list rendering never rescans transcripts that have not changed. import { readSessionTranscriptMessageEventPage, + readSessionTranscriptTitleProbeBatch, readSessionTranscriptWatermark, type SessionTranscriptMessageEvent, type SessionTranscriptReadScope, @@ -32,7 +33,8 @@ type SqliteTitleFieldCacheEntry = ReturnType(); function sqliteTitleFieldCacheKey(target: ResolvedTranscriptReadTarget): string { @@ -157,6 +159,90 @@ function readSqliteTitleFields( return { ...fields }; } +/** Batch-hydrates list title fields once per store, with canonical widening only for misses. */ +export function readSessionTitleFieldsFromTranscriptBatch( + scopes: readonly SessionTranscriptReadScope[], + opts?: { includeInterSession?: boolean }, +): SessionTitleFields[] { + const targets: ResolvedTranscriptReadTarget[] = []; + const variant = opts?.includeInterSession === true ? "includeInterSession" : "default"; + const results = new Map(); + const misses: Array<{ + cacheKey: string; + index: number; + scope: SessionTranscriptReadScope; + target: ResolvedTranscriptReadTarget; + }> = []; + + for (const [index, scope] of scopes.entries()) { + const target = resolveTranscriptReadTarget(scope); + targets.push(target); + const cacheKey = sqliteTitleFieldCacheKey(target); + const cached = sqliteTitleFieldCache.get(cacheKey); + const cachedFields = cached?.fields[variant]; + if (cached && cachedFields) { + // Keep the single-row generation/maxSeq validity contract, but validate only warm rows; + // cold or changed rows still collapse into the one store-batched probe below. + const watermark = readSessionTranscriptWatermark(scope); + if (cached.generation === watermark.generation && cached.maxSeq === watermark.maxSeq) { + setSqliteTitleFieldCache(cacheKey, cached); + results.set(index, { ...cachedFields }); + continue; + } + } + misses.push({ cacheKey, index, scope, target }); + } + + const probes = + misses.length > 0 ? readSessionTranscriptTitleProbeBatch(misses.map((miss) => miss.scope)) : []; + for (const [probeIndex, miss] of misses.entries()) { + const probe = probes[probeIndex]; + if (!probe) { + results.set(miss.index, readSqliteTitleFields(miss.target, opts)); + continue; + } + const cached = sqliteTitleFieldCache.get(miss.cacheKey); + const cachedFields = + cached?.generation === probe.generation && cached.maxSeq === probe.maxSeq + ? cached.fields[variant] + : undefined; + if (cached && cachedFields) { + setSqliteTitleFieldCache(miss.cacheKey, cached); + results.set(miss.index, { ...cachedFields }); + continue; + } + const firstUser = findFirstTitleUserMessage(probe.head, opts?.includeInterSession === true); + const lastText = findLastMessageText(probe.tail); + if (probe.totalMessages > SQLITE_TITLE_PROBE_INITIAL_MESSAGES && (!firstUser || !lastText)) { + results.set(miss.index, readSqliteTitleFields(miss.target, opts)); + continue; + } + const fields = { + firstUserMessage: firstUser ? extractMessageText(firstUser) : null, + lastMessagePreview: lastText, + }; + const fieldsByVariant = + cached?.generation === probe.generation && cached.maxSeq === probe.maxSeq + ? cached.fields + : {}; + fieldsByVariant[variant] = fields; + setSqliteTitleFieldCache(miss.cacheKey, { + generation: probe.generation, + maxSeq: probe.maxSeq, + fields: fieldsByVariant, + }); + results.set(miss.index, { ...fields }); + } + + return targets.map((target, index) => { + const fields = results.get(index); + if (!fields) { + throw new Error(`Missing batched title fields for session ${target.sessionId}`); + } + return fields; + }); +} + /** Reads title and preview text from a transcript through the reader seam. */ export function readSessionTitleFieldsFromTranscript( scope: SessionTranscriptReadScope, diff --git a/src/gateway/session-utils-list.ts b/src/gateway/session-utils-list.ts index 10290b5a916c..5340e6f83707 100644 --- a/src/gateway/session-utils-list.ts +++ b/src/gateway/session-utils-list.ts @@ -23,7 +23,7 @@ import { import { isCronRunSessionKey } from "../sessions/session-key-utils.js"; import { type SessionEntryPair, sortAndLimitSessionEntries } from "./session-list-order.js"; import { resolveStoredSessionKeyForAgentStore } from "./session-store-key.js"; -import { readSessionTitleFieldsFromTranscriptAsync as readScopedSessionTitleFieldsFromTranscriptAsync } from "./session-transcript-title-reader.js"; +import { readSessionTitleFieldsFromTranscriptBatch as readScopedSessionTitleFieldsFromTranscriptBatch } from "./session-transcript-title-reader.js"; import type { SessionActorProfileIdentity, SessionListRowContext, @@ -70,6 +70,7 @@ type ListSessionsFromStoreParams = { storePath: string; store: Record; modelCatalog?: ModelCatalogEntry[]; + lightweightListRows?: boolean; opts: SessionsListParams; }; @@ -492,6 +493,8 @@ export function listSessionsFromStore(params: ListSessionsFromStoreParams): Sess transcriptUsageMaxBytes: SESSIONS_LIST_TRANSCRIPT_USAGE_MAX_BYTES, storeChildSessionsByKey, rowContext: list.rowContext, + skipTranscriptUsageFallback: params.lightweightListRows === true, + lightweightListRow: params.lightweightListRows === true, }); }); return buildSessionsListResult({ cfg, list, modelCatalog: params.modelCatalog, sessions }); @@ -519,6 +522,31 @@ export async function listSessionsFromStoreAsync( const { cfg, store, opts } = params; const list = prepareSessionList(params); const sessions: GatewaySessionRow[] = []; + const transcriptScopes = list.entries + .slice(0, SESSIONS_LIST_TRANSCRIPT_FIELD_ROWS) + .flatMap(([key, entry]) => { + if (!entry.sessionId || (!list.includeDerivedTitles && !list.includeLastMessage)) { + return []; + } + const parsed = parseAgentSessionKey(key); + const agentId = + key === "global" && typeof opts.agentId === "string" + ? normalizeAgentId(opts.agentId) + : parsed?.agentId + ? normalizeAgentId(parsed.agentId) + : resolveDefaultAgentId(cfg); + return [ + { + agentId, + sessionEntry: entry, + sessionId: entry.sessionId, + sessionKey: key, + storePath: list.storePath, + }, + ]; + }); + const transcriptFields = readScopedSessionTitleFieldsFromTranscriptBatch(transcriptScopes); + let transcriptFieldIndex = 0; for (let i = 0; i < list.entries.length; i++) { const [key, entry] = expectDefined(list.entries[i], "entries entry at i"); const includeTranscriptFields = i < SESSIONS_LIST_TRANSCRIPT_FIELD_ROWS; @@ -556,17 +584,11 @@ export async function listSessionsFromStoreAsync( includeTranscriptFields && (list.includeDerivedTitles || list.includeLastMessage) ) { - const parsed = parseAgentSessionKey(key); - const sessionAgentId = - rowAgentId ?? - (parsed?.agentId ? normalizeAgentId(parsed.agentId) : resolveDefaultAgentId(cfg)); - const fields = await readScopedSessionTitleFieldsFromTranscriptAsync({ - agentId: sessionAgentId, - sessionEntry: entry, - sessionId: entry.sessionId, - sessionKey: key, - storePath: list.storePath, - }); + const fields = expectDefined( + transcriptFields[transcriptFieldIndex], + "batched transcript fields at transcriptFieldIndex", + ); + transcriptFieldIndex += 1; if (list.includeDerivedTitles) { row.derivedTitle = deriveSessionTitle(entry, fields.firstUserMessage, row.displayName); } diff --git a/src/gateway/session-utils.perf.test.ts b/src/gateway/session-utils.perf.test.ts index 777647de3ac9..a0cdf3ea22a7 100644 --- a/src/gateway/session-utils.perf.test.ts +++ b/src/gateway/session-utils.perf.test.ts @@ -17,7 +17,8 @@ import { resetPluginRuntimeStateForTest, setActivePluginRegistry } from "../plug import { openOpenClawStateDatabase } from "../state/openclaw-state-db.js"; import { withStateDirEnv } from "../test-helpers/state-dir-env.js"; import * as usageFormat from "../utils/usage-format.js"; -import { listSessionsFromStore } from "./session-utils.js"; +import * as titleReader from "./session-transcript-title-reader.js"; +import { listSessionsFromStore, listSessionsFromStoreAsync } from "./session-utils.js"; /** * Regression smoke for the per-list rowContext resolver cache. The bug we are @@ -263,4 +264,56 @@ describe("listSessionsFromStore resolver cache", () => { } }); }); + + test("batches transcript title hydration once instead of O(rows)", async () => { + await withStateDirEnv("openclaw-perf-title-batch-", async () => { + resetPluginRuntimeStateForTest(); + setActivePluginRegistry(createEmptyPluginRegistry()); + const cfg = { + agents: { defaults: { model: { primary: "openai/gpt-5" } } }, + } as OpenClawConfig; + resetConfigRuntimeState(); + setRuntimeConfigSnapshot(cfg); + const storePath = "/tmp/sessions.json"; + const store: Record = {}; + for (let index = 0; index < 30; index += 1) { + const sessionId = `title-batch-${index}`; + const sessionKey = `agent:main:${sessionId}`; + const entry = { sessionId, updatedAt: 1_000 - index } satisfies SessionEntry; + store[sessionKey] = entry; + } + + const titleBatchSpy = vi + .spyOn(titleReader, "readSessionTitleFieldsFromTranscriptBatch") + .mockImplementation((scopes) => + scopes.map((scope) => ({ + firstUserMessage: `title ${scope.sessionId.slice("title-batch-".length)}`, + lastMessagePreview: `last ${scope.sessionId.slice("title-batch-".length)}`, + })), + ); + try { + const result = await listSessionsFromStoreAsync({ + cfg, + storePath, + store, + opts: { includeDerivedTitles: true, includeLastMessage: true, limit: 30 }, + }); + + expect(result.sessions).toHaveLength(30); + expect(titleBatchSpy).toHaveBeenCalledOnce(); + expect(titleBatchSpy.mock.calls[0]?.[0]).toHaveLength(30); + const sessionsByKey = new Map(result.sessions.map((session) => [session.key, session])); + expect(sessionsByKey.get("agent:main:title-batch-0")).toMatchObject({ + derivedTitle: "title 0", + lastMessagePreview: "last 0", + }); + expect(sessionsByKey.get("agent:main:title-batch-29")).toMatchObject({ + derivedTitle: "title 29", + lastMessagePreview: "last 29", + }); + } finally { + titleBatchSpy.mockRestore(); + } + }); + }); }); diff --git a/src/gateway/session-utils.search.test.ts b/src/gateway/session-utils.search.test.ts index 1a7ba88c792a..28cc8a7a7971 100644 --- a/src/gateway/session-utils.search.test.ts +++ b/src/gateway/session-utils.search.test.ts @@ -16,7 +16,8 @@ import { appendTranscriptMessageSync, replaceSessionEntry, } from "../config/sessions/session-accessor.js"; -import { registerAgentRunContext, resetAgentEventsForTest } from "../infra/agent-events.js"; +import { resetAgentEventsForTest } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { closeOpenClawAgentDatabasesForTest } from "../state/openclaw-agent-db.js"; import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; import { diff --git a/src/gateway/session-utils.subagent.test.ts b/src/gateway/session-utils.subagent.test.ts index f864c623478d..4d638df68584 100644 --- a/src/gateway/session-utils.subagent.test.ts +++ b/src/gateway/session-utils.subagent.test.ts @@ -17,7 +17,8 @@ import type { OpenClawConfig } from "../config/config.js"; import type { SessionEntry } from "../config/sessions.js"; import { canPrewarmCombinedSessionStoresForGateway } from "../config/sessions/combined-store-gateway.js"; import { replaceSessionEntry } from "../config/sessions/session-accessor.js"; -import { registerAgentRunContext, resetAgentEventsForTest } from "../infra/agent-events.js"; +import { resetAgentEventsForTest } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { closeOpenClawAgentDatabasesForTest, resolveIncognitoOpenClawAgentSqlitePath, diff --git a/src/gateway/sessions-resolve.test.ts b/src/gateway/sessions-resolve.test.ts index 87266c6cdc9b..a1dc8aa4673d 100644 --- a/src/gateway/sessions-resolve.test.ts +++ b/src/gateway/sessions-resolve.test.ts @@ -282,6 +282,9 @@ describe("resolveSessionKeyFromResolveParams", () => { expect(hoisted.loadCombinedSessionStoreForGatewayMock).toHaveBeenCalledWith(cfg, { agentId: "main", }); + expect(hoisted.listSessionsFromStoreMock).toHaveBeenCalledWith( + expect.objectContaining({ lightweightListRows: true }), + ); expect(result).toEqual({ ok: false, error: { diff --git a/src/gateway/sessions-resolve.ts b/src/gateway/sessions-resolve.ts index f93dca845433..26d356a74141 100644 --- a/src/gateway/sessions-resolve.ts +++ b/src/gateway/sessions-resolve.ts @@ -199,6 +199,7 @@ export async function resolveSessionKeyFromResolveParams(params: { cfg, storePath, store, + lightweightListRows: true, opts: { includeGlobal: p.includeGlobal === true, includeUnknown: p.includeUnknown === true, @@ -215,7 +216,7 @@ export async function resolveSessionKeyFromResolveParams(params: { }); } if (list.sessions.length > 1) { - const keys = list.sessions.map((s) => s.key).join(", "); + const keys = list.sessions.map((session) => session.key).join(", "); return { ok: false, error: errorShape( @@ -232,6 +233,6 @@ export async function resolveSessionKeyFromResolveParams(params: { } return { ok: true, - key: expectDefined(list.sessions[0], "sessions entry at 0").key, + key: labelKey, }; } diff --git a/src/gateway/worker-environments/live-events.test.ts b/src/gateway/worker-environments/live-events.test.ts index 9d2a6d565496..15db5bd98fab 100644 --- a/src/gateway/worker-environments/live-events.test.ts +++ b/src/gateway/worker-environments/live-events.test.ts @@ -9,16 +9,18 @@ import type { import * as sessions from "../../config/sessions/session-accessor.js"; import type { OpenClawConfig as Config } from "../../config/types.openclaw.js"; import { - claimAgentRunContext, - clearAgentRunContext, emitAgentEvent, getAgentEventLifecycleGeneration, - getAgentRunContext, onAgentRuntimeEvent, - releaseAgentRunContext, - sweepStaleRunContexts, type AgentEventRuntimePayload as Event, } from "../../infra/agent-events.js"; +import { + claimAgentRunContext, + clearAgentRunContext, + getAgentRunContext, + releaseAgentRunContext, + sweepStaleRunContexts, +} from "../../infra/agent-run-registry.js"; import { closeOpenClawAgentDatabasesForTest } from "../../state/openclaw-agent-db.js"; import { loadSqliteTrajectoryRuntimeEventRowsSync } from "../../trajectory/runtime-store.sqlite.js"; import type { WorkerConnectionIdentity as Identity } from "./connection-identity.js"; diff --git a/src/gateway/worker-environments/live-events.ts b/src/gateway/worker-environments/live-events.ts index 70c308d7e3c1..ea3e61ca8186 100644 --- a/src/gateway/worker-environments/live-events.ts +++ b/src/gateway/worker-environments/live-events.ts @@ -7,15 +7,17 @@ import type { import { onSessionIdentityMutation } from "../../config/sessions/session-accessor.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; import { - claimAgentRunContext, emitAgentEventIfCurrent, emitAgentEventForOwner, getAgentEventLifecycleGeneration, +} from "../../infra/agent-events.js"; +import { + claimAgentRunContext, getAgentRunContext, getAgentRunContextOwnerStatus, registerAgentRunContext, releaseAgentRunContext, -} from "../../infra/agent-events.js"; +} from "../../infra/agent-run-registry.js"; import type { WorkerConnectionIdentity } from "./connection-identity.js"; import { createWorkerLiveTrajectoryRecorder, diff --git a/src/infra/agent-events.test.ts b/src/infra/agent-events.test.ts index 80629e980355..97bbcf08f2d8 100644 --- a/src/infra/agent-events.test.ts +++ b/src/infra/agent-events.test.ts @@ -3,35 +3,46 @@ import { beforeEach, describe, expect, test, vi } from "vitest"; import { type AgentEventPayload, captureAgentRunLifecycleGeneration, - claimAgentRunContext, - clearAgentRunContext, emitAgentAuditEvent, emitAgentEvent, emitAgentEventForOwner, emitAgentEventIfCurrent, getAgentEventLifecycleGeneration, - getAgentRunContext, - listAgentRunsForSession, onAgentAuditEvent, onAgentEvent, onAgentRuntimeEvent, - registerAgentRunContext, - releaseAgentRunContext, resetAgentEventsForTest, rotateAgentEventLifecycleGeneration, runOncePerAgentRun, - sweepStaleRunContexts, withAgentRunLifecycleGeneration, } from "./agent-events.js"; +import { + claimAgentRunContext, + clearAgentRunContext, + getAgentRunContext, + listAgentRunsForSession, + readAgentRunIndexVersion, + registerAgentRunContext, + releaseAgentRunContext, + sweepStaleRunContexts, +} from "./agent-run-registry.js"; import { emitAgentRunStatusEvent } from "./agent-run-status-events.js"; import { recordAgentRunOutputTokens } from "./agent-run-usage.js"; -type AgentEventsModule = typeof import("./agent-events.js"); +type AgentEventsModule = { + events: typeof import("./agent-events.js"); + registry: typeof import("./agent-run-registry.js"); +}; const agentEventsModuleUrl = new URL("./agent-events.ts", import.meta.url).href; +const agentRunRegistryModuleUrl = new URL("./agent-run-registry.ts", import.meta.url).href; async function importAgentEventsModule(cacheBust: string): Promise { - return (await import(`${agentEventsModuleUrl}?t=${cacheBust}`)) as AgentEventsModule; + const [events, registry] = await Promise.all([ + import(`${agentEventsModuleUrl}?t=${cacheBust}`), + import(`${agentRunRegistryModuleUrl}?t=${cacheBust}`), + ]); + return { events, registry } as AgentEventsModule; } describe("agent-events sequencing", () => { @@ -66,6 +77,39 @@ describe("agent-events sequencing", () => { expect(getAgentRunContext("run-1")).toBeUndefined(); }); + test("versions active-run projection ownership transitions", () => { + let version = readAgentRunIndexVersion(); + registerAgentRunContext("projected-run", { + projectSessionActive: true, + sessionId: "projected-session-id", + sessionKey: "agent:main:projected", + }); + expect(readAgentRunIndexVersion()).toBeGreaterThan(version); + version = readAgentRunIndexVersion(); + + registerAgentRunContext("projected-run", { verboseLevel: "full" }); + expect(readAgentRunIndexVersion()).toBe(version); + + const claimId = claimAgentRunContext( + "owned-projected-run", + { + projectSessionActive: true, + sessionId: "owned-session-id", + sessionKey: "agent:main:owned", + }, + { ownsContext: true, trackOwner: true }, + ); + expect(readAgentRunIndexVersion()).toBeGreaterThan(version); + version = readAgentRunIndexVersion(); + + releaseAgentRunContext("owned-projected-run", claimId); + expect(readAgentRunIndexVersion()).toBeGreaterThan(version); + version = readAgentRunIndexVersion(); + + clearAgentRunContext("projected-run"); + expect(readAgentRunIndexVersion()).toBeGreaterThan(version); + }); + test("does not let an old execution clear a newer same-id context", () => { registerAgentRunContext("shared-run", { sessionKey: "main", @@ -897,23 +941,23 @@ describe("agent-events sequencing", () => { const first = await importAgentEventsModule(`first-${Date.now()}`); const second = await importAgentEventsModule(`second-${Date.now()}`); - first.resetAgentEventsForTest(); - first.registerAgentRunContext("run-dup", { sessionKey: "session-dup" }); + first.events.resetAgentEventsForTest(); + first.registry.registerAgentRunContext("run-dup", { sessionKey: "session-dup" }); const seen: Array<{ seq: number; sessionKey?: string }> = []; - const stop = first.onAgentEvent((evt) => { + const stop = first.events.onAgentEvent((evt) => { if (evt.runId === "run-dup") { seen.push({ seq: evt.seq, sessionKey: evt.sessionKey }); } }); - second.emitAgentEvent({ + second.events.emitAgentEvent({ runId: "run-dup", stream: "assistant", data: { text: "from second" }, sessionKey: " ", }); - first.emitAgentEvent({ + first.events.emitAgentEvent({ runId: "run-dup", stream: "assistant", data: { text: "from first" }, @@ -922,13 +966,13 @@ describe("agent-events sequencing", () => { stop(); - expect(second.getAgentRunContext("run-dup")?.sessionKey).toBe("session-dup"); + expect(second.registry.getAgentRunContext("run-dup")?.sessionKey).toBe("session-dup"); expect(seen).toEqual([ { seq: 1, sessionKey: "session-dup" }, { seq: 2, sessionKey: "session-dup" }, ]); - first.resetAgentEventsForTest(); + first.events.resetAgentEventsForTest(); }); test("sweeps stale run contexts and clears their sequence state", () => { @@ -944,7 +988,9 @@ describe("agent-events sequencing", () => { emitAgentEvent({ runId: "run-active", stream: "assistant", data: { text: "active" } }); stop.mockReturnValue(1_000); + const versionBeforeSweep = readAgentRunIndexVersion(); expect(sweepStaleRunContexts(500)).toBe(1); + expect(readAgentRunIndexVersion()).toBeGreaterThan(versionBeforeSweep); expect(getAgentRunContext("run-stale")).toBeUndefined(); expect(getAgentRunContext("run-active")?.sessionKey).toBe("session-active"); diff --git a/src/infra/agent-events.ts b/src/infra/agent-events.ts index 57919701fb3b..6fb476dde143 100644 --- a/src/infra/agent-events.ts +++ b/src/infra/agent-events.ts @@ -1,12 +1,17 @@ // Stores and broadcasts agent lifecycle and streaming events. import { AsyncLocalStorage } from "node:async_hooks"; -import { randomUUID } from "node:crypto"; -import type { VerboseLevel } from "../auto-reply/thinking.js"; import { resolveGlobalSingleton } from "../shared/global-singleton.js"; import { notifyListeners, registerListener } from "../shared/listeners.js"; import { hasInvalidLifecycleStartTimestamp } from "./agent-event-lifecycle.js"; import { createAgentRunStaleLifecycleError } from "./agent-lifecycle-error.js"; -import { clearAgentRunUsage, resetAgentRunUsageForTest } from "./agent-run-usage.js"; +import { + getAgentRunContext, + getAgentRunContextOwnership, + getAgentRunLifecycleGeneration, + registerAgentRunSequenceResetHandler, + resetAgentRunRegistryForTest, + rotateAgentRunRegistryLifecycleGeneration, +} from "./agent-run-registry.js"; /** Approval event phase for request/resolution transitions. */ type AgentApprovalEventPhase = "requested" | "resolved"; @@ -75,49 +80,10 @@ export type AgentEventRuntimePayload = AgentEventPayload & { readonly projectSessionLifecycle?: boolean; }; -/** Per-run metadata used to stamp events and gate Control UI visibility. */ -type AgentRunContext = { - sessionKey?: string; - /** Resolved agent owner, including for unscoped session keys. */ - agentId?: string; - /** Owning run's sessionId; stamped onto lifecycle events (see AgentEventPayload.sessionId). */ - sessionId?: string; - /** Gateway lifecycle generation captured when the run was registered. */ - lifecycleGeneration?: string; - /** Producer-owned start captured from this run's accepted lifecycle event. */ - lifecycleStartedAt?: number; - verboseLevel?: VerboseLevel; - isHeartbeat?: boolean; - /** Whether control UI clients should receive chat/agent updates for this run. */ - isControlUiVisible?: boolean; - projectSessionActive?: boolean; - /** Whether lifecycle events may update the shared session row. */ - projectSessionLifecycle?: boolean; - /** Active cadence state by job; admission permits one invocation per job. */ - cronRunsByJobId?: Map; - /** Timestamp when this context was first registered (for TTL-based cleanup). */ - registeredAt?: number; - /** Timestamp of last activity (updated on every emitAgentEvent). */ - lastActiveAt?: number; -}; - type AgentEventState = { seqByRun: Map; listeners: Set<(evt: AgentEventRuntimePayload) => void>; auditListeners: Set<(evt: AgentEventPayload) => void>; - runContextById: Map; - runContextOwnersById?: Map< - string, - { - lifecycleGeneration: string; - claimIds: Set; - preserveAfterRelease: boolean; - clearRequested: boolean; - exclusiveClaimId?: string; - clearListeners?: Map void>; - } - >; - lifecycleGeneration: string; lifecycleRotationHandlers?: Map void>; }; @@ -134,11 +100,13 @@ function getAgentEventState(): AgentEventState { seqByRun: new Map(), listeners: new Set<(evt: AgentEventRuntimePayload) => void>(), auditListeners: new Set<(evt: AgentEventPayload) => void>(), - runContextById: new Map(), - lifecycleGeneration: randomUUID(), })); } +registerAgentRunSequenceResetHandler((runId) => { + getAgentEventState().seqByRun.delete(runId); +}); + function getAgentEventExecutionContext() { return resolveGlobalSingleton>( AGENT_EVENT_EXECUTION_CONTEXT_KEY, @@ -172,11 +140,11 @@ export function runOncePerAgentRun(runId: string, operation: string, run: () } export function getAgentEventLifecycleGeneration(): string { - return getAgentEventState().lifecycleGeneration; + return getAgentRunLifecycleGeneration(); } export function isAgentEventLifecycleGenerationCurrent(lifecycleGeneration: string): boolean { - return lifecycleGeneration === getAgentEventState().lifecycleGeneration; + return lifecycleGeneration === getAgentRunLifecycleGeneration(); } /** Registers process-local state cleanup at the gateway lifecycle boundary. */ @@ -203,385 +171,26 @@ export function assertAgentRunLifecycleGenerationCurrent(lifecycleGeneration: st export function captureAgentRunLifecycleGeneration(runId: string): string { return ( getAgentEventExecutionContext().getStore()?.lifecycleGeneration ?? - getAgentEventState().runContextById.get(runId)?.lifecycleGeneration ?? - getAgentEventState().lifecycleGeneration + getAgentRunContext(runId)?.lifecycleGeneration ?? + getAgentRunLifecycleGeneration() ); } /** Starts a new ownership generation before an in-process gateway restart. */ export function rotateAgentEventLifecycleGeneration(): string { const state = getAgentEventState(); - state.lifecycleGeneration = randomUUID(); + const lifecycleGeneration = rotateAgentRunRegistryLifecycleGeneration(); // Rotation is the liveness choke point: after it returns, no prior-generation // owner is operationally reachable. Recovery and runtime consumers therefore // agree that only current-generation owners can drive or receive work. const errors: unknown[] = []; - notifyListeners( - state.lifecycleRotationHandlers?.values() ?? [], - state.lifecycleGeneration, - (error) => errors.push(error), + notifyListeners(state.lifecycleRotationHandlers?.values() ?? [], lifecycleGeneration, (error) => + errors.push(error), ); if (errors.length > 0) { throw new AggregateError(errors, "Failed to retire stale agent lifecycle owners"); } - return state.lifecycleGeneration; -} - -/** Registers or merges per-run context used by later agent event emissions. */ -export function registerAgentRunContext(runId: string, context: AgentRunContext, claimId?: string) { - if (!runId) { - return; - } - const state = getAgentEventState(); - const lifecycleGeneration = context.lifecycleGeneration ?? state.lifecycleGeneration; - const owners = getAgentRunContextOwners(state).get(runId); - if ( - owners?.lifecycleGeneration === lifecycleGeneration && - owners.exclusiveClaimId && - (owners.exclusiveClaimId !== claimId || owners.clearRequested) - ) { - return; - } - const existing = state.runContextById.get(runId); - if (!existing) { - state.runContextById.set(runId, { - ...context, - lifecycleGeneration: context.lifecycleGeneration ?? state.lifecycleGeneration, - registeredAt: context.registeredAt ?? Date.now(), - }); - return; - } - if ( - context.lifecycleGeneration && - existing.lifecycleGeneration && - context.lifecycleGeneration !== existing.lifecycleGeneration - ) { - return; - } - if (context.sessionKey && existing.sessionKey !== context.sessionKey) { - existing.sessionKey = context.sessionKey; - } - if (context.sessionId && existing.sessionId !== context.sessionId) { - existing.sessionId = context.sessionId; - } - if (context.agentId && existing.agentId !== context.agentId) { - existing.agentId = context.agentId; - } - if (context.verboseLevel && existing.verboseLevel !== context.verboseLevel) { - existing.verboseLevel = context.verboseLevel; - } - if (context.isControlUiVisible !== undefined) { - existing.isControlUiVisible = context.isControlUiVisible; - } - if (context.projectSessionActive !== undefined) { - existing.projectSessionActive = context.projectSessionActive; - } - if (context.projectSessionLifecycle !== undefined) { - existing.projectSessionLifecycle = context.projectSessionLifecycle; - } - if (context.cronRunsByJobId !== undefined) { - existing.cronRunsByJobId ??= new Map(); - for (const [jobId, cronRun] of context.cronRunsByJobId) { - existing.cronRunsByJobId.set(jobId, cronRun); - } - } - if (context.isHeartbeat !== undefined && existing.isHeartbeat !== context.isHeartbeat) { - existing.isHeartbeat = context.isHeartbeat; - } - if (context.registeredAt !== undefined) { - existing.registeredAt = context.registeredAt; - } - if (context.lastActiveAt !== undefined) { - existing.lastActiveAt = context.lastActiveAt; - } -} - -function getAgentRunContextOwners(state = getAgentEventState()) { - state.runContextOwnersById ??= new Map(); - return state.runContextOwnersById; -} - -/** Claims a run id for a newly admitted execution, replacing stale ownership. */ -export function claimAgentRunContext( - runId: string, - context: AgentRunContext, - options: { - /** Adopt a same-generation context only when no tracked execution owns it. */ - adoptExistingUnowned?: boolean; - trackOwner?: boolean; - ownsContext?: boolean; - exclusive?: boolean; - onClearRequested?: (claimId: string) => void; - } = {}, -): string | undefined { - if (!runId) { - return undefined; - } - const state = getAgentEventState(); - const lifecycleGeneration = context.lifecycleGeneration ?? state.lifecycleGeneration; - const existing = state.runContextById.get(runId); - const ownersById = getAgentRunContextOwners(state); - const existingOwners = ownersById.get(runId); - const currentOwners = - existingOwners?.lifecycleGeneration === lifecycleGeneration ? existingOwners : undefined; - const adoptsExistingUnowned = - options.exclusive === true && - options.adoptExistingUnowned === true && - existing?.lifecycleGeneration === lifecycleGeneration && - currentOwners === undefined; - if ( - currentOwners?.exclusiveClaimId || - (options.exclusive && - ((existing?.lifecycleGeneration === lifecycleGeneration && !adoptsExistingUnowned) || - currentOwners !== undefined)) - ) { - return undefined; - } - let claimId: string | undefined; - if (options.trackOwner) { - claimId = randomUUID(); - if (currentOwners) { - currentOwners.claimIds.add(claimId); - if (options.ownsContext) { - currentOwners.preserveAfterRelease = false; - } - if (options.onClearRequested) { - currentOwners.clearListeners ??= new Map(); - currentOwners.clearListeners.set(claimId, options.onClearRequested); - } - } else { - ownersById.set(runId, { - lifecycleGeneration, - claimIds: new Set([claimId]), - preserveAfterRelease: - options.ownsContext !== true && existing?.lifecycleGeneration === lifecycleGeneration, - clearRequested: false, - ...(options.exclusive ? { exclusiveClaimId: claimId } : {}), - ...(options.onClearRequested - ? { clearListeners: new Map([[claimId, options.onClearRequested]]) } - : {}), - }); - } - } else if (existingOwners?.lifecycleGeneration !== lifecycleGeneration) { - // Same-generation untracked claims refresh metadata inside the tracked - // execution. A new lifecycle replaces that ownership outright. - ownersById.delete(runId); - } - if (existing?.lifecycleGeneration === lifecycleGeneration) { - registerAgentRunContext( - runId, - { - ...context, - lifecycleGeneration, - }, - claimId, - ); - return claimId; - } - state.runContextById.set(runId, { - ...context, - lifecycleGeneration, - registeredAt: context.registeredAt ?? Date.now(), - }); - state.seqByRun.delete(runId); - clearAgentRunUsage(runId); - return claimId; -} - -/** Returns the currently registered context for a run, if it has not been cleared or swept. */ -export function getAgentRunContext(runId: string) { - return getAgentEventState().runContextById.get(runId); -} - -/** Records the latest next-check proposal on the matching paced cron run. */ -export function recordCronNextCheckProposal(runId: string, jobId: string, delayMs: number): void { - const context = getAgentEventState().runContextById.get(runId); - const cronRun = context?.cronRunsByJobId?.get(jobId); - if (!cronRun) { - throw new Error("cron next_check is only available to the currently running job"); - } - if (!cronRun.pacingEnabled) { - throw new Error("cron next_check requires pacing on the current job"); - } - cronRun.nextCheckMs = delayMs; -} - -/** Consumes one successful cron run's proposal so it cannot affect a later run. */ -export function consumeCronNextCheckProposal(runId: string, jobId: string): number | undefined { - const context = getAgentEventState().runContextById.get(runId); - const cronRuns = context?.cronRunsByJobId; - const cronRun = cronRuns?.get(jobId); - if (!cronRun) { - return undefined; - } - cronRuns?.delete(jobId); - if (cronRuns?.size === 0 && context) { - delete context.cronRunsByJobId; - } - return cronRun.nextCheckMs; -} - -export function getAgentRunContextOwnerStatus( - runId: string, - claimId: string, - lifecycleGeneration: string, -): "active" | "clear-requested" | undefined { - const state = getAgentEventState(); - const owners = getAgentRunContextOwners(state).get(runId); - if ( - lifecycleGeneration !== state.lifecycleGeneration || - owners?.lifecycleGeneration !== lifecycleGeneration || - !owners.claimIds.has(claimId) - ) { - return undefined; - } - return owners.clearRequested ? "clear-requested" : "active"; -} - -/** Lists active runs bound to one current session identity. */ -export function listAgentRunsForSession(params: { - sessionKey: string; - sessionId?: string; -}): Array<{ runId: string; lifecycleGeneration: string }> { - const currentLifecycleGeneration = getAgentEventState().lifecycleGeneration; - const runs: Array<{ runId: string; lifecycleGeneration: string }> = []; - for (const [runId, context] of getAgentEventState().runContextById) { - const matches = context.sessionId - ? context.sessionId === params.sessionId - : context.sessionKey === params.sessionKey; - if (matches && context.lifecycleGeneration === currentLifecycleGeneration) { - runs.push({ runId, lifecycleGeneration: context.lifecycleGeneration }); - } - } - return runs.toSorted((a, b) => - a.runId === b.runId - ? a.lifecycleGeneration.localeCompare(b.lifecycleGeneration) - : a.runId.localeCompare(b.runId), - ); -} - -export type ProjectedAgentRunIndex = { - sessionKeys: ReadonlySet; - sessionIds: ReadonlySet; -}; - -export function buildProjectedAgentRunIndex(): ProjectedAgentRunIndex { - const state = getAgentEventState(); - const sessionKeys = new Set(); - const sessionIds = new Set(); - for (const context of state.runContextById.values()) { - if ( - context.projectSessionActive !== true || - context.lifecycleGeneration !== state.lifecycleGeneration - ) { - continue; - } - if (context.sessionKey !== undefined) { - sessionKeys.add(context.sessionKey); - } - if (context.sessionId !== undefined) { - sessionIds.add(context.sessionId); - } - } - return { sessionKeys, sessionIds }; -} - -export function hasProjectedAgentRunForSession(params: { - sessionKeys: readonly string[]; - sessionId?: string; - index?: ProjectedAgentRunIndex; -}): boolean { - const index = params.index ?? buildProjectedAgentRunIndex(); - return ( - params.sessionKeys.some((sessionKey) => index.sessionKeys.has(sessionKey)) || - (params.sessionId !== undefined && index.sessionIds.has(params.sessionId)) - ); -} - -/** Clears context and sequence state for a run that has ended or been discarded. */ -export function clearAgentRunContext( - runId: string, - lifecycleGeneration?: string, - claimId?: string, -) { - const state = getAgentEventState(); - const existing = state.runContextById.get(runId); - if (lifecycleGeneration && existing && existing.lifecycleGeneration !== lifecycleGeneration) { - return; - } - const owners = getAgentRunContextOwners(state).get(runId); - if ( - claimId && - (!owners || - (lifecycleGeneration && owners.lifecycleGeneration !== lifecycleGeneration) || - !owners.claimIds.has(claimId)) - ) { - return; - } - // A rejected claimant's cleanup must not evict the exclusive owner. - if (owners?.exclusiveClaimId && owners.exclusiveClaimId !== claimId) { - return; - } - if (owners?.claimIds.size) { - if (!lifecycleGeneration || owners.lifecycleGeneration === lifecycleGeneration) { - owners.clearRequested = true; - for (const [ownerClaimId, listener] of owners.clearListeners ?? []) { - listener(ownerClaimId); - } - } - return; - } - state.runContextById.delete(runId); - state.seqByRun.delete(runId); - clearAgentRunUsage(runId, lifecycleGeneration ?? existing?.lifecycleGeneration); -} - -/** Releases one tracked owner and clears its context after the final owner exits. */ -export function releaseAgentRunContext(runId: string, claimId: string | undefined) { - if (!runId || !claimId) { - return; - } - const state = getAgentEventState(); - const ownersById = getAgentRunContextOwners(state); - const owners = ownersById.get(runId); - if (!owners?.claimIds.delete(claimId)) { - return; - } - owners.clearListeners?.delete(claimId); - if (owners.exclusiveClaimId === claimId) { - owners.exclusiveClaimId = undefined; - } - if (owners.claimIds.size > 0) { - return; - } - ownersById.delete(runId); - if (owners.clearRequested || !owners.preserveAfterRelease) { - clearAgentRunContext(runId, owners.lifecycleGeneration); - } -} - -/** - * Sweep stale run contexts that exceeded the given TTL. - * Guards against orphaned entries when lifecycle "end"/"error" events are missed. - */ -export function sweepStaleRunContexts(maxAgeMs = 30 * 60 * 1000): number { - const state = getAgentEventState(); - const now = Date.now(); - let swept = 0; - for (const [runId, ctx] of state.runContextById.entries()) { - // Use lastActiveAt (refreshed on every event) to avoid sweeping active runs. - // Fall back to registeredAt, then treat missing timestamps as infinitely old. - const lastSeen = ctx.lastActiveAt ?? ctx.registeredAt; - const age = lastSeen ? now - lastSeen : Infinity; - if (age > maxAgeMs) { - state.runContextById.delete(runId); - state.seqByRun.delete(runId); - clearAgentRunUsage(runId, ctx.lifecycleGeneration); - getAgentRunContextOwners(state).delete(runId); - swept++; - } - } - return swept; + return lifecycleGeneration; } function enrichAgentEvent( @@ -589,20 +198,24 @@ function enrichAgentEvent( claimId?: string, ): AgentEventRuntimePayload | undefined { const state = getAgentEventState(); - const owners = getAgentRunContextOwners(state).get(event.runId); + const currentLifecycleGeneration = getAgentRunLifecycleGeneration(); + const owners = getAgentRunContextOwnership(event.runId); if (claimId !== undefined) { if ( - owners?.lifecycleGeneration !== state.lifecycleGeneration || + owners?.lifecycleGeneration !== currentLifecycleGeneration || owners.exclusiveClaimId !== claimId || !owners.claimIds.has(claimId) || owners.clearRequested ) { return undefined; } - } else if (owners?.lifecycleGeneration === state.lifecycleGeneration && owners.exclusiveClaimId) { + } else if ( + owners?.lifecycleGeneration === currentLifecycleGeneration && + owners.exclusiveClaimId + ) { return undefined; } - const context = state.runContextById.get(event.runId); + const context = getAgentRunContext(event.runId); const executionLifecycleGeneration = event.lifecycleGeneration ?? getAgentEventExecutionContext().getStore()?.lifecycleGeneration; const ownedLifecycleGeneration = executionLifecycleGeneration ?? context?.lifecycleGeneration; @@ -613,7 +226,7 @@ function enrichAgentEvent( ) { return undefined; } - if (ownedLifecycleGeneration && ownedLifecycleGeneration !== state.lifecycleGeneration) { + if (ownedLifecycleGeneration && ownedLifecycleGeneration !== currentLifecycleGeneration) { return undefined; } if (hasInvalidLifecycleStartTimestamp(event.stream, event.data)) { @@ -655,7 +268,7 @@ function enrichAgentEvent( event.stream === "lifecycle" ? (event.sessionId ?? context?.sessionId) : event.sessionId; const lifecycleGeneration = event.stream === "lifecycle" - ? (ownedLifecycleGeneration ?? state.lifecycleGeneration) + ? (ownedLifecycleGeneration ?? currentLifecycleGeneration) : ownedLifecycleGeneration; const agentId = event.agentId ?? context?.agentId; const enriched: AgentEventRuntimePayload = { @@ -734,7 +347,7 @@ export function emitAgentAuditEvent(event: Omit if (enriched) { notifyListeners(state.auditListeners, enriched); const phase = event.stream === "lifecycle" ? event.data.phase : undefined; - if ((phase === "end" || phase === "error") && !state.runContextById.has(event.runId)) { + if ((phase === "end" || phase === "error") && !getAgentRunContext(event.runId)) { // Private synthetic runs bypass public terminal cleanup. Release sequence state only // after synchronous audit listeners consume the terminal event and its final ordering. state.seqByRun.delete(event.runId); @@ -762,11 +375,9 @@ export function onAgentAuditEvent(listener: (evt: AgentEventPayload) => void) { export function resetAgentEventsForTest(options?: { preserveListeners?: boolean }) { const state = getAgentEventState(); state.seqByRun.clear(); - resetAgentRunUsageForTest(); + resetAgentRunRegistryForTest(); if (!options?.preserveListeners) { state.listeners.clear(); state.auditListeners.clear(); } - state.runContextById.clear(); - getAgentRunContextOwners(state).clear(); } diff --git a/src/infra/agent-run-registry.ts b/src/infra/agent-run-registry.ts new file mode 100644 index 000000000000..1b2361074e03 --- /dev/null +++ b/src/infra/agent-run-registry.ts @@ -0,0 +1,474 @@ +// Owns process-local agent run context, ownership, and projection state. +import { randomUUID } from "node:crypto"; +import type { VerboseLevel } from "../auto-reply/thinking.js"; +import { resolveGlobalSingleton } from "../shared/global-singleton.js"; +import { clearAgentRunUsage, resetAgentRunUsageForTest } from "./agent-run-usage.js"; + +/** Per-run metadata used to stamp events and gate Control UI visibility. */ +type AgentRunContext = { + sessionKey?: string; + /** Resolved agent owner, including for unscoped session keys. */ + agentId?: string; + /** Owning run's sessionId; stamped onto lifecycle events. */ + sessionId?: string; + /** Gateway lifecycle generation captured when the run was registered. */ + lifecycleGeneration?: string; + /** Producer-owned start captured from this run's accepted lifecycle event. */ + lifecycleStartedAt?: number; + verboseLevel?: VerboseLevel; + isHeartbeat?: boolean; + /** Whether control UI clients should receive chat/agent updates for this run. */ + isControlUiVisible?: boolean; + projectSessionActive?: boolean; + /** Whether lifecycle events may update the shared session row. */ + projectSessionLifecycle?: boolean; + /** Active cadence state by job; admission permits one invocation per job. */ + cronRunsByJobId?: Map; + /** Timestamp when this context was first registered (for TTL-based cleanup). */ + registeredAt?: number; + /** Timestamp of last activity (updated on every emitAgentEvent). */ + lastActiveAt?: number; +}; + +type AgentRunContextOwnership = { + lifecycleGeneration: string; + claimIds: Set; + preserveAfterRelease: boolean; + clearRequested: boolean; + exclusiveClaimId?: string; + clearListeners?: Map void>; +}; + +type AgentRunRegistryState = { + contexts: Map; + owners: Map; + lifecycleGeneration: string; + sequenceResetHandler?: (runId: string) => void; + version: number; +}; + +const AGENT_RUN_REGISTRY_STATE_KEY = Symbol.for("openclaw.agentRunRegistry.state"); + +function getAgentRunRegistryState(): AgentRunRegistryState { + return resolveGlobalSingleton(AGENT_RUN_REGISTRY_STATE_KEY, () => ({ + contexts: new Map(), + owners: new Map(), + lifecycleGeneration: randomUUID(), + version: 0, + })); +} + +function bumpAgentRunIndexVersion(): void { + getAgentRunRegistryState().version += 1; +} + +/** Reads the process-local version of the active-run projection inputs. */ +export function readAgentRunIndexVersion(): number { + return getAgentRunRegistryState().version; +} + +export function getAgentRunLifecycleGeneration(): string { + return getAgentRunRegistryState().lifecycleGeneration; +} + +export function rotateAgentRunRegistryLifecycleGeneration(): string { + const state = getAgentRunRegistryState(); + state.lifecycleGeneration = randomUUID(); + bumpAgentRunIndexVersion(); + return state.lifecycleGeneration; +} + +/** Connects registry cleanup to the event sequencer without reversing ownership. */ +export function registerAgentRunSequenceResetHandler(handler: (runId: string) => void): void { + getAgentRunRegistryState().sequenceResetHandler = handler; +} + +/** Registers or merges per-run context used by later agent event emissions. */ +export function registerAgentRunContext( + runId: string, + context: AgentRunContext, + claimId?: string, +): void { + if (!runId) { + return; + } + const state = getAgentRunRegistryState(); + const lifecycleGeneration = context.lifecycleGeneration ?? state.lifecycleGeneration; + const owners = state.owners.get(runId); + if ( + owners?.lifecycleGeneration === lifecycleGeneration && + owners.exclusiveClaimId && + (owners.exclusiveClaimId !== claimId || owners.clearRequested) + ) { + return; + } + const existing = state.contexts.get(runId); + if (!existing) { + state.contexts.set(runId, { + ...context, + lifecycleGeneration, + registeredAt: context.registeredAt ?? Date.now(), + }); + bumpAgentRunIndexVersion(); + return; + } + if ( + context.lifecycleGeneration && + existing.lifecycleGeneration && + context.lifecycleGeneration !== existing.lifecycleGeneration + ) { + return; + } + let runIndexChanged = false; + if (context.sessionKey && existing.sessionKey !== context.sessionKey) { + existing.sessionKey = context.sessionKey; + runIndexChanged = true; + } + if (context.sessionId && existing.sessionId !== context.sessionId) { + existing.sessionId = context.sessionId; + runIndexChanged = true; + } + if (context.agentId && existing.agentId !== context.agentId) { + existing.agentId = context.agentId; + } + if (context.verboseLevel && existing.verboseLevel !== context.verboseLevel) { + existing.verboseLevel = context.verboseLevel; + } + if (context.isControlUiVisible !== undefined) { + existing.isControlUiVisible = context.isControlUiVisible; + } + if ( + context.projectSessionActive !== undefined && + existing.projectSessionActive !== context.projectSessionActive + ) { + existing.projectSessionActive = context.projectSessionActive; + runIndexChanged = true; + } + if (context.projectSessionLifecycle !== undefined) { + existing.projectSessionLifecycle = context.projectSessionLifecycle; + } + if (context.cronRunsByJobId !== undefined) { + existing.cronRunsByJobId ??= new Map(); + for (const [jobId, cronRun] of context.cronRunsByJobId) { + existing.cronRunsByJobId.set(jobId, cronRun); + } + } + if (context.isHeartbeat !== undefined && existing.isHeartbeat !== context.isHeartbeat) { + existing.isHeartbeat = context.isHeartbeat; + } + if (context.registeredAt !== undefined) { + existing.registeredAt = context.registeredAt; + } + if (context.lastActiveAt !== undefined) { + existing.lastActiveAt = context.lastActiveAt; + } + if (runIndexChanged) { + bumpAgentRunIndexVersion(); + } +} + +/** Claims a run id for a newly admitted execution, replacing stale ownership. */ +export function claimAgentRunContext( + runId: string, + context: AgentRunContext, + options: { + /** Adopt a same-generation context only when no tracked execution owns it. */ + adoptExistingUnowned?: boolean; + trackOwner?: boolean; + ownsContext?: boolean; + exclusive?: boolean; + onClearRequested?: (claimId: string) => void; + } = {}, +): string | undefined { + if (!runId) { + return undefined; + } + const state = getAgentRunRegistryState(); + const lifecycleGeneration = context.lifecycleGeneration ?? state.lifecycleGeneration; + const existing = state.contexts.get(runId); + const existingOwners = state.owners.get(runId); + const currentOwners = + existingOwners?.lifecycleGeneration === lifecycleGeneration ? existingOwners : undefined; + const adoptsExistingUnowned = + options.exclusive === true && + options.adoptExistingUnowned === true && + existing?.lifecycleGeneration === lifecycleGeneration && + currentOwners === undefined; + if ( + currentOwners?.exclusiveClaimId || + (options.exclusive && + ((existing?.lifecycleGeneration === lifecycleGeneration && !adoptsExistingUnowned) || + currentOwners !== undefined)) + ) { + return undefined; + } + let claimId: string | undefined; + if (options.trackOwner) { + claimId = randomUUID(); + if (currentOwners) { + currentOwners.claimIds.add(claimId); + if (options.ownsContext) { + currentOwners.preserveAfterRelease = false; + } + if (options.onClearRequested) { + currentOwners.clearListeners ??= new Map(); + currentOwners.clearListeners.set(claimId, options.onClearRequested); + } + } else { + state.owners.set(runId, { + lifecycleGeneration, + claimIds: new Set([claimId]), + preserveAfterRelease: + options.ownsContext !== true && existing?.lifecycleGeneration === lifecycleGeneration, + clearRequested: false, + ...(options.exclusive ? { exclusiveClaimId: claimId } : {}), + ...(options.onClearRequested + ? { clearListeners: new Map([[claimId, options.onClearRequested]]) } + : {}), + }); + } + } else if (existingOwners?.lifecycleGeneration !== lifecycleGeneration) { + // Same-generation untracked claims refresh metadata inside the tracked + // execution. A new lifecycle replaces that ownership outright. + state.owners.delete(runId); + } + if (existing?.lifecycleGeneration === lifecycleGeneration) { + const versionBeforeRegister = readAgentRunIndexVersion(); + registerAgentRunContext(runId, { ...context, lifecycleGeneration }, claimId); + if (readAgentRunIndexVersion() === versionBeforeRegister) { + bumpAgentRunIndexVersion(); + } + return claimId; + } + state.contexts.set(runId, { + ...context, + lifecycleGeneration, + registeredAt: context.registeredAt ?? Date.now(), + }); + state.sequenceResetHandler?.(runId); + clearAgentRunUsage(runId); + bumpAgentRunIndexVersion(); + return claimId; +} + +/** Returns the currently registered context for a run, if it has not been cleared or swept. */ +export function getAgentRunContext(runId: string): AgentRunContext | undefined { + return getAgentRunRegistryState().contexts.get(runId); +} + +export function getAgentRunContextOwnership(runId: string): AgentRunContextOwnership | undefined { + return getAgentRunRegistryState().owners.get(runId); +} + +/** Records the latest next-check proposal on the matching paced cron run. */ +export function recordCronNextCheckProposal(runId: string, jobId: string, delayMs: number): void { + const context = getAgentRunContext(runId); + const cronRun = context?.cronRunsByJobId?.get(jobId); + if (!cronRun) { + throw new Error("cron next_check is only available to the currently running job"); + } + if (!cronRun.pacingEnabled) { + throw new Error("cron next_check requires pacing on the current job"); + } + cronRun.nextCheckMs = delayMs; +} + +/** Consumes one successful cron run's proposal so it cannot affect a later run. */ +export function consumeCronNextCheckProposal(runId: string, jobId: string): number | undefined { + const context = getAgentRunContext(runId); + const cronRuns = context?.cronRunsByJobId; + const cronRun = cronRuns?.get(jobId); + if (!cronRun) { + return undefined; + } + cronRuns?.delete(jobId); + if (cronRuns?.size === 0 && context) { + delete context.cronRunsByJobId; + } + return cronRun.nextCheckMs; +} + +export function getAgentRunContextOwnerStatus( + runId: string, + claimId: string, + lifecycleGeneration: string, +): "active" | "clear-requested" | undefined { + const state = getAgentRunRegistryState(); + const owners = state.owners.get(runId); + if ( + lifecycleGeneration !== state.lifecycleGeneration || + owners?.lifecycleGeneration !== lifecycleGeneration || + !owners.claimIds.has(claimId) + ) { + return undefined; + } + return owners.clearRequested ? "clear-requested" : "active"; +} + +/** Lists active runs bound to one current session identity. */ +export function listAgentRunsForSession(params: { + sessionKey: string; + sessionId?: string; +}): Array<{ runId: string; lifecycleGeneration: string }> { + const state = getAgentRunRegistryState(); + const runs: Array<{ runId: string; lifecycleGeneration: string }> = []; + for (const [runId, context] of state.contexts) { + const matches = context.sessionId + ? context.sessionId === params.sessionId + : context.sessionKey === params.sessionKey; + if (matches && context.lifecycleGeneration === state.lifecycleGeneration) { + runs.push({ runId, lifecycleGeneration: context.lifecycleGeneration }); + } + } + return runs.toSorted((a, b) => + a.runId === b.runId + ? a.lifecycleGeneration.localeCompare(b.lifecycleGeneration) + : a.runId.localeCompare(b.runId), + ); +} + +export type ProjectedAgentRunIndex = { + sessionKeys: ReadonlySet; + sessionIds: ReadonlySet; +}; + +export function buildProjectedAgentRunIndex(): ProjectedAgentRunIndex { + const state = getAgentRunRegistryState(); + const sessionKeys = new Set(); + const sessionIds = new Set(); + for (const context of state.contexts.values()) { + if ( + context.projectSessionActive !== true || + context.lifecycleGeneration !== state.lifecycleGeneration + ) { + continue; + } + if (context.sessionKey !== undefined) { + sessionKeys.add(context.sessionKey); + } + if (context.sessionId !== undefined) { + sessionIds.add(context.sessionId); + } + } + return { sessionKeys, sessionIds }; +} + +export function hasProjectedAgentRunForSession(params: { + sessionKeys: readonly string[]; + sessionId?: string; + index?: ProjectedAgentRunIndex; +}): boolean { + const index = params.index ?? buildProjectedAgentRunIndex(); + return ( + params.sessionKeys.some((sessionKey) => index.sessionKeys.has(sessionKey)) || + (params.sessionId !== undefined && index.sessionIds.has(params.sessionId)) + ); +} + +/** Clears context state for a run that has ended or been discarded. */ +export function clearAgentRunContext( + runId: string, + lifecycleGeneration?: string, + claimId?: string, +): void { + const state = getAgentRunRegistryState(); + const existing = state.contexts.get(runId); + if (lifecycleGeneration && existing && existing.lifecycleGeneration !== lifecycleGeneration) { + return; + } + const owners = state.owners.get(runId); + if ( + claimId && + (!owners || + (lifecycleGeneration && owners.lifecycleGeneration !== lifecycleGeneration) || + !owners.claimIds.has(claimId)) + ) { + return; + } + // A rejected claimant's cleanup must not evict the exclusive owner. + if (owners?.exclusiveClaimId && owners.exclusiveClaimId !== claimId) { + return; + } + if (owners?.claimIds.size) { + if (!lifecycleGeneration || owners.lifecycleGeneration === lifecycleGeneration) { + const wasClearRequested = owners.clearRequested; + owners.clearRequested = true; + for (const [ownerClaimId, listener] of owners.clearListeners ?? []) { + listener(ownerClaimId); + } + if (!wasClearRequested) { + bumpAgentRunIndexVersion(); + } + } + return; + } + const removed = state.contexts.delete(runId); + state.sequenceResetHandler?.(runId); + clearAgentRunUsage(runId, lifecycleGeneration ?? existing?.lifecycleGeneration); + if (removed) { + bumpAgentRunIndexVersion(); + } +} + +/** Releases one tracked owner and clears its context after the final owner exits. */ +export function releaseAgentRunContext(runId: string, claimId: string | undefined): void { + if (!runId || !claimId) { + return; + } + const state = getAgentRunRegistryState(); + const owners = state.owners.get(runId); + if (!owners?.claimIds.delete(claimId)) { + return; + } + const versionBeforeRelease = readAgentRunIndexVersion(); + owners.clearListeners?.delete(claimId); + if (owners.exclusiveClaimId === claimId) { + owners.exclusiveClaimId = undefined; + } + if (owners.claimIds.size > 0) { + bumpAgentRunIndexVersion(); + return; + } + state.owners.delete(runId); + if (owners.clearRequested || !owners.preserveAfterRelease) { + clearAgentRunContext(runId, owners.lifecycleGeneration); + } + if (readAgentRunIndexVersion() === versionBeforeRelease) { + bumpAgentRunIndexVersion(); + } +} + +/** Sweeps orphaned run contexts that exceeded the given TTL. */ +export function sweepStaleRunContexts(maxAgeMs = 30 * 60 * 1000): number { + const state = getAgentRunRegistryState(); + const now = Date.now(); + let swept = 0; + for (const [runId, context] of state.contexts) { + // Use lastActiveAt (refreshed on every event) to avoid sweeping active runs. + // Fall back to registeredAt, then treat missing timestamps as infinitely old. + const lastSeen = context.lastActiveAt ?? context.registeredAt; + const age = lastSeen ? now - lastSeen : Infinity; + if (age > maxAgeMs) { + state.contexts.delete(runId); + state.sequenceResetHandler?.(runId); + clearAgentRunUsage(runId, context.lifecycleGeneration); + state.owners.delete(runId); + swept += 1; + } + } + if (swept > 0) { + bumpAgentRunIndexVersion(); + } + return swept; +} + +export function resetAgentRunRegistryForTest(): void { + const state = getAgentRunRegistryState(); + const hadRunContexts = state.contexts.size > 0; + resetAgentRunUsageForTest(); + state.contexts.clear(); + state.owners.clear(); + if (hadRunContexts) { + bumpAgentRunIndexVersion(); + } +} diff --git a/src/tasks/task-registry.maintenance.ts b/src/tasks/task-registry.maintenance.ts index 6526f8f6560b..8eb84e717c3b 100644 --- a/src/tasks/task-registry.maintenance.ts +++ b/src/tasks/task-registry.maintenance.ts @@ -21,7 +21,7 @@ import { import type { OpenClawConfig } from "../config/types.openclaw.js"; import { isCronJobActive } from "../cron/active-jobs.js"; import { resolveCronTaskRecordTimestamp } from "../cron/task-run-detail.js"; -import { getAgentRunContext } from "../infra/agent-events.js"; +import { getAgentRunContext } from "../infra/agent-run-registry.js"; import { getSessionBindingService } from "../infra/outbound/session-binding-service.js"; import { createSubsystemLogger } from "../logging/subsystem.js"; import { diff --git a/src/tasks/task-registry.test.ts b/src/tasks/task-registry.test.ts index 9529109bc069..de69120ef084 100644 --- a/src/tasks/task-registry.test.ts +++ b/src/tasks/task-registry.test.ts @@ -4,11 +4,8 @@ import type { AcpSessionStoreEntry } from "../acp/runtime/session-meta.js"; import { startAcpSpawnParentStreamRelay } from "../agents/acp-spawn-parent-stream.js"; import { emitAcpLifecycleStart } from "../agents/command/attempt-execution.js"; import { resetCronActiveJobs } from "../cron/active-jobs.js"; -import { - emitAgentEvent, - registerAgentRunContext, - resetAgentEventsForTest, -} from "../infra/agent-events.js"; +import { emitAgentEvent, resetAgentEventsForTest } from "../infra/agent-events.js"; +import { registerAgentRunContext } from "../infra/agent-run-registry.js"; import { requestHeartbeat, setHeartbeatWakeHandler, diff --git a/src/worker/worker.fault-injection.test.ts b/src/worker/worker.fault-injection.test.ts index 96713974031a..bb2077f91160 100644 --- a/src/worker/worker.fault-injection.test.ts +++ b/src/worker/worker.fault-injection.test.ts @@ -41,13 +41,12 @@ import { } from "../gateway/worker-environments/store.js"; import { createWorkerTranscriptCommitStore } from "../gateway/worker-environments/transcript-commit-store.js"; import { createWorkerTranscriptCommitter } from "../gateway/worker-environments/transcript-commit.js"; +import { getAgentEventLifecycleGeneration, onAgentRuntimeEvent } from "../infra/agent-events.js"; import { claimAgentRunContext, clearAgentRunContext, - getAgentEventLifecycleGeneration, getAgentRunContext, - onAgentRuntimeEvent, -} from "../infra/agent-events.js"; +} from "../infra/agent-run-registry.js"; import { rawDataToString } from "../infra/ws.js"; import type { WorkerProvider, WorkerSshEndpoint } from "../plugins/types.js"; import { diff --git a/test/scripts/audit-seams.test.ts b/test/scripts/audit-seams.test.ts index 34e1c67d9fbd..bc7a4a427d6b 100644 --- a/test/scripts/audit-seams.test.ts +++ b/test/scripts/audit-seams.test.ts @@ -11,7 +11,7 @@ describe("audit-seams cron seam classification", () => { const source = ` import { runCliAgent } from "../../agents/cli-runner.js"; import { runWithModelFallback } from "../../agents/model-fallback-runner.js"; - import { registerAgentRunContext } from "../../infra/agent-events.js"; + import { registerAgentRunContext } from "../../infra/agent-run-registry.js"; import { deliverOutboundPayloads } from "../../infra/outbound/deliver.js"; import { buildOutboundSessionContext } from "../../infra/outbound/session-context.js";