diff --git a/src/plugins/runtime/runtime-task-test-harness.ts b/src/plugins/runtime/runtime-task-test-harness.ts index 1b0562cc81c..72eb9f62c5b 100644 --- a/src/plugins/runtime/runtime-task-test-harness.ts +++ b/src/plugins/runtime/runtime-task-test-harness.ts @@ -1,7 +1,9 @@ import { vi } from "vitest"; import { + resetTaskRegistryControlRuntimeForTests, resetTaskRegistryDeliveryRuntimeForTests, resetTaskRegistryForTests, + setTaskRegistryControlRuntimeForTests, setTaskRegistryDeliveryRuntimeForTests, } from "../../tasks/runtime-internal.js"; import { resetTaskFlowRegistryForTests } from "../../tasks/task-flow-runtime-internal.js"; @@ -12,13 +14,6 @@ const runtimeTaskMocks = vi.hoisted(() => ({ killSubagentRunAdminMock: vi.fn(), })); -vi.mock("../../tasks/task-registry-control.runtime.js", () => ({ - getAcpSessionManager: () => ({ - cancelSession: runtimeTaskMocks.cancelSessionMock, - }), - killSubagentRunAdmin: (params: unknown) => runtimeTaskMocks.killSubagentRunAdminMock(params), -})); - export function getRuntimeTaskMocks() { return runtimeTaskMocks; } @@ -27,11 +22,18 @@ export function installRuntimeTaskDeliveryMock(): void { setTaskRegistryDeliveryRuntimeForTests({ sendMessage: runtimeTaskMocks.sendMessageMock, }); + setTaskRegistryControlRuntimeForTests({ + getAcpSessionManager: () => ({ + cancelSession: runtimeTaskMocks.cancelSessionMock, + }), + killSubagentRunAdmin: (params: unknown) => runtimeTaskMocks.killSubagentRunAdminMock(params), + }); } export function resetRuntimeTaskTestState( taskRegistryOptions?: Parameters[0], ): void { + resetTaskRegistryControlRuntimeForTests(); resetTaskRegistryDeliveryRuntimeForTests(); resetTaskRegistryForTests(taskRegistryOptions); resetTaskFlowRegistryForTests({ persist: false }); diff --git a/src/tasks/runtime-internal.ts b/src/tasks/runtime-internal.ts index d1c7ea55641..c2b245002f8 100644 --- a/src/tasks/runtime-internal.ts +++ b/src/tasks/runtime-internal.ts @@ -3,6 +3,7 @@ export { createTaskRecord, deleteTaskRecordById, ensureTaskRegistryReady, + resetTaskRegistryControlRuntimeForTests, findLatestTaskForOwnerKey, findLatestTaskForFlowId, findLatestTaskForRelatedSessionKey, @@ -25,6 +26,7 @@ export { resolveTaskForLookupToken, resetTaskRegistryForTests, isParentFlowLinkError, + setTaskRegistryControlRuntimeForTests, setTaskRegistryDeliveryRuntimeForTests, setTaskCleanupAfterById, setTaskProgressById, diff --git a/src/tasks/task-registry.ts b/src/tasks/task-registry.ts index 30f51e6edd1..25e82f97913 100644 --- a/src/tasks/task-registry.ts +++ b/src/tasks/task-registry.ts @@ -63,11 +63,19 @@ type TaskRegistryDeliveryRuntime = Pick< typeof import("./task-registry-delivery-runtime.js"), "sendMessage" >; +type TaskRegistryControlRuntime = Pick< + typeof import("./task-registry-control.runtime.js"), + "getAcpSessionManager" | "killSubagentRunAdmin" +>; const TASK_REGISTRY_DELIVERY_RUNTIME_OVERRIDE_KEY = Symbol.for( "openclaw.taskRegistry.deliveryRuntimeOverride", ); +const TASK_REGISTRY_CONTROL_RUNTIME_OVERRIDE_KEY = Symbol.for( + "openclaw.taskRegistry.controlRuntimeOverride", +); type TaskRegistryGlobalWithDeliveryOverride = typeof globalThis & { [TASK_REGISTRY_DELIVERY_RUNTIME_OVERRIDE_KEY]?: TaskRegistryDeliveryRuntime | null; + [TASK_REGISTRY_CONTROL_RUNTIME_OVERRIDE_KEY]?: TaskRegistryControlRuntime | null; }; let deliveryRuntimePromise: Promise | null = null; @@ -379,6 +387,12 @@ function loadTaskRegistryDeliveryRuntime() { } function loadTaskRegistryControlRuntime() { + const controlRuntimeOverride = (globalThis as TaskRegistryGlobalWithDeliveryOverride)[ + TASK_REGISTRY_CONTROL_RUNTIME_OVERRIDE_KEY + ]; + if (controlRuntimeOverride) { + return Promise.resolve(controlRuntimeOverride); + } // Registry reads happen far more often than task cancellation, so keep the ACP/subagent // control graph off the default import path until a cancellation flow actually needs it. controlRuntimePromise ??= import("./task-registry-control.runtime.js"); @@ -1976,3 +1990,17 @@ export function setTaskRegistryDeliveryRuntimeForTests(runtime: TaskRegistryDeli ] = runtime; deliveryRuntimePromise = null; } + +export function resetTaskRegistryControlRuntimeForTests() { + (globalThis as TaskRegistryGlobalWithDeliveryOverride)[ + TASK_REGISTRY_CONTROL_RUNTIME_OVERRIDE_KEY + ] = null; + controlRuntimePromise = null; +} + +export function setTaskRegistryControlRuntimeForTests(runtime: TaskRegistryControlRuntime): void { + (globalThis as TaskRegistryGlobalWithDeliveryOverride)[ + TASK_REGISTRY_CONTROL_RUNTIME_OVERRIDE_KEY + ] = runtime; + controlRuntimePromise = null; +}