@a5c-ai/agent-mux-adapters 0.3.0 → 0.4.1

package/dist/codex-websocket-connection.js

@@ -0,0 +1,733 @@
+import WebSocket from 'ws';
+export class CodexWebSocketConnection {
+    connectionId;
+    connectionType = 'websocket';
+    websocketUrl;
+    endpoint;
+    prompt;
+    cwd;
+    requestedModel;
+    approvalMode;
+    initialSessionId;
+    models;
+    createSocket;
+    ws = null;
+    connected = false;
+    closed = false;
+    requestSeq = 0;
+    pendingResponses = new Map();
+    pendingInteractionRequests = new Map();
+    eventQueue = [];
+    waitingResolver = null;
+    threadId = null;
+    currentTurnId = null;
+    currentModelId;
+    turnIndex = -1;
+    textByItemId = new Map();
+    thinkingByItemId = new Map();
+    constructor(options) {
+        this.websocketUrl = options.websocketUrl;
+        this.connectionId = options.connectionId;
+        this.endpoint = options.websocketUrl;
+        this.prompt = options.prompt;
+        this.cwd = options.cwd;
+        this.requestedModel = options.requestedModel;
+        this.approvalMode = options.approvalMode;
+        this.initialSessionId = options.sessionId;
+        this.models = options.models;
+        this.createSocket = options.createSocket ?? ((url) => new WebSocket(url));
+    }
+    async connect() {
+        const socket = this.createSocket(this.websocketUrl);
+        this.ws = socket;
+        await new Promise((resolve, reject) => {
+            const onOpen = () => {
+                cleanup();
+                resolve();
+            };
+            const onError = (error) => {
+                cleanup();
+                reject(error);
+            };
+            const cleanup = () => {
+                socket.off('open', onOpen);
+                socket.off('error', onError);
+            };
+            socket.once('open', onOpen);
+            socket.once('error', onError);
+        });
+        socket.on('message', (data) => {
+            void this.handleIncomingMessage(data.toString());
+        });
+        socket.on('close', () => {
+            this.connected = false;
+            this.closed = true;
+            if (this.waitingResolver) {
+                const resolve = this.waitingResolver;
+                this.waitingResolver = null;
+                resolve(null);
+            }
+        });
+        socket.on('error', (error) => {
+            this.pushEvent({
+                type: 'error',
+                runId: this.connectionId,
+                agent: 'codex-websocket',
+                timestamp: Date.now(),
+                code: 'INTERNAL',
+                message: error.message,
+                recoverable: false,
+            });
+        });
+        this.connected = true;
+        await this.request('initialize', {
+            clientInfo: { name: 'agent-mux', version: '0.4.0' },
+            capabilities: {
+                experimentalApi: true,
+            },
+        });
+        await this.notify('initialized');
+        if (this.initialSessionId) {
+            const resumed = await this.request('thread/resume', {
+                threadId: this.initialSessionId,
+                cwd: this.cwd,
+                approvalPolicy: this.mapApprovalPolicy(),
+                persistExtendedHistory: false,
+                ...(this.requestedModel ? { model: this.requestedModel } : {}),
+            });
+            this.threadId = resumed.thread.id;
+            this.currentModelId = resumed.model ?? this.requestedModel;
+            this.pushEvent({
+                type: 'session_start',
+                runId: this.connectionId,
+                agent: 'codex-websocket',
+                timestamp: Date.now(),
+                sessionId: resumed.thread.id,
+                resumed: true,
+            });
+        }
+        else {
+            const started = await this.request('thread/start', {
+                cwd: this.cwd,
+                approvalPolicy: this.mapApprovalPolicy(),
+                sandbox: 'danger-full-access',
+                experimentalRawEvents: false,
+                persistExtendedHistory: false,
+                ...(this.requestedModel ? { model: this.requestedModel } : {}),
+            });
+            this.threadId = started.thread.id;
+            this.currentModelId = started.model ?? this.requestedModel;
+            this.pushEvent({
+                type: 'session_start',
+                runId: this.connectionId,
+                agent: 'codex-websocket',
+                timestamp: Date.now(),
+                sessionId: started.thread.id,
+                resumed: false,
+            });
+        }
+        if (this.prompt.trim().length > 0) {
+            await this.startTurn(this.prompt);
+        }
+    }
+    async send(data) {
+        if (!this.connected || this.closed) {
+            throw new Error('Codex app-server connection is not active');
+        }
+        if (this.isUserMessagePayload(data)) {
+            await this.sendUserMessage(data.text);
+            return;
+        }
+        if (this.isInteractionResponsePayload(data)) {
+            await this.sendInteractionResponse(data.interactionId, data.response);
+            return;
+        }
+    }
+    async *receive() {
+        while (!this.closed || this.eventQueue.length > 0) {
+            const event = this.eventQueue.shift();
+            if (event) {
+                yield event;
+                continue;
+            }
+            const nextEvent = await new Promise((resolve) => {
+                this.waitingResolver = resolve;
+            });
+            if (nextEvent) {
+                yield nextEvent;
+            }
+        }
+    }
+    subscribe(_channel) {
+        return this.receive();
+    }
+    async unsubscribe(_channel) { }
+    async close() {
+        this.closed = true;
+        const socket = this.ws;
+        this.ws = null;
+        if (socket && socket.readyState === WebSocket.OPEN) {
+            await new Promise((resolve) => {
+                socket.once('close', () => resolve());
+                socket.close();
+            });
+        }
+        if (this.waitingResolver) {
+            const resolve = this.waitingResolver;
+            this.waitingResolver = null;
+            resolve(null);
+        }
+    }
+    async handleIncomingMessage(raw) {
+        let message;
+        try {
+            message = JSON.parse(raw);
+        }
+        catch {
+            this.pushEvent({
+                type: 'error',
+                runId: this.connectionId,
+                agent: 'codex-websocket',
+                timestamp: Date.now(),
+                code: 'INTERNAL',
+                message: `Invalid Codex app-server message: ${raw}`,
+                recoverable: false,
+            });
+            return;
+        }
+        if (typeof message['method'] === 'string' && message['id'] != null) {
+            this.handleServerRequest({
+                id: message['id'],
+                method: message['method'],
+                params: message['params'] ?? {},
+            });
+            return;
+        }
+        if (typeof message['method'] === 'string') {
+            this.handleNotification({
+                method: message['method'],
+                params: message['params'] ?? {},
+            });
+            return;
+        }
+        if (message['id'] != null) {
+            this.handleResponse({
+                id: message['id'],
+                result: message['result'],
+                error: message['error'],
+            });
+        }
+    }
+    handleResponse(response) {
+        const pending = this.pendingResponses.get(response.id);
+        if (!pending) {
+            return;
+        }
+        this.pendingResponses.delete(response.id);
+        if (response.error) {
+            pending.reject(new Error(response.error.message ?? `Codex app-server request failed (${String(response.id)})`));
+            return;
+        }
+        pending.resolve(response.result);
+    }
+    handleServerRequest(request) {
+        const interactionId = String(request.id);
+        const params = request.params ?? {};
+        this.pendingInteractionRequests.set(interactionId, {
+            id: request.id,
+            method: request.method,
+            params,
+        });
+        const base = this.baseEvent();
+        switch (request.method) {
+            case 'item/commandExecution/requestApproval': {
+                const command = String(params['command'] ?? 'command');
+                const reason = params['reason'];
+                this.pushEvent({
+                    ...base,
+                    type: 'approval_request',
+                    interactionId,
+                    action: 'command_execution',
+                    detail: typeof reason === 'string' && reason.length > 0 ? `${command}\n${reason}` : command,
+                    toolName: command,
+                    riskLevel: 'high',
+                });
+                return;
+            }
+            case 'item/fileChange/requestApproval': {
+                const reason = params['reason'];
+                this.pushEvent({
+                    ...base,
+                    type: 'approval_request',
+                    interactionId,
+                    action: 'file_change',
+                    detail: typeof reason === 'string' && reason.length > 0 ? reason : 'Codex requested file changes',
+                    toolName: 'file_change',
+                    riskLevel: 'medium',
+                });
+                return;
+            }
+            case 'item/permissions/requestApproval': {
+                const reason = params['reason'];
+                this.pushEvent({
+                    ...base,
+                    type: 'approval_request',
+                    interactionId,
+                    action: 'permission_request',
+                    detail: typeof reason === 'string' && reason.length > 0 ? reason : 'Codex requested additional permissions',
+                    toolName: 'permissions',
+                    riskLevel: 'high',
+                });
+                return;
+            }
+            case 'item/tool/requestUserInput': {
+                const questions = params['questions'];
+                const firstQuestion = Array.isArray(questions) ? questions[0] : undefined;
+                this.pushEvent({
+                    ...base,
+                    type: 'input_required',
+                    interactionId,
+                    question: String(firstQuestion?.['question'] ?? 'Codex needs input'),
+                    context: undefined,
+                    source: 'tool',
+                });
+                return;
+            }
+            default: {
+                this.pushEvent({
+                    ...base,
+                    type: 'debug',
+                    level: 'warn',
+                    message: `Unhandled Codex app-server server request: ${request.method}`,
+                });
+            }
+        }
+    }
+    handleNotification(notification) {
+        const params = notification.params ?? {};
+        switch (notification.method) {
+            case 'turn/started': {
+                const turn = params['turn'];
+                this.currentTurnId = typeof turn?.['id'] === 'string' ? turn['id'] : this.currentTurnId;
+                this.turnIndex += 1;
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'turn_start',
+                    turnIndex: this.turnIndex,
+                });
+                return;
+            }
+            case 'turn/completed': {
+                const turn = params['turn'];
+                const durationMs = typeof turn?.['durationMs'] === 'number' ? turn['durationMs'] : undefined;
+                this.currentTurnId = null;
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'turn_end',
+                    turnIndex: this.turnIndex,
+                    ...(durationMs != null ? { cost: undefined } : {}),
+                });
+                return;
+            }
+            case 'item/started': {
+                const item = params['item'];
+                if (!item)
+                    return;
+                this.handleItemStarted(item);
+                return;
+            }
+            case 'item/completed': {
+                const item = params['item'];
+                if (!item)
+                    return;
+                this.handleItemCompleted(item);
+                return;
+            }
+            case 'item/agentMessage/delta': {
+                const itemId = String(params['itemId'] ?? '');
+                const delta = String(params['delta'] ?? '');
+                const accumulated = (this.textByItemId.get(itemId) ?? '') + delta;
+                this.textByItemId.set(itemId, accumulated);
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'text_delta',
+                    delta,
+                    accumulated,
+                });
+                return;
+            }
+            case 'item/reasoning/textDelta':
+            case 'item/reasoning/summaryTextDelta': {
+                const itemId = String(params['itemId'] ?? '');
+                const delta = String(params['delta'] ?? '');
+                const accumulated = (this.thinkingByItemId.get(itemId) ?? '') + delta;
+                this.thinkingByItemId.set(itemId, accumulated);
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'thinking_delta',
+                    delta,
+                    accumulated,
+                });
+                return;
+            }
+            case 'item/commandExecution/outputDelta': {
+                const delta = String(params['delta'] ?? '');
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'shell_stdout_delta',
+                    delta,
+                });
+                return;
+            }
+            case 'thread/tokenUsage/updated': {
+                const tokenUsage = params['tokenUsage'];
+                const total = tokenUsage?.['total'];
+                if (!total)
+                    return;
+                const inputTokens = Number(total['inputTokens'] ?? 0);
+                const outputTokens = Number(total['outputTokens'] ?? 0);
+                const cachedTokens = Number(total['cachedInputTokens'] ?? 0);
+                const thinkingTokens = Number(total['reasoningOutputTokens'] ?? 0);
+                const cost = this.buildCostRecord(inputTokens, outputTokens, thinkingTokens, cachedTokens);
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'token_usage',
+                    inputTokens,
+                    outputTokens,
+                    thinkingTokens,
+                    cachedTokens,
+                });
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'cost',
+                    cost,
+                });
+                return;
+            }
+            case 'error': {
+                const error = params['error'];
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'error',
+                    code: 'INTERNAL',
+                    message: String(error?.['message'] ?? 'Codex app-server error'),
+                    recoverable: Boolean(params['willRetry']),
+                });
+                return;
+            }
+            case 'mcpServer/startupStatus/updated': {
+                const status = String(params['status'] ?? '');
+                if (status === 'failed') {
+                    this.pushEvent({
+                        ...this.baseEvent(),
+                        type: 'debug',
+                        level: 'warn',
+                        message: String(params['error'] ?? 'Codex MCP server failed to start'),
+                    });
+                }
+                return;
+            }
+            default:
+                return;
+        }
+    }
+    handleItemStarted(item) {
+        const itemType = String(item['type'] ?? '');
+        switch (itemType) {
+            case 'agentMessage': {
+                const itemId = String(item['id'] ?? '');
+                this.textByItemId.set(itemId, '');
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'message_start',
+                });
+                return;
+            }
+            case 'reasoning': {
+                const itemId = String(item['id'] ?? '');
+                this.thinkingByItemId.set(itemId, '');
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'thinking_start',
+                    effort: undefined,
+                });
+                return;
+            }
+            case 'commandExecution': {
+                const toolCallId = String(item['id'] ?? '');
+                const command = String(item['command'] ?? 'command');
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'tool_call_start',
+                    toolCallId,
+                    toolName: 'commandExecution',
+                    inputAccumulated: command,
+                });
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'shell_start',
+                    command,
+                    cwd: String(item['cwd'] ?? this.cwd),
+                });
+                return;
+            }
+            case 'mcpToolCall': {
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'mcp_tool_call_start',
+                    toolCallId: String(item['id'] ?? ''),
+                    server: String(item['server'] ?? ''),
+                    toolName: String(item['tool'] ?? ''),
+                    input: item['arguments'] ?? {},
+                });
+                return;
+            }
+            case 'dynamicToolCall': {
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'tool_call_start',
+                    toolCallId: String(item['id'] ?? ''),
+                    toolName: String(item['tool'] ?? ''),
+                    inputAccumulated: JSON.stringify(item['arguments'] ?? {}),
+                });
+                return;
+            }
+            default:
+                return;
+        }
+    }
+    handleItemCompleted(item) {
+        const itemType = String(item['type'] ?? '');
+        switch (itemType) {
+            case 'agentMessage': {
+                const itemId = String(item['id'] ?? '');
+                const text = String(item['text'] ?? this.textByItemId.get(itemId) ?? '');
+                this.textByItemId.delete(itemId);
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'message_stop',
+                    text,
+                });
+                return;
+            }
+            case 'reasoning': {
+                const itemId = String(item['id'] ?? '');
+                const thinking = this.thinkingByItemId.get(itemId) ?? '';
+                this.thinkingByItemId.delete(itemId);
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'thinking_stop',
+                    thinking,
+                });
+                return;
+            }
+            case 'commandExecution': {
+                const toolCallId = String(item['id'] ?? '');
+                const command = String(item['command'] ?? 'command');
+                const output = item['aggregatedOutput'] ?? '';
+                const durationMs = Number(item['durationMs'] ?? 0);
+                const exitCode = Number(item['exitCode'] ?? 0);
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'tool_result',
+                    toolCallId,
+                    toolName: 'commandExecution',
+                    output,
+                    durationMs,
+                });
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'shell_exit',
+                    exitCode,
+                    durationMs,
+                });
+                return;
+            }
+            case 'mcpToolCall': {
+                const toolCallId = String(item['id'] ?? '');
+                const server = String(item['server'] ?? '');
+                const toolName = String(item['tool'] ?? '');
+                if (item['error']) {
+                    this.pushEvent({
+                        ...this.baseEvent(),
+                        type: 'mcp_tool_error',
+                        toolCallId,
+                        server,
+                        toolName,
+                        error: JSON.stringify(item['error']),
+                    });
+                    return;
+                }
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'mcp_tool_result',
+                    toolCallId,
+                    server,
+                    toolName,
+                    output: item['result'] ?? null,
+                });
+                return;
+            }
+            case 'dynamicToolCall': {
+                const toolCallId = String(item['id'] ?? '');
+                const toolName = String(item['tool'] ?? '');
+                if (item['success'] === false) {
+                    this.pushEvent({
+                        ...this.baseEvent(),
+                        type: 'tool_error',
+                        toolCallId,
+                        toolName,
+                        error: JSON.stringify(item['contentItems'] ?? []),
+                    });
+                    return;
+                }
+                this.pushEvent({
+                    ...this.baseEvent(),
+                    type: 'tool_result',
+                    toolCallId,
+                    toolName,
+                    output: item['contentItems'] ?? [],
+                    durationMs: Number(item['durationMs'] ?? 0),
+                });
+                return;
+            }
+            default:
+                return;
+        }
+    }
+    async startTurn(text) {
+        if (!this.threadId) {
+            throw new Error('Codex app-server thread is not initialized');
+        }
+        const response = await this.request('turn/start', {
+            threadId: this.threadId,
+            input: [this.asTextInput(text)],
+            approvalPolicy: this.mapApprovalPolicy(),
+        });
+        this.currentTurnId = response.turn.id;
+    }
+    async sendUserMessage(text) {
+        if (!this.threadId) {
+            throw new Error('Codex app-server thread is not initialized');
+        }
+        if (this.currentTurnId) {
+            await this.request('turn/steer', {
+                threadId: this.threadId,
+                input: [this.asTextInput(text)],
+                expectedTurnId: this.currentTurnId,
+            });
+            return;
+        }
+        await this.startTurn(text);
+    }
+    async sendInteractionResponse(interactionId, response) {
+        const pending = this.pendingInteractionRequests.get(interactionId);
+        if (!pending || !this.ws) {
+            throw new Error(`No pending Codex interaction ${interactionId}`);
+        }
+        this.pendingInteractionRequests.delete(interactionId);
+        const result = this.buildInteractionResult(pending.method, response);
+        this.ws.send(JSON.stringify({ id: pending.id, result }));
+    }
+    buildInteractionResult(method, response) {
+        if (method === 'item/tool/requestUserInput') {
+            return {
+                answers: {
+                    response: {
+                        answers: [response.type === 'text' ? response.text : response.type],
+                    },
+                },
+            };
+        }
+        const approved = response.type === 'approve';
+        switch (method) {
+            case 'item/commandExecution/requestApproval':
+                return { decision: approved ? 'accept' : 'decline' };
+            case 'item/fileChange/requestApproval':
+                return { decision: approved ? 'accept' : 'decline' };
+            case 'item/permissions/requestApproval':
+                return approved
+                    ? { permissions: {}, scope: 'turn' }
+                    : { permissions: {}, scope: 'turn' };
+            default:
+                return {};
+        }
+    }
+    async request(method, params) {
+        if (!this.ws) {
+            throw new Error('Codex app-server socket is not connected');
+        }
+        const id = ++this.requestSeq;
+        const payload = { method, id, params };
+        const promise = new Promise((resolve, reject) => {
+            this.pendingResponses.set(id, { resolve: resolve, reject });
+        });
+        this.ws.send(JSON.stringify(payload));
+        return promise;
+    }
+    async notify(method) {
+        if (!this.ws) {
+            throw new Error('Codex app-server socket is not connected');
+        }
+        this.ws.send(JSON.stringify({ method }));
+    }
+    asTextInput(text) {
+        return {
+            type: 'text',
+            text,
+            text_elements: [],
+        };
+    }
+    mapApprovalPolicy() {
+        return this.approvalMode === 'yolo' ? 'never' : 'on-request';
+    }
+    buildCostRecord(inputTokens, outputTokens, thinkingTokens, cachedTokens) {
+        const model = this.models.find((entry) => entry.modelId === this.currentModelId) ?? this.models[0];
+        const inputPrice = model?.inputPricePerMillion ?? 0;
+        const outputPrice = model?.outputPricePerMillion ?? 0;
+        const thinkingPrice = model?.thinkingPricePerMillion ?? 0;
+        const cachedPrice = model?.cachedInputPricePerMillion ?? inputPrice;
+        const totalUsd = ((inputTokens - cachedTokens) / 1_000_000) * inputPrice +
+            (cachedTokens / 1_000_000) * cachedPrice +
+            (outputTokens / 1_000_000) * outputPrice +
+            (thinkingTokens / 1_000_000) * thinkingPrice;
+        return {
+            totalUsd,
+            inputTokens,
+            outputTokens,
+            thinkingTokens,
+            cachedTokens,
+        };
+    }
+    isUserMessagePayload(data) {
+        if (typeof data !== 'object' || data === null) {
+            return false;
+        }
+        const candidate = data;
+        return candidate.type === 'user_message' && typeof candidate.text === 'string';
+    }
+    isInteractionResponsePayload(data) {
+        if (typeof data !== 'object' || data === null) {
+            return false;
+        }
+        const candidate = data;
+        return candidate.type === 'interaction_response' && typeof candidate.interactionId === 'string' && candidate.response != null;
+    }
+    pushEvent(event) {
+        if (this.waitingResolver) {
+            const resolve = this.waitingResolver;
+            this.waitingResolver = null;
+            resolve(event);
+            return;
+        }
+        this.eventQueue.push(event);
+    }
+    baseEvent() {
+        return {
+            runId: this.connectionId,
+            agent: 'codex-websocket',
+            timestamp: Date.now(),
+        };
+    }
+}
+//# sourceMappingURL=codex-websocket-connection.js.map