diff --git a/src/main.tsx b/src/main.tsx index b80aefaab..f50629bb7 100644 --- a/src/main.tsx +++ b/src/main.tsx @@ -5983,7 +5983,7 @@ async function run(): Promise { }); await sessionManager.init(); - const server = startServer(config, sessionManager); + const server = startServer(config, sessionManager, eventBus); const actualPort = server.port ?? config.port; printBanner(config, undefined, actualPort); diff --git a/src/server/__tests__/sessionHandle.test.ts b/src/server/__tests__/sessionHandle.test.ts index 575a9dfde..9bf10f50d 100644 --- a/src/server/__tests__/sessionHandle.test.ts +++ b/src/server/__tests__/sessionHandle.test.ts @@ -83,7 +83,7 @@ describe('SessionHandle', () => { unsub() }) - test('prompt throws if not running', async () => { + test('prompt throws if stopped', async () => { const bus = createMockEventBus() const handle = new SessionHandle({ sessionId: 'test-id', @@ -92,7 +92,8 @@ describe('SessionHandle', () => { execPath: 'bun', scriptArgs: [], }) - expect(handle.prompt('hello')).rejects.toThrow('not running') + handle.kill() + expect(handle.prompt('hello')).rejects.toThrow('stopped') }) test('getPendingPermissions returns empty initially', () => { @@ -132,6 +133,31 @@ describe('SessionHandle', () => { expect(handle.status).toBe('stopped') }) + test('waitReady throws if stopped', async () => { + const bus = createMockEventBus() + const handle = new SessionHandle({ + sessionId: 'test-id', + cwd: '/tmp', + eventBus: bus, + execPath: 'bun', + scriptArgs: [], + }) + handle.kill() + expect(handle.waitReady(100)).rejects.toThrow('stopped') + }) + + test('ready returns false initially', () => { + const bus = createMockEventBus() + const handle = new SessionHandle({ + sessionId: 'test-id', + cwd: '/tmp', + eventBus: bus, + execPath: 'bun', + scriptArgs: [], + }) + expect(handle.ready).toBe(false) + }) + test('usage returns token counts', () => { const bus = createMockEventBus() const handle = new SessionHandle({ diff --git a/src/server/eventBus.ts b/src/server/eventBus.ts index 3ad404f67..ab4e27506 100644 --- a/src/server/eventBus.ts +++ b/src/server/eventBus.ts @@ -25,13 +25,18 @@ export class EventBus { const client: SSEClient = { id, writer, rawWrite: null, sessionIdFilter } this.clients.set(id, client) - void writer.writeSSE({ + const sendAndCleanup = (opts: { event: string; data: string }) => + writer.writeSSE(opts).catch(() => { + this.clients.delete(id) + }) + + void sendAndCleanup({ event: 'connected', data: JSON.stringify({ type: 'server.connected' }), }) for (const buffered of this.buffer) { - void writer.writeSSE({ + void sendAndCleanup({ event: buffered.event, data: JSON.stringify(buffered.data), }) @@ -51,6 +56,7 @@ export class EventBus { this.buffer.shift() } const payload = JSON.stringify(data) + const deadIds: string[] = [] for (const client of this.clients.values()) { if (client.writer) { if (client.sessionIdFilter) { @@ -59,7 +65,14 @@ export class EventBus { continue } } - void client.writer.writeSSE({ event, data: payload }) + client.writer.writeSSE({ event, data: payload }).catch(() => { + deadIds.push(client.id) + }) + } + } + if (deadIds.length > 0) { + for (const id of deadIds) { + this.clients.delete(id) } } } diff --git a/src/server/routes/permission.ts b/src/server/routes/permission.ts index f4911e680..47168c893 100644 --- a/src/server/routes/permission.ts +++ b/src/server/routes/permission.ts @@ -14,14 +14,23 @@ export function createPermissionRoutes(sessionManager: SessionManager): Hono { if (!found) throw notFound('permission request not found') const body = await c.req.json<{ - behavior: 'allow' | 'deny' + behavior?: 'allow' | 'deny' | 'once' | 'always' | 'reject' updated_input?: Record + updated_permissions?: Record[] message?: string + interrupt?: boolean }>() - found.handle.replyPermission(requestId, body.behavior, { + const mappedBehavior: 'allow' | 'deny' = + body.behavior === 'reject' ? 'deny' : + body.behavior === 'once' || body.behavior === 'always' ? 'allow' : + body.behavior === 'deny' ? 'deny' : 'allow' + + found.handle.replyPermission(requestId, mappedBehavior, { updatedInput: body.updated_input, + updatedPermissions: body.updated_permissions, message: body.message, + interrupt: body.interrupt, }) return c.json({ resolved: true }) diff --git a/src/server/routes/session.ts b/src/server/routes/session.ts index cb798e464..8f9e3f8ab 100644 --- a/src/server/routes/session.ts +++ b/src/server/routes/session.ts @@ -114,34 +114,45 @@ export function createSessionRoutes( const body = await c.req.json<{ cwd?: string permission_mode?: string + permission?: Array<{ permission: string; pattern: string; action: string }> model?: string system_prompt?: string resume_session_id?: string + resume_session_at?: string + hooks?: Record }>() + let permissionMode = body.permission_mode + if (!permissionMode && body.permission) { + const hasDeny = body.permission.some(r => r.action === 'deny') + const hasAsk = body.permission.some(r => r.action === 'ask') + if (!hasDeny && !hasAsk) { + permissionMode = 'bypassPermissions' + } else if (hasAsk) { + permissionMode = 'default' + } + } + try { const handle = await sessionManager.createSession({ cwd: body.cwd, model: body.model, - permissionMode: body.permission_mode, + permissionMode, systemPrompt: body.system_prompt, resumeSessionId: body.resume_session_id, + resumeSessionAt: body.resume_session_at, + hooks: body.hooks, execPath: process.execPath, scriptArgs: getScriptArgsForChild(), }) - const initData = handle.initData - + const info = handle.getInfo() return c.json( { session_id: handle.sessionId, status: handle.status, cwd: handle.cwd, - created_at: handle.getInfo().created_at, - commands: initData?.commands ?? [], - agents: initData?.agents ?? [], - models: initData?.models ?? [], - account: initData?.account ?? {}, + created_at: info.created_at, }, 201, ) @@ -294,6 +305,7 @@ export function createSessionRoutes( parts?: Array<{ type: string; text?: string }> files?: string[] images?: unknown[] + model?: { providerID?: string; modelID?: string } }>() const content = body.content ?? body.parts @@ -302,6 +314,11 @@ export function createSessionRoutes( .join('\n') ?? '' if (!content) throw badRequest('content is required') if (handle.prompting) throw conflict('session is already processing a prompt') + + if (body.model?.modelID) { + try { await handle.setModel(body.model.modelID) } catch {} + } + return ssePrompt(handle, id, content, c) }) .post('/session/:sessionID/prompt_async', async c => { @@ -314,6 +331,7 @@ export function createSessionRoutes( parts?: Array<{ type: string; text?: string }> files?: string[] images?: unknown[] + model?: { providerID?: string; modelID?: string } }>() const content = body.content ?? body.parts @@ -323,7 +341,12 @@ export function createSessionRoutes( if (!content) throw badRequest('content is required') if (handle.prompting) throw conflict('session is already processing a prompt') - handle.prompt(content).catch(() => {}) + void (async () => { + if (body.model?.modelID) { + try { await handle.setModel(body.model.modelID) } catch {} + } + handle.prompt(content).catch(() => {}) + })() return new Response(null, { status: 204 }) }) diff --git a/src/server/server.ts b/src/server/server.ts index 926774e16..0cb61c79b 100644 --- a/src/server/server.ts +++ b/src/server/server.ts @@ -18,8 +18,8 @@ import type { ServerConfig } from './types.js' export function startServer( config: ServerConfig, sessionManager: SessionManager, + eventBus: EventBus, ): BunServer & { port?: number } { - const eventBus = new EventBus() eventBus.startHeartbeat() const app = new Hono() diff --git a/src/server/sessionHandle.ts b/src/server/sessionHandle.ts index 11caf5960..5fb77901c 100644 --- a/src/server/sessionHandle.ts +++ b/src/server/sessionHandle.ts @@ -55,10 +55,14 @@ export type SessionHandleOptions = { permissionMode?: string systemPrompt?: string resumeSessionId?: string + resumeSessionAt?: string + hooks?: Record eventBus: EventBus execPath: string scriptArgs: string[] verbose?: boolean + silent?: boolean + onInit?: (data: InitData) => void } export function getScriptArgsForChild(): string[] { @@ -133,7 +137,9 @@ export class SessionHandle { private _initData: InitData | null = null private initRequestId: string | null = null private initResolve: ((data: InitData) => void) | null = null + private initReject: ((reason: unknown) => void) | null = null private _prompting = false + private _lastMessageUuid: string | null = null get status(): SessionState { return this._status @@ -169,6 +175,36 @@ export class SessionHandle { return this._prompting } + get lastMessageUuid(): string | null { + return this._lastMessageUuid + } + + get ready(): boolean { + return this._status === 'running' + } + + get silent(): boolean { + return this.opts.silent ?? false + } + + async waitReady(timeoutMs = 30000): Promise { + if (this._status === 'running') return + if (this._status === 'stopped') { + throw new Error(`Session ${this.sessionId} is stopped`) + } + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + reject(new Error(`Session ${this.sessionId} init timed out`)) + }, timeoutMs) + const origResolve = this.initResolve + this.initResolve = (data: InitData) => { + clearTimeout(timer) + origResolve?.(data) + resolve() + } + }) + } + constructor(private opts: SessionHandleOptions) { this.sessionId = opts.sessionId this.cwd = opts.cwd @@ -178,12 +214,17 @@ export class SessionHandle { this.verbose = opts.verbose ?? false } + private emitEvent(event: string, data: Record): void { + if (this.opts.silent) return + this.eventBus.publishSessionEvent(this.sessionId, event, data) + } + onMessage(listener: MessageListener): () => void { this.listeners.add(listener) return () => { this.listeners.delete(listener) } } - async start(): Promise { + spawn(): void { const printArgs = [ '--print', '--input-format', @@ -194,11 +235,19 @@ export class SessionHandle { this.sessionId, ...(this.opts.model ? ['--model', this.opts.model] : []), ...(this.opts.permissionMode - ? ['--permission-mode', this.opts.permissionMode] + ? [ + '--permission-mode', + this.opts.permissionMode, + '--permission-prompt-tool', + 'stdio', + ] : []), ...(this.opts.resumeSessionId ? ['--resume', this.opts.resumeSessionId] : []), + ...(this.opts.resumeSessionAt + ? ['--resume-session-at', this.opts.resumeSessionAt] + : []), '--verbose', ] @@ -229,7 +278,7 @@ export class SessionHandle { this.child.on('close', (code, signal) => { if (this._status !== 'stopped') { this._status = 'stopped' - this.eventBus.publishSessionEvent(this.sessionId, 'deleted', { + this.emitEvent('deleted', { status: 'stopped', exit_code: code, signal: signal ?? null, @@ -255,43 +304,58 @@ export class SessionHandle { }) this.initRequestId = crypto.randomUUID() - const initPromise = new Promise((resolve, reject) => { - this.initResolve = resolve - const timeout = setTimeout(() => { - this.initResolve = null - reject(new Error('Init handshake timed out (30s)')) - }, 30000) - const originalResolve = resolve - this.initResolve = (data: InitData) => { - clearTimeout(timeout) - originalResolve(data) - } - }) - this.sendInitialize() - try { - this._initData = await initPromise - } catch (err) { + const timeout = setTimeout(() => { + this.initResolve = null + this.initReject = null this._status = 'stopped' - throw err - } + this.emitEvent('deleted', { + status: 'stopped', + reason: 'init_timeout', + }) + }, 30000) - this._status = 'running' - this.lastActiveAt = Date.now() - this.eventBus.publishSessionEvent(this.sessionId, 'ready', { - status: 'running', - model: this._model, + this.initResolve = (data: InitData) => { + clearTimeout(timeout) + this.initReject = null + } + this.initReject = (_err: unknown) => { + clearTimeout(timeout) + this.initResolve = null + } + } + + async start(): Promise { + if (!this.child) { + this.spawn() + } + if (this._status === 'running') return + return new Promise((resolve, reject) => { + const origResolve = this.initResolve + const origReject = this.initReject + this.initResolve = (data: InitData) => { + origResolve?.(data) + resolve() + } + this.initReject = (err: unknown) => { + origReject?.(err) + reject(err) + } }) } private sendInitialize(): void { + const request: Record = { + subtype: 'initialize', + } + if (this.opts.hooks) { + request.hooks = this.opts.hooks + } const msg = jsonStringify({ type: 'control_request', request_id: this.initRequestId, - request: { - subtype: 'initialize', - }, + request, }) this.writeStdin(msg) } @@ -347,8 +411,17 @@ export class SessionHandle { const first = (initData.models as Array>)[0] this._model = first?.value ?? first?.name } - this.initResolve(initData) + this._status = 'running' + this.lastActiveAt = Date.now() + const resolve = this.initResolve this.initResolve = null + this.initReject = null + resolve?.(initData) + this.emitEvent('ready', { + status: 'running', + model: this._model, + }) + this.opts.onInit?.(initData) return } } @@ -372,14 +445,13 @@ export class SessionHandle { this.emitMessage(msg) switch (msg.type) { - case 'system': { - this.eventBus.publishSessionEvent(this.sessionId, 'message', msg) - break - } case 'assistant': { this.lastActiveAt = Date.now() this._status = 'running' - this.eventBus.publishSessionEvent(this.sessionId, 'message', msg) + if (msg.uuid && typeof msg.uuid === 'string') { + this._lastMessageUuid = msg.uuid + } + this.emitEvent('message', msg) break } case 'result': { @@ -392,7 +464,7 @@ export class SessionHandle { if (cost) this._costUsd += cost if (usage?.input_tokens) this._inputTokens += usage.input_tokens if (usage?.output_tokens) this._outputTokens += usage.output_tokens - this.eventBus.publishSessionEvent(this.sessionId, 'result', msg) + this.emitEvent('result', msg) if (this.promptResolve) { this.promptResolve({ done: true }) this.promptResolve = null @@ -414,11 +486,7 @@ export class SessionHandle { description: `Execute ${request.tool_name}`, } this.pendingPermissions.set(requestId, perm) - this.eventBus.publishSessionEvent( - this.sessionId, - 'control_request', - { request_id: requestId, request }, - ) + this.emitEvent('control_request', { request_id: requestId, request }) } else if (request?.subtype === 'elicitation') { const question: PendingQuestion = { requestId, @@ -430,22 +498,24 @@ export class SessionHandle { (request.requested_schema as Record) ?? {}, } this.pendingQuestions.set(requestId, question) - this.eventBus.publishSessionEvent( - this.sessionId, - 'control_request', - { request_id: requestId, request }, - ) + this.emitEvent('control_request', { request_id: requestId, request }) + } else if (request?.subtype === 'hook_callback') { + this.emitEvent('control_request', { request_id: requestId, request }) } break } + case 'control_cancel_request': { + const cancelRequestId = msg.request_id as string + this.pendingPermissions.delete(cancelRequestId) + this.pendingQuestions.delete(cancelRequestId) + break + } default: { - this.eventBus.publishSessionEvent(this.sessionId, 'message', msg) break } } } - private _outputTokens = 0 private pendingControlResponses = new Map< string, { @@ -503,11 +573,14 @@ export class SessionHandle { } async prompt(content: string): Promise<{ done: boolean; error?: string }> { - if (this._status !== 'running') { + if (this._status === 'stopped') { throw new Error( - `Session ${this.sessionId} is not running (status: ${this._status})`, + `Session ${this.sessionId} is stopped`, ) } + if (this._status !== 'running') { + await this.waitReady() + } if (this._prompting) { throw new Error( `Session ${this.sessionId} is already processing a prompt`, @@ -518,8 +591,16 @@ export class SessionHandle { const userMsg = jsonStringify({ type: 'user', + content, + uuid: crypto.randomUUID(), + session_id: this.sessionId, message: { role: 'user', content }, parent_tool_use_id: null, + }) + + this.emitEvent('message', { + type: 'user', + content, session_id: this.sessionId, }) @@ -565,7 +646,12 @@ export class SessionHandle { replyPermission( requestId: string, behavior: 'allow' | 'deny', - opts?: { updatedInput?: Record; message?: string }, + opts?: { + updatedInput?: Record + updatedPermissions?: Record[] + message?: string + interrupt?: boolean + }, ): void { const response = { type: 'control_response', @@ -574,8 +660,18 @@ export class SessionHandle { request_id: requestId, response: behavior === 'allow' - ? { behavior: 'allow', updatedInput: opts?.updatedInput } - : { behavior: 'deny', message: opts?.message ?? 'Denied' }, + ? { + behavior: 'allow', + updatedInput: opts?.updatedInput, + ...(opts?.updatedPermissions + ? { updatedPermissions: opts.updatedPermissions } + : {}), + } + : { + behavior: 'deny', + message: opts?.message ?? 'Denied', + ...(opts?.interrupt ? { interrupt: true } : {}), + }, }, } this.writeStdin(jsonStringify(response)) @@ -659,12 +755,16 @@ export class SessionHandle { cost_usd: this._costUsd, input_tokens: this._inputTokens, output_tokens: this._outputTokens, + last_message_uuid: this._lastMessageUuid, } } private writeStdin(data: string): void { if (this.child?.stdin && !this.child.stdin.destroyed) { - this.child.stdin.write(data + '\n') + const flushed = this.child.stdin.write(data + '\n') + if (!flushed) { + this.child.stdin.once('drain', () => {}) + } } } diff --git a/src/server/sessionManager.ts b/src/server/sessionManager.ts index fb7a9abf5..7da799c46 100644 --- a/src/server/sessionManager.ts +++ b/src/server/sessionManager.ts @@ -1,6 +1,6 @@ import { readFile, writeFile, mkdir } from 'fs/promises' -import { existsSync } from 'fs' -import { join } from 'path' +import { existsSync, statSync } from 'fs' +import { join, resolve, isAbsolute } from 'path' import { getClaudeConfigHomeDir } from '../utils/envUtils.js' import { logError } from '../utils/log.js' import type { EventBus } from './eventBus.js' @@ -11,6 +11,7 @@ const INDEX_FILE = 'server-sessions.json' export class SessionManager { private sessions = new Map() + private loadedIndex = new Map() private eventBus: EventBus private maxSessions: number private idleTimeoutMs: number @@ -49,8 +50,8 @@ export class SessionManager { try { const raw = await readFile(path, 'utf-8') const index: SessionIndex = JSON.parse(raw) - for (const [, entry] of Object.entries(index)) { - entry + for (const [key, entry] of Object.entries(index)) { + this.loadedIndex.set(key, entry) } } catch (err) { logError(err as Error) @@ -148,17 +149,26 @@ export class SessionManager { permissionMode?: string systemPrompt?: string resumeSessionId?: string + resumeSessionAt?: string + hooks?: Record execPath: string scriptArgs: string[] + silent?: boolean }): Promise { - if (this.sessions.size >= this.maxSessions) { + if (this.maxSessions > 0 && this.sessions.size >= this.maxSessions) { throw new Error( `Maximum concurrent sessions reached (${this.maxSessions})`, ) } const sessionId = crypto.randomUUID() - const cwd = opts.cwd || this.defaultWorkspace || process.cwd() + let cwd = opts.cwd || this.defaultWorkspace || process.cwd() + if (!isAbsolute(cwd)) { + cwd = resolve(process.cwd(), cwd) + } + if (!existsSync(cwd) || !statSync(cwd).isDirectory()) { + throw new Error(`Working directory does not exist: ${cwd}`) + } const handle = new SessionHandle({ sessionId, @@ -167,29 +177,34 @@ export class SessionManager { permissionMode: opts.permissionMode, systemPrompt: opts.systemPrompt, resumeSessionId: opts.resumeSessionId, + resumeSessionAt: opts.resumeSessionAt, + hooks: opts.hooks, eventBus: this.eventBus, execPath: opts.execPath, scriptArgs: opts.scriptArgs, + silent: opts.silent, + onInit: (data) => { + this._cachedInitData = data + this._resolveInitDataReady?.() + }, }) this.sessions.set(sessionId, handle) - this.eventBus.publishSessionEvent(sessionId, 'created', { - status: 'starting', - }) + if (!opts.silent) { + this.eventBus.publishSessionEvent(sessionId, 'created', { + status: 'starting', + }) + } this.scheduleIndexSave() try { - await handle.start() + handle.spawn() } catch (err) { this.sessions.delete(sessionId) this.scheduleIndexSave() throw err } - if (handle.initData) { - this._cachedInitData = handle.initData - } - return handle } @@ -221,10 +236,12 @@ export class SessionManager { cwd: opts.cwd, execPath: opts.execPath, scriptArgs: opts.scriptArgs, + silent: true, }) this._probeHandle = probe - this._probeHandle = null + await probe.start() await this.deleteSession(probe.sessionId) + this._probeHandle = null } catch (err) { this._probeHandle?.forceKill() this._probeHandle = null @@ -244,9 +261,12 @@ export class SessionManager { async deleteSession(id: string): Promise { const handle = this.sessions.get(id) if (!handle) return false + const silent = handle.silent handle.forceKill() this.sessions.delete(id) - this.eventBus.publishSessionEvent(id, 'deleted', { status: 'stopped' }) + if (!silent) { + this.eventBus.publishSessionEvent(id, 'deleted', { status: 'stopped' }) + } this.scheduleIndexSave() return true }