From be78094ddd2f4258eee0ae519c5239b688abdef4 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Wed, 19 Aug 2026 02:03:37 -0700 Subject: [PATCH] fix(agents): retain beta subagent gateway context --- .../announce/subagent-announce-delivery.ts | 2 + .../subagent-announce-direct-delivery.ts | 4 ++ ...ent-announce.requester-settle-wake.test.ts | 12 +++- ...subagent-announce.requester-settle-wake.ts | 2 + .../subagents/announce/subagent-announce.ts | 2 + .../subagent-gateway-context-binding.test.ts | 61 +++++++++++++++++++ ...ent-registry-lifecycle-announce-cleanup.ts | 2 + .../subagent-registry-lifecycle-wake.ts | 5 ++ .../subagent-registry-lifecycle.test.ts | 5 ++ .../registry/subagent-registry-run-launch.ts | 4 ++ .../registry/subagent-registry-run-manager.ts | 2 + .../subagent-registry-run-recovery.ts | 10 +++ .../registry/subagent-registry.test.ts | 15 ++++- .../subagents/spawn/subagent-spawn-deps.ts | 3 + .../subagents/spawn/subagent-spawn.runtime.ts | 1 + .../spawn/subagent-spawn.test-helpers.ts | 2 + .../spawn/subagent-spawn.test-support.ts | 1 + .../subagents/spawn/subagent-spawn.test.ts | 11 +++- src/agents/subagents/spawn/subagent-spawn.ts | 5 +- src/gateway/server-methods/shared-types.ts | 2 + ...-in-process-dispatch.authorization.test.ts | 18 ++++++ .../server-plugin-in-process-dispatch.ts | 8 ++- src/plugins/runtime/gateway-request-scope.ts | 33 ++++++++++ 23 files changed, 203 insertions(+), 7 deletions(-) create mode 100644 src/agents/subagents/registry/subagent-gateway-context-binding.test.ts diff --git a/src/agents/subagents/announce/subagent-announce-delivery.ts b/src/agents/subagents/announce/subagent-announce-delivery.ts index 46fa16f07b29..efcf3b74d7fe 100644 --- a/src/agents/subagents/announce/subagent-announce-delivery.ts +++ b/src/agents/subagents/announce/subagent-announce-delivery.ts @@ -99,6 +99,7 @@ export async function deliverSubagentAnnouncement(params: { directIdempotencyKey: string; onDeliveryResult?: (delivery: SubagentAnnounceDeliveryResult) => void; signal?: AbortSignal; + resolveGatewayContext?: import("../../../gateway/server-methods/types.js").GatewayContextResolver; }): Promise { const sourceOwnerChanged = () => params.isSourceSessionEffectsAllowed?.() === false; if (sourceOwnerChanged()) { @@ -258,6 +259,7 @@ export async function deliverSubagentAnnouncement(params: { onDeliveryResult: params.onDeliveryResult, signal: params.signal, bestEffortDeliver: params.bestEffortDeliver, + resolveGatewayContext: params.resolveGatewayContext, }); }, }); diff --git a/src/agents/subagents/announce/subagent-announce-direct-delivery.ts b/src/agents/subagents/announce/subagent-announce-direct-delivery.ts index 09004aff5e18..618ec9a44c43 100644 --- a/src/agents/subagents/announce/subagent-announce-direct-delivery.ts +++ b/src/agents/subagents/announce/subagent-announce-direct-delivery.ts @@ -73,6 +73,7 @@ async function runAnnounceAgentCall(params: { delegatedToolPolicyHandoff?: SubagentCompletionToolHandoffRegistration; expectFinal?: boolean; timeoutMs?: number; + resolveGatewayContext?: import("../../../gateway/server-methods/types.js").GatewayContextResolver; }): Promise { return await dispatchSubagentAnnounceAgent(params.agentParams, { expectFinal: params.expectFinal, @@ -81,6 +82,7 @@ async function runAnnounceAgentCall(params: { ), delegatedToolPolicyHandoff: params.delegatedToolPolicyHandoff, timeoutMs: params.timeoutMs, + resolveGatewayContext: params.resolveGatewayContext, }); } @@ -105,6 +107,7 @@ export async function sendSubagentAnnounceDirectly(params: { requesterIsSubagent: boolean; onDeliveryResult?: (delivery: SubagentAnnounceDeliveryResult) => void; signal?: AbortSignal; + resolveGatewayContext?: import("../../../gateway/server-methods/types.js").GatewayContextResolver; }): Promise { if (params.signal?.aborted) { return { @@ -369,6 +372,7 @@ export async function sendSubagentAnnounceDirectly(params: { : undefined, expectFinal: true, timeoutMs: announceTimeoutMs, + resolveGatewayContext: params.resolveGatewayContext, }); }, }); diff --git a/src/agents/subagents/announce/subagent-announce.requester-settle-wake.test.ts b/src/agents/subagents/announce/subagent-announce.requester-settle-wake.test.ts index 8c3898896a5a..8fb9d667840b 100644 --- a/src/agents/subagents/announce/subagent-announce.requester-settle-wake.test.ts +++ b/src/agents/subagents/announce/subagent-announce.requester-settle-wake.test.ts @@ -2,6 +2,7 @@ // drain gating, batch idempotency, and the guards that keep the wake out of // nested/cron/single-delivered paths. import { beforeEach, describe, expect, it, vi } from "vitest"; +import { bindGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import type { SubagentRunRecord } from "../registry/subagent-registry.types.js"; import type { SubagentAnnounceDeliveryResult } from "./subagent-announce-dispatch.js"; @@ -181,7 +182,9 @@ describe("maybeWakeRequesterAfterAllChildrenSettled", () => { }); it("wakes the requester once with a batch-stable idempotency key when the fan-out drains", async () => { - registryRuntimeMock.listSubagentRunsForRequester.mockReturnValue([ + const context = { owner: "gateway-a" } as never; + const resolveGatewayContext = () => context; + const children = [ makeSettledChild({ runId: "run-b", completion: { required: true, resultText: "network findings" }, @@ -190,7 +193,11 @@ describe("maybeWakeRequesterAfterAllChildrenSettled", () => { runId: "run-a", completion: { required: true, resultText: "social findings" }, }), - ]); + ]; + for (const child of children) { + bindGatewayContextResolver(child, resolveGatewayContext); + } + registryRuntimeMock.listSubagentRunsForRequester.mockReturnValue(children); const woke = await maybeWakeRequesterAfterAllChildrenSettled(wakeParams()); @@ -203,6 +210,7 @@ describe("maybeWakeRequesterAfterAllChildrenSettled", () => { expect(call.requireDirectDelivery).toBe(true); expect(call.requireVisibleReply).toBeUndefined(); expect(call.directIdempotencyKey).toBe(requesterSettleKey("run-a,run-b")); + expect((call.resolveGatewayContext as (() => typeof context) | undefined)?.()).toBe(context); const message = String(call.triggerMessage); expect(message).toContain("settled"); expect(message).toContain("social findings"); diff --git a/src/agents/subagents/announce/subagent-announce.requester-settle-wake.ts b/src/agents/subagents/announce/subagent-announce.requester-settle-wake.ts index c0961b5e677e..832d5f007beb 100644 --- a/src/agents/subagents/announce/subagent-announce.requester-settle-wake.ts +++ b/src/agents/subagents/announce/subagent-announce.requester-settle-wake.ts @@ -7,6 +7,7 @@ import { SILENT_REPLY_TOKEN } from "../../../auto-reply/tokens.js"; import { getRuntimeConfig } from "../../../config/config.js"; import { logWarn } from "../../../logger.js"; +import { getSharedGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import { isCronSessionKey } from "../../../sessions/session-key-utils.js"; import { type DeliveryContext, @@ -451,6 +452,7 @@ export async function maybeWakeRequesterAfterAllChildrenSettled(params: { attemptIndex === 0 ? wakeKeyBase : `${wakeKeyBase}:retry-${attemptIndex}`, ), signal: params.signal, + resolveGatewayContext: getSharedGatewayContextResolver(settledBatch), }); } catch (error) { // A transport exception can arrive after gateway admission. Replay the diff --git a/src/agents/subagents/announce/subagent-announce.ts b/src/agents/subagents/announce/subagent-announce.ts index 9036078d1e1b..0472571d9479 100644 --- a/src/agents/subagents/announce/subagent-announce.ts +++ b/src/agents/subagents/announce/subagent-announce.ts @@ -193,6 +193,7 @@ export async function runSubagentAnnounceFlow(params: { bestEffortDeliver?: boolean; onDeliveryResult?: (delivery: SubagentAnnounceDeliveryResult) => void; onBeforeDeleteChildSession?: () => boolean; + resolveGatewayContext?: import("../../../gateway/server-methods/types.js").GatewayContextResolver; }): Promise { let announceOutcome: SubagentAnnounceFlowOutcome = "retryable"; const expectsCompletionMessage = params.expectsCompletionMessage === true; @@ -589,6 +590,7 @@ export async function runSubagentAnnounceFlow(params: { directIdempotencyKey, onDeliveryResult: reportDeliveryResult, signal: params.signal, + resolveGatewayContext: params.resolveGatewayContext, }); reportDeliveryResult(delivery); announceOutcome = delivery.disposition ?? (delivery.delivered ? "delivered" : "retryable"); diff --git a/src/agents/subagents/registry/subagent-gateway-context-binding.test.ts b/src/agents/subagents/registry/subagent-gateway-context-binding.test.ts new file mode 100644 index 000000000000..a5b179d5a57c --- /dev/null +++ b/src/agents/subagents/registry/subagent-gateway-context-binding.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it } from "vitest"; +import { + bindGatewayContextResolver, + clearGatewayContextResolver, + getGatewayContextResolver, + getSharedGatewayContextResolver, +} from "../../../plugins/runtime/gateway-request-scope.js"; +import { createSubagentRunRecord } from "../../subagent-test-fixtures.test-helpers.js"; + +describe("subagent Gateway context binding", () => { + it("keeps successor routing private and clears retired owners", () => { + const context = { owner: "gateway-a" } as never; + const resolver = () => context; + const source = createSubagentRunRecord({ runId: "run-source" }); + const successor = createSubagentRunRecord({ runId: "run-successor" }); + const restored = structuredClone(source); + + bindGatewayContextResolver(source, resolver); + bindGatewayContextResolver(successor, getGatewayContextResolver(source)); + + expect(getGatewayContextResolver(successor)?.()).toBe(context); + expect(getGatewayContextResolver(restored)).toBeUndefined(); + clearGatewayContextResolver(source); + expect(getGatewayContextResolver(source)).toBeUndefined(); + }); + + it("refuses to select one Gateway for a mixed-owner settle batch", () => { + const first = createSubagentRunRecord({ runId: "run-first" }); + const second = createSubagentRunRecord({ runId: "run-second" }); + const firstContext = { owner: "gateway-a" } as never; + const secondContext = { owner: "gateway-b" } as never; + bindGatewayContextResolver(first, () => firstContext); + bindGatewayContextResolver(second, () => secondContext); + + expect(getGatewayContextResolver(first)?.()).toBe(firstContext); + expect(getGatewayContextResolver(second)?.()).toBe(secondContext); + const shared = getSharedGatewayContextResolver([first, second]); + expect(shared).toBeTypeOf("function"); + expect(shared?.()).toBeUndefined(); + }); + + it("shares a Gateway context captured by separate sibling resolvers", () => { + const first = createSubagentRunRecord({ runId: "run-first" }); + const second = createSubagentRunRecord({ runId: "run-second" }); + const context = { owner: "gateway-a" } as never; + bindGatewayContextResolver(first, () => context); + bindGatewayContextResolver(second, () => context); + + expect(getSharedGatewayContextResolver([first, second])?.()).toBe(context); + }); + + it("refuses a mixed bound and unbound settle batch", () => { + const bound = createSubagentRunRecord({ runId: "run-bound" }); + const unbound = createSubagentRunRecord({ runId: "run-unbound" }); + bindGatewayContextResolver(bound, () => ({ owner: "gateway-a" }) as never); + + const shared = getSharedGatewayContextResolver([unbound, bound]); + expect(shared).toBeTypeOf("function"); + expect(shared?.()).toBeUndefined(); + }); +}); diff --git a/src/agents/subagents/registry/subagent-registry-lifecycle-announce-cleanup.ts b/src/agents/subagents/registry/subagent-registry-lifecycle-announce-cleanup.ts index c4dd4e124b04..dff2b6363d4d 100644 --- a/src/agents/subagents/registry/subagent-registry-lifecycle-announce-cleanup.ts +++ b/src/agents/subagents/registry/subagent-registry-lifecycle-announce-cleanup.ts @@ -1,3 +1,4 @@ +import { getGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import { defaultRuntime } from "../../../runtime.js"; import { normalizeDeliveryContext } from "../../../utils/delivery-context.shared.js"; import { @@ -606,6 +607,7 @@ export const startSubagentAnnounceCleanupFlow = ( params.persist(runId); } }, + resolveGatewayContext: getGatewayContextResolver(entry), }; runDetachedCleanupAttempt(context, { runId, diff --git a/src/agents/subagents/registry/subagent-registry-lifecycle-wake.ts b/src/agents/subagents/registry/subagent-registry-lifecycle-wake.ts index 70b4bdefd8b6..57b9502f89c1 100644 --- a/src/agents/subagents/registry/subagent-registry-lifecycle-wake.ts +++ b/src/agents/subagents/registry/subagent-registry-lifecycle-wake.ts @@ -1,4 +1,5 @@ import { runWithoutOwnedSessionTranscriptWrites } from "../../../config/sessions/transcript-write-context.js"; +import { clearGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import { runWithGatewayIndependentRootWorkContinuation, runWithGatewayIndependentRootWorkAdmission, @@ -155,6 +156,7 @@ const completeRequesterSettleWakeBatch = ( context.deleteRequesterSettleWakeTimer(runId); } if (entry.requesterSettleWake === undefined || !params.runs.has(runId)) { + clearGatewayContextResolver(entry); params.resumedRuns.delete(runId); params.clearPendingLifecycleError(runId); } @@ -485,6 +487,7 @@ export function completeCleanupBookkeeping( cleanupParams.entry.terminalOwner = previousTerminalOwner; throw error; } + clearGatewayContextResolver(cleanupParams.entry); scheduleCleanupTails({ allowRetiredRow: false, isDeleteCleanup }); retryDeferredCompletedAnnounces(cleanupParams.runId); return; @@ -506,6 +509,7 @@ export function completeCleanupBookkeeping( params.runs.set(cleanupParams.runId, cleanupParams.entry); throw error; } + clearGatewayContextResolver(cleanupParams.entry); scheduleCleanupTails({ allowRetiredRow: true, isDeleteCleanup }); retryDeferredCompletedAnnounces(cleanupParams.runId); return; @@ -549,6 +553,7 @@ export function completeCleanupBookkeeping( cleanupParams.entry.terminalOwner = previousTerminalOwner; throw error; } + clearGatewayContextResolver(cleanupParams.entry); } scheduleCleanupTails({ allowRetiredRow: false, isDeleteCleanup }); retryDeferredCompletedAnnounces(cleanupParams.runId); diff --git a/src/agents/subagents/registry/subagent-registry-lifecycle.test.ts b/src/agents/subagents/registry/subagent-registry-lifecycle.test.ts index 3b5b641ef2db..f44bfe9c0528 100644 --- a/src/agents/subagents/registry/subagent-registry-lifecycle.test.ts +++ b/src/agents/subagents/registry/subagent-registry-lifecycle.test.ts @@ -8,6 +8,7 @@ import { } from "../../../config/sessions/transcript-write-context.js"; import type { CallGatewayOptions } from "../../../gateway/call.js"; import { getAgentEventLifecycleGeneration } from "../../../infra/agent-events.js"; +import { bindGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import { getActiveGatewayRootWorkCount, markGatewayRestartDraining, @@ -2629,6 +2630,9 @@ describe("subagent registry lifecycle hardening", () => { expectsCompletionMessage: true, }); const runSubagentAnnounceFlow = vi.fn(async () => "delivered" as const); + const gatewayContext = { owner: "gateway-a" } as never; + const resolveGatewayContext = () => gatewayContext; + bindGatewayContextResolver(entry, resolveGatewayContext); const controller = createLifecycleController({ entry, persist, runSubagentAnnounceFlow }); @@ -2641,6 +2645,7 @@ describe("subagent registry lifecycle hardening", () => { expect(browserCleanupArg.onWarn).toBeTypeOf("function"); expectFields(firstCallArg(runSubagentAnnounceFlow), { childSessionKey: entry.childSessionKey, + resolveGatewayContext, }); }); diff --git a/src/agents/subagents/registry/subagent-registry-run-launch.ts b/src/agents/subagents/registry/subagent-registry-run-launch.ts index 3aa4a67314cc..555735e25e04 100644 --- a/src/agents/subagents/registry/subagent-registry-run-launch.ts +++ b/src/agents/subagents/registry/subagent-registry-run-launch.ts @@ -1,9 +1,11 @@ /** Owns subagent registration and queued collector launch transitions. */ +import type { GatewayContextResolver } from "../../../gateway/server-methods/types.js"; import { getAgentEventLifecycleGeneration, isAgentEventLifecycleGenerationCurrent, } from "../../../infra/agent-events.js"; import { createSubsystemLogger } from "../../../logging/subsystem.js"; +import { bindGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import { createQueuedTaskRun, createRunningTaskRun, @@ -87,6 +89,7 @@ export type RegisterSubagentRunParams = { /** Required when direct dispatch suppresses Gateway tracking. Out-of-process launches keep Gateway's existing best-effort CLI policy; other callers create a best-effort row here. */ taskRowOwnership?: "required" | "gateway_best_effort"; + gatewayContextResolver?: GatewayContextResolver; }; export class SubagentLaunchManager extends SubagentRecoveryManager { @@ -266,6 +269,7 @@ export class SubagentLaunchManager extends SubagentRecoveryManager { } } } + bindGatewayContextResolver(entry, registerParams.gatewayContextResolver); // Wait through Gateway RPC; the in-process lifecycle listener is the embedded fallback. activateRegistrationLifecycle(); }; diff --git a/src/agents/subagents/registry/subagent-registry-run-manager.ts b/src/agents/subagents/registry/subagent-registry-run-manager.ts index ca0c07527b4f..07f28c525d50 100644 --- a/src/agents/subagents/registry/subagent-registry-run-manager.ts +++ b/src/agents/subagents/registry/subagent-registry-run-manager.ts @@ -8,6 +8,7 @@ import { isAgentEventLifecycleGenerationCurrent, } from "../../../infra/agent-events.js"; import { createSubsystemLogger } from "../../../logging/subsystem.js"; +import { clearGatewayContextResolver } from "../../../plugins/runtime/gateway-request-scope.js"; import { runWithGatewayIndependentRootWorkAdmission } from "../../../process/gateway-work-admission.js"; import { SUBAGENT_KILL_TASK_ERROR } from "../../../tasks/detached-task-runtime-contract.js"; import { finalizeTaskRunByRunId } from "../../../tasks/detached-task-runtime.js"; @@ -45,6 +46,7 @@ class SubagentRunManager extends SubagentLaunchManager { throw error; } this.options.clearPendingLifecycleError(runId); + clearGatewayContextResolver(entry); if (this.shouldDeleteAttachments(entry)) { void safeRemoveAttachmentsDir(entry); } diff --git a/src/agents/subagents/registry/subagent-registry-run-recovery.ts b/src/agents/subagents/registry/subagent-registry-run-recovery.ts index 1c44d4fc8a9f..c4431d08a855 100644 --- a/src/agents/subagents/registry/subagent-registry-run-recovery.ts +++ b/src/agents/subagents/registry/subagent-registry-run-recovery.ts @@ -1,9 +1,14 @@ /** Owns steer replacement and restart-recovery receipt transitions. */ +import type { GatewayContextResolver } from "../../../gateway/server-methods/types.js"; import { getAgentEventLifecycleGeneration, isAgentEventLifecycleGenerationCurrent, } from "../../../infra/agent-events.js"; import { createSubsystemLogger } from "../../../logging/subsystem.js"; +import { + bindGatewayContextResolver, + getGatewayContextResolver, +} from "../../../plugins/runtime/gateway-request-scope.js"; import { finalizeTaskRunByRunId } from "../../../tasks/detached-task-runtime.js"; import { removeInternalSessionEffectsSession } from "../../internal-session-effects.js"; import type { AgentRunSessionTarget } from "../../run-session-target.js"; @@ -140,6 +145,7 @@ export class SubagentRecoveryManager extends SubagentWaitManager { restartRecovery?: SubagentRestartRecoveryReceipt; lifecycleGeneration?: string; persistenceFailure?: "return-false" | "throw"; + gatewayContextResolver?: GatewayContextResolver; }): boolean => { const previousRunId = replaceParams.previousRunId.trim(); const nextRunId = replaceParams.nextRunId.trim(); @@ -270,6 +276,10 @@ export class SubagentRecoveryManager extends SubagentWaitManager { archiveAtMs: undefined, runTimeoutSeconds, }); + bindGatewayContextResolver( + next, + replaceParams.gatewayContextResolver ?? getGatewayContextResolver(source), + ); clearDeliveryState(next); if (previousRunId !== nextRunId) { diff --git a/src/agents/subagents/registry/subagent-registry.test.ts b/src/agents/subagents/registry/subagent-registry.test.ts index 91951be567e1..546c0126ec12 100644 --- a/src/agents/subagents/registry/subagent-registry.test.ts +++ b/src/agents/subagents/registry/subagent-registry.test.ts @@ -19,7 +19,11 @@ import type { OpenClawConfig } from "../../../config/types.openclaw.js"; import type { GatewayRecoveryRuntime } from "../../../gateway/server-instance-runtime.types.js"; import type { AgentEventPayload } from "../../../infra/agent-events.js"; import { createEmptyPluginRegistry } from "../../../plugins/registry-empty.js"; -import { getPluginRuntimeGatewayRequestScope } from "../../../plugins/runtime/gateway-request-scope.js"; +import { + bindGatewayContextResolver, + getGatewayContextResolver, + getPluginRuntimeGatewayRequestScope, +} from "../../../plugins/runtime/gateway-request-scope.js"; import { getActiveGatewayRootWorkCount, markGatewayRestartDraining, @@ -6461,7 +6465,7 @@ describe("subagent registry seam flow", () => { }, }; }); - mod.addSubagentRunForTests({ + const entry = createSubagentRunRecord({ runId: "run-release-context-engine", childSessionKey: "agent:main:session:child", controllerSessionKey: "agent:main:session:parent", @@ -6479,8 +6483,15 @@ describe("subagent registry seam flow", () => { accumulatedRuntimeMs: 0, cleanupHandled: false, }); + mod.addSubagentRunForTests(entry); + const registeredEntry = mod + .listSubagentRunsForRequester("agent:main:session:parent") + .find((run) => run.runId === "run-release-context-engine"); + expect(registeredEntry).toBeDefined(); + bindGatewayContextResolver(registeredEntry!, () => ({ owner: "gateway-a" }) as never); mod.releaseSubagentRun("run-release-context-engine"); + expect(getGatewayContextResolver(registeredEntry!)).toBeUndefined(); await waitForFast(() => { expect(mocks.onSubagentEnded).toHaveBeenCalledWith({ diff --git a/src/agents/subagents/spawn/subagent-spawn-deps.ts b/src/agents/subagents/spawn/subagent-spawn-deps.ts index e5e6b5c44d14..f7b1590a8357 100644 --- a/src/agents/subagents/spawn/subagent-spawn-deps.ts +++ b/src/agents/subagents/spawn/subagent-spawn-deps.ts @@ -5,6 +5,7 @@ import { ensureContextEnginesInitialized, forkSessionEntryFromParent, getGlobalHookRunner, + getInProcessGatewayRequestContext, getRuntimeConfig, hasInProcessGatewayContext, loadPreparedModelCatalog, @@ -16,6 +17,7 @@ type SubagentSpawnDeps = { dispatchGatewayMethodInProcess: typeof dispatchGatewayMethodInProcess; forkSessionEntryFromParent: typeof forkSessionEntryFromParent; getGlobalHookRunner: () => SubagentLifecycleHookRunner | null; + getInProcessGatewayRequestContext: typeof getInProcessGatewayRequestContext; getRuntimeConfig: typeof getRuntimeConfig; hasInProcessGatewayContext: typeof hasInProcessGatewayContext; ensureContextEnginesInitialized: typeof ensureContextEnginesInitialized; @@ -28,6 +30,7 @@ const defaultSubagentSpawnDeps: SubagentSpawnDeps = { dispatchGatewayMethodInProcess, forkSessionEntryFromParent, getGlobalHookRunner, + getInProcessGatewayRequestContext, getRuntimeConfig, hasInProcessGatewayContext, ensureContextEnginesInitialized, diff --git a/src/agents/subagents/spawn/subagent-spawn.runtime.ts b/src/agents/subagents/spawn/subagent-spawn.runtime.ts index 65b9059dc477..be0605a31634 100644 --- a/src/agents/subagents/spawn/subagent-spawn.runtime.ts +++ b/src/agents/subagents/spawn/subagent-spawn.runtime.ts @@ -14,6 +14,7 @@ export { resolveContextEngine } from "../../../context-engine/registry.js"; export { callGateway } from "../../../gateway/call.js"; export { dispatchGatewayMethodInProcess, + getInProcessGatewayRequestContext, hasInProcessGatewayContext, } from "../../../gateway/server-plugins.js"; export { diff --git a/src/agents/subagents/spawn/subagent-spawn.test-helpers.ts b/src/agents/subagents/spawn/subagent-spawn.test-helpers.ts index 06e256f2b1cf..003a6c744ab4 100644 --- a/src/agents/subagents/spawn/subagent-spawn.test-helpers.ts +++ b/src/agents/subagents/spawn/subagent-spawn.test-helpers.ts @@ -130,6 +130,7 @@ export function expectPersistedRuntimeModel(params: { export async function loadSubagentSpawnModuleForTest(params: { callGatewayMock: MockFn; dispatchGatewayMethodInProcessMock?: MockFn; + getInProcessGatewayRequestContextMock?: MockFn; hasInProcessGatewayContextMock?: MockFn; getRuntimeConfig?: () => Record; loadSessionStoreMock?: MockFn; @@ -219,6 +220,7 @@ export async function loadSubagentSpawnModuleForTest(params: { callGateway: (opts: unknown) => params.callGatewayMock(opts), dispatchGatewayMethodInProcess: (...args: unknown[]) => params.dispatchGatewayMethodInProcessMock?.(...args), + getInProcessGatewayRequestContext: () => params.getInProcessGatewayRequestContextMock?.(), hasInProcessGatewayContext: () => Boolean(params.hasInProcessGatewayContextMock?.()), buildSubagentSystemPrompt: () => "system-prompt", forkSessionEntryFromParent: diff --git a/src/agents/subagents/spawn/subagent-spawn.test-support.ts b/src/agents/subagents/spawn/subagent-spawn.test-support.ts index 72a3d303bb9a..02702ae48d1c 100644 --- a/src/agents/subagents/spawn/subagent-spawn.test-support.ts +++ b/src/agents/subagents/spawn/subagent-spawn.test-support.ts @@ -9,6 +9,7 @@ type SpawnDeps = Omit< | "ensureContextEnginesInitialized" | "forkSessionEntryFromParent" | "getGlobalHookRunner" + | "getInProcessGatewayRequestContext" | "getRuntimeConfig" | "hasInProcessGatewayContext" | "loadPreparedModelCatalog" diff --git a/src/agents/subagents/spawn/subagent-spawn.test.ts b/src/agents/subagents/spawn/subagent-spawn.test.ts index 1dbd645ea0c3..84025aa3909d 100644 --- a/src/agents/subagents/spawn/subagent-spawn.test.ts +++ b/src/agents/subagents/spawn/subagent-spawn.test.ts @@ -28,6 +28,7 @@ const hoisted = vi.hoisted(() => ({ completeCollectorLaunchCleanupMock: vi.fn(), emitSessionLifecycleEventMock: vi.fn(), dispatchGatewayMethodInProcessMock: vi.fn(), + getInProcessGatewayRequestContextMock: vi.fn(), hasInProcessGatewayContextMock: vi.fn(), resolveAgentConfigMock: vi.fn(), resolveContextEngineMock: vi.fn(), @@ -160,6 +161,7 @@ describe("spawnSubagentDirect seam flow", () => { ({ resetSubagentRegistryForTests, spawnSubagentDirect } = await loadSubagentSpawnModuleForTest({ callGatewayMock: hoisted.callGatewayMock, dispatchGatewayMethodInProcessMock: hoisted.dispatchGatewayMethodInProcessMock, + getInProcessGatewayRequestContextMock: hoisted.getInProcessGatewayRequestContextMock, hasInProcessGatewayContextMock: hoisted.hasInProcessGatewayContextMock, getRuntimeConfig: () => hoisted.configOverride, loadSessionStoreMock: hoisted.loadSessionStoreMock, @@ -194,6 +196,7 @@ describe("spawnSubagentDirect seam flow", () => { hoisted.completeCollectorLaunchCleanupMock.mockReset(); hoisted.emitSessionLifecycleEventMock.mockReset(); hoisted.dispatchGatewayMethodInProcessMock.mockReset(); + hoisted.getInProcessGatewayRequestContextMock.mockReset(); hoisted.hasInProcessGatewayContextMock.mockReset().mockReturnValue(false); hoisted.resolveAgentConfigMock.mockReset(); hoisted.resolveContextEngineMock.mockReset().mockResolvedValue({}); @@ -1233,6 +1236,8 @@ describe("spawnSubagentDirect seam flow", () => { }); it("dispatches spawned agent runs in process when a gateway context is available", async () => { + const gatewayContext = { owner: "gateway-a" } as never; + hoisted.getInProcessGatewayRequestContextMock.mockReturnValue(gatewayContext); hoisted.hasInProcessGatewayContextMock.mockReturnValue(true); hoisted.callGatewayMock.mockRejectedValue(new Error("unexpected websocket gateway call")); hoisted.dispatchGatewayMethodInProcessMock.mockImplementation(async (method: string) => { @@ -1274,7 +1279,11 @@ describe("spawnSubagentDirect seam flow", () => { expect(agentOptions.allowSyntheticModelOverride).toBeUndefined(); // In-process dispatch claims the task row directly, unlike ACP's best-effort // registration (see acp-spawn.test.ts). - expect(firstRegisteredSubagentRun().taskRowOwnership).toBe("required"); + const registration = firstRegisteredSubagentRun(); + expect(registration.taskRowOwnership).toBe("required"); + expect( + (registration.gatewayContextResolver as (() => typeof gatewayContext) | undefined)?.(), + ).toBe(gatewayContext); }); it("authorizes explicit model overrides for in-process child launches", async () => { diff --git a/src/agents/subagents/spawn/subagent-spawn.ts b/src/agents/subagents/spawn/subagent-spawn.ts index 3e9569e2b132..c694d667f947 100644 --- a/src/agents/subagents/spawn/subagent-spawn.ts +++ b/src/agents/subagents/spawn/subagent-spawn.ts @@ -51,7 +51,7 @@ import type { SpawnSubagentParams, SpawnSubagentResult, } from "./subagent-spawn-contract.js"; -import { setSubagentSpawnDepsForTest } from "./subagent-spawn-deps.js"; +import { getSubagentSpawnDeps, setSubagentSpawnDepsForTest } from "./subagent-spawn-deps.js"; import { callNativeSubagentGateway, readGatewayRunId } from "./subagent-spawn-gateway.js"; import { buildSubagentLaunchRequest } from "./subagent-spawn-launch-request.js"; import { createSubagentSpawnLifecycleEmitter } from "./subagent-spawn-lifecycle.js"; @@ -105,6 +105,8 @@ export async function spawnSubagentDirect( const requestThreadBinding = params.thread === true; const sandboxMode = params.sandbox === "require" ? "require" : "inherit"; const requesterSessionKey = ctx.agentSessionKey; + const gatewayContext = getSubagentSpawnDeps().getInProcessGatewayRequestContext(); + const gatewayContextResolver = gatewayContext ? () => gatewayContext : undefined; let requestedAgentId = params.agentId?.trim(); const requestResolution = resolveSubagentSpawnRequest(params, ctx, { initial: requestedAgentId, @@ -546,6 +548,7 @@ export async function spawnSubagentDirect( queuedLaunch, queued: params.collect === true, taskRowOwnership, + ...(gatewayContextResolver ? { gatewayContextResolver } : {}), attachmentsDir: attachmentAbsDir, attachmentsRootDir: attachmentRootDir, retainAttachmentsOnKeep: retainOnSessionKeep, diff --git a/src/gateway/server-methods/shared-types.ts b/src/gateway/server-methods/shared-types.ts index 680e6eb210b9..f3c673f78fe8 100644 --- a/src/gateway/server-methods/shared-types.ts +++ b/src/gateway/server-methods/shared-types.ts @@ -386,6 +386,8 @@ type GatewayResidentBridgeContext = { ) => void; }; +export type GatewayContextResolver = () => GatewayRequestContext | undefined; + /** Complete runtime context available to gateway request handlers. */ export type GatewayRequestContext = GatewayKernelContext & GatewayTransportContext & diff --git a/src/gateway/server-plugin-in-process-dispatch.authorization.test.ts b/src/gateway/server-plugin-in-process-dispatch.authorization.test.ts index 7fd53117583c..e4305d9d41c3 100644 --- a/src/gateway/server-plugin-in-process-dispatch.authorization.test.ts +++ b/src/gateway/server-plugin-in-process-dispatch.authorization.test.ts @@ -9,6 +9,10 @@ import { upsertSessionEntryCore } from "../config/sessions/session-accessor.js"; import { withPluginRuntimeGatewayRequestScope } from "../plugins/runtime/gateway-request-scope.js"; import { withOpenClawTestState } from "../test-utils/openclaw-test-state.js"; import type { GatewayRequestContext, GatewayRequestOptions } from "./server-methods/types.js"; +import { + clearFallbackGatewayContext, + setFallbackGatewayContext, +} from "./server-plugin-fallback-context.js"; import { dispatchGatewayMethodInProcess } from "./server-plugin-in-process-dispatch.js"; const startTurn = vi.hoisted(() => vi.fn()); @@ -100,10 +104,24 @@ async function dispatchScopedMethod(params: { describe("typed in-process agent authorization", () => { beforeEach(() => { + clearFallbackGatewayContext(); startTurn.mockReset(); waitForTurn.mockReset(); }); + it("fails closed when an explicit Gateway resolver loses its owner", async () => { + setFallbackGatewayContext(createContext()); + + await expect( + dispatchGatewayMethodInProcess( + "agent", + { message: "detached completion", idempotencyKey: "detached-completion" }, + { resolveGatewayContext: () => undefined }, + ), + ).rejects.toThrow("No scope set and no fallback context available"); + expect(startTurn).not.toHaveBeenCalled(); + }); + it("rejects a scoped agent turn without operator.write", async () => { await expect( dispatchScopedAgent({ diff --git a/src/gateway/server-plugin-in-process-dispatch.ts b/src/gateway/server-plugin-in-process-dispatch.ts index 88f560dcee6b..2b84eadc9657 100644 --- a/src/gateway/server-plugin-in-process-dispatch.ts +++ b/src/gateway/server-plugin-in-process-dispatch.ts @@ -13,6 +13,7 @@ import type { AgentRunRequest } from "./server-methods/agent-request-types.js"; import type { TrustedSessionCreation } from "./server-methods/session-creation-provenance.js"; import type { GatewayAgentRunTaskOwner, + GatewayContextResolver, GatewayRequestContext, GatewayRequestOptions, TrustedAgentToolCaller, @@ -52,6 +53,7 @@ type DispatchGatewayMethodInProcessOptions = { syntheticScopes?: string[]; timeoutMs?: number; signal?: AbortSignal; + resolveGatewayContext?: GatewayContextResolver; }; type ResolvedInProcessGatewayDispatch = { @@ -66,7 +68,11 @@ function resolveInProcessGatewayDispatch( options?: DispatchGatewayMethodInProcessOptions, ): ResolvedInProcessGatewayDispatch { const scope = getPluginRuntimeGatewayRequestScope(); - const context = scope?.context ?? getFallbackGatewayContext(); + const context = + scope?.context ?? + (options?.resolveGatewayContext + ? options.resolveGatewayContext() + : getFallbackGatewayContext()); const isWebchatConnect = scope?.isWebchatConnect ?? (() => false); if (!context) { throw new Error( diff --git a/src/plugins/runtime/gateway-request-scope.ts b/src/plugins/runtime/gateway-request-scope.ts index 53eceaa1b48b..5acfa1798594 100644 --- a/src/plugins/runtime/gateway-request-scope.ts +++ b/src/plugins/runtime/gateway-request-scope.ts @@ -1,6 +1,7 @@ // Gateway request scope tracks request-local plugin runtime context across async work. import { AsyncLocalStorage } from "node:async_hooks"; import type { + GatewayContextResolver, GatewayRequestContext, GatewayRequestOptions, } from "../../gateway/server-methods/types.js"; @@ -37,6 +38,38 @@ const pluginRuntimeGatewayRequestScope = resolveGlobalSingleton< PLUGIN_RUNTIME_GATEWAY_REQUEST_SCOPE_KEY, () => new AsyncLocalStorage(), ); +const gatewayContextResolvers = new WeakMap(); +const resolveNoGatewayContext: GatewayContextResolver = () => undefined; + +export function bindGatewayContextResolver( + owner: object, + resolver: GatewayContextResolver | undefined, +): void { + if (resolver) { + gatewayContextResolvers.set(owner, resolver); + } +} + +export const getGatewayContextResolver = (owner: object) => gatewayContextResolvers.get(owner); + +export const clearGatewayContextResolver = (owner: object) => gatewayContextResolvers.delete(owner); + +export function getSharedGatewayContextResolver( + owners: readonly object[], +): GatewayContextResolver | undefined { + const resolvers = owners.map((owner) => gatewayContextResolvers.get(owner)); + if (resolvers.every((resolver) => resolver === undefined)) { + return undefined; + } + if (resolvers.some((resolver) => resolver === undefined)) { + return resolveNoGatewayContext; + } + const contexts = resolvers.map((resolver) => resolver?.()); + const first = contexts[0]; + return first && contexts.every((context) => context === first) + ? () => first + : resolveNoGatewayContext; +} /** * Runs plugin gateway handlers with request-scoped context that runtime helpers can read.