test/devin-connect-stream-lifecycle.test.js
Rama master · réplica de lectura del repositorio original.
import { afterEach, beforeEach, describe, it, mock } from 'node:test';
import assert from 'node:assert/strict';
import { addAccountByKey, removeAccount, getAccountInternal, __resetReloginState, __setReloginDeps } from '../src/auth.js';
import { handleChatCompletions, __setConnectDeps, __resetConnectDeps } from '../src/handlers/chat.js';
import { activeSseCount, abortActiveSse } from '../src/sse-registry.js';
// DEVIN_CONNECT stream lifecycle (audit finding): the connect streaming handler
// gained the same abort/heartbeat/SSE-registry wiring the Cascade path already
// had. These exercise the real handler with the connect network call mocked, so
// no token/socket is touched.
// 1. client disconnect mid-stream aborts and does NOT hop to a fresh account
// 2. a server-level context.signal abort likewise stops account-hopping
// 3. heartbeat `: ping` comments are emitted while the stream is silent
// 4. the SSE controller is registered during the stream and unregistered after
// 5. graceful-shutdown drain (abortActiveSse) tears the stream down
const createdIds = [];
const prevEnv = {};
const HEARTBEAT_MS = 15_000; // mirrors the constant in chat.js
function seed(label) {
const key = `devin-session-token$lc-${label}-${Math.random().toString(36).slice(2)}`;
const acct = addAccountByKey(key, label);
createdIds.push(acct.id);
return acct;
}
function rateLimited() {
// Pre-emit account dry-well: WITHOUT the abort guard this triggers a
// cross-account failover hop. It is the control that proves the guard fires.
return Object.assign(new Error('rate limited'), { code: 'RATE_LIMITED' });
}
// Fake res mirroring the SSE write/end/on contract the handler relies on, plus a
// `disconnect()` that fires 'close' WITHOUT marking writableEnded — exactly how
// a real client hang-up (and messages.js's _clientDisconnected) surfaces, so the
// handler's close listener takes the abort path.
function fakeRes() {
const listeners = new Map();
return {
body: '', writableEnded: false,
write(chunk) { this.body += String(chunk); return true; },
end(chunk) { if (chunk) this.write(chunk); this.writableEnded = true; for (const cb of listeners.get('close') || []) cb(); },
on(event, cb) { if (!listeners.has(event)) listeners.set(event, []); listeners.get(event).push(cb); return this; },
disconnect() { for (const cb of listeners.get('close') || []) cb(); },
};
}
function parseFrames(raw) {
return raw.split('\n\n').filter(Boolean).filter(f => !f.startsWith(':')).map(f => {
const d = f.split('\n').find(l => l.startsWith('data: '))?.slice(6) || '';
return d === '[DONE]' ? '[DONE]' : JSON.parse(d);
});
}
beforeEach(() => {
prevEnv.DEVIN_CONNECT = process.env.DEVIN_CONNECT;
prevEnv.MAX = process.env.DEVIN_CONNECT_FAILOVER_MAX;
process.env.DEVIN_CONNECT = '1';
delete process.env.DEVIN_CONNECT_FAILOVER_MAX;
});
afterEach(() => {
mock.timers.reset();
__resetConnectDeps();
__resetReloginState();
__setReloginDeps(null);
if (prevEnv.DEVIN_CONNECT === undefined) delete process.env.DEVIN_CONNECT;
else process.env.DEVIN_CONNECT = prevEnv.DEVIN_CONNECT;
if (prevEnv.MAX === undefined) delete process.env.DEVIN_CONNECT_FAILOVER_MAX;
else process.env.DEVIN_CONNECT_FAILOVER_MAX = prevEnv.MAX;
while (createdIds.length) removeAccount(createdIds.pop());
});
describe('DEVIN_CONNECT stream — client disconnect stops account-hopping', () => {
it('aborts on client disconnect and does NOT fail over to a fresh account', async () => {
const a = seed('dc-1');
const b = seed('dc-2');
let res; // assigned before handler() runs; referenced by the mock closure
const seen = [];
__setConnectDeps({
streamChatCompletion: async (params) => {
seen.push(params.token);
// Client hangs up while the first upstream call is in flight, THEN the
// upstream returns a normally-failover-eligible RATE_LIMITED. The abort
// guard must win: no hop to account b.
res.disconnect();
throw rateLimited();
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
res = fakeRes();
await result.handler(res);
assert.equal(seen.length, 1, 'stopped after the first account — no failover to a dead socket');
assert.ok(!seen.includes(b.apiKey), 'never touched the second pooled account');
assert.ok(seen.includes(a.apiKey), 'the first account was the one attempted');
// No error frame is spun up for a client that already left.
const frames = parseFrames(res.body);
assert.ok(!frames.some(f => f !== '[DONE]' && f.error), 'no error frame emitted to the gone client');
});
it('does not replay a transport error and aborts cleanly when the client disconnects', async () => {
const a = seed('replay-abort');
const b = seed('replay-abort-2');
let res;
const seen = [];
__setConnectDeps({
streamChatCompletion: async (params) => {
seen.push(params.token);
// Transport fault before any byte. The handler must NOT replay it (the
// adapter owns the one bounded retry); with the client already gone it
// must abort without hopping or writing an error frame.
res.disconnect();
throw Object.assign(new Error('socket hang up'), { code: 'ECONNRESET' });
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
res = fakeRes();
await result.handler(res);
assert.equal(seen.length, 1, 'no handler-side replay; the adapter owns the bounded retry');
assert.equal(seen[0], a.apiKey, 'the first account was the one attempted');
assert.ok(!seen.includes(b.apiKey), 'never hopped to second account');
const frames = parseFrames(res.body);
assert.ok(!frames.some(f => f !== '[DONE]' && f.error), 'no error frame emitted to the gone client');
assert.ok(!frames.includes('[DONE]'), 'no [DONE] written to a disconnected socket');
});
it('aborts on a server-level context.signal and does NOT fail over', async () => {
const a = seed('sig-1');
const b = seed('sig-2');
const ctrl = new AbortController();
const seen = [];
__setConnectDeps({
streamChatCompletion: async (params) => {
seen.push(params.token);
ctrl.abort(); // e.g. server drain / request teardown from server.js wiring
throw rateLimited();
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '', signal: ctrl.signal },
);
const res = fakeRes();
await result.handler(res);
assert.equal(seen.length, 1, 'context.signal abort stops the failover loop');
assert.ok(!seen.includes(b.apiKey), 'no hop to a second account after the caller aborted');
assert.ok(seen.includes(a.apiKey));
});
it('still fails over normally when the client stays connected (guard is abort-gated)', async () => {
// Control: without an abort, a pre-emit RATE_LIMITED must still hop — proving
// the new guard only trips on a real abort, not on every failover.
seed('ok-1');
seed('ok-2');
const seen = [];
let call = 0;
__setConnectDeps({
streamChatCompletion: async (params, send) => {
seen.push(params.token);
// First account (whichever the pool picks) is dry → must hop; the second
// streams clean. No abort, so the guard stays out of the way.
if (call++ === 0) throw rateLimited();
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: { content: 'OK' }, finish_reason: null }] });
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] });
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
const res = fakeRes();
await result.handler(res);
assert.equal(seen.length, 2, 'connected client still gets the failover hop');
const frames = parseFrames(res.body);
assert.ok(frames.some(f => f !== '[DONE]' && f.choices?.[0]?.delta?.content === 'OK'), 'streamed from the healthy account');
});
});
describe('DEVIN_CONNECT stream — inflight release on pre-attempt abort (audit #2/#3)', () => {
// REF-3: the top-of-loop abort break at chat.js:2570 fires when a client
// disconnects between failover hops — AFTER an account was acquired but
// BEFORE attemptStream (→ finalizeConnectAccount → releaseAccountById) ran for
// it. Without the release in that break, the acquired account's _inflight slot
// leaks until the acquire reservation ages out (≤15min), shrinking the pool
// under disconnect churn. This asserts the slot is returned to 0.
it('no account is left in-flight when a disconnect interrupts mid-failover', async () => {
// Account A fails over eligibly, but the disconnect fires during A's attempt.
// acquireConnectFailover then short-circuits on the aborted signal (returns
// null → no B acquired), and the post-attempt abort check breaks. A was
// already finalized (released) by attemptStream. This asserts the whole
// failover teardown leaves ZERO in-flight slots — the pool-integrity
// invariant the leak would violate.
const a = seed('inflight-a');
const b = seed('inflight-b');
let res;
const seen = [];
__setConnectDeps({
streamChatCompletion: async (params) => {
seen.push(params.token);
res.disconnect(); // client hangs up during the (only) attempt
throw rateLimited();
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
res = fakeRes();
await result.handler(res);
assert.equal(seen.length, 1, 'no failover hop to a dead socket');
assert.equal(getAccountInternal(a.id)?._inflight || 0, 0, 'attempted account released');
assert.equal(getAccountInternal(b.id)?._inflight || 0, 0, 'untouched account holds no slot');
});
it('leaves no in-flight slot when the client aborts before the very first attempt', async () => {
// hop 0 edge: ccAcct was acquired before the loop; if the abort is already
// visible at loop entry, the break must still release ccAcct.
const a = seed('inflight-first');
const ctrl = new AbortController();
ctrl.abort(); // already aborted before the handler streams
const seen = [];
__setConnectDeps({
streamChatCompletion: async (params) => { seen.push(params.token); throw rateLimited(); },
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '', signal: ctrl.signal },
);
const res = fakeRes();
await result.handler(res);
assert.equal(seen.length, 0, 'pre-aborted request never attempts an upstream call');
assert.equal(getAccountInternal(a.id)?._inflight || 0, 0, 'ccAcct released at the entry abort break (no leak)');
});
});
describe('DEVIN_CONNECT stream — heartbeat', () => {
it('emits `: ping` comments while the upstream is silent', async () => {
seed('hb-1');
mock.timers.enable({ apis: ['setInterval'] });
let release;
const gate = new Promise((r) => { release = r; });
__setConnectDeps({
streamChatCompletion: async (params, send) => {
// Hold the stream open with no bytes so the heartbeat interval is the
// only thing writing — the idle-timeout risk the fix addresses.
await gate;
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: { content: 'DONE_NOW' }, finish_reason: null }] });
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] });
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
const res = fakeRes();
const done = result.handler(res);
// Let the handler park at `await gate` (interval already armed), then fire it.
await Promise.resolve();
mock.timers.tick(HEARTBEAT_MS);
mock.timers.tick(HEARTBEAT_MS);
release();
await done;
const pings = (res.body.match(/: ping/g) || []).length;
assert.ok(pings >= 2, `expected heartbeat pings while idle, got ${pings}`);
const frames = parseFrames(res.body);
assert.ok(frames.some(f => f !== '[DONE]' && f.choices?.[0]?.delta?.content === 'DONE_NOW'), 'content still streamed after the idle window');
assert.equal(frames.at(-1), '[DONE]');
});
it('stops the heartbeat after the stream ends (no interval leak)', async () => {
seed('hb-2');
mock.timers.enable({ apis: ['setInterval'] });
__setConnectDeps({
streamChatCompletion: async (params, send) => {
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: { content: 'X' }, finish_reason: null }] });
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] });
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
const res = fakeRes();
await result.handler(res);
const before = res.body.length;
// The stream is over and the interval should be cleared: further ticks add
// no more pings.
mock.timers.tick(HEARTBEAT_MS * 5);
assert.equal(res.body.length, before, 'no pings after stopHeartbeat()');
});
});
describe('DEVIN_CONNECT stream — SSE registry', () => {
it('registers a controller during the stream and unregisters when it ends', async () => {
seed('reg-1');
const before = activeSseCount();
let midCount = -1;
__setConnectDeps({
streamChatCompletion: async (params, send) => {
midCount = activeSseCount();
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: { content: 'OK' }, finish_reason: null }] });
send({ id: 's', object: 'chat.completion.chunk', choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] });
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
const res = fakeRes();
await result.handler(res);
assert.equal(midCount, before + 1, 'controller was registered while streaming');
assert.equal(activeSseCount(), before, 'controller unregistered in the finally after the stream ended');
});
it('graceful-shutdown drain (abortActiveSse) tears the stream down and stops hopping', async () => {
const a = seed('drain-1');
const b = seed('drain-2');
let release;
const gate = new Promise((r) => { release = r; });
const seen = [];
__setConnectDeps({
streamChatCompletion: async (params) => {
seen.push(params.token);
await gate; // park so the drain lands mid-stream
throw rateLimited();
},
});
const result = await handleChatCompletions(
{ model: 'swe-1-6-slow', stream: true, messages: [{ role: 'user', content: 'hi' }] },
{ callerKey: '' },
);
const res = fakeRes();
const done = result.handler(res);
await Promise.resolve();
const drained = abortActiveSse('server shutting down');
release();
await done;
assert.ok(drained >= 1, 'the connect stream was visible to the shutdown drain');
assert.equal(seen.length, 1, 'drain aborted the stream — no failover hop');
assert.ok(!seen.includes(b.apiKey), 'never hopped to the second account during drain');
assert.ok(seen.includes(a.apiKey));
const frames = parseFrames(res.body);
assert.ok(frames.some(f => f !== '[DONE]' && f.error?.code === 'server_shutdown'), 'emitted a server_shutdown frame');
});
});