diff --git a/src/main/providers/ssh-filesystem-provider-stream.test.ts b/src/main/providers/ssh-filesystem-provider-stream.test.ts new file mode 100644 index 000000000..1932af4ce --- /dev/null +++ b/src/main/providers/ssh-filesystem-provider-stream.test.ts @@ -0,0 +1,182 @@ +import { describe, expect, it, vi, beforeEach } from 'vitest' +import { SshFilesystemProvider } from './ssh-filesystem-provider' + +type MockMultiplexer = { + request: ReturnType + notify: ReturnType + onNotification: ReturnType + onNotificationByMethod: ReturnType + onDispose: ReturnType + dispose: ReturnType + isDisposed: ReturnType + _emitMethod: (method: string, params: Record) => void +} + +function createMockMux(): MockMultiplexer { + const methodHandlers = new Map) => void>>() + return { + request: vi.fn().mockResolvedValue(undefined), + notify: vi.fn(), + onNotification: vi.fn(), + onNotificationByMethod: vi.fn( + (method: string, handler: (params: Record) => void) => { + let set = methodHandlers.get(method) + if (!set) { + set = new Set() + methodHandlers.set(method, set) + } + set.add(handler) + return () => set!.delete(handler) + } + ), + onDispose: vi.fn(() => () => {}), + dispose: vi.fn(), + isDisposed: vi.fn().mockReturnValue(false), + _emitMethod: (method, params) => { + const set = methodHandlers.get(method) + if (set) { + for (const handler of Array.from(set)) { + handler(params) + } + } + } + } +} + +describe('SshFilesystemProvider readFile streaming', () => { + let mux: MockMultiplexer + let provider: SshFilesystemProvider + + beforeEach(() => { + mux = createMockMux() + provider = new SshFilesystemProvider('conn-1', mux as never) + }) + + it('streams via fs.readFileStream and reassembles utf-8 text', async () => { + const text = 'hello world' + const totalSize = Buffer.byteLength(text, 'utf-8') + mux.request.mockImplementation(async (method: string) => { + if (method !== 'fs.readFileStream') { + throw new Error(`unexpected method ${method}`) + } + // Why: setImmediate fires after the metadata-resolution .then has set + // streamIdRef, ensuring subscribed handlers see a matching streamId. + setImmediate(() => { + mux._emitMethod('fs.streamChunk', { + streamId: 1, + seq: 0, + data: Buffer.from(text, 'utf-8').toString('base64') + }) + mux._emitMethod('fs.streamEnd', { streamId: 1 }) + }) + return { + streamId: 1, + totalSize, + isBinary: false, + chunkEncoding: 'base64', + resultEncoding: 'utf-8' + } + }) + + const result = await provider.readFile('/home/user/file.txt') + expect(mux.request).toHaveBeenCalledWith('fs.readFileStream', { + filePath: '/home/user/file.txt' + }) + expect(result).toEqual({ content: text, isBinary: false }) + }) + + it('falls back to legacy fs.readFile on -32601 method-not-found', async () => { + const legacyResult = { content: 'legacy', isBinary: false } + mux.request.mockImplementation(async (method: string) => { + if (method === 'fs.readFileStream') { + const err = new Error('Method not found') as Error & { code: number } + err.code = -32601 + throw err + } + if (method === 'fs.readFile') { + return legacyResult + } + throw new Error(`unexpected method ${method}`) + }) + + const result = await provider.readFile('/home/user/file.txt') + expect(mux.request).toHaveBeenCalledWith('fs.readFile', { filePath: '/home/user/file.txt' }) + expect(result).toEqual(legacyResult) + }) + + it('rejects when chunk arrives out of order', async () => { + const totalSize = 256 * 1024 * 2 + mux.request.mockImplementation(async () => { + setImmediate(() => { + mux._emitMethod('fs.streamChunk', { + streamId: 1, + seq: 1, + data: Buffer.alloc(256 * 1024).toString('base64') + }) + }) + return { + streamId: 1, + totalSize, + isBinary: true, + chunkEncoding: 'base64', + resultEncoding: 'base64' + } + }) + await expect(provider.readFile('/home/x.bin')).rejects.toThrow(/out-of-order/i) + }) + + it('rejects when totalSize exceeds client cap without allocating', async () => { + mux.request.mockResolvedValue({ + streamId: 1, + totalSize: 51 * 1024 * 1024, + isBinary: true, + chunkEncoding: 'base64', + resultEncoding: 'base64' + }) + await expect(provider.readFile('/home/x.bin')).rejects.toThrow(/exceeds client cap/i) + expect(mux.notify).toHaveBeenCalledWith('fs.cancelStream', { streamId: 1 }) + }) + + it('rejects on fs.streamError notification', async () => { + const totalSize = 1024 + mux.request.mockImplementation(async () => { + setImmediate(() => { + mux._emitMethod('fs.streamError', { + streamId: 7, + code: 'ENOENT', + message: 'gone' + }) + }) + return { + streamId: 7, + totalSize, + isBinary: false, + chunkEncoding: 'base64', + resultEncoding: 'utf-8' + } + }) + await expect(provider.readFile('/home/x.txt')).rejects.toThrow(/gone/) + }) + + it('rejects on chunk count mismatch at streamEnd', async () => { + const totalSize = 256 * 1024 * 3 + mux.request.mockImplementation(async () => { + setImmediate(() => { + mux._emitMethod('fs.streamChunk', { + streamId: 1, + seq: 0, + data: Buffer.alloc(256 * 1024).toString('base64') + }) + mux._emitMethod('fs.streamEnd', { streamId: 1 }) + }) + return { + streamId: 1, + totalSize, + isBinary: true, + chunkEncoding: 'base64', + resultEncoding: 'base64' + } + }) + await expect(provider.readFile('/home/x.bin')).rejects.toThrow(/count mismatch/i) + }) +}) diff --git a/src/main/providers/ssh-filesystem-provider.test.ts b/src/main/providers/ssh-filesystem-provider.test.ts index e83ac485e..dca8cf06b 100644 --- a/src/main/providers/ssh-filesystem-provider.test.ts +++ b/src/main/providers/ssh-filesystem-provider.test.ts @@ -5,17 +5,43 @@ type MockMultiplexer = { request: ReturnType notify: ReturnType onNotification: ReturnType + onNotificationByMethod: ReturnType + onDispose: ReturnType dispose: ReturnType isDisposed: ReturnType + _methodHandlers: Map) => void>> + _emitMethod: (method: string, params: Record) => void } function createMockMux(): MockMultiplexer { + const methodHandlers = new Map) => void>>() return { request: vi.fn().mockResolvedValue(undefined), notify: vi.fn(), onNotification: vi.fn(), + onNotificationByMethod: vi.fn( + (method: string, handler: (params: Record) => void) => { + let set = methodHandlers.get(method) + if (!set) { + set = new Set() + methodHandlers.set(method, set) + } + set.add(handler) + return () => set!.delete(handler) + } + ), + onDispose: vi.fn(() => () => {}), dispose: vi.fn(), - isDisposed: vi.fn().mockReturnValue(false) + isDisposed: vi.fn().mockReturnValue(false), + _methodHandlers: methodHandlers, + _emitMethod: (method, params) => { + const set = methodHandlers.get(method) + if (set) { + for (const handler of Array.from(set)) { + handler(params) + } + } + } } } @@ -47,13 +73,10 @@ describe('SshFilesystemProvider', () => { }) describe('readFile', () => { - it('sends fs.readFile request', async () => { - const fileResult = { content: 'hello world', isBinary: false } - mux.request.mockResolvedValue(fileResult) - - const result = await provider.readFile('/home/user/file.txt') - expect(mux.request).toHaveBeenCalledWith('fs.readFile', { filePath: '/home/user/file.txt' }) - expect(result).toEqual(fileResult) + it('short-circuits on empty:true metadata without subscribing to chunks', async () => { + mux.request.mockResolvedValue({ totalSize: 0, isBinary: false, empty: true }) + const result = await provider.readFile('/home/user/empty.txt') + expect(result).toEqual({ content: '', isBinary: false }) }) }) diff --git a/src/main/providers/ssh-filesystem-provider.ts b/src/main/providers/ssh-filesystem-provider.ts index 903dc71e7..d0c545d68 100644 --- a/src/main/providers/ssh-filesystem-provider.ts +++ b/src/main/providers/ssh-filesystem-provider.ts @@ -1,4 +1,5 @@ import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer' +import { isMethodNotFoundError, readFileViaStream } from '../ssh/ssh-filesystem-stream-reader' import type { IFilesystemProvider, FileStat, FileReadResult } from './types' import type { DirEntry, FsChangeEvent, SearchOptions, SearchResult } from '../../shared/types' @@ -13,6 +14,10 @@ export class SshFilesystemProvider implements IFilesystemProvider { // multiplexer. Without this, notification callbacks keep firing after // the provider is torn down on disconnect, routing events to stale state. private unsubscribeNotifications: (() => void) | null = null + // Why: relays from a previous build may not implement fs.readFileStream. + // We log the fallback once per session at warn level so users on stale + // relays get diagnosed quickly without per-read log spam. + private loggedStreamFallback = false constructor(connectionId: string, mux: SshChannelMultiplexer) { this.connectionId = connectionId @@ -48,7 +53,25 @@ export class SshFilesystemProvider implements IFilesystemProvider { } async readFile(filePath: string): Promise { - return (await this.mux.request('fs.readFile', { filePath })) as FileReadResult + // Why: streaming is the default path so previews above the legacy single- + // frame budget (~12 MB after base64) don't hit MAX_MESSAGE_SIZE. Old relays + // that don't implement fs.readFileStream surface as MethodNotFound; we fall + // back to the legacy single-shot fs.readFile (which retains the old 10 MB + // cap on those hosts). + try { + return await readFileViaStream(this.mux, filePath) + } catch (err) { + if (isMethodNotFoundError(err)) { + if (!this.loggedStreamFallback) { + this.loggedStreamFallback = true + console.warn( + '[ssh-fs] Relay does not implement fs.readFileStream; falling back to fs.readFile (10 MB cap)' + ) + } + return (await this.mux.request('fs.readFile', { filePath })) as FileReadResult + } + throw err + } } async writeFile(filePath: string, content: string): Promise { diff --git a/src/main/ssh/relay-protocol.ts b/src/main/ssh/relay-protocol.ts index 5a134fc7f..52a0e1508 100644 --- a/src/main/ssh/relay-protocol.ts +++ b/src/main/ssh/relay-protocol.ts @@ -36,9 +36,27 @@ export const RelayErrorCode = { PermissionDenied: -33002, PathNotFound: -33003, PtyAllocationFailed: -33004, - DiskFull: -33005 + DiskFull: -33005, + TooManyStreams: -33006, + StreamProtocolError: -33007 } as const +export const JsonRpcErrorCode = { + MethodNotFound: -32601 +} as const + +// ── Streaming constants (see docs/relay-file-stream-design.md) ───── + +/** Per-chunk payload size for fs.readFileStream. Mirrors VS Code's + * `bufferSize: 256 * 1024` (vs/platform/files/node/diskFileSystemProvider.ts). + * 256KB raw → ~340KB base64, well under MAX_MESSAGE_SIZE. */ +export const STREAM_CHUNK_SIZE = 256 * 1024 + +/** Cap on concurrent in-flight streams per relay; mirrors fs.watch's + * 20-watcher cap idiom. Prevents file-descriptor exhaustion from a buggy + * client. */ +export const MAX_CONCURRENT_STREAMS = 16 + // ── JSON-RPC types ────────────────────────────────────────────────── export type JsonRpcRequest = { diff --git a/src/main/ssh/ssh-channel-multiplexer.test.ts b/src/main/ssh/ssh-channel-multiplexer.test.ts index b547d5f51..f8ba76209 100644 --- a/src/main/ssh/ssh-channel-multiplexer.test.ts +++ b/src/main/ssh/ssh-channel-multiplexer.test.ts @@ -167,6 +167,40 @@ describe('SshChannelMultiplexer', () => { expect(handler).toHaveBeenCalledWith('pty.exit', { id: 'pty-1', code: 0 }) }) + + it('typed dispatcher only fires for its method', () => { + const chunkHandler = vi.fn() + const otherHandler = vi.fn() + const generic = vi.fn() + mux.onNotificationByMethod('fs.streamChunk', chunkHandler) + mux.onNotificationByMethod('fs.streamEnd', otherHandler) + mux.onNotification(generic) + + transport.dataCallbacks[0]( + makeNotificationFrame('fs.streamChunk', { streamId: 1, seq: 0, data: 'aGk=' }, 1) + ) + + expect(chunkHandler).toHaveBeenCalledWith({ streamId: 1, seq: 0, data: 'aGk=' }) + expect(otherHandler).not.toHaveBeenCalled() + expect(generic).toHaveBeenCalledWith('fs.streamChunk', { + streamId: 1, + seq: 0, + data: 'aGk=' + }) + }) + + it('typed dispatcher unsubscribe removes only that handler', () => { + const a = vi.fn() + const b = vi.fn() + const unsubA = mux.onNotificationByMethod('fs.streamEnd', a) + mux.onNotificationByMethod('fs.streamEnd', b) + unsubA() + + transport.dataCallbacks[0](makeNotificationFrame('fs.streamEnd', { streamId: 7 }, 1)) + + expect(a).not.toHaveBeenCalled() + expect(b).toHaveBeenCalledWith({ streamId: 7 }) + }) }) describe('keepalive', () => { diff --git a/src/main/ssh/ssh-channel-multiplexer.ts b/src/main/ssh/ssh-channel-multiplexer.ts index 8d1a963ee..11a18bebe 100644 --- a/src/main/ssh/ssh-channel-multiplexer.ts +++ b/src/main/ssh/ssh-channel-multiplexer.ts @@ -27,6 +27,7 @@ type PendingRequest = { } export type NotificationHandler = (method: string, params: Record) => void +export type MethodNotificationHandler = (params: Record) => void const REQUEST_TIMEOUT_MS = 30_000 @@ -40,6 +41,10 @@ export class SshChannelMultiplexer { private lastReceivedAt = Date.now() private pendingRequests = new Map() private notificationHandlers: NotificationHandler[] = [] + // Why: per-method dispatch map keeps streaming consumers (fs.streamChunk, + // fs.streamEnd, fs.streamError) from accreting string-match logic in the + // generic notification listener that already serves fs.changed. + private methodNotificationHandlers = new Map>() private disposeHandlers: ((reason: 'shutdown' | 'connection_lost') => void)[] = [] private keepaliveTimer: ReturnType | null = null private timeoutTimer: ReturnType | null = null @@ -85,6 +90,25 @@ export class SshChannelMultiplexer { } } + onNotificationByMethod(method: string, handler: MethodNotificationHandler): () => void { + let set = this.methodNotificationHandlers.get(method) + if (!set) { + set = new Set() + this.methodNotificationHandlers.set(method, set) + } + set.add(handler) + return () => { + const current = this.methodNotificationHandlers.get(method) + if (!current) { + return + } + current.delete(handler) + if (current.size === 0) { + this.methodNotificationHandlers.delete(method) + } + } + } + // Why: the session needs to know when the relay channel dies so it can // auto-reconnect. Without this, a relay channel close (e.g. --connect // bridge exits) leaves the session in 'ready' state with a dead mux @@ -178,6 +202,7 @@ export class SshChannelMultiplexer { } this.unackedTimestamps.clear() + this.methodNotificationHandlers.clear() this.decoder.reset() this.transport.close?.() @@ -286,12 +311,19 @@ export class SshChannelMultiplexer { private handleNotification(msg: JsonRpcNotification): void { const params = msg.params ?? {} // Why: handlers may unsubscribe during iteration (via the returned disposer - // from onNotification), which splices the live array and skips the next handler. - // Iterating a snapshot prevents that. + // from onNotification / onNotificationByMethod), which mutates the live + // collection and skips the next handler. Iterating a snapshot prevents that. const snapshot = Array.from(this.notificationHandlers) for (const handler of snapshot) { handler(msg.method, params) } + const methodHandlers = this.methodNotificationHandlers.get(msg.method) + if (methodHandlers && methodHandlers.size > 0) { + const methodSnapshot = Array.from(methodHandlers) + for (const handler of methodSnapshot) { + handler(params) + } + } } private startKeepalive(): void { diff --git a/src/main/ssh/ssh-filesystem-stream-reader.ts b/src/main/ssh/ssh-filesystem-stream-reader.ts new file mode 100644 index 000000000..b06ab52ed --- /dev/null +++ b/src/main/ssh/ssh-filesystem-stream-reader.ts @@ -0,0 +1,302 @@ +import type { SshChannelMultiplexer } from './ssh-channel-multiplexer' +import { STREAM_CHUNK_SIZE, JsonRpcErrorCode, RelayErrorCode } from './relay-protocol' +import type { FileReadResult } from '../providers/types' + +const RESULT_ENCODING_BASE64 = 'base64' +const SENTINEL_STREAM_ID = -1 + +const MAX_PREVIEWABLE_BINARY_SIZE = 50 * 1024 * 1024 +const MAX_TEXT_FILE_SIZE = 10 * 1024 * 1024 + +type StreamMetadataResponse = { + streamId?: number + totalSize: number + isBinary: boolean + isImage?: boolean + mimeType?: string + resultEncoding?: 'base64' | 'utf-8' + empty?: boolean +} + +export function isMethodNotFoundError(err: unknown): boolean { + if (!err || typeof err !== 'object') { + return false + } + const code = (err as { code?: unknown }).code + return code === JsonRpcErrorCode.MethodNotFound +} + +export class StreamProtocolError extends Error { + readonly code = RelayErrorCode.StreamProtocolError + constructor(message: string) { + super(message) + } +} + +export async function readFileViaStream( + mux: SshChannelMultiplexer, + filePath: string +): Promise { + // Why: subscribe BEFORE awaiting the metadata response so a chunk arriving + // immediately after the response cannot beat the listener registration. + // streamIdRef stays at SENTINEL_STREAM_ID until metadata resolves; chunk + // handlers compare against it and drop unmatched ids cleanly. + const streamIdRef = { current: SENTINEL_STREAM_ID } + const unsubscribers: (() => void)[] = [] + const cleanup = (): void => { + while (unsubscribers.length > 0) { + const fn = unsubscribers.pop() + try { + fn?.() + } catch { + // Best-effort cleanup + } + } + } + + return new Promise((resolve, reject) => { + let buffer: Buffer | null = null + let resultEncoding: 'base64' | 'utf-8' = RESULT_ENCODING_BASE64 + let isBinary = false + let isImage: boolean | undefined + let mimeType: string | undefined + let totalSize = 0 + let expectedSeq = 0 + let receivedChunks = 0 + let totalChunks = 0 + let settled = false + + // Why: chunk/end/error frames may arrive in the same dispatch tick as the + // metadata response. Queue them until streamIdRef is set, then drain. + type PendingFrame = + | { kind: 'chunk'; params: Record } + | { kind: 'end'; params: Record } + | { kind: 'error'; params: Record } + const pending: PendingFrame[] = [] + let metadataReady = false + + const cancel = (): void => { + if (streamIdRef.current !== SENTINEL_STREAM_ID && !mux.isDisposed()) { + try { + mux.notify('fs.cancelStream', { streamId: streamIdRef.current }) + } catch { + // Best-effort + } + } + } + + const fail = (err: Error): void => { + if (settled) { + return + } + settled = true + cancel() + cleanup() + reject(err) + } + + const succeed = (value: FileReadResult): void => { + if (settled) { + return + } + settled = true + cleanup() + resolve(value) + } + + const handleChunk = (params: Record): void => { + if (settled) { + return + } + const id = params.streamId as number | undefined + if (id !== streamIdRef.current) { + return + } + const seq = params.seq as number + const data = params.data as string + if (typeof seq !== 'number' || typeof data !== 'string') { + fail(new StreamProtocolError(`Malformed chunk for stream ${id}`)) + return + } + if (seq !== expectedSeq) { + fail( + new StreamProtocolError( + `Out-of-order chunk for stream ${id}: expected ${expectedSeq}, got ${seq}` + ) + ) + return + } + const offset = seq * STREAM_CHUNK_SIZE + const decoded = Buffer.from(data, 'base64') + if (offset + decoded.length > totalSize) { + fail( + new StreamProtocolError( + `Chunk overflows declared totalSize: offset=${offset} len=${decoded.length} total=${totalSize}` + ) + ) + return + } + if (!buffer) { + fail(new StreamProtocolError(`Chunk arrived before metadata for stream ${id}`)) + return + } + decoded.copy(buffer, offset) + expectedSeq += 1 + receivedChunks += 1 + } + + const handleEnd = (params: Record): void => { + if (settled) { + return + } + const id = params.streamId as number | undefined + if (id !== streamIdRef.current) { + return + } + if (receivedChunks !== totalChunks) { + fail( + new StreamProtocolError( + `Chunk count mismatch for stream ${id}: expected ${totalChunks}, received ${receivedChunks}` + ) + ) + return + } + if (!buffer) { + fail(new StreamProtocolError(`Stream end before metadata for stream ${id}`)) + return + } + const content = + resultEncoding === RESULT_ENCODING_BASE64 + ? buffer.toString('base64') + : buffer.toString('utf-8') + succeed({ + content, + isBinary, + ...(isImage !== undefined ? { isImage } : {}), + ...(mimeType !== undefined ? { mimeType } : {}) + }) + } + + const handleStreamError = (params: Record): void => { + if (settled) { + return + } + const id = params.streamId as number | undefined + if (id !== streamIdRef.current) { + return + } + const message = (params.message as string | undefined) ?? 'stream error' + const code = (params.code as string | undefined) ?? 'ESTREAMERROR' + const err = new Error(message) as Error & { code: string } + err.code = code + fail(err) + } + + const drainPending = (): void => { + while (!settled && pending.length > 0) { + const frame = pending.shift()! + if (frame.kind === 'chunk') { + handleChunk(frame.params) + } else if (frame.kind === 'end') { + handleEnd(frame.params) + } else { + handleStreamError(frame.params) + } + } + } + + unsubscribers.push( + mux.onNotificationByMethod('fs.streamChunk', (params) => { + if (!metadataReady) { + pending.push({ kind: 'chunk', params }) + return + } + handleChunk(params) + }) + ) + unsubscribers.push( + mux.onNotificationByMethod('fs.streamEnd', (params) => { + if (!metadataReady) { + pending.push({ kind: 'end', params }) + return + } + handleEnd(params) + }) + ) + unsubscribers.push( + mux.onNotificationByMethod('fs.streamError', (params) => { + if (!metadataReady) { + pending.push({ kind: 'error', params }) + return + } + handleStreamError(params) + }) + ) + + const onDispose = mux.onDispose((reason) => { + const message = + reason === 'connection_lost' + ? 'SSH connection lost, reconnecting...' + : 'Multiplexer disposed' + const err = new Error(message) as Error & { code: string } + err.code = reason === 'connection_lost' ? 'CONNECTION_LOST' : 'DISPOSED' + fail(err) + }) + unsubscribers.push(onDispose) + + void mux + .request('fs.readFileStream', { filePath }) + .then((rawMetadata) => { + if (settled) { + return + } + const metadata = rawMetadata as StreamMetadataResponse + isBinary = metadata.isBinary + isImage = metadata.isImage + mimeType = metadata.mimeType + resultEncoding = metadata.resultEncoding ?? RESULT_ENCODING_BASE64 + + if (metadata.empty) { + succeed({ + content: '', + isBinary: metadata.isBinary, + ...(metadata.isImage !== undefined ? { isImage: metadata.isImage } : {}), + ...(metadata.mimeType !== undefined ? { mimeType: metadata.mimeType } : {}) + }) + return + } + + if (typeof metadata.streamId !== 'number') { + fail(new StreamProtocolError('Metadata missing streamId for non-empty stream')) + return + } + + const cap = metadata.isBinary ? MAX_PREVIEWABLE_BINARY_SIZE : MAX_TEXT_FILE_SIZE + if (metadata.totalSize < 0 || metadata.totalSize > cap) { + streamIdRef.current = metadata.streamId + fail( + new StreamProtocolError( + `Reported totalSize ${metadata.totalSize} exceeds client cap ${cap}` + ) + ) + return + } + + totalSize = metadata.totalSize + totalChunks = totalSize === 0 ? 0 : Math.ceil(totalSize / STREAM_CHUNK_SIZE) + try { + buffer = Buffer.alloc(totalSize) + } catch (err) { + streamIdRef.current = metadata.streamId + fail(new Error(`Failed to allocate ${totalSize} bytes: ${(err as Error).message}`)) + return + } + streamIdRef.current = metadata.streamId + metadataReady = true + drainPending() + }) + .catch((err) => { + fail(err as Error) + }) + }) +} diff --git a/src/main/ssh/ssh-relay-deploy.test.ts b/src/main/ssh/ssh-relay-deploy.test.ts index a6b0ca8d5..362b60ac3 100644 --- a/src/main/ssh/ssh-relay-deploy.test.ts +++ b/src/main/ssh/ssh-relay-deploy.test.ts @@ -87,6 +87,7 @@ describe('deployAndLaunchRelay', () => { const mockExecCommand = vi.mocked(execCommand) mockExecCommand.mockResolvedValueOnce('Linux x86_64') // uname -sm mockExecCommand.mockResolvedValueOnce('/home/user') // echo $HOME + mockExecCommand.mockResolvedValueOnce('ORCA-NATIVE-DEPS-OK') // native deps probe mockExecCommand.mockResolvedValueOnce('DEAD') // socket probe mockExecCommand.mockResolvedValueOnce('READY') // socket poll @@ -100,6 +101,7 @@ describe('deployAndLaunchRelay', () => { const mockExecCommand = vi.mocked(execCommand) mockExecCommand.mockResolvedValueOnce('Linux x86_64') mockExecCommand.mockResolvedValueOnce('/home/user') + mockExecCommand.mockResolvedValueOnce('ORCA-NATIVE-DEPS-OK') // native deps probe mockExecCommand.mockResolvedValueOnce('DEAD') // socket probe mockExecCommand.mockResolvedValueOnce('READY') // socket poll @@ -115,6 +117,7 @@ describe('deployAndLaunchRelay', () => { const mockExecCommand = vi.mocked(execCommand) mockExecCommand.mockResolvedValueOnce('Linux x86_64') mockExecCommand.mockResolvedValueOnce('/home/user') + mockExecCommand.mockResolvedValueOnce('ORCA-NATIVE-DEPS-OK') mockExecCommand.mockResolvedValueOnce('DEAD') mockExecCommand.mockResolvedValueOnce('READY') @@ -159,10 +162,12 @@ describe('deployAndLaunchRelay', () => { mockExecCommand .mockResolvedValueOnce('Linux x86_64') // uname A .mockResolvedValueOnce('/home/user') // $HOME A + .mockResolvedValueOnce('ORCA-NATIVE-DEPS-OK') // native deps probe A .mockResolvedValueOnce('DEAD') // probe A .mockResolvedValueOnce('READY') // poll A .mockResolvedValueOnce('Linux x86_64') // uname B .mockResolvedValueOnce('/home/user') // $HOME B + .mockResolvedValueOnce('ORCA-NATIVE-DEPS-OK') // native deps probe B .mockResolvedValueOnce('DEAD') // probe B .mockResolvedValueOnce('READY') // poll B diff --git a/src/main/ssh/ssh-relay-deploy.ts b/src/main/ssh/ssh-relay-deploy.ts index a166d2d7b..03dbeaf03 100644 --- a/src/main/ssh/ssh-relay-deploy.ts +++ b/src/main/ssh/ssh-relay-deploy.ts @@ -117,7 +117,9 @@ async function deployAndLaunchRelayInner( const alreadyInstalled = await isRelayAlreadyInstalled(conn, remoteRelayDir) console.log(`[ssh-relay] Already installed at ${fullVersion}: ${alreadyInstalled}`) - if (!alreadyInstalled) { + if (alreadyInstalled) { + await repairInstalledNativeDeps(conn, remoteRelayDir, platform) + } else { // Why: serialize concurrent first-installs of the same version against // each other via an atomic mkdir lock. The losing caller polls and either // re-checks `alreadyInstalled` (now true) or steals a stale lock. @@ -132,7 +134,7 @@ async function deployAndLaunchRelayInner( console.log('[ssh-relay] Upload complete') onProgress?.('Installing native dependencies...') - console.log('[ssh-relay] Installing node-pty...') + console.log('[ssh-relay] Installing native dependencies...') await installNativeDeps(conn, remoteRelayDir, platform) console.log('[ssh-relay] Native deps installed') @@ -219,10 +221,54 @@ async function uploadRelay( } } -// Why: node-pty is a native addon that can't be bundled by esbuild. It must -// be compiled on the remote host against its Node.js version and OS. We -// write a minimal package.json + run `npm install node-pty` in the relay -// directory so `require('node-pty')` resolves to the local node_modules. +const RELAY_NATIVE_DEPS = ['node-pty', '@parcel/watcher'] as const + +async function hasRequiredNativeDeps(conn: SshConnection, remoteDir: string): Promise { + const nodePath = await resolveRemoteNodePath(conn) + const nodeBinDir = nodePath.replace(/\/node$/, '') + const escapedDir = shellEscape(remoteDir) + const escapedBinDir = shellEscape(nodeBinDir) + const escapedNode = shellEscape(nodePath) + try { + const probe = await execCommand( + conn, + `export PATH=${escapedBinDir}:$PATH && cd ${escapedDir} && (${escapedNode} -e 'require.resolve("node-pty"); require.resolve("@parcel/watcher"); console.log("ORCA-NATIVE-DEPS-OK")' 2>/dev/null || echo MISSING)` + ) + return probe.includes('ORCA-NATIVE-DEPS-OK') + } catch { + return false + } +} + +async function repairInstalledNativeDeps( + conn: SshConnection, + remoteDir: string, + platform: RelayPlatform +): Promise { + if (await hasRequiredNativeDeps(conn, remoteDir)) { + return + } + + console.warn(`[ssh-relay] Repairing missing native deps at ${remoteDir}`) + await acquireInstallLock(conn, remoteDir) + try { + // Why: older complete relay dirs were created before @parcel/watcher was + // installed. Re-probe under the lock so only one reconnect mutates the dir. + if (!(await hasRequiredNativeDeps(conn, remoteDir))) { + await installNativeDeps(conn, remoteDir, platform) + await finalizeInstall(conn, remoteDir) + } else { + await abandonInstall(conn, remoteDir) + } + } catch (err) { + await abandonInstall(conn, remoteDir) + throw err + } +} + +// Why: node-pty and @parcel/watcher are native addons that can't be bundled by +// esbuild. They must be installed on the remote host against its Node.js version +// and OS so dynamic imports/require calls resolve from the relay dir. // // TODO(#1693): VS Code ships per-platform tarballs with node-pty pre-built // from CI and skips `npm install` on the remote entirely. That approach @@ -251,7 +297,8 @@ async function installNativeDeps( name: 'orca-relay', version: '1.0.0', private: true, - type: 'commonjs' + type: 'commonjs', + dependencies: Object.fromEntries(RELAY_NATIVE_DEPS.map((name) => [name, '*'])) })}\n` const sftpPkg = await conn.sftp() try { @@ -269,9 +316,10 @@ async function installNativeDeps( } try { + const installArgs = RELAY_NATIVE_DEPS.map((dep) => shellEscape(dep)).join(' ') await execCommand( conn, - `export PATH=${escapedBinDir}:$PATH && cd ${escapedDir} && npm install node-pty 2>&1` + `export PATH=${escapedBinDir}:$PATH && cd ${escapedDir} && npm install ${installArgs} 2>&1` ) } catch (err) { // Don't write .install-complete on hard fail; reconnect retries on a @@ -279,7 +327,7 @@ async function installNativeDeps( // searchable. const msg = (err as Error).message console.warn( - `[ssh-relay][NPTY-INSTALL-FAIL] npm install node-pty failed at ${remoteDir} (${platform}): ${msg}` + `[ssh-relay][NATIVE-DEPS-INSTALL-FAIL] npm install native deps failed at ${remoteDir} (${platform}): ${msg}` ) throw err } diff --git a/src/main/ssh/ssh-relay-native-deps-install.test.ts b/src/main/ssh/ssh-relay-native-deps-install.test.ts index a5d381f9a..cb4c6b6ec 100644 --- a/src/main/ssh/ssh-relay-native-deps-install.test.ts +++ b/src/main/ssh/ssh-relay-native-deps-install.test.ts @@ -57,6 +57,7 @@ import { deployAndLaunchRelay } from './ssh-relay-deploy' import { execCommand } from './ssh-relay-deploy-helpers' import { parseUnameToRelayPlatform } from './relay-protocol' import { + acquireInstallLock, abandonInstall, finalizeInstall, isRelayAlreadyInstalled @@ -233,9 +234,12 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { // Why: pin commonjs so a future Node default flip doesn't silently // break `require('node-pty')`. expect(parsed.type).toBe('commonjs') + expect(parsed.dependencies).toEqual({ '@parcel/watcher': '*', 'node-pty': '*' }) const execCalls = vi.mocked(execCommand).mock.calls.map(([, c]) => c) - const npmInstallIdx = execCalls.findIndex((c) => c.includes('npm install node-pty')) + const npmInstallIdx = execCalls.findIndex( + (c) => c.includes('npm install') && c.includes('node-pty') && c.includes('@parcel/watcher') + ) expect(npmInstallIdx).toBeGreaterThanOrEqual(0) // Pin actual ordering: number of execCommand calls observed at the moment // ws.end() ran for package.json must be < the index of `npm install`. @@ -261,7 +265,7 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { expect(vi.mocked(finalizeInstall)).not.toHaveBeenCalled() const warnMessages = warnSpy.mock.calls.map((args) => String(args[0] ?? '')) - expect(warnMessages.some((m) => m.includes('[ssh-relay][NPTY-INSTALL-FAIL]'))).toBe(true) + expect(warnMessages.some((m) => m.includes('[ssh-relay][NATIVE-DEPS-INSTALL-FAIL]'))).toBe(true) }) it('warns clearly when node-pty installs but require() fails (built-but-unloadable)', async () => { @@ -298,7 +302,9 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { // test pass while exercising a different failure path. const execCalls = vi.mocked(execCommand).mock.calls.map(([, c]) => c) const probeCallIdx = execCalls.findIndex((c) => c.includes('require("node-pty")')) - const npmInstallIdx = execCalls.findIndex((c) => c.includes('npm install node-pty')) + const npmInstallIdx = execCalls.findIndex( + (c) => c.includes('npm install') && c.includes('node-pty') && c.includes('@parcel/watcher') + ) expect(probeCallIdx, 'probe must have been invoked').toBeGreaterThanOrEqual(0) // Probe must come strictly AFTER `npm install` — otherwise we'd be // probing into an empty install dir and this whole failure mode @@ -309,7 +315,9 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { // Channel failure must NOT be conflated with "node-pty missing" or with // "npm install failed". expect(warnMessages.some((m) => m.includes('[ssh-relay][NPTY-MISSING]'))).toBe(false) - expect(warnMessages.some((m) => m.includes('[ssh-relay][NPTY-INSTALL-FAIL]'))).toBe(false) + expect(warnMessages.some((m) => m.includes('[ssh-relay][NATIVE-DEPS-INSTALL-FAIL]'))).toBe( + false + ) expect(vi.mocked(finalizeInstall)).not.toHaveBeenCalled() // Lock must be released so a future reconnect can retry. @@ -333,7 +341,9 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { // reason. const execCalls = vi.mocked(execCommand).mock.calls.map(([, c]) => c) const probeIdx = execCalls.findIndex((c) => c.includes('require("node-pty")')) - const npmInstallIdx = execCalls.findIndex((c) => c.includes('npm install node-pty')) + const npmInstallIdx = execCalls.findIndex( + (c) => c.includes('npm install') && c.includes('node-pty') && c.includes('@parcel/watcher') + ) expect(probeIdx).toBeGreaterThan(npmInstallIdx) const warnMessages = warnSpy.mock.calls.map((args) => String(args[0] ?? '')) @@ -365,7 +375,9 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { // probe would silently break spawn-helper bits; one that probes before // npm install would test an empty dir. const all = vi.mocked(execCommand).mock.calls.map(([, c]) => c) - const npmIdx = all.findIndex((c) => c.includes('npm install node-pty')) + const npmIdx = all.findIndex( + (c) => c.includes('npm install') && c.includes('node-pty') && c.includes('@parcel/watcher') + ) const chmodPrebuildsIdx = all.findIndex( (c) => c.includes('spawn-helper') && c.includes('chmod +x') ) @@ -420,7 +432,7 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { expect(vi.mocked(abandonInstall)).not.toHaveBeenCalled() }) - it('includes the platform tuple in NPTY-MISSING and NPTY-INSTALL-FAIL logs', async () => { + it('includes the platform tuple in NPTY-MISSING and native install failure logs', async () => { // Platform tuple lets bug reports be triaged for prebuild availability // without asking the user to dig out their arch. const conn = makeMockConnection(sftpCapture) @@ -459,4 +471,45 @@ describe('installNativeDeps (via deployAndLaunchRelay)', () => { expect(second).toBe(first) }) + + it('repairs an existing complete relay dir that is missing @parcel/watcher', async () => { + vi.mocked(isRelayAlreadyInstalled).mockResolvedValue(true) + const conn = makeMockConnection(sftpCapture) + feed([ + 'Linux x86_64', + '/home/u', + 'MISSING', // first native-deps probe before lock + 'MISSING', // re-probe after lock + '', // npm install native deps + '', // chmod prebuilds + 'ORCA-NPTY-PROBE-OK\n', + '', // rm probe stderr + 'DEAD', + 'READY' + ]) + + await deployAndLaunchRelay(conn) + + expect(vi.mocked(acquireInstallLock)).toHaveBeenCalledTimes(1) + expect(vi.mocked(finalizeInstall)).toHaveBeenCalledTimes(1) + const execCalls = vi.mocked(execCommand).mock.calls.map(([, c]) => c) + expect( + execCalls.some( + (c) => c.includes('npm install') && c.includes('node-pty') && c.includes('@parcel/watcher') + ) + ).toBe(true) + }) + + it('does not mutate an existing relay dir when required native deps are present', async () => { + vi.mocked(isRelayAlreadyInstalled).mockResolvedValue(true) + const conn = makeMockConnection(sftpCapture) + feed(['Linux x86_64', '/home/u', 'ORCA-NATIVE-DEPS-OK', 'DEAD', 'READY']) + + await deployAndLaunchRelay(conn) + + expect(vi.mocked(acquireInstallLock)).not.toHaveBeenCalled() + expect(vi.mocked(finalizeInstall)).not.toHaveBeenCalled() + const execCalls = vi.mocked(execCommand).mock.calls.map(([, c]) => c) + expect(execCalls.some((c) => c.includes('npm install'))).toBe(false) + }) }) diff --git a/src/relay/fs-handler-file-read.ts b/src/relay/fs-handler-file-read.ts index d114f5837..d8f2f9077 100644 --- a/src/relay/fs-handler-file-read.ts +++ b/src/relay/fs-handler-file-read.ts @@ -1,5 +1,8 @@ -import { readFile, stat } from 'fs/promises' +import { open, readFile, stat } from 'fs/promises' import { extname } from 'path' +import type { RelayDispatcher, RequestContext } from './dispatcher' +import { STREAM_CHUNK_SIZE, RelayErrorCode } from './protocol' +import type { RelayStreamRegistry, TooManyStreamsError } from './fs-stream-registry' import { BINARY_PROBE_BYTES, IMAGE_MIME_TYPES, @@ -24,8 +27,6 @@ export async function readRelayFileContent(filePath: string) { return { content: buffer.toString('base64'), isBinary: true, isImage: true, mimeType } } - // Why: SSH reads serialize through bounded relay frames; probing large - // unknown files prevents binary archives from consuming frame budget. if (stats.size > BINARY_PROBE_BYTES && (await isBinaryFilePrefix(filePath))) { return { content: '', isBinary: true } } @@ -36,3 +37,171 @@ export async function readRelayFileContent(filePath: string) { } return { content: buffer.toString('utf-8'), isBinary: false } } + +export type StreamMetadata = { + streamId?: number + totalSize: number + isBinary: boolean + isImage?: boolean + mimeType?: string + /** On-the-wire encoding of each chunk's `data` field. Always 'base64'. */ + chunkEncoding?: 'base64' + /** Encoding of the assembled FileReadResult.content. */ + resultEncoding?: 'base64' | 'utf-8' + /** True for empty files and binary archives that short-circuit without pumping. */ + empty?: boolean +} + +export async function readRelayFileStreamMetadata( + filePath: string, + dispatcher: RelayDispatcher, + registry: RelayStreamRegistry, + context: RequestContext +): Promise { + const stats = await stat(filePath) + const mimeType = IMAGE_MIME_TYPES[extname(filePath).toLowerCase()] + const sizeLimit = mimeType ? MAX_PREVIEWABLE_BINARY_SIZE : MAX_TEXT_FILE_SIZE + if (stats.size > sizeLimit) { + throw new Error( + `File too large: ${(stats.size / 1024 / 1024).toFixed(1)}MB exceeds ${sizeLimit / 1024 / 1024}MB limit` + ) + } + + if (stats.size === 0) { + return { + totalSize: 0, + isBinary: !!mimeType, + mimeType, + isImage: mimeType ? true : undefined, + empty: true + } + } + // Why: unlike the legacy single-shot path, streaming does not read the full + // buffer before classifying content. Probe every unknown file so small binary + // files do not get decoded as UTF-8 text over SSH. + if (!mimeType && (await isBinaryFilePrefix(filePath))) { + return { totalSize: 0, isBinary: true, empty: true } + } + + const handle = await open(filePath, 'r') + let streamId: number + try { + streamId = registry.register(handle) + } catch (err) { + await handle.close() + throw err + } + + process.stderr.write(`[relay] stream start id=${streamId} size=${stats.size}\n`) + + // Why: pumpChunks owns its own try/finally for handle release; the outer + // setImmediate kicks the pump off the metadata-response task so the client + // sees the response before the first chunk frame. + setImmediate(() => { + void pumpChunks(streamId, stats.size, dispatcher, registry, context) + }) + + return { + streamId, + totalSize: stats.size, + isBinary: !!mimeType, + isImage: mimeType ? true : undefined, + mimeType, + chunkEncoding: 'base64', + resultEncoding: mimeType ? 'base64' : 'utf-8' + } +} + +async function pumpChunks( + streamId: number, + totalSize: number, + dispatcher: RelayDispatcher, + registry: RelayStreamRegistry, + context: RequestContext +): Promise { + const entry = registry.get(streamId) + if (!entry) { + return + } + const buffer = Buffer.allocUnsafe(STREAM_CHUNK_SIZE) + let offset = 0 + let seq = 0 + let endReason: 'end' | 'aborted' | 'stale' | 'error' = 'end' + let errorCode: string | null = null + let errorMessage: string | null = null + + try { + try { + while (offset < totalSize) { + if (context.isStale()) { + endReason = 'stale' + break + } + if (registry.isAborted(streamId)) { + endReason = 'aborted' + break + } + const want = Math.min(STREAM_CHUNK_SIZE, totalSize - offset) + const { bytesRead } = await entry.handle.read(buffer, 0, want, offset) + if (bytesRead === 0) { + endReason = 'error' + errorCode = 'ESTREAMTRUNCATED' + errorMessage = `File truncated mid-stream: expected ${totalSize}, got ${offset}` + break + } + if (context.isStale()) { + endReason = 'stale' + break + } + if (registry.isAborted(streamId)) { + endReason = 'aborted' + break + } + const data = buffer.subarray(0, bytesRead).toString('base64') + dispatcher.notify('fs.streamChunk', { streamId, seq, data }) + offset += bytesRead + seq += 1 + } + } catch (err) { + // Why: a read() rejection that races with disposeAll surfaces as EBADF; + // treat as aborted so we don't emit a spurious streamError to a client + // that is already gone. + const code = (err as { code?: string }).code + if (code === 'EBADF' && registry.isAborted(streamId)) { + endReason = 'aborted' + } else { + endReason = 'error' + errorCode = code ?? 'ESTREAMREAD' + errorMessage = err instanceof Error ? err.message : String(err) + } + } + + try { + if (endReason === 'end') { + dispatcher.notify('fs.streamEnd', { streamId }) + process.stderr.write(`[relay] stream end id=${streamId}\n`) + } else if (endReason === 'error') { + dispatcher.notify('fs.streamError', { + streamId, + code: errorCode ?? 'ESTREAMERROR', + message: errorMessage ?? 'stream error' + }) + process.stderr.write(`[relay] stream error id=${streamId} code=${errorCode}\n`) + } else if (endReason === 'aborted') { + process.stderr.write(`[relay] stream cancel id=${streamId}\n`) + } else { + process.stderr.write(`[relay] stream stale id=${streamId}\n`) + } + } catch (err) { + process.stderr.write( + `[relay] stream notify failed id=${streamId}: ${err instanceof Error ? err.message : String(err)}\n` + ) + } + } finally { + await registry.release(streamId) + } +} + +export function isTooManyStreamsError(err: unknown): err is TooManyStreamsError { + return err instanceof Error && (err as { code?: number }).code === RelayErrorCode.TooManyStreams +} diff --git a/src/relay/fs-handler-stream.test.ts b/src/relay/fs-handler-stream.test.ts new file mode 100644 index 000000000..25db03417 --- /dev/null +++ b/src/relay/fs-handler-stream.test.ts @@ -0,0 +1,253 @@ +import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest' +import { FsHandler } from './fs-handler' +import { RelayContext } from './context' +import type { RelayDispatcher } from './dispatcher' +import * as fs from 'fs/promises' +import * as path from 'path' +import { mkdtempSync, writeFileSync } from 'fs' +import { tmpdir } from 'os' + +vi.mock('@parcel/watcher', () => ({ subscribe: vi.fn() })) + +type Notification = { method: string; params?: Record } + +function createMockDispatcher() { + const requestHandlers = new Map< + string, + (params: Record, context?: { isStale: () => boolean }) => Promise + >() + const notificationHandlers = new Map) => void>() + const notifications: Notification[] = [] + return { + onRequest: vi.fn( + ( + method: string, + handler: typeof requestHandlers extends Map ? H : never + ) => { + requestHandlers.set(method, handler as never) + } + ), + onNotification: vi.fn((method: string, handler: (params: Record) => void) => { + notificationHandlers.set(method, handler) + }), + notify: vi.fn((method: string, params?: Record) => { + notifications.push({ method, params }) + }), + _notifications: notifications, + callRequest( + method: string, + params: Record = {}, + context?: { isStale: () => boolean } + ) { + const handler = requestHandlers.get(method) + if (!handler) { + throw new Error(`No handler for ${method}`) + } + return handler(params, context) + }, + callNotification(method: string, params: Record = {}) { + const handler = notificationHandlers.get(method) + if (!handler) { + throw new Error(`No handler for ${method}`) + } + handler(params) + } + } +} + +type StreamOutcome = { + chunks: { seq: number; data: string }[] + end: { streamId: number } | null + err: { code: string; message: string } | null +} + +function collectStream(d: ReturnType): StreamOutcome { + const chunks: { seq: number; data: string }[] = [] + let end: { streamId: number } | null = null + let err: { code: string; message: string } | null = null + for (const n of d._notifications) { + if (n.method === 'fs.streamChunk') { + chunks.push({ seq: n.params!.seq as number, data: n.params!.data as string }) + } else if (n.method === 'fs.streamEnd') { + end = { streamId: n.params!.streamId as number } + } else if (n.method === 'fs.streamError') { + err = { + code: n.params!.code as string, + message: n.params!.message as string + } + } + } + return { chunks, end, err } +} + +async function flush(times = 5): Promise { + for (let i = 0; i < times; i++) { + await new Promise((r) => setImmediate(r)) + } +} + +async function waitFor(predicate: () => boolean, timeoutMs = 5000): Promise { + const deadline = Date.now() + timeoutMs + while (!predicate()) { + if (Date.now() > deadline) { + throw new Error('waitFor: predicate did not become true in time') + } + await new Promise((r) => setImmediate(r)) + } +} + +describe('FsHandler readFileStream', () => { + let dispatcher: ReturnType + let handler: FsHandler + let tmpDir: string + + beforeEach(() => { + tmpDir = mkdtempSync(path.join(tmpdir(), 'relay-stream-')) + dispatcher = createMockDispatcher() + handler = new FsHandler(dispatcher as unknown as RelayDispatcher, new RelayContext()) + }) + + afterEach(async () => { + handler.dispose() + await fs.rm(tmpDir, { recursive: true, force: true }) + }) + + it('streams a binary file in chunked notifications', async () => { + const filePath = path.join(tmpDir, 'image.png') + const content = Buffer.alloc(300 * 1024, 0x42) + writeFileSync(filePath, content) + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath }, + { isStale: () => false } + )) as { streamId: number; totalSize: number; resultEncoding: string } + expect(meta.streamId).toBeDefined() + expect(meta.totalSize).toBe(content.length) + expect(meta.resultEncoding).toBe('base64') + + await waitFor(() => collectStream(dispatcher).end !== null) + const { chunks, end, err } = collectStream(dispatcher) + expect(err).toBeNull() + expect(end).toEqual({ streamId: meta.streamId }) + const reassembled = Buffer.concat(chunks.map((c) => Buffer.from(c.data, 'base64'))) + expect(reassembled.equals(content)).toBe(true) + }) + + it('returns empty:true for 0-byte files without opening a handle', async () => { + const filePath = path.join(tmpDir, 'empty.txt') + writeFileSync(filePath, '') + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath }, + { isStale: () => false } + )) as { totalSize: number; empty: boolean; streamId?: number } + expect(meta.empty).toBe(true) + expect(meta.totalSize).toBe(0) + expect(meta.streamId).toBeUndefined() + + await flush() + const { chunks } = collectStream(dispatcher) + expect(chunks).toHaveLength(0) + }) + + it('returns empty:true for binary archives over the probe threshold', async () => { + const filePath = path.join(tmpDir, 'archive.bin') + const content = Buffer.alloc(20 * 1024, 0x61) + content[0] = 0x00 + writeFileSync(filePath, content) + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath }, + { isStale: () => false } + )) as { totalSize: number; empty: boolean; isBinary: boolean } + expect(meta.empty).toBe(true) + expect(meta.isBinary).toBe(true) + }) + + it('returns empty:true for small binary files under the probe threshold', async () => { + const filePath = path.join(tmpDir, 'small.bin') + const content = Buffer.from([0x41, 0x00, 0x42]) + writeFileSync(filePath, content) + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath }, + { isStale: () => false } + )) as { totalSize: number; empty: boolean; isBinary: boolean; streamId?: number } + expect(meta.empty).toBe(true) + expect(meta.isBinary).toBe(true) + expect(meta.totalSize).toBe(0) + expect(meta.streamId).toBeUndefined() + }) + + it('rejects when totalSize exceeds the binary cap', async () => { + const filePath = path.join(tmpDir, 'huge.png') + writeFileSync(filePath, Buffer.alloc(51 * 1024 * 1024)) + + await expect( + dispatcher.callRequest('fs.readFileStream', { filePath }, { isStale: () => false }) + ).rejects.toThrow(/File too large/) + }) + + it('exits the pump and emits no further chunks when isStale flips', async () => { + const filePath = path.join(tmpDir, 'big.png') + const content = Buffer.alloc(800 * 1024, 0x42) + writeFileSync(filePath, content) + + let stale = false + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath }, + { isStale: () => stale } + )) as { streamId: number } + + await new Promise((r) => setImmediate(r)) + stale = true + await flush(10) + + const { end, err } = collectStream(dispatcher) + expect(end).toBeNull() + expect(err).toBeNull() + expect(meta.streamId).toBeGreaterThan(0) + }) + + it('honors fs.cancelStream by stopping the pump and emitting no end frame', async () => { + const filePath = path.join(tmpDir, 'cancel.png') + writeFileSync(filePath, Buffer.alloc(2 * 1024 * 1024, 0x42)) + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath }, + { isStale: () => false } + )) as { streamId: number } + + await new Promise((r) => setImmediate(r)) + dispatcher.callNotification('fs.cancelStream', { streamId: meta.streamId }) + await flush(10) + + const { end, err } = collectStream(dispatcher) + expect(end).toBeNull() + expect(err).toBeNull() + }) + + it('rejects the 17th concurrent stream with TooManyStreams', async () => { + const paths: string[] = [] + for (let i = 0; i < 17; i++) { + const p = path.join(tmpDir, `s${i}.png`) + writeFileSync(p, Buffer.alloc(8 * 1024 * 1024, 0x42)) + paths.push(p) + } + + const isStale = () => false + for (let i = 0; i < 16; i++) { + await dispatcher.callRequest('fs.readFileStream', { filePath: paths[i] }, { isStale }) + } + await expect( + dispatcher.callRequest('fs.readFileStream', { filePath: paths[16] }, { isStale }) + ).rejects.toThrow(/Too many concurrent streams/) + await flush(50) + }, 20_000) +}) diff --git a/src/relay/fs-handler-utils.ts b/src/relay/fs-handler-utils.ts index d49d082b0..0ea0fa179 100644 --- a/src/relay/fs-handler-utils.ts +++ b/src/relay/fs-handler-utils.ts @@ -22,8 +22,10 @@ import type { SearchResult as SharedSearchResult } from '../shared/types' // the old 5MB search cap would block common JSON/log files before Monaco's // large-file optimizations can handle them. export const MAX_TEXT_FILE_SIZE = 10 * 1024 * 1024 -// 10MB for relayed binaries (base64 → ~13.3MB frame payload at 16MB relay cap) -export const MAX_PREVIEWABLE_BINARY_SIZE = 10 * 1024 * 1024 +// Why: matches the local cap (src/main/ipc/filesystem.ts MAX_PREVIEWABLE_BINARY_SIZE). +// Reads above the legacy 16MB single-frame budget go through fs.readFileStream, +// which chunks at STREAM_CHUNK_SIZE; see docs/relay-file-stream-design.md. +export const MAX_PREVIEWABLE_BINARY_SIZE = 50 * 1024 * 1024 export const BINARY_PROBE_BYTES = 8192 export const SEARCH_TIMEOUT_MS = SHARED_SEARCH_TIMEOUT_MS export const DEFAULT_MAX_RESULTS = 2000 diff --git a/src/relay/fs-handler.ts b/src/relay/fs-handler.ts index 89070b883..c7ecebe8b 100644 --- a/src/relay/fs-handler.ts +++ b/src/relay/fs-handler.ts @@ -15,7 +15,8 @@ import { listFilesWithGit, searchWithGitGrep } from './fs-handler-git-fallback' import { listFilesWithReaddir } from './fs-handler-readdir-fallback' import { buildExcludePathPrefixes } from '../shared/quick-open-filter' import { buildInstallRgMessage } from './fs-handler-install-rg' -import { readRelayFileContent } from './fs-handler-file-read' +import { readRelayFileContent, readRelayFileStreamMetadata } from './fs-handler-file-read' +import { RelayStreamRegistry } from './fs-stream-registry' type WatchState = { rootPath: string @@ -27,6 +28,7 @@ type WatchState = { export class FsHandler { private dispatcher: RelayDispatcher private watches = new Map() + private streamRegistry = new RelayStreamRegistry() constructor(dispatcher: RelayDispatcher, _context: RelayContext) { this.dispatcher = dispatcher @@ -36,6 +38,7 @@ export class FsHandler { private registerHandlers(): void { this.dispatcher.onRequest('fs.readDir', (p) => this.readDir(p)) this.dispatcher.onRequest('fs.readFile', (p) => this.readFile(p)) + this.dispatcher.onRequest('fs.readFileStream', (p, c) => this.readFileStream(p, c)) this.dispatcher.onRequest('fs.writeFile', (p) => this.writeFile(p)) this.dispatcher.onRequest('fs.stat', (p) => this.stat(p)) this.dispatcher.onRequest('fs.deletePath', (p) => this.deletePath(p)) @@ -48,6 +51,7 @@ export class FsHandler { this.dispatcher.onRequest('fs.listFiles', (p) => this.listFiles(p)) this.dispatcher.onRequest('fs.watch', (p, context) => this.watch(p, context)) this.dispatcher.onNotification('fs.unwatch', (p) => this.unwatch(p)) + this.dispatcher.onNotification('fs.cancelStream', (p) => this.cancelStream(p)) } private async readDir(params: Record) { @@ -72,6 +76,22 @@ export class FsHandler { return readRelayFileContent(filePath) } + private async readFileStream( + params: Record, + context?: { isStale: () => boolean } + ) { + const filePath = expandTilde(params.filePath as string) + const ctx = context ?? { isStale: () => false } + return readRelayFileStreamMetadata(filePath, this.dispatcher, this.streamRegistry, ctx) + } + + private cancelStream(params: Record): void { + const streamId = params.streamId as number | undefined + if (typeof streamId === 'number') { + this.streamRegistry.abort(streamId) + } + } + private async writeFile(params: Record) { const filePath = expandTilde(params.filePath as string) const content = params.content as string @@ -296,5 +316,6 @@ export class FsHandler { state.unwatchFn?.() } this.watches.clear() + void this.streamRegistry.disposeAll() } } diff --git a/src/relay/fs-stream-registry.ts b/src/relay/fs-stream-registry.ts new file mode 100644 index 000000000..b1ac9fb1f --- /dev/null +++ b/src/relay/fs-stream-registry.ts @@ -0,0 +1,72 @@ +import type { FileHandle } from 'fs/promises' +import { MAX_CONCURRENT_STREAMS, RelayErrorCode } from './protocol' + +type StreamEntry = { + handle: FileHandle + aborted: boolean +} + +export class TooManyStreamsError extends Error { + readonly code = RelayErrorCode.TooManyStreams + constructor() { + super(`Too many concurrent streams (max ${MAX_CONCURRENT_STREAMS})`) + } +} + +export class RelayStreamRegistry { + private streams = new Map() + private nextId = 1 + + register(handle: FileHandle): number { + if (this.streams.size >= MAX_CONCURRENT_STREAMS) { + throw new TooManyStreamsError() + } + const streamId = this.nextId++ + this.streams.set(streamId, { handle, aborted: false }) + return streamId + } + + abort(streamId: number): void { + const entry = this.streams.get(streamId) + if (entry) { + entry.aborted = true + } + } + + isAborted(streamId: number): boolean { + return this.streams.get(streamId)?.aborted ?? true + } + + get(streamId: number): StreamEntry | undefined { + return this.streams.get(streamId) + } + + async release(streamId: number): Promise { + const entry = this.streams.get(streamId) + if (!entry) { + return + } + this.streams.delete(streamId) + try { + await entry.handle.close() + } catch { + // release runs from multiple exit paths (pump, cancel, dispose); a + // second close throws EBADF — swallow it. + } + } + + size(): number { + return this.streams.size + } + + async disposeAll(): Promise { + // Why: flag every stream as aborted so any in-flight pump exits its loop + // cleanly on the next iteration boundary instead of seeing EBADF when + // release closes the handle out from under an in-flight read. + for (const id of this.streams.keys()) { + this.abort(id) + } + const ids = Array.from(this.streams.keys()) + await Promise.all(ids.map((id) => this.release(id))) + } +} diff --git a/src/relay/integration.test.ts b/src/relay/integration.test.ts index 1bb3f9b30..1e69f9556 100644 --- a/src/relay/integration.test.ts +++ b/src/relay/integration.test.ts @@ -12,6 +12,7 @@ import { rm, readFile, stat } from 'fs/promises' import * as path from 'path' import { tmpdir } from 'os' import { execFileSync } from 'child_process' +import { randomBytes } from 'crypto' import { SshChannelMultiplexer, @@ -118,6 +119,15 @@ describe('Integration: Client Mux ↔ Relay Dispatcher', () => { expect(result.isBinary).toBe(false) }) + it('readFileStream round-trip preserves a 12 MB binary file', async () => { + const filePath = path.join(tmpDir, 'big.png') + const original = randomBytes(12 * 1024 * 1024) + writeFileSync(filePath, original) + const { readFileViaStream } = await import('../main/ssh/ssh-filesystem-stream-reader') + const { content } = await readFileViaStream(mux, filePath) + expect(Buffer.from(content, 'base64').equals(original)).toBe(true) + }, 30_000) + it('writeFile creates/overwrites file content', async () => { const filePath = path.join(tmpDir, 'output.txt') @@ -129,13 +139,10 @@ describe('Integration: Client Mux ↔ Relay Dispatcher', () => { it('stat returns file metadata', async () => { writeFileSync(path.join(tmpDir, 'sized.txt'), 'abcdef') - const result = (await mux.request('fs.stat', { filePath: path.join(tmpDir, 'sized.txt') })) as { size: number; type: string; mtime: number } - - expect(result.type).toBe('file') - expect(result.size).toBe(6) + expect(result).toMatchObject({ type: 'file', size: 6 }) expect(typeof result.mtime).toBe('number') }) diff --git a/src/relay/protocol.ts b/src/relay/protocol.ts index e555289bc..66818efaf 100644 --- a/src/relay/protocol.ts +++ b/src/relay/protocol.ts @@ -43,6 +43,16 @@ export function parseHandshakeMessage(payload: Buffer): HandshakeMessage { export const KEEPALIVE_SEND_MS = 5_000 export const TIMEOUT_MS = 20_000 +// ── Streaming constants (see docs/relay-file-stream-design.md) ───── + +export const STREAM_CHUNK_SIZE = 256 * 1024 +export const MAX_CONCURRENT_STREAMS = 16 + +export const RelayErrorCode = { + TooManyStreams: -33006, + StreamProtocolError: -33007 +} as const + export type JsonRpcRequest = { jsonrpc: '2.0' id: number diff --git a/src/renderer/src/hooks/useEditorExternalWatch.test.ts b/src/renderer/src/hooks/useEditorExternalWatch.test.ts index d8f080e71..f6f809ff4 100644 --- a/src/renderer/src/hooks/useEditorExternalWatch.test.ts +++ b/src/renderer/src/hooks/useEditorExternalWatch.test.ts @@ -16,10 +16,29 @@ vi.mock('@/components/editor/editor-autosave', () => ({ import { createExternalWatchEventHandler, - getOverflowExternalReloadTargets + getOverflowExternalReloadTargets, + getWatchedTargetKey } from './useEditorExternalWatch' import { useAppStore } from '@/store' +describe('getWatchedTargetKey', () => { + it('changes when a worktree gains an SSH connection id', () => { + expect( + getWatchedTargetKey({ + worktreeId: 'wt-1', + worktreePath: '/repo', + connectionId: undefined + }) + ).not.toBe( + getWatchedTargetKey({ + worktreeId: 'wt-1', + worktreePath: '/repo', + connectionId: 'conn-1' + }) + ) + }) +}) + describe('getOverflowExternalReloadTargets', () => { const setExternalMutation = vi.fn() diff --git a/src/renderer/src/hooks/useEditorExternalWatch.ts b/src/renderer/src/hooks/useEditorExternalWatch.ts index 416ef1913..26e97953c 100644 --- a/src/renderer/src/hooks/useEditorExternalWatch.ts +++ b/src/renderer/src/hooks/useEditorExternalWatch.ts @@ -4,7 +4,6 @@ readable in one file. */ import { useEffect, useMemo, useRef } from 'react' import { useAppStore } from '@/store' -import { getConnectionId } from '@/lib/connection-context' import { basename, joinPath } from '@/lib/path' import { normalizeAbsolutePath } from '@/components/right-sidebar/file-explorer-paths' import { getExternalFileChangeRelativePath } from '@/components/right-sidebar/useFileExplorerWatch' @@ -66,6 +65,13 @@ type ExternalWatchNotification = { relativePath: string } +export function getWatchedTargetKey(target: WatchedTarget): string { + // Why: SSH worktrees can exist in the store before their remote filesystem + // provider is ready. Include connectionId so a local/unknown placeholder + // watch is replaced by the real SSH watch when the repo metadata hydrates. + return `${target.worktreeId}::${target.worktreePath}::${target.connectionId ?? 'local'}` +} + // Why: macOS atomic writes (Claude Code Edit, vim :w, VSCode save) deliver a // delete event immediately followed by a create event for the same path. When // those two land in separate fs:changed payloads a few ms apart, the tab @@ -98,6 +104,7 @@ type PendingDeleteTimer = { export function useEditorExternalWatch(): void { const openFiles = useAppStore((s) => s.openFiles) const worktreesByRepo = useAppStore((s) => s.worktreesByRepo) + const repos = useAppStore((s) => s.repos) const activeWorktreeId = useAppStore((s) => s.activeWorktreeId) // Why: unify the target computation and the dependency key into one memo so @@ -123,15 +130,17 @@ export function useEditorExternalWatch(): void { if (!wt) { continue } - nextTargets.push({ + const repo = repos.find((r) => r.id === wt.repoId) + const target = { worktreeId: id, worktreePath: wt.path, - connectionId: getConnectionId(id) ?? undefined - }) - parts.push(`${id}::${wt.path}`) + connectionId: repo?.connectionId ?? undefined + } + nextTargets.push(target) + parts.push(getWatchedTargetKey(target)) } return { targets: nextTargets, targetsKey: parts.join('|') } - }, [openFiles, worktreesByRepo, activeWorktreeId]) + }, [openFiles, worktreesByRepo, repos, activeWorktreeId]) const targetsRef = useRef([]) const latestTargetsRef = useRef(targets) @@ -144,10 +153,10 @@ export function useEditorExternalWatch(): void { useEffect(() => { const nextTargets = latestTargetsRef.current const prev = targetsRef.current - const prevIds = new Set(prev.map((t) => t.worktreeId)) - const nextIds = new Set(nextTargets.map((t) => t.worktreeId)) - const removed = prev.filter((t) => !nextIds.has(t.worktreeId)) - const added = nextTargets.filter((t) => !prevIds.has(t.worktreeId)) + const prevKeys = new Set(prev.map(getWatchedTargetKey)) + const nextKeys = new Set(nextTargets.map(getWatchedTargetKey)) + const removed = prev.filter((t) => !nextKeys.has(getWatchedTargetKey(t))) + const added = nextTargets.filter((t) => !prevKeys.has(getWatchedTargetKey(t))) for (const target of removed) { void window.api.fs.unwatchWorktree({