fix: stabilize cloud computer use mcp

This commit is contained in:
Simos Mikelatos
2026-06-19 20:47:53 +00:00
parent 077baee5f2
commit 1c05fe0905
7 changed files with 146 additions and 33 deletions

View File

@@ -3,6 +3,7 @@ import fs from 'node:fs/promises';
import path from 'node:path'; import path from 'node:path';
const IPC_PREFIX = '@@CUAGENT@@'; const IPC_PREFIX = '@@CUAGENT@@';
const TARGET_STATUS_TIMEOUT_MS = 5000;
function getDesktopPath() { function getDesktopPath() {
const currentPath = process.env.PATH || ''; 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 * Keeps a Computer Use desktop agent connected to running cloud environments
* while desktop access is enabled. * while desktop access is enabled.
@@ -102,7 +135,8 @@ export class ComputerAgentController {
async sync() { async sync() {
const targets = this.settings.enabled ? (this.getRunningEnvironmentUrls?.() || []) : []; 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 = const sameTargets =
wsTargets.length === this.currentTargets.length && wsTargets.length === this.currentTargets.length &&
@@ -209,6 +243,12 @@ export class ComputerAgentController {
this.connectedUrls.delete(payload.url); this.connectedUrls.delete(payload.url);
this.lastEvent = 'disconnected'; this.lastEvent = 'disconnected';
this.onChange?.(); 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; break;
case 'starting': case 'starting':
this.lastEvent = 'starting'; this.lastEvent = 'starting';

View File

@@ -245,8 +245,12 @@ function connect(url: string): void {
} }
}); });
const scheduleReconnect = () => { const scheduleReconnect = (code?: number, reason?: Buffer) => {
emitToParent({ type: 'disconnected', url }); const reasonText = reason?.toString() || '';
emitToParent({ type: 'disconnected', url, code, reason: reasonText });
if (code === 1008 && /computer use.*disabled/i.test(reasonText)) {
return;
}
setTimeout(open, reconnectMs); setTimeout(open, reconnectMs);
reconnectMs = Math.min(reconnectMs * 2, RECONNECT_MAX_MS); reconnectMs = Math.min(reconnectMs * 2, RECONNECT_MAX_MS);
}; };

View File

@@ -470,19 +470,25 @@ async function handleMessage(message: JsonRpcRequest) {
throw new Error(`Unsupported method: ${message.method}`); throw new Error(`Unsupported method: ${message.method}`);
} }
function writeMessage(message: Record<string, unknown>) { type MessageFraming = 'content-length' | 'line';
function writeMessage(message: Record<string, unknown>, framing: MessageFraming) {
const payload = JSON.stringify(message); 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}`); 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) { if (id === undefined) {
return; 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) { if (id === undefined) {
return; return;
} }
@@ -493,28 +499,69 @@ function sendError(id: string | number | null | undefined, error: unknown) {
code: -32000, code: -32000,
message: error instanceof Error ? error.message : String(error), message: error instanceof Error ? error.message : String(error),
}, },
}); }, framing);
} }
let buffer = Buffer.alloc(0); 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) => { process.stdin.on('data', (chunk) => {
buffer = Buffer.concat([buffer, chunk]); buffer = Buffer.concat([buffer, chunk]);
while (true) { while (true) {
const headerEnd = buffer.indexOf('\r\n\r\n'); const headerEnd = findHeaderEnd(buffer);
if (headerEnd === -1) { if (!headerEnd) {
return; 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); const lengthMatch = /Content-Length:\s*(\d+)/i.exec(header);
if (!lengthMatch) { if (!lengthMatch) {
buffer = buffer.slice(headerEnd + 4); buffer = buffer.slice(headerEnd.index + headerEnd.length);
continue; continue;
} }
const length = Number.parseInt(lengthMatch[1], 10); const length = Number.parseInt(lengthMatch[1], 10);
const messageStart = headerEnd + 4; const messageStart = headerEnd.index + headerEnd.length;
const messageEnd = messageStart + length; const messageEnd = messageStart + length;
if (buffer.length < messageEnd) { if (buffer.length < messageEnd) {
return; return;
@@ -522,16 +569,6 @@ process.stdin.on('data', (chunk) => {
const rawMessage = buffer.slice(messageStart, messageEnd).toString('utf8'); const rawMessage = buffer.slice(messageStart, messageEnd).toString('utf8');
buffer = buffer.slice(messageEnd); buffer = buffer.slice(messageEnd);
handleRawMessage(rawMessage, 'content-length');
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);
}
})();
} }
}); });

View File

@@ -22,7 +22,7 @@ try {
} }
}); });
} catch (e) { } 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 // Keep the default database in a stable user-level location so rebuilding dist-server

View File

@@ -465,6 +465,7 @@ export const computerUseService = {
await this.registerAgentMcp(); await this.registerAgentMcp();
} else { } else {
await this.unregisterAgentMcp(); await this.unregisterAgentMcp();
desktopAgentRelay.disconnectAll('Computer Use was disabled in this environment.');
stopSessions('settings:disabled', 'Computer Use was disabled in settings.'); stopSessions('settings:disabled', 'Computer Use was disabled in settings.');
} }
return next; return next;
@@ -909,6 +910,7 @@ export const computerUseService = {
// Drive cloud MCP exposure + session teardown off desktop-agent connectivity. // Drive cloud MCP exposure + session teardown off desktop-agent connectivity.
desktopAgentRelay.setHooks({ desktopAgentRelay.setHooks({
canAcceptConnection: () => getRuntime() === 'cloud' && readSettings().enabled,
onFirstConnect: () => computerUseService.onDesktopAgentConnected(), onFirstConnect: () => computerUseService.onDesktopAgentConnected(),
onLastDisconnect: () => computerUseService.onDesktopAgentDisconnected(), onLastDisconnect: () => computerUseService.onDesktopAgentDisconnected(),
}); });

View File

@@ -18,6 +18,7 @@ type ConnectedAgent = {
}; };
type RelayLifecycleHooks = { type RelayLifecycleHooks = {
canAcceptConnection?: () => boolean;
onFirstConnect?: () => void | Promise<void>; onFirstConnect?: () => void | Promise<void>;
onLastDisconnect?: () => void | Promise<void>; onLastDisconnect?: () => void | Promise<void>;
}; };
@@ -54,15 +55,25 @@ export const desktopAgentRelay = {
hooks = next; 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; const wasEmpty = pickAgent() === undefined;
agents.set(ws, { ws, label, registeredAt: new Date().toISOString() }); agents.set(ws, { ws, label, registeredAt: new Date().toISOString() });
console.log(`[DesktopAgent] Registered (${label}); ${agents.size} connected.`); console.log(`[DesktopAgent] Registered (${label}); ${agents.size} connected.`);
ws.on('close', () => { ws.on('close', () => {
agents.delete(ws); const wasRegistered = agents.delete(ws);
console.log(`[DesktopAgent] Disconnected (${label}); ${agents.size} remain.`); console.log(`[DesktopAgent] Disconnected (${label}); ${agents.size} remain.`);
if (pickAgent() === undefined) { if (wasRegistered && pickAgent() === undefined) {
rejectAllPending('Desktop agent disconnected.'); rejectAllPending('Desktop agent disconnected.');
void hooks.onLastDisconnect?.(); void hooks.onLastDisconnect?.();
} }
@@ -71,6 +82,24 @@ export const desktopAgentRelay = {
if (wasEmpty) { if (wasEmpty) {
void hooks.onFirstConnect?.(); 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. */ /** Resolves a pending relay call with the desktop agent's reply. */

View File

@@ -102,6 +102,7 @@ export default function ComputerUseSettingsTab() {
const isCloud = status?.runtime === 'cloud'; const isCloud = status?.runtime === 'cloud';
const effectiveEnabled = isCloud ? status?.enabled === true : settings.enabled; const effectiveEnabled = isCloud ? status?.enabled === true : settings.enabled;
const showCloudDesktopAccess = Boolean(isCloud && effectiveEnabled);
const needsRuntime = Boolean(effectiveEnabled && !isCloud && status && (!status.nutInstalled || !status.screenshotInstalled)); const needsRuntime = Boolean(effectiveEnabled && !isCloud && status && (!status.nutInstalled || !status.screenshotInstalled));
const desktopAgentCount = status?.desktopAgentCount ?? (status?.desktopAgentConnected ? 1 : 0); const desktopAgentCount = status?.desktopAgentCount ?? (status?.desktopAgentConnected ? 1 : 0);
const modeDescription = isCloud const modeDescription = isCloud
@@ -144,7 +145,7 @@ export default function ComputerUseSettingsTab() {
/> />
</SettingsRow> </SettingsRow>
{isCloud && ( {showCloudDesktopAccess && (
<SettingsRow <SettingsRow
label="Cloud desktop access" label="Cloud desktop access"
description={status?.desktopAgentConnected description={status?.desktopAgentConnected
@@ -176,9 +177,9 @@ export default function ComputerUseSettingsTab() {
</SettingsRow> </SettingsRow>
)} )}
{(needsRuntime || isCloud || error) && ( {(needsRuntime || showCloudDesktopAccess || error) && (
<div className="space-y-4 px-4 py-4"> <div className="space-y-4 px-4 py-4">
{isCloud && !status?.desktopAgentConnected && ( {showCloudDesktopAccess && !status?.desktopAgentConnected && (
<div className="rounded-md border border-border bg-muted/40 px-3 py-3 text-sm text-muted-foreground"> <div className="rounded-md border border-border bg-muted/40 px-3 py-3 text-sm text-muted-foreground">
<div className="font-medium text-foreground">To link this computer</div> <div className="font-medium text-foreground">To link this computer</div>
<ol className="mt-2 list-decimal space-y-1 pl-5"> <ol className="mt-2 list-decimal space-y-1 pl-5">
@@ -190,7 +191,7 @@ export default function ComputerUseSettingsTab() {
</div> </div>
)} )}
{isCloud && status?.desktopAgentConnected && ( {showCloudDesktopAccess && status?.desktopAgentConnected && (
<div className="rounded-md border border-border bg-muted/40 px-3 py-2 text-sm text-muted-foreground"> <div className="rounded-md border border-border bg-muted/40 px-3 py-2 text-sm text-muted-foreground">
{desktopAgentCount > 1 {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.` ? `${desktopAgentCount} desktops are linked. Agents will use one available desktop; stop Computer Use on any desktop you do not want agents to control.`