test(tasks): add control runtime override seam

Author: vincentkoc

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<typeof resetTaskRegistryForTests>[0],
 ): void {
+  resetTaskRegistryControlRuntimeForTests();
   resetTaskRegistryDeliveryRuntimeForTests();
   resetTaskRegistryForTests(taskRegistryOptions);
   resetTaskFlowRegistryForTests({ persist: false });

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,

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<typeof import("./task-registry-delivery-runtime.js")> | 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;
+}