From 09e88f7dcedef3b0ab183008dfbd1dfb50d6bb27 Mon Sep 17 00:00:00 2001 From: Amp Date: Fri, 21 Aug 2026 13:16:13 +0000 Subject: [PATCH] fix(gateway): enforce probe timeout budget --- src/commands/gateway-status.test.ts | 71 ++++++++++++++++++++-- src/commands/gateway-status.ts | 20 ++++--- src/commands/gateway-status/probe-run.ts | 49 +++++++++++++--- src/infra/bonjour-discovery.test.ts | 75 ++++++++++++++++++++++++ src/infra/bonjour-discovery.ts | 38 +++++++++--- 5 files changed, 226 insertions(+), 27 deletions(-) diff --git a/src/commands/gateway-status.test.ts b/src/commands/gateway-status.test.ts index 09773649e015..5cc14d8fb682 100644 --- a/src/commands/gateway-status.test.ts +++ b/src/commands/gateway-status.test.ts @@ -390,6 +390,59 @@ describe("gateway-status command", () => { requireRecord(firstTarget.summary, "first target summary"); }); + it("stops probe phases when discovery exhausts the overall deadline", async () => { + let now = 10_000; + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); + readBestEffortConfig.mockResolvedValueOnce({ + gateway: { mode: "local", auth: { token: "ltok" } }, + } as never); + discoverGatewayBeacons.mockImplementationOnce(async (opts: unknown) => { + const deadlineMs = (opts as { deadlineMs?: number }).deadlineMs; + expect(deadlineMs).toBe(11_000); + now = deadlineMs ?? now; + return []; + }); + const { runtime } = createRuntimeCapture(); + + try { + await expect(runGatewayStatus(runtime, { timeout: "1000", json: true })).rejects.toThrow( + "__exit__:1", + ); + } finally { + nowSpy.mockRestore(); + } + + expect(probeGateway).not.toHaveBeenCalled(); + expect(startSshPortForward).not.toHaveBeenCalled(); + }); + + it("passes positive remaining budgets to tunnel setup and probes", async () => { + let now = 20_000; + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); + resolveSshConfig.mockImplementationOnce(async () => { + now += 200; + return null; + }); + discoverGatewayBeacons.mockImplementationOnce(async () => { + now += 100; + return []; + }); + const { runtime } = createRuntimeCapture(); + + try { + await runGatewayStatus(runtime, { + timeout: "1000", + json: true, + ssh: "me@studio", + }); + } finally { + nowSpy.mockRestore(); + } + + expect(requireSshForwardCall().timeoutMs).toBe(700); + expect(requireProbeCall("ws://127.0.0.1:18789").timeoutMs).toBe(700); + }); + it("does not run Windows LAN firewall diagnostics during fast gateway status", async () => { readBestEffortConfig.mockResolvedValueOnce({ gateway: { @@ -1085,7 +1138,8 @@ describe("gateway-status command", () => { expect(requireProbeCall("ws://127.0.0.1:18789").timeoutMs).toBe(15_000); }); - it("uses --port for the local loopback probe target", async () => { + it("uses --port for the local loopback probe target within the remaining budget", async () => { + const nowSpy = vi.spyOn(Date, "now").mockReturnValue(30_000); const { runtime, runtimeLogs, runtimeErrors } = createRuntimeCapture(); probeGateway.mockClear(); readBestEffortConfig.mockResolvedValueOnce({ @@ -1096,7 +1150,11 @@ describe("gateway-status command", () => { }, } as never); - await runGatewayStatus(runtime, { timeout: "15000", json: true, port: "19080" }); + try { + await runGatewayStatus(runtime, { timeout: "15000", json: true, port: "19080" }); + } finally { + nowSpy.mockRestore(); + } expect(runtimeErrors).toHaveLength(0); expect(requireProbeCall("ws://127.0.0.1:19080").timeoutMs).toBe(15_000); @@ -1147,11 +1205,16 @@ describe("gateway-status command", () => { ]); }); - it("passes the full caller timeout through to active configured remote probes", async () => { + it("bounds active configured remote probes by the remaining caller timeout", async () => { + const nowSpy = vi.spyOn(Date, "now").mockReturnValue(40_000); const { runtime } = createRuntimeCapture(); probeGateway.mockClear(); - await runGatewayStatus(runtime, { timeout: "15000", json: true }); + try { + await runGatewayStatus(runtime, { timeout: "15000", json: true }); + } finally { + nowSpy.mockRestore(); + } const remoteProbeCall = requireProbeCall("wss://remote.example:18789"); expect(remoteProbeCall.timeoutMs).toBe(15_000); diff --git a/src/commands/gateway-status.ts b/src/commands/gateway-status.ts index d6d835bc517a..dbbd9a596813 100644 --- a/src/commands/gateway-status.ts +++ b/src/commands/gateway-status.ts @@ -60,6 +60,7 @@ export async function gatewayStatusCommand( const rich = isRich() && opts.json !== true; const defaultTimeoutMs = 3000; const overallTimeoutMs = parseTimeoutMsWithFallback(opts.timeout, defaultTimeoutMs); + const deadlineMs = Date.now() + overallTimeoutMs; const portOverride = parseGatewayPortOption(opts.port); const wideAreaDomain = resolveWideAreaDiscoveryDomain({ configDomain: cfg.discovery?.wideArea?.domain, @@ -85,13 +86,17 @@ export async function gatewayStatusCommand( } if (sshTarget) { - const resolved = await resolveSshTarget({ - rawTarget: sshTarget, - identity: sshIdentity, - overallTimeoutMs, - loadSshConfigModule, - loadSshTunnelModule, - }); + const remainingMs = Math.floor(deadlineMs - Date.now()); + const resolved = + remainingMs > 0 + ? await resolveSshTarget({ + rawTarget: sshTarget, + identity: sshIdentity, + overallTimeoutMs: remainingMs, + loadSshConfigModule, + loadSshTunnelModule, + }) + : null; if (resolved) { sshTarget = resolved.target; if (!sshIdentity && resolved.identity) { @@ -118,6 +123,7 @@ export async function gatewayStatusCommand( cfg, opts, overallTimeoutMs, + deadlineMs, discoveryTimeoutMs, wideAreaDomain, baseTargets, diff --git a/src/commands/gateway-status/probe-run.ts b/src/commands/gateway-status/probe-run.ts index 675190757128..aa834ef55f7d 100644 --- a/src/commands/gateway-status/probe-run.ts +++ b/src/commands/gateway-status/probe-run.ts @@ -38,6 +38,7 @@ export async function runGatewayStatusProbePass(params: { sshAuto?: boolean; }; overallTimeoutMs: number; + deadlineMs: number; discoveryTimeoutMs: number; wideAreaDomain?: string | null; baseTargets: GatewayStatusTarget[]; @@ -46,6 +47,7 @@ export async function runGatewayStatusProbePass(params: { sshIdentity: string | null; loadSshTunnelModule: () => Promise; localTlsFingerprint?: string; + now?: () => number; }): Promise<{ discovery: GatewayBonjourBeacon[]; probed: GatewayStatusProbedTarget[]; @@ -53,27 +55,39 @@ export async function runGatewayStatusProbePass(params: { sshTunnelStarted: boolean; sshTunnelError: string | null; }> { - const discoveryPromise = discoverGatewayBeacons({ - timeoutMs: params.discoveryTimeoutMs, - wideAreaDomain: params.wideAreaDomain, - }); + const now = params.now ?? Date.now; + const remainingMs = () => Math.floor(params.deadlineMs - now()); + const discoveryBudget = Math.min(params.discoveryTimeoutMs, remainingMs()); + const discoveryPromise = + discoveryBudget > 0 + ? discoverGatewayBeacons({ + timeoutMs: discoveryBudget, + deadlineMs: params.deadlineMs, + wideAreaDomain: params.wideAreaDomain, + now, + }) + : Promise.resolve([]); let sshTarget = params.sshTarget; let sshTunnelError: string | null = null; let sshTunnelStarted = false; const tryStartTunnel = async () => { - if (!sshTarget) { + if (!sshTarget || remainingMs() <= 0) { return null; } try { const { startSshPortForward } = await params.loadSshTunnelModule(); + const tunnelBudget = Math.min(1500, remainingMs()); + if (tunnelBudget <= 0) { + return null; + } const tunnel = await startSshPortForward({ target: sshTarget, identity: params.sshIdentity ?? undefined, localPortPreferred: params.remotePort, remotePort: params.remotePort, - timeoutMs: Math.min(1500, params.overallTimeoutMs), + timeoutMs: tunnelBudget, }); sshTunnelStarted = true; return tunnel; @@ -88,6 +102,15 @@ export async function runGatewayStatusProbePass(params: { const [discovery, tunnelFirst] = await Promise.all([discoveryTask, tunnelTask]); if (!sshTarget && params.opts.sshAuto) { + if (remainingMs() <= 0) { + return { + discovery, + probed: [], + sshTarget, + sshTunnelStarted, + sshTunnelError, + }; + } const { parseSshTarget } = await params.loadSshTunnelModule(); sshTarget = pickAutoSshTargetFromDiscovery({ discovery, @@ -125,10 +148,20 @@ export async function runGatewayStatusProbePass(params: { try { const probed = await Promise.all( targets.map(async (target) => { + if (remainingMs() <= 0) { + return null; + } const authResolution = await resolveAuthForTarget(params.cfg, target, { token: readStringValue(params.opts.token), password: readStringValue(params.opts.password), }); + const probeBudget = Math.min( + resolveProbeBudgetMs(params.overallTimeoutMs, target), + remainingMs(), + ); + if (probeBudget <= 0) { + return null; + } const probe = await probeGateway({ url: target.url, config: params.cfg, @@ -147,7 +180,7 @@ export async function runGatewayStatusProbePass(params: { target.kind === "localLoopback" && target.url.startsWith("wss://") ? params.localTlsFingerprint : undefined, - timeoutMs: resolveProbeBudgetMs(params.overallTimeoutMs, target), + timeoutMs: probeBudget, }); return { target, @@ -161,7 +194,7 @@ export async function runGatewayStatusProbePass(params: { return { discovery, - probed, + probed: probed.filter((entry): entry is GatewayStatusProbedTarget => entry !== null), sshTarget, sshTunnelStarted, sshTunnelError, diff --git a/src/infra/bonjour-discovery.test.ts b/src/infra/bonjour-discovery.test.ts index 43230160c547..f6b3465b7aba 100644 --- a/src/infra/bonjour-discovery.test.ts +++ b/src/infra/bonjour-discovery.test.ts @@ -134,6 +134,8 @@ describe("bonjour-discovery", () => { expect(browseCalls.map((c) => c.argv[3])).toContain("local."); expect(browseCalls.map((c) => c.argv[3])).toContain(WIDE_AREA_DOMAIN); expect([...new Set(browseCalls.map((c) => c.timeoutMs))]).toEqual([1234]); + const resolveCalls = calls.filter((c) => c.argv[0] === "dns-sd" && c.argv[1] === "-L"); + expect([...new Set(resolveCalls.map((c) => c.timeoutMs))]).toEqual([1234]); }); it("decodes dns-sd octal escapes in TXT displayName", async () => { @@ -190,6 +192,79 @@ describe("bonjour-discovery", () => { expect(beacon.txt?.displayName).toBe("Peter’s Mac Studio"); }); + it("shrinks Darwin resolve budgets against an aggregate deadline", async () => { + let now = 1_000; + const resolveBudgets: number[] = []; + const run = vi.fn(async (argv: string[], options: { timeoutMs: number }) => { + if (argv[1] === "-B") { + now += 100; + return { + stdout: [ + "Add 2 3 local. _openclaw-gw._tcp. First Gateway", + "Add 2 3 local. _openclaw-gw._tcp. Second Gateway", + ].join("\n"), + stderr: "", + code: 0, + signal: null, + killed: false, + }; + } + resolveBudgets.push(options.timeoutMs); + now += 300; + return { + stdout: `${argv[2]}._openclaw-gw._tcp. can be reached at gateway.local:18789`, + stderr: "", + code: 0, + signal: null, + killed: false, + }; + }); + + await discoverGatewayBeacons({ + platform: "darwin", + timeoutMs: 1_000, + deadlineMs: 2_000, + domains: ["local."], + run: run as unknown as typeof runCommandWithTimeout, + now: () => now, + }); + + expect(resolveBudgets).toEqual([900, 600]); + }); + + it("does not start another Darwin resolve after the aggregate deadline", async () => { + let now = 1_000; + const resolvedInstances: string[] = []; + const run = vi.fn(async (argv: string[]) => { + if (argv[1] === "-B") { + return { + stdout: [ + "Add 2 3 local. _openclaw-gw._tcp. First Gateway", + "Add 2 3 local. _openclaw-gw._tcp. Second Gateway", + ].join("\n"), + stderr: "", + code: 0, + signal: null, + killed: false, + }; + } + resolvedInstances.push(argv[2] ?? ""); + now = 2_000; + return { stdout: "", stderr: "", code: 0, signal: null, killed: false }; + }); + + await discoverGatewayBeacons({ + platform: "darwin", + timeoutMs: 1_000, + deadlineMs: 2_000, + domains: ["local."], + run: run as unknown as typeof runCommandWithTimeout, + now: () => now, + }); + + expect(resolvedInstances).toEqual(["First Gateway"]); + }); + it("rejects malformed and out-of-range advertised ports", async () => { const run = vi.fn(async (argv: string[]) => { const domain = argv[3] ?? ""; diff --git a/src/infra/bonjour-discovery.ts b/src/infra/bonjour-discovery.ts index 07c141bea3ca..b3c8ff269db7 100644 --- a/src/infra/bonjour-discovery.ts +++ b/src/infra/bonjour-discovery.ts @@ -66,10 +66,12 @@ export function resolveGatewayDiscoveryEndpoint( type GatewayBonjourDiscoverOpts = { timeoutMs?: number; + deadlineMs?: number; domains?: string[]; wideAreaDomain?: string | null; platform?: NodeJS.Platform; run?: typeof runCommandWithTimeout; + now?: () => number; }; const DEFAULT_TIMEOUT_MS = 2000; @@ -307,15 +309,24 @@ async function discoverViaDnsSd( domain: string, timeoutMs: number, run: typeof runCommandWithTimeout, + remainingMs?: () => number, ): Promise { + const browseBudget = remainingMs ? Math.min(timeoutMs, remainingMs()) : timeoutMs; + if (browseBudget <= 0) { + return []; + } const browse = await run(["dns-sd", "-B", GATEWAY_SERVICE_TYPE, domain], { - timeoutMs, + timeoutMs: browseBudget, }); const instances = parseDnsSdBrowse(browse.stdout); const results: GatewayBonjourBeacon[] = []; for (const instance of instances) { + const resolveBudget = remainingMs ? Math.min(timeoutMs, remainingMs()) : timeoutMs; + if (resolveBudget <= 0) { + break; + } const resolved = await run(["dns-sd", "-L", instance, GATEWAY_SERVICE_TYPE, domain], { - timeoutMs, + timeoutMs: resolveBudget, }); const parsed = parseDnsSdResolve(resolved.stdout, instance); if (parsed) { @@ -329,12 +340,13 @@ async function discoverWideAreaViaTailnetDns( domain: string, timeoutMs: number, run: typeof runCommandWithTimeout, + now: () => number = Date.now, ): Promise { if (!domain || domain === "local.") { return []; } - const startedAt = Date.now(); - const remainingMs = () => timeoutMs - (Date.now() - startedAt); + const startedAt = now(); + const remainingMs = () => timeoutMs - (now() - startedAt); const tailscaleCandidates = ["tailscale", "/Applications/Tailscale.app/Contents/MacOS/Tailscale"]; let ips: string[] = []; @@ -575,6 +587,9 @@ export async function discoverGatewayBeacons( const timeoutMs = opts.timeoutMs ?? DEFAULT_TIMEOUT_MS; const platform = opts.platform ?? process.platform; const run = opts.run ?? runCommandWithTimeout; + const now = opts.now ?? Date.now; + const deadlineMs = opts.deadlineMs; + const remainingMs = deadlineMs === undefined ? undefined : () => Math.floor(deadlineMs - now()); const wideAreaDomain = resolveWideAreaDiscoveryDomain({ configDomain: opts.wideAreaDomain }); const domainsRaw = Array.isArray(opts.domains) ? opts.domains : []; const defaultDomains = ["local.", ...(wideAreaDomain ? [wideAreaDomain] : [])]; @@ -585,7 +600,7 @@ export async function discoverGatewayBeacons( try { if (platform === "darwin") { const perDomain = await Promise.allSettled( - domains.map(async (domain) => await discoverViaDnsSd(domain, timeoutMs, run)), + domains.map(async (domain) => await discoverViaDnsSd(domain, timeoutMs, run, remainingMs)), ); const discovered = perDomain.flatMap((r) => (r.status === "fulfilled" ? r.value : [])); @@ -595,9 +610,16 @@ export async function discoverGatewayBeacons( : false; if (wantsWideArea && !hasWideArea && wideAreaDomain) { - const fallback = await discoverWideAreaViaTailnetDns(wideAreaDomain, timeoutMs, run).catch( - () => [], - ); + const fallbackBudget = remainingMs ? Math.min(timeoutMs, remainingMs()) : timeoutMs; + if (fallbackBudget <= 0) { + return discovered; + } + const fallback = await discoverWideAreaViaTailnetDns( + wideAreaDomain, + fallbackBudget, + run, + now, + ).catch(() => []); return [...discovered, ...fallback]; }