From 1c05fe0905d1cbc85edccff0d95aa6776d20978b Mon Sep 17 00:00:00 2001 From: Simos Mikelatos Date: Fri, 19 Jun 2026 20:47:53 +0000 Subject: [PATCH] fix: stabilize cloud computer use mcp --- electron/computerAgent.js | 42 +++++++++- server/computer-use-agent.ts | 8 +- server/computer-use-mcp.ts | 81 ++++++++++++++----- server/load-env.js | 2 +- .../computer-use/computer-use.service.ts | 2 + .../desktop-agent-relay.service.ts | 35 +++++++- .../ComputerUseSettingsTab.tsx | 9 ++- 7 files changed, 146 insertions(+), 33 deletions(-) diff --git a/electron/computerAgent.js b/electron/computerAgent.js index c06ab26a..c72a73ec 100644 --- a/electron/computerAgent.js +++ b/electron/computerAgent.js @@ -3,6 +3,7 @@ import fs from 'node:fs/promises'; import path from 'node:path'; const IPC_PREFIX = '@@CUAGENT@@'; +const TARGET_STATUS_TIMEOUT_MS = 5000; function getDesktopPath() { const currentPath = process.env.PATH || ''; @@ -35,6 +36,38 @@ function toAgentWsUrl(httpUrl) { } } +async function isComputerUseEnabledTarget(httpUrl, apiKey) { + let statusUrl; + try { + statusUrl = new URL('/api/computer-use/status', httpUrl).toString(); + } catch { + return false; + } + + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), TARGET_STATUS_TIMEOUT_MS); + try { + const response = await fetch(statusUrl, { + signal: controller.signal, + headers: apiKey ? { 'X-API-Key': apiKey } : undefined, + }); + const body = await response.json().catch(() => null); + return response.ok && body?.success !== false && body?.data?.enabled === true; + } catch { + return false; + } finally { + clearTimeout(timeout); + } +} + +async function filterEnabledComputerUseTargets(targets, apiKey) { + const checks = await Promise.all(targets.map(async (target) => ({ + target, + enabled: await isComputerUseEnabledTarget(target, apiKey), + }))); + return checks.filter((item) => item.enabled).map((item) => item.target); +} + /** * Keeps a Computer Use desktop agent connected to running cloud environments * while desktop access is enabled. @@ -102,7 +135,8 @@ export class ComputerAgentController { async sync() { const targets = this.settings.enabled ? (this.getRunningEnvironmentUrls?.() || []) : []; - const wsTargets = targets.map(toAgentWsUrl).filter(Boolean); + const enabledTargets = this.settings.enabled ? await filterEnabledComputerUseTargets(targets, this.getApiKey?.() || '') : []; + const wsTargets = enabledTargets.map(toAgentWsUrl).filter(Boolean); const sameTargets = wsTargets.length === this.currentTargets.length && @@ -209,6 +243,12 @@ export class ComputerAgentController { this.connectedUrls.delete(payload.url); this.lastEvent = 'disconnected'; this.onChange?.(); + if (payload.reason && /computer use.*disabled/i.test(payload.reason)) { + void this.sync().catch((error) => { + this.lastError = error instanceof Error ? error.message : 'Failed to sync Computer Use targets.'; + this.onChange?.(); + }); + } break; case 'starting': this.lastEvent = 'starting'; diff --git a/server/computer-use-agent.ts b/server/computer-use-agent.ts index 5402b2b6..d81e6c2b 100644 --- a/server/computer-use-agent.ts +++ b/server/computer-use-agent.ts @@ -245,8 +245,12 @@ function connect(url: string): void { } }); - const scheduleReconnect = () => { - emitToParent({ type: 'disconnected', url }); + const scheduleReconnect = (code?: number, reason?: Buffer) => { + const reasonText = reason?.toString() || ''; + emitToParent({ type: 'disconnected', url, code, reason: reasonText }); + if (code === 1008 && /computer use.*disabled/i.test(reasonText)) { + return; + } setTimeout(open, reconnectMs); reconnectMs = Math.min(reconnectMs * 2, RECONNECT_MAX_MS); }; diff --git a/server/computer-use-mcp.ts b/server/computer-use-mcp.ts index c6e717ba..ee998a33 100644 --- a/server/computer-use-mcp.ts +++ b/server/computer-use-mcp.ts @@ -470,19 +470,25 @@ async function handleMessage(message: JsonRpcRequest) { throw new Error(`Unsupported method: ${message.method}`); } -function writeMessage(message: Record) { +type MessageFraming = 'content-length' | 'line'; + +function writeMessage(message: Record, framing: MessageFraming) { const payload = JSON.stringify(message); + if (framing === 'line') { + process.stdout.write(`${payload}\n`); + return; + } process.stdout.write(`Content-Length: ${Buffer.byteLength(payload, 'utf8')}\r\n\r\n${payload}`); } -function sendResult(id: string | number | null | undefined, result: unknown) { +function sendResult(id: string | number | null | undefined, result: unknown, framing: MessageFraming) { if (id === undefined) { return; } - writeMessage({ jsonrpc: '2.0', id, result }); + writeMessage({ jsonrpc: '2.0', id, result }, framing); } -function sendError(id: string | number | null | undefined, error: unknown) { +function sendError(id: string | number | null | undefined, error: unknown, framing: MessageFraming) { if (id === undefined) { return; } @@ -493,28 +499,69 @@ function sendError(id: string | number | null | undefined, error: unknown) { code: -32000, message: error instanceof Error ? error.message : String(error), }, - }); + }, framing); } let buffer = Buffer.alloc(0); +function handleRawMessage(rawMessage: string, framing: MessageFraming) { + void (async () => { + let request: JsonRpcRequest | null = null; + try { + request = JSON.parse(rawMessage) as JsonRpcRequest; + const result = await handleMessage(request); + sendResult(request.id, result, framing); + } catch (error) { + sendError(request?.id ?? null, error, framing); + } + })(); +} + +function findHeaderEnd(input: Buffer): { index: number; length: number } | null { + const crlf = input.indexOf('\r\n\r\n'); + if (crlf !== -1) { + return { index: crlf, length: 4 }; + } + const lf = input.indexOf('\n\n'); + if (lf !== -1) { + return { index: lf, length: 2 }; + } + return null; +} + process.stdin.on('data', (chunk) => { buffer = Buffer.concat([buffer, chunk]); while (true) { - const headerEnd = buffer.indexOf('\r\n\r\n'); - if (headerEnd === -1) { - return; + const headerEnd = findHeaderEnd(buffer); + if (!headerEnd) { + if (/^Content-Length:/i.test(buffer.toString('utf8', 0, Math.min(buffer.length, 32)))) { + return; + } + + const newline = buffer.indexOf('\n'); + if (newline === -1) { + return; + } + + const rawLine = buffer.slice(0, newline).toString('utf8').trim(); + buffer = buffer.slice(newline + 1); + if (!rawLine) { + continue; + } + + handleRawMessage(rawLine, 'line'); + continue; } - const header = buffer.slice(0, headerEnd).toString('utf8'); + const header = buffer.slice(0, headerEnd.index).toString('utf8'); const lengthMatch = /Content-Length:\s*(\d+)/i.exec(header); if (!lengthMatch) { - buffer = buffer.slice(headerEnd + 4); + buffer = buffer.slice(headerEnd.index + headerEnd.length); continue; } const length = Number.parseInt(lengthMatch[1], 10); - const messageStart = headerEnd + 4; + const messageStart = headerEnd.index + headerEnd.length; const messageEnd = messageStart + length; if (buffer.length < messageEnd) { return; @@ -522,16 +569,6 @@ process.stdin.on('data', (chunk) => { const rawMessage = buffer.slice(messageStart, messageEnd).toString('utf8'); buffer = buffer.slice(messageEnd); - - void (async () => { - let request: JsonRpcRequest | null = null; - try { - request = JSON.parse(rawMessage) as JsonRpcRequest; - const result = await handleMessage(request); - sendResult(request.id, result); - } catch (error) { - sendError(request?.id ?? null, error); - } - })(); + handleRawMessage(rawMessage, 'content-length'); } }); diff --git a/server/load-env.js b/server/load-env.js index 2c889534..8d2954b6 100644 --- a/server/load-env.js +++ b/server/load-env.js @@ -22,7 +22,7 @@ try { } }); } catch (e) { - console.log('No .env file found or error reading it:', e.message); + console.error('No .env file found or error reading it:', e.message); } // Keep the default database in a stable user-level location so rebuilding dist-server diff --git a/server/modules/computer-use/computer-use.service.ts b/server/modules/computer-use/computer-use.service.ts index 86094e6e..8608395a 100644 --- a/server/modules/computer-use/computer-use.service.ts +++ b/server/modules/computer-use/computer-use.service.ts @@ -465,6 +465,7 @@ export const computerUseService = { await this.registerAgentMcp(); } else { await this.unregisterAgentMcp(); + desktopAgentRelay.disconnectAll('Computer Use was disabled in this environment.'); stopSessions('settings:disabled', 'Computer Use was disabled in settings.'); } return next; @@ -909,6 +910,7 @@ export const computerUseService = { // Drive cloud MCP exposure + session teardown off desktop-agent connectivity. desktopAgentRelay.setHooks({ + canAcceptConnection: () => getRuntime() === 'cloud' && readSettings().enabled, onFirstConnect: () => computerUseService.onDesktopAgentConnected(), onLastDisconnect: () => computerUseService.onDesktopAgentDisconnected(), }); diff --git a/server/modules/computer-use/desktop-agent-relay.service.ts b/server/modules/computer-use/desktop-agent-relay.service.ts index fc3c62cd..aa1171f0 100644 --- a/server/modules/computer-use/desktop-agent-relay.service.ts +++ b/server/modules/computer-use/desktop-agent-relay.service.ts @@ -18,6 +18,7 @@ type ConnectedAgent = { }; type RelayLifecycleHooks = { + canAcceptConnection?: () => boolean; onFirstConnect?: () => void | Promise; onLastDisconnect?: () => void | Promise; }; @@ -54,15 +55,25 @@ export const desktopAgentRelay = { hooks = next; }, - register(ws: WebSocket, label = 'desktop-agent'): void { + register(ws: WebSocket, label = 'desktop-agent'): boolean { + if (hooks.canAcceptConnection && !hooks.canAcceptConnection()) { + console.log(`[DesktopAgent] Rejected (${label}); Computer Use is disabled.`); + try { + ws.close(1008, 'Computer Use is disabled in this environment.'); + } catch { + // ignore close failures + } + return false; + } + const wasEmpty = pickAgent() === undefined; agents.set(ws, { ws, label, registeredAt: new Date().toISOString() }); console.log(`[DesktopAgent] Registered (${label}); ${agents.size} connected.`); ws.on('close', () => { - agents.delete(ws); + const wasRegistered = agents.delete(ws); console.log(`[DesktopAgent] Disconnected (${label}); ${agents.size} remain.`); - if (pickAgent() === undefined) { + if (wasRegistered && pickAgent() === undefined) { rejectAllPending('Desktop agent disconnected.'); void hooks.onLastDisconnect?.(); } @@ -71,6 +82,24 @@ export const desktopAgentRelay = { if (wasEmpty) { void hooks.onFirstConnect?.(); } + return true; + }, + + disconnectAll(reason = 'Desktop agent disconnected.'): void { + const hadAgent = pickAgent() !== undefined; + const sockets = [...agents.keys()]; + agents.clear(); + for (const ws of sockets) { + try { + ws.close(1008, reason); + } catch { + // ignore close failures + } + } + rejectAllPending(reason); + if (hadAgent) { + void hooks.onLastDisconnect?.(); + } }, /** Resolves a pending relay call with the desktop agent's reply. */ diff --git a/src/components/settings/view/tabs/computer-use-settings/ComputerUseSettingsTab.tsx b/src/components/settings/view/tabs/computer-use-settings/ComputerUseSettingsTab.tsx index 798eaa3a..7338459f 100644 --- a/src/components/settings/view/tabs/computer-use-settings/ComputerUseSettingsTab.tsx +++ b/src/components/settings/view/tabs/computer-use-settings/ComputerUseSettingsTab.tsx @@ -102,6 +102,7 @@ export default function ComputerUseSettingsTab() { const isCloud = status?.runtime === 'cloud'; const effectiveEnabled = isCloud ? status?.enabled === true : settings.enabled; + const showCloudDesktopAccess = Boolean(isCloud && effectiveEnabled); const needsRuntime = Boolean(effectiveEnabled && !isCloud && status && (!status.nutInstalled || !status.screenshotInstalled)); const desktopAgentCount = status?.desktopAgentCount ?? (status?.desktopAgentConnected ? 1 : 0); const modeDescription = isCloud @@ -144,7 +145,7 @@ export default function ComputerUseSettingsTab() { /> - {isCloud && ( + {showCloudDesktopAccess && ( )} - {(needsRuntime || isCloud || error) && ( + {(needsRuntime || showCloudDesktopAccess || error) && (
- {isCloud && !status?.desktopAgentConnected && ( + {showCloudDesktopAccess && !status?.desktopAgentConnected && (
To link this computer
    @@ -190,7 +191,7 @@ export default function ComputerUseSettingsTab() {
)} - {isCloud && status?.desktopAgentConnected && ( + {showCloudDesktopAccess && status?.desktopAgentConnected && (
{desktopAgentCount > 1 ? `${desktopAgentCount} desktops are linked. Agents will use one available desktop; stop Computer Use on any desktop you do not want agents to control.`