From 120728c6ed26218b41cfb5e94597a4bfc038ec31 Mon Sep 17 00:00:00 2001 From: gfwangjie Date: Thu, 9 Apr 2026 13:29:49 +0800 Subject: [PATCH] fix(server): harden graph streaming --- .../test/unit/server-connection.test.ts | 88 ++++++++++ gitnexus/src/server/api.ts | 162 ++++++++++++++++-- .../test/unit/api-graph-streaming.test.ts | 128 ++++++++++++++ 3 files changed, 365 insertions(+), 13 deletions(-) create mode 100644 gitnexus/test/unit/api-graph-streaming.test.ts diff --git a/gitnexus-web/test/unit/server-connection.test.ts b/gitnexus-web/test/unit/server-connection.test.ts index b704f983f..ccbecd307 100644 --- a/gitnexus-web/test/unit/server-connection.test.ts +++ b/gitnexus-web/test/unit/server-connection.test.ts @@ -41,6 +41,27 @@ afterEach(() => { }); describe('fetchGraph', () => { + it('requests streamed graph responses from the backend', async () => { + setBackendUrl('http://localhost:4747'); + + const fetchMock = vi.fn().mockResolvedValue( + new Response('{"nodes":[],"relationships":[]}', { + status: 200, + headers: { + 'Content-Type': 'application/json', + }, + }), + ); + vi.stubGlobal('fetch', fetchMock); + + await fetchGraph('big-repo'); + + expect(fetchMock).toHaveBeenCalledWith( + expect.stringContaining('/api/graph?repo=big-repo&stream=true'), + expect.any(Object), + ); + }); + it('parses NDJSON graph streams incrementally', async () => { setBackendUrl('http://localhost:4747'); @@ -80,4 +101,71 @@ describe('fetchGraph', () => { expect(result.relationships[0].type).toBe('CONTAINS'); expect(progress).toHaveBeenCalled(); }); + + it('parses NDJSON graph lines split across chunks', async () => { + setBackendUrl('http://localhost:4747'); + + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode( + '{"type":"node","data":{"id":"File:src/app.ts","label":"File","properties":{"name":"app.ts"', + ), + ); + controller.enqueue( + encoder.encode( + ',"filePath":"src/app.ts"}}}\n{"type":"relationship","data":{"id":"File:src/app.ts_CONTAINS_Function:src/app.ts:main","type":"CONTAINS","sourceId":"File:src/app.ts","targetId":"Function:src/app.ts:main"}}\n', + ), + ); + controller.close(); + }, + }); + + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue( + new Response(stream, { + status: 200, + headers: { + 'Content-Type': 'application/x-ndjson', + }, + }), + ), + ); + + const result = await fetchGraph('big-repo'); + + expect(result.nodes).toHaveLength(1); + expect(result.relationships).toHaveLength(1); + expect(result.nodes[0].properties.filePath).toBe('src/app.ts'); + }); + + it('throws backend errors emitted in the NDJSON stream', async () => { + setBackendUrl('http://localhost:4747'); + + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode('{"type":"error","error":"stream failed"}\n')); + controller.close(); + }, + }); + + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue( + new Response(stream, { + status: 200, + headers: { + 'Content-Type': 'application/x-ndjson', + }, + }), + ), + ); + + await expect(fetchGraph('big-repo')).rejects.toMatchObject({ + message: 'stream failed', + }); + }); }); diff --git a/gitnexus/src/server/api.ts b/gitnexus/src/server/api.ts index efbdb7faf..c0e34e927 100644 --- a/gitnexus/src/server/api.ts +++ b/gitnexus/src/server/api.ts @@ -106,6 +106,104 @@ export const isAllowedOrigin = (origin: string | undefined): boolean => { return false; }; +type GraphStreamRecord = + | { type: 'node'; data: GraphNode } + | { type: 'relationship'; data: GraphRelationship } + | { type: 'error'; error: string }; + +export class ClientDisconnectedError extends Error { + constructor() { + super('Client disconnected during graph stream'); + this.name = 'ClientDisconnectedError'; + } +} + +export const isIgnorableGraphQueryError = (err: unknown): boolean => { + const message = err instanceof Error ? err.message : String(err); + return ( + message.includes('does not exist') || + message.includes('not found') || + message.includes('No table named') + ); +}; + +const ensureStreamIsWritable = ( + res: express.Response, + signal?: AbortSignal, +): void => { + if (signal?.aborted || res.destroyed || res.writableEnded) { + throw new ClientDisconnectedError(); + } +}; + +const waitForDrain = async ( + res: express.Response, + signal?: AbortSignal, +): Promise => { + ensureStreamIsWritable(res, signal); + + await new Promise((resolve, reject) => { + const cleanup = () => { + res.off('drain', onDrain); + res.off('close', onClose); + signal?.removeEventListener('abort', onAbort); + }; + + const onDrain = () => { + cleanup(); + resolve(); + }; + const onClose = () => { + cleanup(); + reject(new ClientDisconnectedError()); + }; + const onAbort = () => { + cleanup(); + reject(new ClientDisconnectedError()); + }; + + res.once('drain', onDrain); + res.once('close', onClose); + signal?.addEventListener('abort', onAbort, { once: true }); + + if (signal?.aborted || res.destroyed || res.writableEnded) { + onAbort(); + } + }); + + ensureStreamIsWritable(res, signal); +}; + +const isClientDisconnectWriteError = (err: unknown): boolean => { + if (!(err instanceof Error)) return false; + return ( + (err as NodeJS.ErrnoException).code === 'ERR_STREAM_DESTROYED' || + (err as NodeJS.ErrnoException).code === 'EPIPE' || + (err as NodeJS.ErrnoException).code === 'ECONNRESET' || + err.message.includes('write after end') + ); +}; + +export const writeNdjsonRecord = async ( + res: express.Response, + record: GraphStreamRecord, + signal?: AbortSignal, +): Promise => { + ensureStreamIsWritable(res, signal); + + try { + const canContinue = res.write(JSON.stringify(record) + '\n'); + if (!canContinue) { + await waitForDrain(res, signal); + } + } catch (err) { + if (isClientDisconnectWriteError(err)) { + throw new ClientDisconnectedError(); + } + throw err; + } +}; + const buildGraph = async ( includeContent = false, ): Promise<{ nodes: GraphNode[]; relationships: GraphRelationship[] }> => { @@ -116,8 +214,10 @@ const buildGraph = async ( for (const row of rows) { nodes.push(mapGraphNodeRow(table, row, includeContent)); } - } catch { - // ignore empty tables + } catch (err) { + if (!isIgnorableGraphQueryError(err)) { + throw err; + } } } @@ -184,31 +284,38 @@ const mapGraphRelationshipRow = (row: any): GraphRelationship => ({ step: row.step, }); -const streamGraphNdjson = async ( +export const streamGraphNdjson = async ( res: express.Response, includeContent = false, + signal?: AbortSignal, ): Promise => { for (const table of NODE_TABLES) { try { await streamQuery(getNodeQuery(table, includeContent), async (row) => { - res.write( - JSON.stringify({ + await writeNdjsonRecord( + res, + { type: 'node', data: mapGraphNodeRow(table, row, includeContent), - }) + '\n', + }, + signal, ); }); - } catch { - // ignore empty tables + } catch (err) { + if (!isIgnorableGraphQueryError(err)) { + throw err; + } } } await streamQuery(GRAPH_RELATIONSHIP_QUERY, async (row) => { - res.write( - JSON.stringify({ + await writeNdjsonRecord( + res, + { type: 'relationship', data: mapGraphRelationshipRow(row), - }) + '\n', + }, + signal, ); }); }; @@ -506,17 +613,46 @@ export const createServer = async (port: number, host: string = '127.0.0.1') => const stream = req.query.stream === 'true'; if (stream) { + const abortController = new AbortController(); + let responseFinished = false; + const markFinished = () => { + responseFinished = true; + }; + const abortStreaming = () => { + if (!responseFinished) { + abortController.abort(); + } + }; + res.setHeader('Content-Type', 'application/x-ndjson; charset=utf-8'); res.setHeader('Cache-Control', 'no-cache'); res.flushHeaders(); - await withLbugDb(lbugPath, async () => streamGraphNdjson(res, includeContent)); - res.end(); + + req.once('aborted', abortStreaming); + res.once('finish', markFinished); + res.once('close', abortStreaming); + + try { + await withLbugDb(lbugPath, async () => + streamGraphNdjson(res, includeContent, abortController.signal), + ); + if (!abortController.signal.aborted && !res.writableEnded) { + res.end(); + } + } finally { + req.off('aborted', abortStreaming); + res.off('finish', markFinished); + res.off('close', abortStreaming); + } return; } const graph = await withLbugDb(lbugPath, async () => buildGraph(includeContent)); res.json(graph); } catch (err: any) { + if (err instanceof ClientDisconnectedError) { + return; + } const message = err.message || 'Failed to build graph'; if (res.headersSent) { try { diff --git a/gitnexus/test/unit/api-graph-streaming.test.ts b/gitnexus/test/unit/api-graph-streaming.test.ts new file mode 100644 index 000000000..2d22b5ad9 --- /dev/null +++ b/gitnexus/test/unit/api-graph-streaming.test.ts @@ -0,0 +1,128 @@ +import { EventEmitter } from 'node:events'; +import { describe, expect, it, vi, beforeEach } from 'vitest'; + +const { lbugMocks } = vi.hoisted(() => ({ + lbugMocks: { + streamQuery: vi.fn(), + }, +})); + +vi.mock('../../src/core/lbug/lbug-adapter.js', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, ...lbugMocks }; +}); + +import { + ClientDisconnectedError, + streamGraphNdjson, +} from '../../src/server/api.js'; + +const createMockResponse = (writeImpl?: (chunk: string) => boolean) => { + const response = new EventEmitter() as any; + response.writableEnded = false; + response.destroyed = false; + response.write = vi.fn((chunk: string) => (writeImpl ? writeImpl(chunk) : true)); + return response; +}; + +describe('streamGraphNdjson', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('waits for drain when writes hit backpressure', async () => { + lbugMocks.streamQuery.mockImplementation(async (query: string, onRow: (row: any) => Promise) => { + if (query.includes('MATCH (n:File)')) { + await onRow({ id: 'File:src/app.ts', name: 'app.ts', filePath: 'src/app.ts' }); + return 1; + } + if (query.includes('CodeRelation')) { + await onRow({ + sourceId: 'File:src/app.ts', + targetId: 'Function:src/app.ts:main', + type: 'CONTAINS', + }); + return 1; + } + return 0; + }); + + const writes: string[] = []; + let firstWrite = true; + const response = createMockResponse((chunk) => { + writes.push(chunk); + if (firstWrite) { + firstWrite = false; + return false; + } + return true; + }); + + let settled = false; + const pending = streamGraphNdjson(response, false).then(() => { + settled = true; + }); + + await Promise.resolve(); + expect(writes).toHaveLength(1); + expect(settled).toBe(false); + + response.emit('drain'); + await pending; + + expect(writes).toHaveLength(2); + }); + + it('stops streaming when the client disconnects', async () => { + const controller = new AbortController(); + lbugMocks.streamQuery.mockImplementation(async (query: string, onRow: (row: any) => Promise) => { + if (!query.includes('MATCH (n:File)')) { + return 0; + } + await onRow({ id: 'File:src/app.ts', name: 'app.ts', filePath: 'src/app.ts' }); + controller.abort(); + await onRow({ id: 'File:src/other.ts', name: 'other.ts', filePath: 'src/other.ts' }); + return 2; + }); + + const response = createMockResponse(); + + await expect(streamGraphNdjson(response, false, controller.signal)).rejects.toBeInstanceOf( + ClientDisconnectedError, + ); + expect(response.write).toHaveBeenCalledTimes(1); + }); + + it('rethrows non-missing table errors', async () => { + lbugMocks.streamQuery.mockImplementation(async (query: string) => { + if (query.includes('MATCH (n:File)')) { + throw new Error('database unavailable'); + } + return 0; + }); + + const response = createMockResponse(); + await expect(streamGraphNdjson(response, false)).rejects.toThrow('database unavailable'); + }); + + it('ignores missing-table errors while continuing the stream', async () => { + lbugMocks.streamQuery.mockImplementation(async (query: string, onRow: (row: any) => Promise) => { + if (query.includes('MATCH (n:File)')) { + throw new Error('Table File does not exist'); + } + if (query.includes('CodeRelation')) { + await onRow({ + sourceId: 'File:src/app.ts', + targetId: 'Function:src/app.ts:main', + type: 'CONTAINS', + }); + return 1; + } + return 0; + }); + + const response = createMockResponse(); + await expect(streamGraphNdjson(response, false)).resolves.toBeUndefined(); + expect(response.write).toHaveBeenCalledTimes(1); + }); +});