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);