Support multiple bridge clients per server

copilotj/server/bridge.py

@@ -182,6 +182,20 @@ async def _on_message_text(self, msg: aiohttp.WSMessage) -> None:
 
     async def _on_event(self, response: _BridgeEvent) -> None:
         match response.event:
+            case "auth":
+                token = response.data
+                if isinstance(token, dict):
+                    token = token.get("access_token")
+                if token and isinstance(token, str):
+                    self._bridge._token_to_client[token] = self.id
+                    _log.info("Client %s authenticated with access token", self.id)
+                # Send auth acknowledgement
+                auth_payload = _BridgeEvent(
+                    id=uuid.uuid4(), event_id=response.event_id, event="auth_ok", data=None, err=None
+                )
+                await self._text_message_queue.put(auth_payload.model_dump_json())
+                return
+
             case "query_id":
                 data = _IdChangedEventData(id=self.id)
                 payload = _BridgeEvent(
@@ -216,6 +230,7 @@ def __init__(self) -> None:
         self._clients = dict[uuid.UUID, _Client]()
         self._close_event: asyncio.Event | None = None  # TODO
         self._used_client_ids = dict[uuid.UUID, datetime]()
+        self._token_to_client = dict[str, uuid.UUID]()
 
     async def on_plugin_connect(self, request: web.Request) -> web.StreamResponse:
         client = _Client(self)
@@ -224,6 +239,12 @@ async def on_plugin_connect(self, request: web.Request) -> web.StreamResponse:
         _log.info("Client WebSocket connection established: %s", client.id)
 
         ws = await client.run(request)
+
+        # Clean up token mappings for this client
+        tokens_to_remove = [t for t, cid in self._token_to_client.items() if cid == client.id]
+        for t in tokens_to_remove:
+            del self._token_to_client[t]
+
         self._clients.pop(client.id, None)
         return ws
 
@@ -262,6 +283,15 @@ async def close(self) -> None:
             await c.close()
 
         self._clients.clear()
+        self._token_to_client.clear()
+
+    def get_client_id_by_token(self, token: str) -> uuid.UUID | None:
+        """Resolve an access token to a client ID."""
+        return self._token_to_client.get(token)
+
+    def list_clients(self) -> list[dict[str, str]]:
+        """Return info about connected clients."""
+        return [{"id": str(cid)} for cid in self._clients]
 
     def _negotiate_id(self, client: _Client, new_id: uuid.UUID) -> str | None:
         if new_id == client.id:  # not changed

copilotj/server/server.py

@@ -21,7 +21,7 @@ class Server:
     def __init__(self):
         super().__init__()
         self._bridge = Bridge()
-        self._threads = Threads()
+        self._threads = Threads(self._bridge)
         self._app = self._create_app()
 
     def add_background_task(self, task: asyncio.Task) -> None:
@@ -44,6 +44,7 @@ def _create_app(self) -> web.Application:
 
         r = app.router
         r.add_get("/api/ping", _on_ping)
+        r.add_get("/api/clients", _on_list_clients)
         r.add_get("/api/plugins", self._bridge.on_plugin_connect)
         r.add_post("/api/plugins/events", self._bridge.on_forward_event)
         r.add_post("/api/threads", self._threads.new_thread)
@@ -72,8 +73,17 @@ async def on_shutdown(app: web.Application) -> None:
             await asyncio.gather(self._threads.close(), self._bridge.close())
 
         app.on_shutdown.append(on_shutdown)
+        app["server"] = self
         return app
 
 
 async def _on_ping(request: web.Request) -> web.Response:
     return web.Response(text="pong")
+
+
+async def _on_list_clients(request: web.Request) -> web.Response:
+    """List connected bridge clients."""
+    import json
+
+    clients = request.app["server"]._bridge.list_clients()
+    return web.Response(text=json.dumps(clients), content_type="application/json")

copilotj/server/threads.py

@@ -8,7 +8,7 @@
 import uuid
 from asyncio import Future
 from contextlib import suppress
-from typing import AsyncGenerator, Literal, override
+from typing import TYPE_CHECKING, AsyncGenerator, Literal, override
 
 import aiohttp.web as web
 import langfuse
@@ -18,7 +18,10 @@
 from copilotj.core.config import get_llm_and_key
 from copilotj.core.ui import UIEventContentMarkdown
 from copilotj.multiagent.leader_multiagent import LeaderDriven
-from copilotj.plugin.api import HTTPPluginAPI, PluginAPI
+from copilotj.plugin.api import DEV_CLIENT_ID, HTTPPluginAPI, PluginAPI
+
+if TYPE_CHECKING:
+    from copilotj.server.bridge import Bridge
 
 __all__ = ["Threads"]
 
@@ -61,14 +64,22 @@ class _OptimizePrompt(pydantic.BaseModel):
 
 
 class _Thread(UI):
-    def __init__(self, thread_id: str, *, config: _ConfigQuery | None = None, trace_context: langfuse.Langfuse):
+    def __init__(
+        self,
+        thread_id: str,
+        *,
+        client_id: uuid.UUID | None = None,
+        config: _ConfigQuery | None = None,
+        trace_context: langfuse.Langfuse,
+    ):
         self.thread_id = thread_id
         self._trace_ctx = trace_context
 
         self._mailbox = asyncio.Queue[UIEvent | _Signal]()
 
         self._apis: PluginAPI = HTTPPluginAPI("http://127.0.0.1:8786")  # TODO: make configurable
-        client_apis = self._apis.attach_dev_client()  # TODO: should be removed
+        cid = client_id or DEV_CLIENT_ID
+        client_apis = self._apis.with_client(cid)
 
         config_model = config and config.model
         self._agent = LeaderDriven(
@@ -239,12 +250,22 @@ async def request_user_manipulate(self, role: str, message: str | None = None) -
 
 
 class Threads:
-    def __init__(self):
+    def __init__(self, bridge: "Bridge"):
         super().__init__()
+        self._bridge = bridge
         self._threads: dict[str, tuple[_Thread, threading.Lock]] = {}
         self._threads_lock = threading.Lock()
         self._trace_ctx = langfuse.Langfuse()
 
+    def _resolve_client_id(self, request: web.Request) -> uuid.UUID | None:
+        """Extract access token from Authorization header and resolve to a client ID."""
+        auth_header = request.headers.get("Authorization", "")
+        if auth_header.startswith("Bearer "):
+            token = auth_header[7:].strip()
+            if token:
+                return self._bridge.get_client_id_by_token(token)
+        return None
+
     async def new_thread(self, request: web.Request) -> web.Response:
         try:
             data = await request.json()
@@ -267,8 +288,9 @@ async def new_thread(self, request: web.Request) -> web.Response:
                     text="No model configured. Please click the Settings gear icon in the sidebar to set up a model and API key.",
                 )
 
+        client_id = self._resolve_client_id(request)
         thread_id = str(uuid.uuid4())
-        thread = _Thread(thread_id, config=config, trace_context=self._trace_ctx)
+        thread = _Thread(thread_id, client_id=client_id, config=config, trace_context=self._trace_ctx)
         thread_lock = threading.Lock()
         with self._threads_lock:
             self._threads[thread_id] = (thread, thread_lock)
@@ -428,7 +450,8 @@ async def optimize_prompt_standalone(self, request: web.Request) -> web.Response
             # Create a temporary agent instance for optimization
             # Use default model from settings
             apis = HTTPPluginAPI("http://127.0.0.1:8786")
-            client = apis.attach_dev_client()  # TODO: should be removed or made configurable
+            client_id = self._resolve_client_id(request)
+            client = apis.with_client(client_id or DEV_CLIENT_ID)
             temp_agent = LeaderDriven(apis=client)
             # Optimize the prompt (without chat history context)
             optimized = await temp_agent.optimize_prompt(prompt_data.prompt)

plugin/src/main/java/copilotj/Connection.java

@@ -42,13 +42,15 @@ public interface ConnectionStateListener {
   private Timer timer;
 
   private final String server;
+  private final String accessToken;
 
   public Connection(final String serverURL, final EventHandler handler, final LogService log,
-      final int maxRetryWaitSecond) {
+      final int maxRetryWaitSecond, final String accessToken) {
     this.server = serverURL;
     this.handler = handler;
     this.log = log;
     this.maxRetryWaitSecond = maxRetryWaitSecond;
+    this.accessToken = accessToken;
     Connection.notifyStateChange(this, State.DISCONNECTED, "Initialized");
   }
 
@@ -109,9 +111,17 @@ public void onOpen(final ServerHandshake handshakedata) {
           // send a ping to check if the connection is alive
           webSocketClient.sendPing();
 
-          // Notify the handler of the new connection
+          // Send ID negotiation first so the server can assign/display a client ID
           final String event = handler.newConnectedEvent();
           this.send(event);
+
+          // Then send auth event if access token is configured
+          if (accessToken != null && !accessToken.isEmpty()) {
+            final String authEvent = handler.newAuthEvent(accessToken);
+            if (authEvent != null) {
+              this.send(authEvent);
+            }
+          }
         }
 
         @Override

plugin/src/main/java/copilotj/CopilotJBridgeDialog.java

@@ -15,6 +15,7 @@
 import javax.swing.JFrame;
 import javax.swing.JLabel;
 import javax.swing.JPanel;
+import javax.swing.JPasswordField;
 import javax.swing.JTextField;
 import javax.swing.SwingUtilities;
 
@@ -55,10 +56,14 @@ public void run() {
     final JFrame frame = new JFrame("CopilotJ Bridge Config");
     opened = frame;
 
-    frame.setSize(400, 200);
+    frame.setSize(400, 230);
     frame.setLayout(new BorderLayout());
 
     final JTextField urlField = new JTextField(service.getServerUrl());
+    final JPasswordField tokenField = new JPasswordField();
+    if (service.getAccessToken() != null) {
+      tokenField.setText(service.getAccessToken());
+    }
     final Connection existingConn = service.getConnection();
     final boolean isActive = existingConn != null &&
         existingConn.getState() != Connection.State.DISCONNECTED;
@@ -75,8 +80,9 @@ public void run() {
       }
 
       final String newUrl = urlField.getText();
+      final String token = new String(tokenField.getPassword());
       logService.info("Attempting to connect to: " + newUrl);
-      service.start(newUrl); // This will create a new connection or restart the existing one
+      service.start(newUrl, token.isEmpty() ? null : token); // This will create a new connection or restart the existing one
       connectButton.setText("Disconnect");
 
       // Register listeners when new connection is created
@@ -101,13 +107,21 @@ public void run() {
     inputPanel.add(urlField, BorderLayout.CENTER);
     inputPanel.add(connectButton, BorderLayout.EAST);
 
+    final JPanel tokenPanel = new JPanel(new BorderLayout(5, 5)); // Panel for access token
+    tokenPanel.add(new JLabel("Access Token:"), BorderLayout.WEST);
+    tokenPanel.add(tokenField, BorderLayout.CENTER);
+
+    final JPanel fieldsPanel = new JPanel(new BorderLayout(5, 5));
+    fieldsPanel.add(inputPanel, BorderLayout.NORTH);
+    fieldsPanel.add(tokenPanel, BorderLayout.CENTER);
+
     final JPanel statusPanel = new JPanel(new BorderLayout(5, 5)); // Panel for status and ID
     statusPanel.add(statusLabel, BorderLayout.NORTH);
     statusPanel.add(idLabel, BorderLayout.SOUTH);
 
     final JPanel mainPanel = new JPanel(new BorderLayout(10, 10)); // Main content panel
     mainPanel.setBorder(javax.swing.BorderFactory.createEmptyBorder(10, 10, 10, 10)); // Add padding
-    mainPanel.add(inputPanel, BorderLayout.NORTH);
+    mainPanel.add(fieldsPanel, BorderLayout.NORTH);
     mainPanel.add(statusPanel, BorderLayout.CENTER);
 
     // Register ID listener

plugin/src/main/java/copilotj/CopilotJBridgeService.java

@@ -14,7 +14,9 @@ public interface CopilotJBridgeService extends SciJavaService {
 
   public String getServerUrl();
 
+  public String getAccessToken();
+
   public Connection getConnection();
 
-  public void start(final String serverURL);
+  public void start(final String serverURL, final String accessToken);
 }

plugin/src/main/java/copilotj/DefaultCopilotJBridgeService.java

@@ -57,6 +57,7 @@ public static void main(final String... args) throws Exception {
   private LogService log;
 
   private final String serverUrl = "http://127.0.0.1:8786";
+  private String accessToken = null;
   private EventHandler eventHandler;
   private Connection connection;
 
@@ -73,6 +74,11 @@ public String getServerUrl() {
     return serverUrl;
   }
 
+  @Override
+  public String getAccessToken() {
+    return accessToken;
+  }
+
   @Override
   public Connection getConnection() {
     return connection;
@@ -85,9 +91,14 @@ public void initialize() {
     final boolean debug = Boolean.getBoolean("ij.debug");
     this.eventHandler = new EventHandler(context, log, debug);
 
+    final String token = System.getProperty("copilotj.accessToken");
+    if (token != null && !token.isEmpty()) {
+      this.accessToken = token;
+    }
+
     if (debug) {
       log.info("Automatically starting connection to " + serverUrl);
-      this.start(serverUrl);
+      this.start(serverUrl, accessToken);
     }
 
     if (ui.isVisible()) {
@@ -98,15 +109,16 @@ public void initialize() {
   }
 
   @Override
-  public void start(final String serverURL) {
+  public void start(final String serverURL, final String accessToken) {
     if (this.connection != null) {
       this.connection.close();
       this.connection = null;
     }
 
     log.info("copilotj: Start connection " + serverURL);
     final int maxRetryWaitSecond = Integer.getInteger("copilotj.maxRetryWaitSecond", 64);
-    this.connection = new Connection(serverURL, eventHandler, log, maxRetryWaitSecond);
+    this.accessToken = accessToken;
+    this.connection = new Connection(serverURL, eventHandler, log, maxRetryWaitSecond, accessToken);
     this.connection.connect();
   }
 

plugin/src/main/java/copilotj/EventHandler.java

@@ -122,6 +122,18 @@ public void removeListener(final IdListener listener) {
     }
   }
 
+  public String newAuthEvent(final String accessToken) {
+    final ObjectNode data = objectMapper.createObjectNode();
+    data.put("access_token", accessToken);
+    final Payload payload = new Payload("auth", data);
+    try {
+      return objectMapper.writeValueAsString(payload);
+    } catch (final IOException e) {
+      log.error("Error serializing auth event: " + e.getMessage());
+      return null;
+    }
+  }
+
   public String newConnectedEvent() {
     final Payload payload = this.id != null && !this.id.isEmpty()
         ? new Payload("negotiate_id", objectMapper.valueToTree(new IdChanged(this.id)))
@@ -214,6 +226,11 @@ private JsonNode handle(final Payload payload) {
         return null;
       }
 
+      case "auth_ok": {
+        log.info("Authentication successful");
+        return null;
+      }
+
       case "capture_screen": {
         final ScreenCapturer.ScreenPreviews image = screenCapturer.captureScreen();
         return objectMapper.valueToTree(image);

web/src/apis/clients.ts

@@ -0,0 +1,20 @@
+/**
+ * SPDX-FileCopyrightText: Copyright contributors to the CopilotJ project.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { baseUrl } from "./base.ts";
+
+export interface ClientInfo {
+  id: string;
+}
+
+export async function listClients(): Promise<ClientInfo[]> {
+  const url = `${baseUrl}/clients`;
+  const response = await fetch(url);
+  if (!response.ok) {
+    throw new Error(`HTTP error! status: ${response.status}, message: ${await response.text()}`);
+  }
+  return response.json();
+}