From 077baee5f27f5c93ac40fa90557b910ace4dbd4d Mon Sep 17 00:00:00 2001
From: Simos Mikelatos <simosmik@gmail.com>
Date: Fri, 19 Jun 2026 15:52:23 +0000
Subject: [PATCH] fix: authenticate desktop agent websocket

---
 electron/computerAgent.js                       |  4 +++-
 electron/main.js                                |  1 +
 server/computer-use-agent.ts                    |  5 ++---
 .../services/desktop-agent-websocket.service.ts | 17 +++++++++++++----
 4 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/electron/computerAgent.js b/electron/computerAgent.js
index b28cea6e..c06ab26a 100644
--- a/electron/computerAgent.js
+++ b/electron/computerAgent.js
@@ -40,11 +40,12 @@ function toAgentWsUrl(httpUrl) {
  * while desktop access is enabled.
  */
 export class ComputerAgentController {
-  constructor({ appRoot, settingsPath, isPackaged = false, getRunningEnvironmentUrls, promptConsent, onChange }) {
+  constructor({ appRoot, settingsPath, isPackaged = false, getRunningEnvironmentUrls, getApiKey, promptConsent, onChange }) {
     this.appRoot = appRoot;
     this.settingsPath = settingsPath;
     this.isPackaged = isPackaged;
     this.getRunningEnvironmentUrls = getRunningEnvironmentUrls;
+    this.getApiKey = getApiKey;
     this.promptConsent = promptConsent;
     this.onChange = onChange;
     this.settings = { enabled: false, consentMode: 'ask' };
@@ -138,6 +139,7 @@ export class ComputerAgentController {
         ...runtime.env,
         PATH: getDesktopPath(),
         CLOUDCLI_DESKTOP_AGENT_URLS: wsTargets.join(','),
+        CLOUDCLI_DESKTOP_AGENT_API_KEY: this.getApiKey?.() || '',
         CLOUDCLI_COMPUTER_USE_CONSENT_MODE: this.settings.consentMode,
       },
       stdio: ['pipe', 'pipe', 'pipe'],
diff --git a/electron/main.js b/electron/main.js
index 2e647ae3..451ba68f 100644
--- a/electron/main.js
+++ b/electron/main.js
@@ -909,6 +909,7 @@ async function bootstrap() {
     settingsPath: getComputerUseSettingsPath(),
     isPackaged: app.isPackaged,
     getRunningEnvironmentUrls,
+    getApiKey: () => cloud.getAccount()?.apiKey || '',
     promptConsent: promptComputerUseConsent,
     onChange: syncDesktopState,
   });
diff --git a/server/computer-use-agent.ts b/server/computer-use-agent.ts
index 53fc772f..5402b2b6 100644
--- a/server/computer-use-agent.ts
+++ b/server/computer-use-agent.ts
@@ -42,6 +42,7 @@ const RECONNECT_MAX_MS = 30_000;
 
 const consentMode: ConsentMode = process.env.CLOUDCLI_COMPUTER_USE_CONSENT_MODE === 'auto' ? 'auto' : 'ask';
 const agentLabel = process.env.CLOUDCLI_DESKTOP_AGENT_LABEL || 'cloudcli-desktop';
+const desktopAgentApiKey = process.env.CLOUDCLI_DESKTOP_AGENT_API_KEY || '';
 
 function parseTargets(): string[] {
   const raw =
@@ -195,9 +196,7 @@ function connect(url: string): void {
 
   const open = () => {
     socket = new WebSocket(url, {
-      headers: process.env.CLOUDCLI_DESKTOP_AGENT_TOKEN
-        ? { 'x-cloudcli-agent-token': process.env.CLOUDCLI_DESKTOP_AGENT_TOKEN }
-        : undefined,
+      headers: desktopAgentApiKey ? { 'X-API-Key': desktopAgentApiKey } : undefined,
     });
 
     socket.on('open', () => {
diff --git a/server/modules/websocket/services/desktop-agent-websocket.service.ts b/server/modules/websocket/services/desktop-agent-websocket.service.ts
index 3f9a7614..aadd129b 100644
--- a/server/modules/websocket/services/desktop-agent-websocket.service.ts
+++ b/server/modules/websocket/services/desktop-agent-websocket.service.ts
@@ -15,9 +15,7 @@ export function handleDesktopAgentConnection(
   ws: WebSocket,
   request: AuthenticatedWebSocketRequest
 ): void {
-  const label = request.user?.username ? `desktop:${request.user.username}` : 'desktop-agent';
-  console.log('[INFO] Desktop agent websocket connected:', label);
-  desktopAgentRelay.register(ws, label);
+  let registered = false;
 
   ws.on('message', (rawMessage) => {
     const data = parseIncomingJsonObject(rawMessage);
@@ -25,6 +23,17 @@ export function handleDesktopAgentConnection(
       return;
     }
     const kind = typeof data.kind === 'string' ? data.kind : typeof data.type === 'string' ? data.type : '';
+    if (kind === 'register' && !registered) {
+      const label = typeof data.label === 'string' && data.label.trim()
+        ? data.label.trim()
+        : request.user?.username
+          ? `desktop:${request.user.username}`
+          : 'desktop-agent';
+      registered = true;
+      console.log('[INFO] Desktop agent websocket registered:', label);
+      desktopAgentRelay.register(ws, label);
+      return;
+    }
     if (kind === 'computer_relay_result' && typeof data.id === 'string') {
       desktopAgentRelay.handleResult(
         data.id,
@@ -37,6 +46,6 @@ export function handleDesktopAgentConnection(
   });
 
   ws.on('close', () => {
-    console.log('[INFO] Desktop agent websocket disconnected:', label);
+    console.log('[INFO] Desktop agent websocket disconnected');
   });
 }