diff --git a/src/main/broker-proxy-service.test.ts b/src/main/broker-proxy-service.test.ts index 7b6ea1c..f6df427 100644 --- a/src/main/broker-proxy-service.test.ts +++ b/src/main/broker-proxy-service.test.ts @@ -450,3 +450,268 @@ describe('extractGitHubJobInfo', () => { expect(info.githubWorkflow).toBe('integration'); }); }); + +describe('message routing', () => { + interface Instance { sessionId?: string; runner: { agentName: string } } + interface RoutingInternals { + targets: Map }>; + messageQueues: Map; + pendingTargetAssignments: string[]; + localSessions: Map; + handleRequest(req: unknown, res: unknown): Promise; + processMessage(state: unknown, instance: unknown, body: string): Promise; + } + + // Broker messages as GitHub delivers them: the inner body is a JSON string. + const jobMessage = JSON.stringify({ + messageId: 2, + messageType: 'RunnerJobRequest', + body: JSON.stringify({ runner_request_id: 'req-1' }), + }); + const refreshMessage = JSON.stringify({ + messageId: 1, + messageType: 'RunnerRefreshConfig', + body: JSON.stringify({ config_type: 'runner' }), + }); + const cancelMessage = JSON.stringify({ + messageId: 3, + messageType: 'JobCancellation', + body: JSON.stringify({ jobId: 'req-1' }), + }); + + const fakeRequest = (method: string, url: string, body = '') => { + const req = new EventEmitter() as EventEmitter & { + method: string; url: string; headers: Record; + [Symbol.asyncIterator]: () => AsyncGenerator; + }; + req.method = method; + req.url = url; + req.headers = {}; + req[Symbol.asyncIterator] = async function* () { if (body) yield Buffer.from(body); }; + return req; + }; + + const fakeResponse = () => { + let ended = false; + const res = new EventEmitter() as EventEmitter & { + statusCode: number; body: string; writableEnded: boolean; + writeHead: (code: number) => unknown; end: (chunk?: unknown) => unknown; + }; + Object.defineProperty(res, 'writableEnded', { get: () => ended }); + res.writeHead = (code) => { res.statusCode = code; return res; }; + res.end = (chunk) => { res.body = chunk ? String(chunk) : ''; ended = true; res.emit('close'); return res; }; + return res; + }; + + let service: BrokerProxyService; + let internals: RoutingInternals; + + const request = async (method: string, url: string, body?: string) => { + const res = fakeResponse(); + await internals.handleRequest(fakeRequest(method, url, body), res); + return res; + }; + + const addTargetWithRunner = (id: string, agentName: string) => { + const target = createMockTarget({ id, displayName: id }); + const cred = createMockInstanceCredentials(1); + service.addTarget(target, [{ ...cred, runner: { ...cred.runner, agentName } }]); + // An upstream session already exists, so /session doesn't try to create one. + internals.targets.get(id)!.instances.get(1)!.sessionId = `upstream-${id}`; + return target; + }; + + const createSession = async (body?: string): Promise => { + const res = await request('POST', '/session', body); + expect(res.statusCode).toBe(201); + return JSON.parse(res.body).sessionId; + }; + + beforeEach(() => { + service = new BrokerProxyService(8787); + internals = service as unknown as RoutingInternals; + // Upstream calls (acknowledge) succeed silently. + mockHttpsRequest.mockImplementation((...args: unknown[]) => { + const callback = args[1] as (res: EventEmitter) => void; + const req = new EventEmitter() as EventEmitter & { setTimeout: () => void; write: () => void; end: () => void }; + req.setTimeout = () => {}; + req.write = () => {}; + req.end = () => { + const res = new EventEmitter() as EventEmitter & { statusCode: number }; + res.statusCode = 200; + callback(res); + res.emit('end'); + }; + return req; + }); + }); + + it('delivers the queued job before a stale RunnerRefreshConfig', async () => { + // GitHub pushes RunnerRefreshConfig between jobs. If the next worker's first + // poll returns that instead of its job, the runner rewrites its config and + // restarts its session, and the job is never delivered. + const target = addTargetWithRunner('target-a', 'runner-a.1'); + internals.messageQueues.set(target.id, [refreshMessage, jobMessage]); + internals.pendingTargetAssignments.push(target.id); + + const sessionId = await createSession(); + const res = await request('GET', `/message?sessionId=${sessionId}`); + + expect(res.statusCode).toBe(200); + expect(JSON.parse(res.body).messageType).toBe('RunnerJobRequest'); + }); + + it('drops RunnerRefreshConfig instead of queuing it for the next worker', async () => { + const target = addTargetWithRunner('target-a', 'runner-a.1'); + const state = internals.targets.get(target.id)!; + + await internals.processMessage(state, state.instances.get(1), refreshMessage); + + expect(internals.messageQueues.get(target.id) ?? []).toEqual([]); + }); + + it('queues a JobCancellation behind the job it is for', async () => { + const target = addTargetWithRunner('target-a', 'runner-a.1'); + internals.messageQueues.set(target.id, [jobMessage]); + const state = internals.targets.get(target.id)!; + + await internals.processMessage(state, state.instances.get(1), cancelMessage); + + expect(internals.messageQueues.get(target.id)).toEqual([jobMessage, cancelMessage]); + }); + + it('queues a JobCancellation for a job a worker is running', async () => { + const target = addTargetWithRunner('target-a', 'runner-a.1'); + internals.messageQueues.set(target.id, [jobMessage]); + internals.pendingTargetAssignments.push(target.id); + const sessionId = await createSession(); + await request('GET', `/message?sessionId=${sessionId}`); + const state = internals.targets.get(target.id)!; + + await internals.processMessage(state, state.instances.get(1), cancelMessage); + + expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]); + }); + + it('does not hand a cancellation to a worker that holds no job', async () => { + // With no job queued, taking the head of the queue gave a jobless worker a + // cancellation meant for whoever runs that job - and marked its session as + // holding a job it never had. + const target = addTargetWithRunner('target-a', 'runner-a.1'); + internals.messageQueues.set(target.id, [cancelMessage]); + internals.pendingTargetAssignments.push(target.id); + const sessionId = await createSession(); + + // Nothing is deliverable, so the handler long-polls rather than answering - + // which is the point. Let it poll, assert the cancellation stayed put, then + // end the poll the way it really ends, by a job arriving: that both cleans + // the request up and shows the cancellation was skipped rather than eaten. + const pending = request('GET', `/message?sessionId=${sessionId}`); + await new Promise((r) => setTimeout(r, 60)); + + expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]); + expect(internals.localSessions.get(sessionId)?.currentJobId).toBeUndefined(); + + internals.messageQueues.get(target.id)!.push(jobMessage); + const res = await pending; + + expect(res.body).toBe(jobMessage); + expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]); + }); + + it('does hand the cancellation to the worker actually running that job', async () => { + const target = addTargetWithRunner('target-a', 'runner-a.1'); + internals.messageQueues.set(target.id, [jobMessage]); + internals.pendingTargetAssignments.push(target.id); + const sessionId = await createSession(); + await request('GET', `/message?sessionId=${sessionId}`); // takes the job + internals.messageQueues.set(target.id, [cancelMessage]); + + const res = await request('GET', `/message?sessionId=${sessionId}`); + + expect(res.body).toContain('Cancel'); + expect(internals.messageQueues.get(target.id)).toEqual([]); + }); + + it('drops a JobCancellation for a job that is neither queued nor running', async () => { + // GitHub redelivers a cancellation for a while after the job ends. Queued + // for a target with no worker, it would be handed to the next worker. + const target = addTargetWithRunner('target-a', 'runner-a.1'); + const state = internals.targets.get(target.id)!; + + await internals.processMessage(state, state.instances.get(1), cancelMessage); + + expect(internals.messageQueues.get(target.id) ?? []).toEqual([]); + }); + + it('does not let a queued cancellation keep itself alive across redeliveries', async () => { + // GitHub redelivers a cancellation while the job is unfinished. A queued + // cancellation names the job, so counting it as evidence the job is live + // made every redelivery queue another copy - unbounded, and the next + // worker to poll gets a cancellation instead of a job. + const target = addTargetWithRunner('target-a', 'runner-a.1'); + const state = internals.targets.get(target.id)!; + internals.messageQueues.set(target.id, [cancelMessage]); + + await internals.processMessage(state, state.instances.get(1), cancelMessage); + + expect(internals.messageQueues.get(target.id)).toEqual([cancelMessage]); + }); + + it("answers the runner's own acknowledge locally", async () => { + const res = await request('POST', '/acknowledge?sessionId=abc', JSON.stringify({ runnerRequestId: 'req-1' })); + + expect(res.statusCode).toBe(200); + expect(res.body).toBe('{}'); + }); + + it('rejects an oversized session request instead of buffering it', async () => { + const res = await request('POST', '/session', 'x'.repeat(65 * 1024)); + + expect(res.statusCode).toBe(413); + }); + + describe('session to target binding', () => { + it('leaves a session unbound when its target has no job waiting', async () => { + // A listener spawned ahead of any job runs in the generic sandbox, not + // the repository's approved policy. Binding it would let it win the next + // job and fail it (seen live: cargo denied reading ~/.rustup). + addTargetWithRunner('target-a', 'runner-a.1'); + addTargetWithRunner('target-b', 'runner-b.1'); + + const sessionId = await createSession(JSON.stringify({ agent: { name: 'runner-b.1' } })); + + expect(internals.localSessions.get(sessionId)?.targetId).toBeUndefined(); + }); + + it("never hands a named session another target's pending assignment", async () => { + const targetA = addTargetWithRunner('target-a', 'runner-a.1'); + addTargetWithRunner('target-b', 'runner-b.1'); + internals.pendingTargetAssignments.push(targetA.id); + + const sessionId = await createSession(JSON.stringify({ agent: { name: 'runner-b.1' } })); + + expect(internals.localSessions.get(sessionId)?.targetId).toBeUndefined(); + expect(internals.pendingTargetAssignments).toEqual([targetA.id]); + }); + + it('consumes the matching pending assignment, not the first one', async () => { + const targetA = addTargetWithRunner('target-a', 'runner-a.1'); + const targetB = addTargetWithRunner('target-b', 'runner-b.1'); + internals.pendingTargetAssignments.push(targetA.id, targetB.id); + + await createSession(JSON.stringify({ agent: { name: 'runner-b.1' } })); + + expect(internals.pendingTargetAssignments).toEqual([targetA.id]); + }); + + it('falls back to the pending assignment when the request names no runner', async () => { + const targetA = addTargetWithRunner('target-a', 'runner-a.1'); + internals.pendingTargetAssignments.push(targetA.id); + + const sessionId = await createSession(); + + expect(internals.localSessions.get(sessionId)?.targetId).toBe(targetA.id); + }); + }); +}); diff --git a/src/main/broker-proxy-service.ts b/src/main/broker-proxy-service.ts index 147570e..d8a7341 100644 --- a/src/main/broker-proxy-service.ts +++ b/src/main/broker-proxy-service.ts @@ -173,6 +173,69 @@ function createJWT(clientId: string, authorizationUrl: string, privateKey: crypt return `${signingInput}.${signature.toString('base64url')}`; } +// ============================================================================ +// Broker Message Helpers +// ============================================================================ + +/** + * True for a job assignment. JobCancellation is a signal about a job, not a + * job; RunnerRefreshConfig and AgentRefresh are runner housekeeping. + */ +function isJobAssignmentMessage(message: string): boolean { + try { + const messageType = String(JSON.parse(message).messageType || '').toLowerCase(); + return messageType.includes('job') && !messageType.includes('cancel'); + } catch { + return false; + } +} + +/** The job a broker message is about, under either name GitHub uses for it. */ +function jobIdFromMessage(message: string): string | undefined { + try { + const parsed = JSON.parse(message); + const innerBody = typeof parsed.body === 'string' ? JSON.parse(parsed.body) : parsed.body; + return innerBody?.jobId || innerBody?.runner_request_id || undefined; + } catch { + return undefined; + } +} + +/** The runner name a session request carries; the runner sends it as `agent.name`. */ +function agentNameFromSessionRequest(body: string): string | undefined { + try { + const name = JSON.parse(body)?.agent?.name; + return typeof name === 'string' && name ? name : undefined; + } catch { + return undefined; + } +} + +/** Largest request body the proxy reads; a runner's are a few hundred bytes. */ +const MAX_REQUEST_BODY_BYTES = 64 * 1024; + +class RequestBodyTooLargeError extends Error { + constructor() { + super('request body too large'); + } +} + +/** + * Read a request body in full, or throw once it exceeds the limit. The stream + * is drained either way: leaving the loop early destroys the socket, and the + * 413 would never reach the client. + */ +async function readRequestBody(req: http.IncomingMessage): Promise { + const chunks: Buffer[] = []; + let size = 0; + for await (const chunk of req) { + size += (chunk as Buffer).length; + if (size <= MAX_REQUEST_BODY_BYTES) chunks.push(chunk as Buffer); + } + if (size > MAX_REQUEST_BODY_BYTES) throw new RequestBodyTooLargeError(); + return Buffer.concat(chunks).toString(); +} + // ============================================================================ // Broker Proxy Service // ============================================================================ @@ -554,18 +617,30 @@ export class BrokerProxyService extends EventEmitter { this.emit('job-received', state.target.id, jobId, instance.runner.agentName, githubInfo); this.emitStatusUpdate(); } else { - // Non-job messages (including cancel signals) must also be forwarded to the runner - // Log more details for debugging cancel signal delivery - const isCancelLike = messageType.toLowerCase().includes('cancel') || + const isCancelLike = messageTypeLower.includes('cancel') || (innerBody && typeof innerBody === 'object' && JSON.stringify(innerBody).toLowerCase().includes('cancel')); - if (isCancelLike) { - log()?.info(`[BrokerProxy] CANCEL signal detected! messageType=${messageType}, target=${state.target.displayName}/${instance.instanceNum}`); - log()?.debug(`[BrokerProxy] Cancel message received (${body.length} bytes)`); - // Acknowledge cancel messages immediately so they don't block job requests - // Cancel signals are only relevant if there's an active job to cancel + if (!isCancelLike) { + // Only jobs and cancellations are forwarded. Anything else, such as + // RunnerRefreshConfig or AgentRefresh, makes the runner rewrite its + // config and restart its session. Queued between jobs it reaches the + // next worker before its job, which then never arrives. localmost + // owns the runner registration, so nothing is lost by dropping it. + log()?.warn(`[BrokerProxy] Dropping ${messageType} from ${state.target.displayName}/${instance.instanceNum}: not forwarded to runners`); await this.acknowledgeMessageUpstream(state, instance, messageId); - } else { - log()?.info(`[BrokerProxy] Non-job message (${messageType}) received from ${state.target.displayName}/${instance.instanceNum}, forwarding to runner`); + return; + } + + log()?.info(`[BrokerProxy] CANCEL signal detected! messageType=${messageType}, target=${state.target.displayName}/${instance.instanceNum}`); + log()?.debug(`[BrokerProxy] Cancel message received (${body.length} bytes)`); + // Acknowledge cancel messages immediately so they don't block job requests + await this.acknowledgeMessageUpstream(state, instance, messageId); + + // GitHub redelivers a cancellation until the job ends, and for a while + // after. One for a job no worker here holds or will hold has nothing to + // cancel; queued, it would only be handed to the next worker. + if (jobId && !this.isJobLive(jobId)) { + log()?.info(`[BrokerProxy] Dropping cancellation for ${jobId}: not queued or running here`); + return; } if (!this.messageQueues.has(targetId)) { this.messageQueues.set(targetId, []); @@ -574,6 +649,25 @@ export class BrokerProxyService extends EventEmitter { } } + /** + * Whether a job is queued for a worker here or held by a live worker session. + * + * Only a queued job *assignment* counts. A queued cancellation names the same + * job, so counting it made a cancellation evidence of its own liveness: every + * redelivery - and GitHub redelivers while a job is unfinished - queued + * another copy, unbounded, and the next worker to poll would be handed a + * cancellation instead of a job. Which is the stall this class exists to stop. + */ + private isJobLive(jobId: string): boolean { + for (const session of this.localSessions.values()) { + if (session.currentJobId === jobId) return true; + } + for (const queue of this.messageQueues.values()) { + if (queue.some(message => isJobAssignmentMessage(message) && jobIdFromMessage(message) === jobId)) return true; + } + return false; + } + /** * Stop the polling loop. */ @@ -1134,14 +1228,21 @@ export class BrokerProxyService extends EventEmitter { try { if (method === 'POST' && url.pathname === '/session') { - await this.handleSessionCreate(res); + await this.handleSessionCreate(req, res); } else if (method === 'GET' && url.pathname === '/message') { await this.handleMessagePoll(res, url); } else if (method === 'DELETE' && url.pathname === '/session') { await this.handleSessionDelete(res, url); } else if (method === 'POST' && url.pathname === '/acknowledge') { // Handle acknowledge locally - the broker proxy already received the message - // when it polled GitHub, so workers don't need to acknowledge upstream + // when it polled GitHub, so workers don't need to acknowledge upstream. + // The request is logged because it shows the shape GitHub's acknowledge + // endpoint expects, which acknowledgeMessageUpstream does not yet match. + const ackBody = await readRequestBody(req); + // JSON-encoded, so a body containing CR/LF cannot forge log lines: the + // log file writes messages verbatim. At debug because it is one line + // per acknowledge, which is one per message the runner receives. + log()?.debug(`[BrokerProxy] Runner acknowledge${url.search}: ${JSON.stringify(ackBody.slice(0, 300))}`); res.writeHead(200, { 'Content-Type': 'application/json' }); res.end('{}'); } else if (method === 'POST' && url.pathname === '/acquirejob') { @@ -1152,17 +1253,49 @@ export class BrokerProxyService extends EventEmitter { await this.handleForward(req, res, url); } } catch (error) { + if (error instanceof RequestBodyTooLargeError) { + res.writeHead(413, { 'Content-Type': 'text/plain' }); + res.end('request body too large'); + return; + } log()?.error( `[BrokerProxy] Error handling ${method} ${url.pathname}: ${(error as Error).message}`); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: (error as Error).message })); } } - private async handleSessionCreate(res: http.ServerResponse): Promise { + /** + * The target a new worker session belongs to, or none. + * + * A worker spawned for a job carries the repository's approved sandbox + * policy; a listener spawned ahead of any job runs in the generic sandbox + * and must not be handed one. The runner names itself in the request, so a + * session is bound only when the target it is registered under has a job + * waiting, and it takes that entry rather than whichever is first. An + * unnamed request falls back to the positional queue. + */ + private resolveSessionTarget(agentName: string | undefined): string | undefined { + if (!agentName) return this.pendingTargetAssignments.shift(); + for (const state of this.targets.values()) { + for (const instance of state.instances.values()) { + if (instance.runner.agentName !== agentName) continue; + const pending = this.pendingTargetAssignments.indexOf(state.target.id); + if (pending < 0) return undefined; + this.pendingTargetAssignments.splice(pending, 1); + return state.target.id; + } + } + log()?.warn(`[BrokerProxy] Session request names unknown runner ${agentName}; leaving it unbound`); + return undefined; + } + + private async handleSessionCreate(req: http.IncomingMessage, res: http.ServerResponse): Promise { const sessionId = crypto.randomUUID(); - // Assign this session to a target (from pending assignments queue) - const targetId = this.pendingTargetAssignments.shift(); + const requestBody = await readRequestBody(req); + const agentName = agentNameFromSessionRequest(requestBody); + log()?.info(`[BrokerProxy] Session request from ${agentName ?? 'unnamed runner'}`); + const targetId = this.resolveSessionTarget(agentName); log()?.debug(`[BrokerProxy] Creating local session ${sessionId} for target ${targetId || 'unknown'}`); // Only create upstream sessions for instances that don't already have them @@ -1207,18 +1340,6 @@ export class BrokerProxyService extends EventEmitter { const session = this.localSessions.get(sessionId)!; const targetId = session.targetId; - // Helper to check if a message is a job assignment (vs cancel or other signal) - const isJobAssignmentMessage = (message: string): boolean => { - try { - const parsed = JSON.parse(message); - const messageType = (parsed.messageType || '').toLowerCase(); - // JobCancellation is NOT a job assignment - it's a signal to cancel - return messageType.includes('job') && !messageType.includes('cancel'); - } catch { - return false; - } - }; - // If this worker already has a job, only deliver non-job messages (like cancel signals). // Don't give them another job message. if (session.currentJobId) { @@ -1293,24 +1414,28 @@ export class BrokerProxyService extends EventEmitter { return undefined; } const queue = this.messageQueues.get(targetId); - return queue?.shift(); + if (!queue || queue.length === 0) return undefined; + // The job goes first. A cancellation queued ahead of it is for a job + // this worker doesn't hold yet; it follows on the next poll. + const jobIndex = queue.findIndex(isJobAssignmentMessage); + if (jobIndex >= 0) return queue.splice(jobIndex, 1)[0]; + + // No job queued. Taking the head anyway handed a worker holding no job + // somebody else's cancellation - stealing it from the worker that runs + // that job, and marking this session as holding a job it never had. A + // cancellation goes only to the worker whose job it names. + const held = session.currentJobId; + if (!held) return undefined; + const mine = queue.findIndex(message => jobIdFromMessage(message) === held); + return mine >= 0 ? queue.splice(mine, 1)[0] : undefined; }; // Helper to extract job ID from message and mark session const markSessionWithJob = (message: string): void => { - try { - const parsed = JSON.parse(message); - let innerBody = parsed.body; - if (typeof innerBody === 'string') { - innerBody = JSON.parse(innerBody); - } - const jobId = innerBody?.jobId || innerBody?.runner_request_id; - if (jobId) { - session.currentJobId = jobId; - log()?.debug(`[BrokerProxy] Marked session ${sessionId} with job ${jobId}`); - } - } catch { - // Could not parse, still deliver the message + const jobId = jobIdFromMessage(message); + if (jobId) { + session.currentJobId = jobId; + log()?.debug(`[BrokerProxy] Marked session ${sessionId} with job ${jobId}`); } };