Skip to content
50 changes: 49 additions & 1 deletion packages/cli/src/cli/commands/core.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1196,9 +1196,16 @@ describe('registerCoreCommands', () => {
});
void runCommand(program, ['up']);

// SIGINT/SIGTERM are now registered before the broker starts (so a
// signal arriving during startup is handled gracefully too), so
// registration alone no longer implies `relay` is set. Wait for the
// broker to actually be up before firing the signal.
for (
let i = 0;
i < 10 && (deps.onSignal as unknown as { mock: { calls: unknown[][] } }).mock.calls.length === 0;
i < 20 &&
!(deps.log as unknown as { mock: { calls: unknown[][] } }).mock.calls.some(
Comment thread
coderabbitai[bot] marked this conversation as resolved.
(call) => call[0] === 'Broker started.'
);
i += 1
) {
await Promise.resolve();
Expand All @@ -1223,6 +1230,47 @@ describe('registerCoreCommands', () => {
expect(logCalls.filter((call) => call[0] === '\nStopping...')).toHaveLength(1);
});

it('up shuts down the in-flight broker candidate when SIGTERM arrives before the status check resolves', async () => {
let resolveStatus: (() => void) | undefined;
const relay = createRelayMock({
getStatus: vi.fn(
() =>
new Promise((resolve) => {
resolveStatus = () => resolve({ agent_count: 0, pending_delivery_count: 0 });
})
),
});
const { program, deps } = createHarness({ relay });
void runCommand(program, ['up']);

// Wait until the status check has actually started. By this point
// `startBrokerWithPortFallback`'s `onCandidateReady` callback has
// already assigned the outer `relay` -- well before the check itself
// resolves. This is exactly the window where `relay` used to still be
// null and a signal would leak the broker child instead of shutting it
// down.
for (
let i = 0;
i < 20 && (relay.getStatus as unknown as { mock: { calls: unknown[][] } }).mock.calls.length === 0;
i += 1
) {
await Promise.resolve();
}
expect(relay.getStatus).toHaveBeenCalled();
expect(relay.shutdown).not.toHaveBeenCalled();

const onSignalMock = deps.onSignal as unknown as { mock: { calls: unknown[][] } };
const sigtermHandler = onSignalMock.mock.calls.find((call) => call[0] === 'SIGTERM')?.[1] as
| (() => Promise<void>)
| undefined;
expect(sigtermHandler).toBeDefined();

await expect((sigtermHandler as () => Promise<void>)()).rejects.toMatchObject({ code: 0 });

expect(relay.shutdown).toHaveBeenCalledTimes(1);
resolveStatus?.();
});

it('down stops broker and cleans stale files', async () => {
const connectionPath = '/tmp/project/.agentworkforce/relay/connection.json';
const relaySockPath = '/tmp/project/.agentworkforce/relay/relay.sock';
Expand Down
74 changes: 73 additions & 1 deletion packages/cli/src/cli/lib/broker-lifecycle.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,13 @@ import {
classifyBrokerStartError,
classifyBrokerStartStage,
describeErrorWithCause,
getBrokerStatusWithRetry,
isBundledBunExecutableEntrypoint,
readNodeDeliveryStatus,
resolveNodeIdentityFromSession,
waitForNodeDelivery,
} from './broker-lifecycle.js';
import type { CoreDependencies } from '../commands/core.js';
import type { CoreDependencies, CoreRelay } from '../commands/core.js';

describe('isBundledBunExecutableEntrypoint', () => {
it.each(['/$bunfs/root/agent-relay', 'B:/~BUN/root/agent-relay.exe', 'B:\\~BUN\\root\\agent-relay.exe'])(
Expand Down Expand Up @@ -128,6 +129,77 @@ describe('classifyBrokerStartStage', () => {
});
});

describe('getBrokerStatusWithRetry', () => {
function createDeps(sleep = vi.fn(async () => undefined)): CoreDependencies {
return { log: vi.fn(), sleep } as unknown as CoreDependencies;
}

it('returns the status on the first successful attempt without sleeping', async () => {
const candidate: Pick<CoreRelay, 'getStatus'> = {
getStatus: vi.fn(async () => ({ agent_count: 0, pending_delivery_count: 0 })),
};
const deps = createDeps();

const result = await getBrokerStatusWithRetry(candidate, deps);

expect(result).toEqual({ agent_count: 0, pending_delivery_count: 0 });
expect(candidate.getStatus).toHaveBeenCalledTimes(1);
expect(deps.sleep).not.toHaveBeenCalled();
});

it('retries a transient connect failure and returns the status once the broker responds', async () => {
let attempt = 0;
const candidate: Pick<CoreRelay, 'getStatus'> = {
getStatus: vi.fn(async () => {
attempt += 1;
if (attempt < 3) {
throw new TypeError('Unable to connect. Is the computer able to access the url?');
}
return { agent_count: 0, pending_delivery_count: 0 };
}),
};
const deps = createDeps();

const result = await getBrokerStatusWithRetry(candidate, deps, true);

expect(result).toEqual({ agent_count: 0, pending_delivery_count: 0 });
expect(candidate.getStatus).toHaveBeenCalledTimes(3);
expect(deps.sleep).toHaveBeenCalledTimes(2);
expect(deps.log).toHaveBeenCalledWith(
expect.stringContaining('Broker status check failed (attempt 1/4), retrying in 300ms...')
);
});

it('exhausts its retry budget and throws the last error when the broker never responds', async () => {
const err = new TypeError('Unable to connect. Is the computer able to access the url?');
const candidate: Pick<CoreRelay, 'getStatus'> = {
getStatus: vi.fn(async () => {
throw err;
}),
};
const deps = createDeps();

await expect(getBrokerStatusWithRetry(candidate, deps)).rejects.toBe(err);
// 4 total attempts: the initial try plus 3 retries.
expect(candidate.getStatus).toHaveBeenCalledTimes(4);
expect(deps.sleep).toHaveBeenCalledTimes(3);
});

it('does not retry a non-connect failure -- fails after a single attempt', async () => {
const err = new Error('unauthorized');
const candidate: Pick<CoreRelay, 'getStatus'> = {
getStatus: vi.fn(async () => {
throw err;
}),
};
const deps = createDeps();

await expect(getBrokerStatusWithRetry(candidate, deps)).rejects.toBe(err);
expect(candidate.getStatus).toHaveBeenCalledTimes(1);
expect(deps.sleep).not.toHaveBeenCalled();
});
});

describe('readNodeDeliveryStatus', () => {
it('reads the canonical snake_case broker status shape', () => {
expect(
Expand Down
Loading
Loading