Repository Public
1import { randomBytes } from 'node:crypto';2import net from 'node:net';3import { afterAll, afterEach, beforeAll, describe, expect, it } from 'vitest';4import type WebSocket from 'ws';5import { AosEvent } from '@aos/shared';6import { Client, FAKE_ANTHROPIC_KEY, FAKE_OPENAI_KEY, ORIGIN, PASSWORD, createHarness, insertOtherUser, login, openWs, seedMission, setupAdmin, until, upgradeStatus, type Harness } from '../testkit/harness.js';78let h: Harness;9let admin: Client;10let secret: string;11let addr: string;1213const closed = (ws: WebSocket): Promise<number> => new Promise((resolve) => ws.once('close', (code) => resolve(code)));14/** Opens the events socket and collects every frame as a validated AosEvent. */15const opened: WebSocket[] = [];16/** Everything emitted so far has gone through the kernel feed and the hub: a new socket starts from a quiet state. */17async function settle(): Promise<void> {18 await h.kernel.drained();19 await h.server.ctx.events.idle();20}21async function connect(query = '', client = admin): Promise<{ ws: WebSocket; events: AosEvent[]; raw: string[] }> {22 await settle();23 const ws = openWs(addr, `/events${query}`, { cookie: client.cookieHeader });24 opened.push(ws);25 const events: AosEvent[] = [];26 const raw: string[] = [];27 ws.on('message', (data, isBinary) => {28 expect(isBinary).toBe(false);29 raw.push(String(data));30 events.push(AosEvent.parse(JSON.parse(String(data))));31 });32 expect(await upgradeStatus(ws)).toBe(101);33 return { ws, events, raw };34}35const persisted = (events: AosEvent[]): AosEvent[] => events.filter((e) => e.id >= 0);36const idsOf = (events: AosEvent[]): number[] => persisted(events).map((e) => e.id);37const headId = async (): Promise<number> => (await h.db.one<{ id: number | null }>('SELECT max(id) AS id FROM events'))?.id ?? 0;3839beforeAll(async () => {40 h = await createHarness({ tweak: (c) => (c.wsAuthRecheckMs = 150) });41 ({ client: admin, secret } = await setupAdmin(h));42 addr = await h.listen();43});44afterEach(async () => {45 for (const ws of opened.splice(0)) ws.terminate();46 await until(() => h.server.ctx.events.size === 0);47 ({ client: admin } = { client: await login(h, secret) }); // some tests end sessions48});49afterAll(async () => {50 await h.close();51});5253describe('upgrade authorisation', () => {54 it('refuses without a cookie (401), with a forged cookie (401), a wrong or missing Origin (403), before the second factor (403)', async () => {55 expect(await upgradeStatus(openWs(addr, '/events'))).toBe(401);56 expect(await upgradeStatus(openWs(addr, '/events', { cookie: 'aos_sid=' + 'B'.repeat(43) }))).toBe(401);57 expect(await upgradeStatus(openWs(addr, '/events', { cookie: admin.cookieHeader, origin: 'https://evil.example' }))).toBe(403);58 expect(await upgradeStatus(openWs(addr, '/events', { cookie: admin.cookieHeader, origin: null }))).toBe(403);59 const half = new Client(h);60 await half.post('/auth/login', { username: 'admin', password: PASSWORD });61 expect(await upgradeStatus(openWs(addr, '/events', { cookie: half.cookieHeader }))).toBe(403);62 expect(h.server.ctx.events.size).toBe(0);63 });64 it('validates the query before upgrading: bad cursor / mission id ⇒ 400, foreign or unknown mission ⇒ 404', async () => {65 for (const q of ['?cursor=abc', '?cursor=-1', '?cursor=1.5', '?missionId=nope', `?cursor=${'9'.repeat(30)}`]) {66 expect([q, await upgradeStatus(openWs(addr, `/events${q}`, { cookie: admin.cookieHeader }))]).toEqual([q, 400]);67 }68 expect(await upgradeStatus(openWs(addr, '/events?missionId=00000000-0000-4000-8000-000000000000', { cookie: admin.cookieHeader }))).toBe(404);69 const other = await insertOtherUser(h, 'ws-other');70 const m = await h.db.one<{ id: string }>("INSERT INTO missions (owner_id, title, objective, budget_usd) VALUES ($1, 't', 'o', 1) RETURNING id", [other]);71 expect(await upgradeStatus(openWs(addr, `/events?missionId=${m?.id}`, { cookie: admin.cookieHeader }))).toBe(404);72 });73});7475describe('live stream', () => {76 it('delivers persisted events in id order as AosEvent frames, and transient deltas with id -1', async () => {77 const c = await connect();78 const { missionId, runId } = await seedMission(h, admin);79 await until(() => persisted(c.events).length >= 2);80 expect(c.events.slice(0, 2).map((e) => [e.kind, e.missionId, e.runId])).toEqual([['mission.created', missionId, null], ['run.created', missionId, runId]]);81 expect(c.events[0]?.id).toBeLessThan(c.events[1]?.id as number);82 expect(Number.isNaN(Date.parse(c.events[0]?.ts as string))).toBe(false);83 h.kernel.emitTransient(missionId, runId, { type: 'text', text: 'Hel' });84 h.kernel.emitTransient(missionId, runId, { type: 'text', text: 'lo' });85 await until(() => c.events.filter((e) => e.kind === 'model.delta').length === 2);86 expect(c.events.filter((e) => e.kind === 'model.delta').map((e) => [e.id, e.runId, e.payload.text])).toEqual([[-1, runId, 'Hel'], [-1, runId, 'lo']]);87 expect(await h.db.q("SELECT 1 FROM events WHERE kind = 'model.delta'")).toEqual([]); // never persisted88 const id = await h.kernel.emit(missionId, runId, 'step.completed', { name: 'read_file' }, 'system', 3);89 await until(() => c.events.some((e) => e.id === id));90 expect(c.events.find((e) => e.id === id)).toMatchObject({ kind: 'step.completed', seq: 3, trust: 'system', payload: { name: 'read_file' } });91 // a model step's event: the normalised result is forwarded, the provider-native transcript is not92 const mid = await h.kernel.emit(missionId, runId, 'step.completed', { kind: 'model_call', name: 'm', result: { blocks: [{ type: 'reasoning', summary: 'plan' }], stop: 'end_turn', native: [{ type: 'reasoning', encrypted_content: 'ENCRYPTED-BLOB' }] } }, 'system', 4);93 await until(() => c.events.some((e) => e.id === mid));94 expect(c.events.find((e) => e.id === mid)?.payload).toEqual({ kind: 'model_call', name: 'm', result: { blocks: [{ type: 'reasoning', summary: 'plan' }], stop: 'end_turn' } });95 c.ws.close();96 await closed(c.ws);97 await until(() => h.server.ctx.events.size === 0);98 });99 it('without a cursor nothing is replayed; frames sent by the browser are ignored', async () => {100 const { missionId } = await seedMission(h, admin);101 const c = await connect();102 c.ws.send('{"type":"subscribe","everything":true}');103 c.ws.send(Buffer.from([1, 2, 3]));104 const id = await h.kernel.emit(missionId, null, 'system.notice', { code: 'probe' });105 await until(() => c.events.length > 0);106 expect(idsOf(c.events)).toEqual([id]);107 c.ws.close();108 });109 it('events of another owner and of other missions (with missionId) are not delivered', async () => {110 const other = await insertOtherUser(h, 'ws-foreign');111 const foreign = await h.db.one<{ id: string }>("INSERT INTO missions (owner_id, title, objective, budget_usd) VALUES ($1, 't', 'o', 1) RETURNING id", [other]);112 const a = await seedMission(h, admin);113 const b = await seedMission(h, admin);114 const all = await connect();115 const onlyA = await connect(`?missionId=${a.missionId}`);116 await h.kernel.emit(foreign?.id as string, null, 'system.notice', { code: 'FOREIGN' });117 h.kernel.emitTransient(foreign?.id as string, '00000000-0000-4000-8000-000000000000', { text: 'FOREIGN-DELTA' });118 const idB = await h.kernel.emit(b.missionId, null, 'system.notice', { code: 'b' });119 h.kernel.emitTransient(b.missionId, b.runId, { text: 'delta-b' });120 const idA = await h.kernel.emit(a.missionId, null, 'system.notice', { code: 'a' });121 await until(() => all.events.some((e) => e.id === idA) && onlyA.events.some((e) => e.id === idA));122 expect(idsOf(all.events)).toEqual([idB, idA]);123 expect(all.events.filter((e) => e.id === -1).map((e) => e.payload.text)).toEqual(['delta-b']);124 expect(onlyA.events.map((e) => e.id)).toEqual([idA]);125 expect(all.raw.join('') + onlyA.raw.join('')).not.toContain('FOREIGN');126 all.ws.close();127 onlyA.ws.close();128 });129});130131describe('cursor replay', () => {132 it('reconnecting with the last id replays exactly what was missed, then continues live without duplicates', async () => {133 const { missionId, runId } = await seedMission(h, admin);134 const first = await connect(`?missionId=${missionId}`);135 const e1 = await h.kernel.emit(missionId, runId, 'step.started', { n: 1 });136 await until(() => first.events.some((e) => e.id === e1));137 first.ws.close();138 await closed(first.ws);139 const missed = [await h.kernel.emit(missionId, runId, 'step.completed', { n: 2 }), await h.kernel.emit(missionId, runId, 'run.usage', { n: 3 }), await h.kernel.emit(missionId, runId, 'step.started', { n: 4 })];140 const second = await connect(`?cursor=${e1}&missionId=${missionId}`);141 await until(() => second.events.length >= 3);142 expect(idsOf(second.events)).toEqual(missed);143 expect(second.events.map((e) => e.payload.n)).toEqual([2, 3, 4]);144 const live = await h.kernel.emit(missionId, runId, 'step.completed', { n: 5 });145 await until(() => second.events.some((e) => e.id === live));146 expect(idsOf(second.events)).toEqual([...missed, live]);147 second.ws.close();148 });149 it('cursor=0 replays the whole history of the owner; a cursor at the head replays nothing', async () => {150 const mine = (await h.db.q<{ id: number }>("SELECT e.id FROM events e JOIN missions m ON m.id = e.mission_id JOIN users u ON u.id = m.owner_id WHERE u.username = 'admin' ORDER BY e.id")).map((r) => r.id);151 const all = await connect('?cursor=0');152 await until(() => all.events.length >= mine.length);153 expect(idsOf(all.events)).toEqual(mine);154 all.ws.close();155 const head = await connect(`?cursor=${await headId()}`);156 const { missionId } = await seedMission(h, admin);157 await until(() => head.events.length >= 2);158 expect(head.events.every((e) => e.missionId === missionId)).toBe(true);159 head.ws.close();160 });161 it('no gap and no duplicate while events keep arriving during the replay hand-over', async () => {162 const { missionId, runId } = await seedMission(h, admin);163 const start = await headId();164 let emitted = 0;165 let stop = false;166 const producer = (async () => {167 while (!stop) {168 await h.kernel.emit(missionId, runId, 'run.usage', { n: emitted++ });169 if (emitted % 5 === 0) await new Promise((r) => setTimeout(r, 1));170 }171 })();172 const sockets: Awaited<ReturnType<typeof connect>>[] = [];173 for (let i = 0; i < 6; i++) {174 await new Promise((r) => setTimeout(r, 25));175 sockets.push(await connect(`?cursor=${start}&missionId=${missionId}`));176 }177 await new Promise((r) => setTimeout(r, 150));178 stop = true;179 await producer;180 const expected = (await h.db.q<{ id: number }>('SELECT id FROM events WHERE mission_id = $1 AND id > $2 ORDER BY id', [missionId, start])).map((r) => r.id);181 expect(expected.length).toBeGreaterThan(50);182 for (const s of sockets) {183 await until(() => persisted(s.events).length >= expected.length);184 expect(idsOf(s.events)).toEqual(expected); // same ids, same order, each exactly once185 s.ws.close();186 }187 });188 it('more than the replay bound is missing ⇒ the bound, then a system.notice "gap, refetch", then live', async () => {189 const { missionId, runId } = await seedMission(h, admin);190 const start = await headId();191 const ids: number[] = [];192 for (let i = 0; i < 9; i++) ids.push(await h.kernel.emit(missionId, runId, 'run.usage', { n: i }));193 const before = h.config.eventsReplayMax;194 h.config.eventsReplayMax = 5;195 try {196 const c = await connect(`?cursor=${start}&missionId=${missionId}`);197 await until(() => c.events.some((e) => e.kind === 'system.notice'));198 expect(idsOf(c.events)).toEqual(ids.slice(0, 5));199 const gap = c.events[5] as AosEvent;200 expect(gap).toMatchObject({ id: -1, kind: 'system.notice', trust: 'system', payload: { code: 'gap', message: 'gap, refetch', latestId: ids[8] } });201 const live = await h.kernel.emit(missionId, runId, 'run.usage', { n: 'live' });202 await until(() => c.events.some((e) => e.id === live));203 c.ws.close();204 // exactly the bound is not a gap205 h.config.eventsReplayMax = 10;206 const exact = await connect(`?cursor=${start}&missionId=${missionId}`);207 await until(() => exact.events.length >= 10);208 expect(exact.events.some((e) => e.kind === 'system.notice')).toBe(false);209 exact.ws.close();210 } finally {211 h.config.eventsReplayMax = before;212 }213 });214});215216describe('slow consumers', () => {217 it('a browser that stops reading is dropped; the others keep receiving', async () => {218 const { missionId, runId } = await seedMission(h, admin);219 const healthy = await connect(`?missionId=${missionId}`);220 // A raw socket that completes the WebSocket handshake and then never reads.221 const [host, port] = addr.split(':') as [string, string];222 const stuck = net.connect(Number(port), host);223 await new Promise<void>((r) => stuck.once('connect', () => r()));224 stuck.write(225 `GET /api/v1/events?missionId=${missionId} HTTP/1.1\r\nHost: ${addr}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n` +226 `Sec-WebSocket-Key: ${randomBytes(16).toString('base64')}\r\nSec-WebSocket-Version: 13\r\nOrigin: ${ORIGIN}\r\nCookie: ${admin.cookieHeader}\r\n\r\n`,227 );228 await new Promise<void>((r) => stuck.once('data', () => r())); // 101 Switching Protocols229 stuck.pause();230 await until(() => h.server.ctx.events.size === 2);231 const before = h.config.eventsSlowBytes;232 h.config.eventsSlowBytes = 64 * 1024;233 try {234 const chunk = 'x'.repeat(16 * 1024);235 for (let i = 0; i < 3000 && h.server.ctx.events.size === 2; i++) {236 h.kernel.emitTransient(missionId, runId, { text: chunk, i });237 if (i % 50 === 0) await new Promise((r) => setTimeout(r, 5));238 }239 await until(() => h.server.ctx.events.size === 1, 10_000);240 const last = await h.kernel.emit(missionId, runId, 'step.completed', { after: 'drop' });241 await until(() => healthy.events.some((e) => e.id === last), 10_000);242 } finally {243 h.config.eventsSlowBytes = before;244 stuck.destroy();245 healthy.ws.close();246 }247 });248});249250describe('auth session lifetime on open sockets', () => {251 it('logout sends system.notice auth.expired, then closes with 4401', async () => {252 const c = await login(h, secret);253 const s = await connect('', c);254 const done = closed(s.ws);255 expect((await c.post('/auth/logout')).statusCode).toBe(200);256 expect(await done).toBe(4401);257 expect(s.events.at(-1)).toMatchObject({ id: -1, kind: 'system.notice', payload: { code: 'auth.expired' } });258 await until(() => h.server.ctx.events.size === 0);259 expect(await upgradeStatus(openWs(addr, '/events', { cookie: c.cookieHeader }))).toBe(401);260 });261 it('a session revoked or expired behind our back is caught by the periodic re-check; nothing is delivered afterwards', async () => {262 const c = await login(h, secret);263 const { missionId } = await seedMission(h, c);264 const s = await connect('', c);265 const done = closed(s.ws);266 await h.db.q("UPDATE auth_sessions SET expires_at = now() - interval '1 second' WHERE token_hash <> '' AND id = (SELECT id FROM auth_sessions WHERE revoked_at IS NULL ORDER BY created_at DESC LIMIT 1)");267 expect(await done).toBe(4401);268 const n = s.events.length;269 await h.kernel.emit(missionId, null, 'system.notice', { code: 'after-expiry' });270 await new Promise((r) => setTimeout(r, 100));271 expect(s.events.length).toBe(n);272 expect(s.raw.join('')).not.toContain('after-expiry');273 });274 it('changing the password closes the sockets of every other session', async () => {275 const a = await login(h, secret);276 const b = await login(h, secret);277 const sb = await connect('', b);278 const done = closed(sb.ws);279 expect((await a.post('/auth/password', { currentPassword: PASSWORD, newPassword: PASSWORD })).statusCode).toBe(200);280 expect(await done).toBe(4401);281 });282});283284describe('secrets and notifications', () => {285 it('no frame ever contains a configured model API key, replayed or live, persisted or transient', async () => {286 const { missionId, runId } = await seedMission(h, admin);287 const start = await headId();288 // What a careless tool or provider error could put in the journal:289 await h.kernel.emit(missionId, runId, 'step.failed', { error: `401 invalid x-api-key: ${FAKE_ANTHROPIC_KEY}` });290 await h.kernel.emit(missionId, runId, 'step.completed', { output: `OPENAI_API_KEY=${FAKE_OPENAI_KEY}\nPATH=/usr/bin` }, 'tool');291 const c = await connect(`?cursor=${start}`);292 await until(() => persisted(c.events).length >= 2);293 await h.kernel.emit(missionId, runId, 'owner.message', { text: `my key is ${FAKE_ANTHROPIC_KEY}` }, 'owner');294 h.kernel.emitTransient(missionId, runId, { type: 'text', text: `echo ${FAKE_OPENAI_KEY}${FAKE_ANTHROPIC_KEY}` });295 await until(() => c.events.length >= 4);296 const all = c.raw.join('\n');297 expect(all).not.toContain(FAKE_ANTHROPIC_KEY);298 expect(all).not.toContain(FAKE_OPENAI_KEY);299 expect(all).toContain('[redacted]');300 expect(c.events.find((e) => e.kind === 'step.completed')?.payload.output).toBe('OPENAI_API_KEY=[redacted]\nPATH=/usr/bin');301 c.ws.close();302 });303 it('a wait and the end of a mission become in-app notifications (titles only)', async () => {304 await h.db.q('DELETE FROM notifications');305 const { missionId, runId } = await seedMission(h, admin, { mission: { title: 'Ship it' } });306 await h.kernel.seedWait(runId, 1, 'approval', { tool: 'git_push', input: { secretish: 'TOOLINPUT' } });307 await h.kernel.seedWait(runId, 2, 'question', { text: 'QUESTIONTEXT?' });308 await h.kernel.seedWait(runId, 3, 'timer', { ms: 10 });309 await h.kernel.emit(missionId, runId, 'run.status', { status: 'done' });310 await h.kernel.emit(missionId, null, 'mission.status', { status: 'running' });311 await h.kernel.emit(missionId, null, 'mission.status', { status: 'budget_exhausted' });312 await until(async () => (await h.db.q('SELECT 1 FROM notifications')).length >= 3);313 await new Promise((r) => setTimeout(r, 60));314 const list = (await admin.get('/notifications')).json<{ kind: string; title: string; body: string; link: string }[]>();315 expect(list.map((n) => [n.kind, n.title, n.body, n.link]).reverse()).toEqual([316 ['wait.approval', 'Approval needed', 'Ship it', `/mission/?id=${missionId}`],317 ['wait.question', 'The agent has a question', 'Ship it', `/mission/?id=${missionId}`],318 ['mission.finished', 'Mission budget exhausted', 'Ship it', `/mission/?id=${missionId}`],319 ]);320 expect(JSON.stringify(list)).not.toMatch(/TOOLINPUT|QUESTIONTEXT/);321 });322});323