From 1584ff101ef7888efabd169f48e65dd08b43eebd Mon Sep 17 00:00:00 2001 From: Zexin Yuan Date: Wed, 22 Apr 2026 23:24:23 +0800 Subject: [PATCH] Add per-bridge access token for multi-bridge support Bridges authenticate by sending an "auth" WebSocket message after ID negotiation. The server maps access tokens to client IDs so the web frontend can route requests to specific bridges. The Settings UI now has a Bridge tab with a client selector dropdown and access token input. --- copilotj/server/bridge.py | 30 ++++++++ copilotj/server/server.py | 12 +++- copilotj/server/threads.py | 37 ++++++++-- plugin/src/main/java/copilotj/Connection.java | 14 +++- .../java/copilotj/CopilotJBridgeDialog.java | 20 +++++- .../java/copilotj/CopilotJBridgeService.java | 4 +- .../DefaultCopilotJBridgeService.java | 18 ++++- .../src/main/java/copilotj/EventHandler.java | 17 +++++ web/src/apis/clients.ts | 20 ++++++ web/src/apis/index.ts | 1 + web/src/apis/threads.ts | 24 +++++-- web/src/components/Settings.vue | 70 ++++++++++++++++++- web/src/store/settings.ts | 25 ++++++- 13 files changed, 265 insertions(+), 27 deletions(-) create mode 100644 web/src/apis/clients.ts diff --git a/copilotj/server/bridge.py b/copilotj/server/bridge.py index f08e3d16..c28cf936 100644 --- a/copilotj/server/bridge.py +++ b/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 diff --git a/copilotj/server/server.py b/copilotj/server/server.py index 7aab8283..cbacfb96 100644 --- a/copilotj/server/server.py +++ b/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") diff --git a/copilotj/server/threads.py b/copilotj/server/threads.py index 46737e66..b9dc6287 100644 --- a/copilotj/server/threads.py +++ b/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) diff --git a/plugin/src/main/java/copilotj/Connection.java b/plugin/src/main/java/copilotj/Connection.java index ec2041fe..c15c1aeb 100644 --- a/plugin/src/main/java/copilotj/Connection.java +++ b/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 diff --git a/plugin/src/main/java/copilotj/CopilotJBridgeDialog.java b/plugin/src/main/java/copilotj/CopilotJBridgeDialog.java index 40eb3ec5..e3d636e9 100644 --- a/plugin/src/main/java/copilotj/CopilotJBridgeDialog.java +++ b/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 diff --git a/plugin/src/main/java/copilotj/CopilotJBridgeService.java b/plugin/src/main/java/copilotj/CopilotJBridgeService.java index 4822efa7..c8a05472 100644 --- a/plugin/src/main/java/copilotj/CopilotJBridgeService.java +++ b/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); } diff --git a/plugin/src/main/java/copilotj/DefaultCopilotJBridgeService.java b/plugin/src/main/java/copilotj/DefaultCopilotJBridgeService.java index f81848e9..87d0bb3f 100644 --- a/plugin/src/main/java/copilotj/DefaultCopilotJBridgeService.java +++ b/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,7 +109,7 @@ 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; @@ -106,7 +117,8 @@ public void start(final String serverURL) { 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(); } diff --git a/plugin/src/main/java/copilotj/EventHandler.java b/plugin/src/main/java/copilotj/EventHandler.java index 7838edde..0d7e892d 100644 --- a/plugin/src/main/java/copilotj/EventHandler.java +++ b/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); diff --git a/web/src/apis/clients.ts b/web/src/apis/clients.ts new file mode 100644 index 00000000..54311831 --- /dev/null +++ b/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 { + 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(); +} diff --git a/web/src/apis/index.ts b/web/src/apis/index.ts index a3157ba9..1bb2d79e 100644 --- a/web/src/apis/index.ts +++ b/web/src/apis/index.ts @@ -5,3 +5,4 @@ */ export * from "./threads.ts"; +export * from "./clients.ts"; diff --git a/web/src/apis/threads.ts b/web/src/apis/threads.ts index 10a93ef1..1b56ebfc 100644 --- a/web/src/apis/threads.ts +++ b/web/src/apis/threads.ts @@ -5,6 +5,7 @@ */ import { baseUrl } from "./base.ts"; +import { useSettings } from "../store"; interface Payload { type: T; @@ -93,11 +94,20 @@ export interface OptimizePromptResponse { optimized: string; } +function authHeaders(): Record { + const token = useSettings().accessToken; + const headers: Record = { "Content-Type": "application/json" }; + if (token) { + headers["Authorization"] = `Bearer ${token}`; + } + return headers; +} + export async function newThread(config: ThreadConfigQuery): Promise { const url = `${baseUrl}/threads`; const options = { method: "POST", - headers: { "Content-Type": "application/json" }, + headers: authHeaders(), body: JSON.stringify({ config }), }; const response = await fetch(url, options); @@ -115,7 +125,7 @@ export async function* newThreadPost( const url = `${baseUrl}/threads/${threadId}/posts`; const options: RequestInit = { method: "POST", - headers: { "Content-Type": "application/json" }, + headers: authHeaders(), body: JSON.stringify({ prompt }), signal, }; @@ -124,7 +134,7 @@ export async function* newThreadPost( export async function getThreadConfig(threadId: string): Promise { const url = `${baseUrl}/threads/${threadId}/config`; - const options = { method: "GET", headers: { "Content-Type": "application/json" } }; + const options = { method: "GET", headers: authHeaders() }; const response = await fetch(url, options); if (!response.ok) { @@ -137,7 +147,7 @@ export async function updateThreadConfig(threadId: string, config: ThreadConfigQ const url = `${baseUrl}/threads/${threadId}/config`; const options = { method: "POST", - headers: { "Content-Type": "application/json" }, + headers: authHeaders(), body: JSON.stringify(config), }; @@ -150,7 +160,7 @@ export async function updateThreadConfig(threadId: string, config: ThreadConfigQ export async function deleteThread(threadId: string): Promise { const url = `${baseUrl}/threads/${threadId}`; - const options = { method: "DELETE" }; + const options = { method: "DELETE", headers: authHeaders() }; const response = await fetch(url, options); if (!response.ok) { throw new Error(`HTTP error! status: ${response.status}, message: ${await response.text()}`); @@ -162,7 +172,7 @@ export async function optimizePrompt(threadId: string, prompt: string): Promise< const url = `${baseUrl}/threads/${threadId}/optimize-prompt`; const options = { method: "POST", - headers: { "Content-Type": "application/json" }, + headers: authHeaders(), body: JSON.stringify({ prompt }), }; const response = await fetch(url, options); @@ -176,7 +186,7 @@ export async function optimizePromptBeforeThread(prompt: string): Promise import { IconRefresh, IconUpload } from "@tabler/icons-vue"; -import { ref } from "vue"; -import type { ThreadConfigModel, ThreadConfigQuery } from "../apis"; +import { computed, onMounted, ref } from "vue"; +import type { ThreadConfigModel, ThreadConfigQuery, ClientInfo } from "../apis"; +import { listClients } from "../apis"; import { useSettings } from "../store"; import SettingModel from "./SettingModel.vue"; @@ -18,6 +19,8 @@ const emit = defineEmits<{ const settings = useSettings(); const fileInput = ref(null); +const accessTokenInput = ref(settings.accessToken); +const clients = ref([]); function handleFileUpload(event: Event) { const input = event.target as HTMLInputElement; @@ -30,12 +33,38 @@ function sumbmitModel(model: ThreadConfigModel | null) { settings.setModel(model); emit("submit", { model: settings.model }); } + +function submitAccessToken() { + settings.setAccessToken(accessTokenInput.value); +} + +async function refreshClients() { + try { + clients.value = await listClients(); + } catch { + clients.value = []; + } +} + +const clientOptions = computed(() => + clients.value.map((c) => ({ + label: c.id.substring(0, 8) + "...", + value: c.id, + })), +); + +function onClientSelected() { + settings.setSelectedClientId(settings.selectedClientId); +} + +onMounted(refreshClients);