From 2038ea595716c65256ff3538c0e4381f7e230614 Mon Sep 17 00:00:00 2001 From: Jerome Date: Thu, 12 Jun 2025 13:38:20 +0100 Subject: [PATCH] Implement auth token persistence and MCP connection improvements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add OpenRouter API key persistence to localStorage with auto-restore - Add automatic provider restoration on app startup for seamless login - Fix MCP OAuth token persistence by preserving connection IDs across refreshes - Add manual reconnect button for failed/disconnected MCP servers - Add auto-reconnect for MCP servers on app startup - Maintain backward compatibility with old connection storage format Resolves auth token loss on page refresh for both inference and MCP providers. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- src/components/MCPStatus.tsx | 12 ++++++- src/contexts/InferenceContext.tsx | 25 ++++++++++++++- src/contexts/MCPContext.tsx | 40 +++++++++++++++++++----- src/providers/openrouter/api-provider.ts | 29 +++++++++++++++++ 4 files changed, 96 insertions(+), 10 deletions(-) diff --git a/src/components/MCPStatus.tsx b/src/components/MCPStatus.tsx index 819d34b..6c2bc54 100644 --- a/src/components/MCPStatus.tsx +++ b/src/components/MCPStatus.tsx @@ -4,7 +4,7 @@ import React from 'react'; import { useMCP } from '@/contexts/MCPContext'; export function MCPStatus() { - const { connections, getAllTools } = useMCP(); + const { connections, getAllTools, reconnectServer } = useMCP(); const connectedServers = connections.filter(conn => conn.status === 'connected'); const allTools = getAllTools(); @@ -89,6 +89,16 @@ export function MCPStatus() {

)} + + {(connection.status === 'failed' || connection.status === 'disconnected') && ( + + )} {connection.tools.length > 0 && ( diff --git a/src/contexts/InferenceContext.tsx b/src/contexts/InferenceContext.tsx index c7246c8..04cf371 100644 --- a/src/contexts/InferenceContext.tsx +++ b/src/contexts/InferenceContext.tsx @@ -1,12 +1,14 @@ // React context for inference provider management -import React, { createContext, useContext, useState, useCallback, ReactNode } from 'react'; +import React, { createContext, useContext, useState, useCallback, useEffect, ReactNode } from 'react'; import type { InferenceProvider, InferenceRequest, InferenceResponse, Model, } from '@/types/inference'; +import { OpenRouterApiProvider } from '@/providers/openrouter/api-provider'; +import { OpenRouterOAuthProvider } from '@/providers/openrouter/oauth-provider'; interface InferenceContextValue { // Current provider state @@ -64,6 +66,27 @@ export function InferenceProvider({ children }: InferenceProviderProps) { setAuthStateVersion(prev => prev + 1); }, []); + // Auto-restore provider with stored credentials on mount + useEffect(() => { + const tryRestoreProvider = async () => { + // Try API provider first + const apiProvider = new OpenRouterApiProvider(); + if (apiProvider.isAuthenticated) { + setProvider(apiProvider); + return; + } + + // Try OAuth provider + const oauthProvider = new OpenRouterOAuthProvider(); + if (oauthProvider.isAuthenticated) { + setProvider(oauthProvider); + return; + } + }; + + tryRestoreProvider().catch(console.error); + }, [setProvider]); + const generateResponse = useCallback(async (request: InferenceRequest): Promise => { if (!provider) { throw new Error('No inference provider configured'); diff --git a/src/contexts/MCPContext.tsx b/src/contexts/MCPContext.tsx index 7794688..3f6b47b 100644 --- a/src/contexts/MCPContext.tsx +++ b/src/contexts/MCPContext.tsx @@ -39,22 +39,42 @@ export function MCPProvider({ children }: MCPProviderProps) { try { const persistedData = localStorage.getItem('mcp_connections'); if (persistedData) { - const persistedConfigs: MCPServerConfig[] = JSON.parse(persistedData); + const persistedData_parsed = JSON.parse(persistedData); - // Restore connections directly without using addMcpServer to avoid loops - for (const config of persistedConfigs) { + // Handle both old format (array of configs) and new format (array of {id, config}) + const persistedConnections = Array.isArray(persistedData_parsed) && persistedData_parsed.length > 0 + ? (typeof persistedData_parsed[0] === 'object' && 'config' in persistedData_parsed[0] + ? persistedData_parsed as {id: string, config: MCPServerConfig}[] + : persistedData_parsed.map((config: MCPServerConfig) => ({id: uuidv4(), config}))) + : []; + + // Restore connections and auto-reconnect + for (const {id: connectionId, config} of persistedConnections) { try { - const connectionId = uuidv4(); const manager = new MCPConnectionManager(connectionId, config); + // Set up callback for connection state updates + manager.setConnectionUpdateCallback(() => { + setConnections(prev => + prev.map(conn => + conn.id === connectionId ? manager.getConnection() : conn + ) + ); + }); + // Add to managers map setManagers(prev => new Map(prev).set(connectionId, manager)); // Add initial connection state setConnections(prev => [...prev, manager.getConnection()]); - // Don't auto-connect during restoration - let user manually connect - console.log(`Restored connection config for ${config.name}`); + // Auto-connect on restoration + try { + await manager.connect(); + console.log(`Auto-reconnected to ${config.name}`); + } catch (error) { + console.warn(`Failed to auto-reconnect to ${config.name}:`, error); + } } catch (error) { console.warn(`Failed to restore connection to ${config.name}:`, error); } @@ -70,8 +90,12 @@ export function MCPProvider({ children }: MCPProviderProps) { const persistConnections = useCallback(() => { try { - const configs = connections.map(conn => conn.config); - localStorage.setItem('mcp_connections', JSON.stringify(configs)); + // Store both config and connection ID to maintain OAuth token association + const connectionData = connections.map(conn => ({ + id: conn.id, + config: conn.config + })); + localStorage.setItem('mcp_connections', JSON.stringify(connectionData)); } catch (error) { console.error('Failed to persist MCP connections:', error); } diff --git a/src/providers/openrouter/api-provider.ts b/src/providers/openrouter/api-provider.ts index ee444f8..34c98bf 100644 --- a/src/providers/openrouter/api-provider.ts +++ b/src/providers/openrouter/api-provider.ts @@ -31,6 +31,9 @@ export class OpenRouterApiProvider extends InferenceProvider { if (config?.defaultModel) { // We'll set this after loading models } + + // Load stored API key + this.loadStoredApiKey(); } get isAuthenticated(): boolean { @@ -60,6 +63,9 @@ export class OpenRouterApiProvider extends InferenceProvider { try { // Test the API key by loading models await this.loadModels(); + + // Store API key after successful validation + this.storeApiKey(); } catch (error) { this._authError = error instanceof Error ? error.message : 'Authentication failed'; this.apiKey = undefined; @@ -72,6 +78,7 @@ export class OpenRouterApiProvider extends InferenceProvider { this._authError = undefined; this._models = []; this._selectedModel = undefined; + this.clearStoredApiKey(); } async loadModels(): Promise { @@ -152,4 +159,26 @@ export class OpenRouterApiProvider extends InferenceProvider { private isInferenceError(error: any): error is InferenceError { return error && typeof error === 'object' && 'type' in error && 'message' in error; } + + private storeApiKey(): void { + if (this.apiKey) { + localStorage.setItem('openrouter_api_key', this.apiKey); + } + } + + private loadStoredApiKey(): void { + const storedKey = localStorage.getItem('openrouter_api_key'); + if (storedKey) { + this.apiKey = storedKey; + // Try to load models to verify the key is still valid + this.loadModels().catch(() => { + // If loading fails, clear invalid key + this.logout(); + }); + } + } + + private clearStoredApiKey(): void { + localStorage.removeItem('openrouter_api_key'); + } } \ No newline at end of file