fix(gateway): enforce probe timeout budget

This commit is contained in:
Amp
2026-08-21 13:16:13 +00:00
parent 75c44b2b98
commit 09e88f7dce
5 changed files with 226 additions and 27 deletions
+67 -4
View File
@@ -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);
+13 -7
View File
@@ -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,
+41 -8
View File
@@ -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<typeof import("../../infra/ssh-tunnel.js")>;
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,
+75
View File
@@ -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("Peters 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] ?? "";
+30 -8
View File
@@ -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<GatewayBonjourBeacon[]> {
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<GatewayBonjourBeacon[]> {
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];
}