Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion cli/src/api/apiMachine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -562,10 +562,14 @@ export class ApiMachineClient {
this.stopKeepAlive()
})

this.socket.on('rpc-request', async (data: { method: string; params: string }, callback: (response: string) => void) => {
this.socket.on('rpc-request', async (data: { method: string; params: string; requestId?: string }, callback: (response: string) => void) => {
callback(await this.rpcHandlerManager.handleRequest(data))
})

this.socket.on('rpc-cancel', ({ requestId }) => {
this.rpcHandlerManager.cancelRequest(requestId)
})

this.socket.on('update', (data: Update) => {
if (data.body.t !== 'update-machine') {
return
Expand Down
6 changes: 5 additions & 1 deletion cli/src/api/apiSession.ts
Original file line number Diff line number Diff line change
Expand Up @@ -341,10 +341,14 @@ export class ApiSessionClient extends EventEmitter {
})
})

this.socket.on('rpc-request', async (data: { method: string; params: string }, callback: (response: string) => void) => {
this.socket.on('rpc-request', async (data: { method: string; params: string; requestId?: string }, callback: (response: string) => void) => {
callback(await this.rpcHandlerManager.handleRequest(data))
})

this.socket.on('rpc-cancel', ({ requestId }) => {
this.rpcHandlerManager.cancelRequest(requestId)
})

this.socket.on('disconnect', (reason) => {
logger.debug('[API] Socket disconnected:', reason)
this.rpcHandlerManager.onSocketDisconnect()
Expand Down
68 changes: 68 additions & 0 deletions cli/src/api/rpc/RpcHandlerManager.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
import { describe, expect, it } from 'vitest'
import { RpcHandlerManager } from './RpcHandlerManager'

describe('RpcHandlerManager cancellation', () => {
it('aborts an in-flight request by request id', async () => {
const manager = new RpcHandlerManager({ scopePrefix: 'session-1' })
manager.registerHandler('long-operation', async (_data, signal) => {
return await new Promise<{ cancelled: boolean }>((resolve) => {
signal?.addEventListener('abort', () => resolve({ cancelled: true }), { once: true })
})
})

const request = manager.handleRequest({
method: 'session-1:long-operation',
params: '{}',
requestId: 'request-1',
})

expect(manager.cancelRequest('request-1')).toBe(true)
await expect(request).resolves.toBe(JSON.stringify({ cancelled: true }))
expect(manager.cancelRequest('request-1')).toBe(false)
})

it('aborts active requests when the socket disconnects', async () => {
const manager = new RpcHandlerManager({ scopePrefix: 'session-1' })
manager.registerHandler('long-operation', async (_data, signal) => {
return await new Promise<{ cancelled: boolean }>((resolve) => {
signal?.addEventListener('abort', () => resolve({ cancelled: true }), { once: true })
})
})

const request = manager.handleRequest({
method: 'session-1:long-operation',
params: '{}',
requestId: 'request-2',
})

manager.onSocketDisconnect()
await expect(request).resolves.toBe(JSON.stringify({ cancelled: true }))
})

it('does not log expected abort errors', async () => {
const logs: unknown[] = []
const manager = new RpcHandlerManager({
scopePrefix: 'session-1',
logger: (...args) => logs.push(args),
})
manager.registerHandler('long-operation', async (_data, signal) => {
return await new Promise<never>((_resolve, reject) => {
signal?.addEventListener('abort', () => {
const error = new Error('Request aborted')
error.name = 'AbortError'
reject(error)
}, { once: true })
})
})

const request = manager.handleRequest({
method: 'session-1:long-operation',
params: '{}',
requestId: 'request-3',
})

expect(manager.cancelRequest('request-3')).toBe(true)
await expect(request).resolves.toBe(JSON.stringify({ error: 'Request aborted' }))
expect(logs).toEqual([])
})
})
36 changes: 35 additions & 1 deletion cli/src/api/rpc/RpcHandlerManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,16 @@ function safeJsonParse(value: string): unknown {
}
}

function isAbortError(error: unknown): boolean {
return Boolean(error && typeof error === 'object' && (error as { name?: unknown }).name === 'AbortError')
}

export class RpcHandlerManager {
private handlers: RpcHandlerMap = new Map()
private readonly scopePrefix: string
private readonly logger: (message: string, data?: any) => void
private socket: Socket | null = null
private readonly inFlightRequests = new Map<string, AbortController>()

constructor(config: RpcHandlerConfig) {
this.scopePrefix = config.scopePrefix
Expand All @@ -40,6 +45,13 @@ export class RpcHandlerManager {
}

async handleRequest(request: RpcRequest): Promise<string> {
const requestId = request.requestId
const abortController = requestId ? new AbortController() : null
if (requestId && abortController) {
this.inFlightRequests.get(requestId)?.abort()
this.inFlightRequests.set(requestId, abortController)
}

try {
const handler = this.handlers.get(request.method)
if (!handler) {
Expand All @@ -48,17 +60,35 @@ export class RpcHandlerManager {
}

const params = safeJsonParse(request.params)
const result = await handler(params as any)
const result = await handler(params as any, abortController?.signal)
return JSON.stringify(result)
} catch (error) {
if (isAbortError(error)) {
return JSON.stringify({ error: 'Request aborted' })
}

const details = error instanceof Error
? { message: error.message, stack: error.stack }
: { error: String(error) }
this.logger('[RPC] [ERROR] Error handling request', details)
return JSON.stringify({
error: error instanceof Error ? error.message : 'Unknown error'
})
} finally {
if (requestId && abortController && this.inFlightRequests.get(requestId) === abortController) {
this.inFlightRequests.delete(requestId)
}
}
}

cancelRequest(requestId: string): boolean {
const controller = this.inFlightRequests.get(requestId)
if (!controller) {
return false
}

controller.abort()
return true
}

onSocketConnect(socket: Socket): void {
Expand All @@ -70,6 +100,10 @@ export class RpcHandlerManager {

onSocketDisconnect(): void {
this.socket = null
for (const controller of this.inFlightRequests.values()) {
controller.abort()
}
this.inFlightRequests.clear()
}

getHandlerCount(): number {
Expand Down
4 changes: 3 additions & 1 deletion cli/src/api/rpc/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,8 @@
* @template TResponse - The response data type
*/
export type RpcHandler<TRequest = any, TResponse = any> = (
data: TRequest
data: TRequest,
signal?: AbortSignal
) => TResponse | Promise<TResponse>;

/**
Expand All @@ -22,6 +23,7 @@ export type RpcHandlerMap = Map<string, RpcHandler>;
export interface RpcRequest {
method: string;
params: string; // JSON string
requestId?: string;
}

/**
Expand Down
54 changes: 53 additions & 1 deletion cli/src/modules/common/handlers/directories.test.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,17 @@
import { beforeEach, describe, expect, it } from 'vitest'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { mkdir, rm, symlink, writeFile } from 'fs/promises'
import { join } from 'path'
import { tmpdir } from 'os'
import { RpcHandlerManager } from '../../../api/rpc/RpcHandlerManager'
import { registerDirectoryHandlers } from './directories'

const { statMock } = vi.hoisted(() => ({ statMock: vi.fn() }))

vi.mock('fs/promises', async () => {
const actual = await vi.importActual<typeof import('fs/promises')>('fs/promises')
return { ...actual, stat: statMock }
})

async function createTempDir(prefix: string): Promise<string> {
const base = tmpdir()
const path = join(base, `${prefix}-${Date.now()}-${Math.random().toString(16).slice(2)}`)
Expand All @@ -17,6 +24,10 @@ describe('directory RPC handlers', () => {
let rpc: RpcHandlerManager

beforeEach(async () => {
const actual = await vi.importActual<typeof import('fs/promises')>('fs/promises')
statMock.mockReset()
statMock.mockImplementation(actual.stat)

if (rootDir) {
await rm(rootDir, { recursive: true, force: true })
}
Expand Down Expand Up @@ -81,6 +92,47 @@ describe('directory RPC handlers', () => {
expect(parsed.entries?.[2]).toEqual({ path: 'missing.txt' })
})

it('stops starting later stat batches after cancellation', async () => {
const paths = Array.from({ length: 32 }, (_, index) => `file-${index}.txt`)
await Promise.all(paths.map((path) => writeFile(join(rootDir, path), path)))

let firstStatStarted!: () => void
const firstStat = new Promise<void>((resolve) => {
firstStatStarted = resolve
})
let releaseFirstStat!: () => void
const firstStatRelease = new Promise<void>((resolve) => {
releaseFirstStat = resolve
})
const originalStat = (await vi.importActual<typeof import('fs/promises')>('fs/promises')).stat
let statCallCount = 0
statMock.mockImplementation(async (path: Parameters<typeof originalStat>[0]) => {
statCallCount += 1
if (statCallCount === 1) {
firstStatStarted()
await firstStatRelease
}
return await originalStat(path)
})

try {
const request = rpc.handleRequest({
method: 'session-test:statFiles',
params: JSON.stringify({ paths }),
requestId: 'stat-files-cancel'
})

await firstStat
expect(rpc.cancelRequest('stat-files-cancel')).toBe(true)
releaseFirstStat()

await expect(request).resolves.toBe(JSON.stringify({ error: 'Request aborted' }))
expect(statCallCount).toBe(16)
} finally {
statMock.mockReset()
}
})

it('rejects stat paths outside the session working directory', async () => {
const response = await rpc.handleRequest({
method: 'session-test:statFiles',
Expand Down
40 changes: 26 additions & 14 deletions cli/src/modules/common/handlers/directories.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,24 @@ interface StatFilesRequest {
paths: string[]
}

const STAT_FILES_BATCH_SIZE = 16

type StatFileEntry = NonNullable<StatFilesResponse['entries']>[number]

async function statFile(path: string, workingDirectory: string): Promise<StatFileEntry> {
try {
const stats = await stat(resolve(workingDirectory, path))
return {
path,
size: stats.size,
modified: stats.mtime.getTime()
}
} catch (error) {
logger.debug(`Failed to stat ${path}:`, error)
return { path }
}
}

interface TreeNode {
name: string
path: string
Expand Down Expand Up @@ -97,7 +115,7 @@ export function registerDirectoryHandlers(rpcHandlerManager: RpcHandlerManager,
}
})

rpcHandlerManager.registerHandler<StatFilesRequest, StatFilesResponse>(RPC_METHODS.StatFiles, async (data) => {
rpcHandlerManager.registerHandler<StatFilesRequest, StatFilesResponse>(RPC_METHODS.StatFiles, async (data, signal) => {
if (!Array.isArray(data.paths) || data.paths.length > 500) {
return rpcError('Invalid file paths')
}
Expand All @@ -109,19 +127,13 @@ export function registerDirectoryHandlers(rpcHandlerManager: RpcHandlerManager,
}
}

const entries = await Promise.all(data.paths.map(async (path) => {
try {
const stats = await stat(resolve(workingDirectory, path))
return {
path,
size: stats.size,
modified: stats.mtime.getTime()
}
} catch (error) {
logger.debug(`Failed to stat ${path}:`, error)
return { path }
}
}))
const entries: StatFileEntry[] = []
for (let index = 0; index < data.paths.length; index += STAT_FILES_BATCH_SIZE) {
signal?.throwIfAborted()
const batch = data.paths.slice(index, index + STAT_FILES_BATCH_SIZE)
entries.push(...await Promise.all(batch.map((path) => statFile(path, workingDirectory))))
}
signal?.throwIfAborted()

return { success: true, entries }
})
Expand Down
38 changes: 38 additions & 0 deletions cli/src/modules/common/handlers/ripgrep.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
import { describe, expect, it, vi } from 'vitest'
import { RpcHandlerManager } from '../../../api/rpc/RpcHandlerManager'
import { registerRipgrepHandlers } from './ripgrep'

const { runFileSearchMock } = vi.hoisted(() => ({ runFileSearchMock: vi.fn() }))

vi.mock('@/modules/ripgrep/index', () => ({
run: vi.fn(),
runFileSearch: runFileSearchMock
}))

describe('ripgrep RPC handlers', () => {
it('passes AbortError through without logging it as a ripgrep failure', async () => {
const logs: unknown[] = []
const manager = new RpcHandlerManager({
scopePrefix: 'session-test',
logger: (...args) => logs.push(args)
})
registerRipgrepHandlers(manager, '/workspace')

const abortError = new Error('Request aborted')
abortError.name = 'AbortError'
runFileSearchMock.mockRejectedValueOnce(abortError)

const response = await manager.handleRequest({
method: 'session-test:ripgrep',
params: JSON.stringify({
args: ['--files'],
cwd: '/workspace',
fileSearch: { query: 'src', limit: 1 }
}),
requestId: 'ripgrep-cancel'
})

expect(response).toBe(JSON.stringify({ error: 'Request aborted' }))
expect(logs).toEqual([])
})
})
Loading
Loading