From db219052918f3673fed60c55136a1c2972a0f3e1 Mon Sep 17 00:00:00 2001 From: Bryan Marty Date: Wed, 4 Mar 2026 17:53:36 +0000 Subject: [PATCH] refactor: address code review feedback on restart sentinel fix - Extract parseDeliveryContextFromParams() into restart-request.ts and import it in both config.ts and update.ts, eliminating the duplicated inline IIFE parsing in update.ts - Add comment in gateway-tool.ts explaining why agentThreadId is intentionally excluded from liveDeliveryContextForRpc: threadId is reliably derived server-side from the session key via parseSessionThreadInfo() and is not subject to heartbeat contamination - Add beforeEach(vi.clearAllMocks) to server-restart-sentinel.test.ts and remove ad-hoc mockClear() calls from individual tests to prevent mock state from leaking between test cases --- src/agents/tools/gateway-tool.ts | 6 ++++ src/gateway/server-methods/config.ts | 27 +-------------- src/gateway/server-methods/restart-request.ts | 33 +++++++++++++++++++ src/gateway/server-methods/update.ts | 17 ++-------- src/gateway/server-restart-sentinel.test.ts | 10 +++--- 5 files changed, 47 insertions(+), 46 deletions(-) diff --git a/src/agents/tools/gateway-tool.ts b/src/agents/tools/gateway-tool.ts index d4b7b707dd6..48f924f0ae4 100644 --- a/src/agents/tools/gateway-tool.ts +++ b/src/agents/tools/gateway-tool.ts @@ -157,6 +157,12 @@ export function createGatewayTool(opts?: { // accurate sentinel without reading the (potentially stale) session // store. The store is frequently overwritten by heartbeat runs to // { channel: "webchat", to: "heartbeat" }. See #18612. + // + // Note: agentThreadId is intentionally excluded here. threadId is + // reliably derived server-side from the session key (via + // parseSessionThreadInfo), which encodes it as :thread:N or :topic:N. + // That parsing is not subject to heartbeat contamination, so there is + // no need to forward it through the RPC params. const liveDeliveryContextForRpc = opts?.agentChannel != null && String(opts.agentChannel).trim() ? { diff --git a/src/gateway/server-methods/config.ts b/src/gateway/server-methods/config.ts index c6a55a87315..20cb8e37da9 100644 --- a/src/gateway/server-methods/config.ts +++ b/src/gateway/server-methods/config.ts @@ -52,7 +52,7 @@ import { validateConfigSetParams, } from "../protocol/index.js"; import { resolveBaseHashParam } from "./base-hash.js"; -import { parseRestartRequestParams } from "./restart-request.js"; +import { parseDeliveryContextFromParams, parseRestartRequestParams } from "./restart-request.js"; import type { GatewayRequestHandlers, RespondFn } from "./types.js"; import { assertValidParams } from "./validation.js"; @@ -212,31 +212,6 @@ function resolveConfigRestartRequest(params: unknown): { }; } -function parseDeliveryContextFromParams( - params: unknown, -): { channel?: string; to?: string; accountId?: string } | undefined { - const raw = (params as { deliveryContext?: unknown }).deliveryContext; - if (!raw || typeof raw !== "object") { - return undefined; - } - const channel = - typeof (raw as { channel?: unknown }).channel === "string" - ? (raw as { channel: string }).channel.trim() || undefined - : undefined; - const to = - typeof (raw as { to?: unknown }).to === "string" - ? (raw as { to: string }).to.trim() || undefined - : undefined; - const accountId = - typeof (raw as { accountId?: unknown }).accountId === "string" - ? (raw as { accountId: string }).accountId.trim() || undefined - : undefined; - if (!channel && !to) { - return undefined; - } - return { channel, to, accountId }; -} - function buildConfigRestartSentinelPayload(params: { kind: RestartSentinelPayload["kind"]; mode: string; diff --git a/src/gateway/server-methods/restart-request.ts b/src/gateway/server-methods/restart-request.ts index f8b2ddb8c0d..e07395803aa 100644 --- a/src/gateway/server-methods/restart-request.ts +++ b/src/gateway/server-methods/restart-request.ts @@ -1,3 +1,36 @@ +/** + * Parse the live deliveryContext passed by gateway-tool clients. + * + * Clients capture delivery context from the active agent run and forward it + * so server-side handlers can write an accurate sentinel without reading the + * persisted session store, which heartbeat runs frequently overwrite to + * { channel: "webchat", to: "heartbeat" }. See #18612. + */ +export function parseDeliveryContextFromParams( + params: unknown, +): { channel?: string; to?: string; accountId?: string } | undefined { + const raw = (params as { deliveryContext?: unknown }).deliveryContext; + if (!raw || typeof raw !== "object") { + return undefined; + } + const channel = + typeof (raw as { channel?: unknown }).channel === "string" + ? (raw as { channel: string }).channel.trim() || undefined + : undefined; + const to = + typeof (raw as { to?: unknown }).to === "string" + ? (raw as { to: string }).to.trim() || undefined + : undefined; + const accountId = + typeof (raw as { accountId?: unknown }).accountId === "string" + ? (raw as { accountId: string }).accountId.trim() || undefined + : undefined; + if (!channel && !to) { + return undefined; + } + return { channel, to, accountId }; +} + export function parseRestartRequestParams(params: unknown): { sessionKey: string | undefined; note: string | undefined; diff --git a/src/gateway/server-methods/update.ts b/src/gateway/server-methods/update.ts index 148b80af1e2..4b6a483bb99 100644 --- a/src/gateway/server-methods/update.ts +++ b/src/gateway/server-methods/update.ts @@ -11,7 +11,7 @@ import { normalizeUpdateChannel } from "../../infra/update-channels.js"; import { runGatewayUpdate } from "../../infra/update-runner.js"; import { formatControlPlaneActor, resolveControlPlaneActor } from "../control-plane-audit.js"; import { validateUpdateRunParams } from "../protocol/index.js"; -import { parseRestartRequestParams } from "./restart-request.js"; +import { parseDeliveryContextFromParams, parseRestartRequestParams } from "./restart-request.js"; import type { GatewayRequestHandlers } from "./types.js"; import { assertValidParams } from "./validation.js"; @@ -23,21 +23,8 @@ export const updateHandlers: GatewayRequestHandlers = { const actor = resolveControlPlaneActor(client); const { sessionKey, note, restartDelayMs } = parseRestartRequestParams(params); // Prefer live deliveryContext from params over extractDeliveryInfo() (see #18612). - const paramsDeliveryContextRaw = (params as { deliveryContext?: unknown }).deliveryContext; - const paramsDeliveryContext = - paramsDeliveryContextRaw && typeof paramsDeliveryContextRaw === "object" - ? (() => { - const dc = paramsDeliveryContextRaw as Record; - const channel = - typeof dc.channel === "string" ? dc.channel.trim() || undefined : undefined; - const to = typeof dc.to === "string" ? dc.to.trim() || undefined : undefined; - const accountId = - typeof dc.accountId === "string" ? dc.accountId.trim() || undefined : undefined; - return channel || to ? { channel, to, accountId } : undefined; - })() - : undefined; const { deliveryContext: extractedDeliveryContext, threadId } = extractDeliveryInfo(sessionKey); - const deliveryContext = paramsDeliveryContext ?? extractedDeliveryContext; + const deliveryContext = parseDeliveryContextFromParams(params) ?? extractedDeliveryContext; const timeoutMsRaw = (params as { timeoutMs?: unknown }).timeoutMs; const timeoutMs = typeof timeoutMsRaw === "number" && Number.isFinite(timeoutMsRaw) diff --git a/src/gateway/server-restart-sentinel.test.ts b/src/gateway/server-restart-sentinel.test.ts index bcb7a57ebd7..c247e8d0d6a 100644 --- a/src/gateway/server-restart-sentinel.test.ts +++ b/src/gateway/server-restart-sentinel.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ resolveSessionAgentId: vi.fn(() => "agent-from-key"), @@ -84,6 +84,10 @@ vi.mock("../infra/system-events.js", () => ({ const { scheduleRestartSentinelWake } = await import("./server-restart-sentinel.js"); describe("scheduleRestartSentinelWake", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + it("calls agentCommand with resolved channel, to, and sessionKey after restart", async () => { await scheduleRestartSentinelWake({ deps: {} as never }); @@ -106,7 +110,6 @@ describe("scheduleRestartSentinelWake", () => { it("falls back to enqueueSystemEvent when agentCommand throws", async () => { mocks.agentCommand.mockRejectedValueOnce(new Error("agent failed")); - mocks.enqueueSystemEvent.mockClear(); await scheduleRestartSentinelWake({ deps: {} as never }); @@ -121,8 +124,6 @@ describe("scheduleRestartSentinelWake", () => { ok: false, error: new Error("no-target"), } as never); - mocks.agentCommand.mockClear(); - mocks.enqueueSystemEvent.mockClear(); await scheduleRestartSentinelWake({ deps: {} as never }); @@ -134,7 +135,6 @@ describe("scheduleRestartSentinelWake", () => { it("falls back to enqueueSystemEvent on main session key when sentinel has no sessionKey", async () => { mocks.consumeRestartSentinel.mockResolvedValueOnce({ payload: { sessionKey: "" } } as never); - mocks.enqueueSystemEvent.mockClear(); await scheduleRestartSentinelWake({ deps: {} as never });