diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index 3332d7f7..bf75b086 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -11,7 +11,13 @@ import * as path from 'path' import * as os from 'os' import { execFileSync } from 'node:child_process' import { fileURLToPath } from 'node:url' -import { ConversationService, ConversationStartupError, conversationService } from '../services/conversationService.js' +import { + ConversationService, + ConversationStartupError, + MAX_CAPTURED_SDK_MESSAGE_BYTES, + MAX_CAPTURED_SDK_TOTAL_BYTES, + conversationService, +} from '../services/conversationService.js' import { SessionService, sessionService } from '../services/sessionService.js' import { ProviderService } from '../services/providerService.js' import { resetTerminalShellEnvironmentCacheForTests } from '../../utils/terminalShellEnvironment.js' @@ -631,6 +637,80 @@ describe('ConversationService', () => { }) }) + it('should bound retained SDK payload bytes without truncating live callbacks', () => { + const svc = new ConversationService() + const liveMessages: any[] = [] + const oversizedText = 'x'.repeat(MAX_CAPTURED_SDK_MESSAGE_BYTES * 4) + + ;(svc as any).sessions.set('session-sdk-byte-limit', { + proc: { pid: 1 }, + outputCallbacks: [(message: any) => liveMessages.push(message)], + workDir: process.cwd(), + permissionMode: 'default', + sdkToken: 'token', + sdkSocket: null, + pendingOutbound: [], + stderrLines: [], + sdkMessages: [], + initMessage: null, + pendingPermissionRequests: new Map(), + }) + + ;(svc as any).handleSdkPayload('session-sdk-byte-limit', JSON.stringify({ + type: 'assistant', + message: { + content: [{ type: 'text', text: oversizedText }], + }, + })) + + expect(liveMessages).toHaveLength(1) + expect(liveMessages[0].message.content[0].text).toBe(oversizedText) + const retained = svc.getRecentSdkMessages('session-sdk-byte-limit') + expect(retained).toHaveLength(1) + expect(retained[0]).toMatchObject({ + type: 'assistant', + truncated: true, + message: { + content: [{ type: 'text' }], + }, + }) + expect(retained[0].message.content[0].text).toEndWith('[truncated]') + expect(Buffer.byteLength(JSON.stringify(retained[0]), 'utf-8')) + .toBeLessThanOrEqual(MAX_CAPTURED_SDK_MESSAGE_BYTES) + }) + + it('should cap the total byte size of recent SDK diagnostics', () => { + const svc = new ConversationService() + let callbackCount = 0 + + ;(svc as any).sessions.set('session-sdk-total-byte-limit', { + proc: { pid: 1 }, + outputCallbacks: [() => { callbackCount += 1 }], + workDir: process.cwd(), + permissionMode: 'default', + sdkToken: 'token', + sdkSocket: null, + pendingOutbound: [], + stderrLines: [], + sdkMessages: [], + initMessage: null, + pendingPermissionRequests: new Map(), + }) + + for (let index = 0; index < 40; index++) { + ;(svc as any).handleSdkPayload('session-sdk-total-byte-limit', JSON.stringify({ + type: 'stream_event', + index, + output: `${index}:${'y'.repeat(32 * 1024)}`, + })) + } + + const retainedBytes = svc.getRecentSdkMessages('session-sdk-total-byte-limit') + .reduce((total, message) => total + Buffer.byteLength(JSON.stringify(message), 'utf-8'), 0) + expect(callbackCount).toBe(40) + expect(retainedBytes).toBeLessThanOrEqual(MAX_CAPTURED_SDK_TOTAL_BYTES) + }) + it('should expose live SDK permission requests for reconnecting clients', () => { const svc = new ConversationService() diff --git a/src/server/services/conversationService.ts b/src/server/services/conversationService.ts index 2dd134fb..1a3b39d1 100644 --- a/src/server/services/conversationService.ts +++ b/src/server/services/conversationService.ts @@ -56,6 +56,9 @@ import { const MAX_CAPTURED_PROCESS_LINES = 80 const MAX_CAPTURED_SDK_MESSAGES = 40 const MAX_CAPTURED_SDK_SUMMARY = 20 +export const MAX_CAPTURED_SDK_MESSAGE_BYTES = 64 * 1024 +export const MAX_CAPTURED_SDK_TOTAL_BYTES = 512 * 1024 +const MAX_CAPTURED_SDK_DIAGNOSTIC_TEXT_BYTES = 4 * 1024 const CONTROL_READY_POLL_MS = 50 const AUTO_MEMORY_DIRNAME = 'memory' export const DESKTOP_CLI_GRACEFUL_SHUTDOWN_TIMEOUT_MS = 6_000 @@ -143,6 +146,7 @@ type SessionProcess = { stderrLines: string[] outputDrain: Promise sdkMessages: any[] + sdkMessageBytes?: number initMessage: any | null usesOfficialOAuth: boolean officialOAuthToken: string | null @@ -407,6 +411,7 @@ export class ConversationService { stderrLines: [], outputDrain: Promise.resolve(), sdkMessages: [], + sdkMessageBytes: 0, initMessage: null, usesOfficialOAuth, officialOAuthToken: childEnv.CLAUDE_CODE_OAUTH_TOKEN ?? null, @@ -887,10 +892,7 @@ export class ConversationService { for (const line of lines) { try { const msg = JSON.parse(line) - session.sdkMessages.push(msg) - if (session.sdkMessages.length > MAX_CAPTURED_SDK_MESSAGES) { - session.sdkMessages.splice(0, session.sdkMessages.length - MAX_CAPTURED_SDK_MESSAGES) - } + this.retainSdkMessage(session, msg, Buffer.byteLength(line, 'utf-8')) const sdkError = this.extractSdkErrorEvent(msg) if (sdkError) { void diagnosticsService.recordEvent({ @@ -970,6 +972,95 @@ export class ConversationService { } } + private retainSdkMessage( + session: SessionProcess, + message: any, + rawBytes: number, + ): void { + const retainedMessage = rawBytes <= MAX_CAPTURED_SDK_MESSAGE_BYTES + ? message + : this.compactSdkMessageForRetention(message, rawBytes) + const retainedBytes = this.capturedSdkMessageBytes(retainedMessage) + let totalBytes = session.sdkMessageBytes + ?? session.sdkMessages.reduce( + (total, existing) => total + this.capturedSdkMessageBytes(existing), + 0, + ) + + session.sdkMessages.push(retainedMessage) + totalBytes += retainedBytes + while ( + session.sdkMessages.length > MAX_CAPTURED_SDK_MESSAGES + || totalBytes > MAX_CAPTURED_SDK_TOTAL_BYTES + ) { + const removed = session.sdkMessages.shift() + if (removed === undefined) break + totalBytes -= this.capturedSdkMessageBytes(removed) + } + session.sdkMessageBytes = Math.max(0, totalBytes) + } + + private capturedSdkMessageBytes(message: any): number { + return Buffer.byteLength(JSON.stringify(message), 'utf-8') + } + + private compactSdkMessageForRetention(message: any, originalBytes: number): Record { + if (!message || typeof message !== 'object') { + return { type: 'unknown', truncated: true, originalBytes } + } + + const compact: Record = { + type: typeof message.type === 'string' ? message.type : 'unknown', + truncated: true, + originalBytes, + } + if (typeof message.subtype === 'string') compact.subtype = message.subtype + if (typeof message.is_error === 'boolean') compact.is_error = message.is_error + if (typeof message.isApiErrorMessage === 'boolean') { + compact.isApiErrorMessage = message.isApiErrorMessage + } + for (const field of ['status', 'error', 'result'] as const) { + const value = this.truncateSdkDiagnosticText(message[field]) + if (value !== undefined) compact[field] = value + } + if (Array.isArray(message.errors)) { + compact.errors = message.errors + .slice(0, 5) + .map((value: unknown) => this.truncateSdkDiagnosticText(value)) + .filter((value: string | undefined): value is string => value !== undefined) + } + + const assistantText = this.extractAssistantText(message) + if (assistantText) { + compact.message = { + content: [{ + type: 'text', + text: this.truncateSdkDiagnosticText(assistantText), + }], + } + } else { + const messageText = this.truncateSdkDiagnosticText(message.message) + if (messageText !== undefined) compact.message = messageText + } + return compact + } + + private truncateSdkDiagnosticText(value: unknown): string | undefined { + if (typeof value !== 'string') return undefined + if (Buffer.byteLength(value, 'utf-8') <= MAX_CAPTURED_SDK_DIAGNOSTIC_TEXT_BYTES) { + return value + } + const prefixBytes = Buffer.from( + value.slice(0, MAX_CAPTURED_SDK_DIAGNOSTIC_TEXT_BYTES), + 'utf-8', + ) + const truncated = prefixBytes + .subarray(0, MAX_CAPTURED_SDK_DIAGNOSTIC_TEXT_BYTES) + .toString('utf-8') + .replace(/\uFFFD$/, '') + return `${truncated}\n[truncated]` + } + stopSession(sessionId: string): void { const session = this.sessions.get(sessionId) if (!session) return