From 3bdcfaf980d5b3a703c3fbed5e3335e0c3e496fc Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Mon, 13 Jul 2026 21:50:21 +0200 Subject: [PATCH] refactor(agents): move heartbeat wrapper to its owner --- .../run/attempt-tool-catalog.ts | 37 +----------------- .../run/tool-activity-heartbeat.test.ts | 2 +- .../run/tool-activity-heartbeat.ts | 39 ++++++++++++++++++- 3 files changed, 40 insertions(+), 38 deletions(-) diff --git a/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts b/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts index ece289e112d7..eeadbdf0b34d 100644 --- a/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts +++ b/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts @@ -2,11 +2,7 @@ * Prepares the attempt-local tool catalog, schema projection, and diagnostics. */ import type { DiagnosticTraceContext } from "../../../infra/diagnostic-trace-context.js"; -import { copyPluginToolMeta } from "../../../plugins/tools.js"; -import { copyBeforeToolCallHookMarker } from "../../agent-tools.before-tool-call.js"; import { resolveToolLoopDetectionConfig } from "../../agent-tools.js"; -import { copyChannelAgentToolMeta } from "../../channel-tools.js"; -import { copyCodeModeControlToolIdentity } from "../../code-mode-control-tools.js"; import { applyCodeModeCatalog, CODE_MODE_EXEC_TOOL_NAME, @@ -30,49 +26,18 @@ import { TOOL_SEARCH_RAW_TOOL_NAME, type ToolSearchCatalogToolExecutor, } from "../../tool-search.js"; -import { copyToolTerminalPresentation } from "../../tool-terminal-presentation.js"; -import type { AnyAgentTool } from "../../tools/common.js"; import { log } from "../logger.js"; import type { prepareEmbeddedAttemptBundleTools } from "./attempt-bundle-tools.js"; import { collectAttemptExplicitToolAllowlistSources } from "./attempt-tool-allowlist.js"; import type { prepareEmbeddedAttemptToolBase } from "./attempt-tool-base-prepare.js"; import { buildToolSearchRunPlan } from "./attempt.tool-search-run-plan.js"; -import { notifyToolActivity } from "./tool-activity-heartbeat.js"; +import { wrapEmbeddedAttemptToolWithActivity } from "./tool-activity-heartbeat.js"; import type { EmbeddedRunAttemptParams } from "./types.js"; type PreparedToolBase = ReturnType; type PreparedBundleTools = Awaited>; type ProviderRuntimeHandle = Parameters[0]["runtimeHandle"]; -export function wrapEmbeddedAttemptToolWithActivity( - tool: T, - runId: string, -): T { - const originalExecute = tool.execute; - const wrappedTool = { - ...tool, - execute: (async (...args: Parameters) => { - // Long-running tools keep the attempt's idle watchdog alive. - const interval = setInterval(() => notifyToolActivity(runId), 60_000); - interval.unref?.(); - try { - notifyToolActivity(runId); - return await originalExecute(...args); - } finally { - clearInterval(interval); - notifyToolActivity(runId); - } - }) as typeof originalExecute, - } as T; - // Tool metadata lives in identity-keyed WeakMaps, so object spread is insufficient. - copyPluginToolMeta(tool, wrappedTool); - copyChannelAgentToolMeta(tool, wrappedTool); - copyBeforeToolCallHookMarker(tool, wrappedTool); - copyToolTerminalPresentation(tool, wrappedTool); - copyCodeModeControlToolIdentity(tool, wrappedTool); - return wrappedTool; -} - export function prepareEmbeddedAttemptToolCatalog(input: { attempt: EmbeddedRunAttemptParams; preparedToolBase: PreparedToolBase; diff --git a/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.test.ts b/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.test.ts index 9b32b16bd92e..996ea9f14d41 100644 --- a/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.test.ts +++ b/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.test.ts @@ -6,12 +6,12 @@ import { getToolTerminalPresentation, setToolTerminalPresentation, } from "../../tool-terminal-presentation.js"; -import { wrapEmbeddedAttemptToolWithActivity } from "./attempt-tool-catalog.js"; import { clearToolActivityRun, getLastToolActivityMs, notifyToolActivity, onToolActivity, + wrapEmbeddedAttemptToolWithActivity, } from "./tool-activity-heartbeat.js"; const RUN = "test-run"; diff --git a/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.ts b/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.ts index 4f1790fc7376..10a50a2ffc84 100644 --- a/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.ts +++ b/src/agents/embedded-agent-runner/run/tool-activity-heartbeat.ts @@ -1,6 +1,43 @@ -export { +import { copyPluginToolMeta } from "../../../plugins/tools.js"; +import { clearToolActivityRun, getLastToolActivityMs, notifyToolActivity, onToolActivity, } from "../../../shared/tool-activity-heartbeat.js"; +import { copyBeforeToolCallHookMarker } from "../../agent-tools.before-tool-call.js"; +import { copyChannelAgentToolMeta } from "../../channel-tools.js"; +import { copyCodeModeControlToolIdentity } from "../../code-mode-control-tools.js"; +import { copyToolTerminalPresentation } from "../../tool-terminal-presentation.js"; +import type { AnyAgentTool } from "../../tools/common.js"; + +export { clearToolActivityRun, getLastToolActivityMs, notifyToolActivity, onToolActivity }; + +export function wrapEmbeddedAttemptToolWithActivity( + tool: T, + runId: string, +): T { + const originalExecute = tool.execute; + const wrappedTool = { + ...tool, + execute: (async (...args: Parameters) => { + // Long-running tools keep the attempt's idle watchdog alive. + const interval = setInterval(() => notifyToolActivity(runId), 60_000); + interval.unref?.(); + try { + notifyToolActivity(runId); + return await originalExecute(...args); + } finally { + clearInterval(interval); + notifyToolActivity(runId); + } + }) as typeof originalExecute, + } as T; + // Tool metadata lives in identity-keyed WeakMaps, so object spread is insufficient. + copyPluginToolMeta(tool, wrappedTool); + copyChannelAgentToolMeta(tool, wrappedTool); + copyBeforeToolCallHookMarker(tool, wrappedTool); + copyToolTerminalPresentation(tool, wrappedTool); + copyCodeModeControlToolIdentity(tool, wrappedTool); + return wrappedTool; +}