Implement auth token persistence and MCP connection improvements

- 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 <noreply@anthropic.com>
This commit is contained in:
Jerome 2025-06-12 13:38:20 +01:00
commit 2038ea5957
No known key found for this signature in database
4 changed files with 96 additions and 10 deletions

View file

@ -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() {
</p>
)}
</div>
{(connection.status === 'failed' || connection.status === 'disconnected') && (
<button
onClick={() => reconnectServer(connection.id)}
className="ml-2 px-2 py-1 text-xs bg-blue-100 dark:bg-blue-900/50 text-blue-800 dark:text-blue-200 rounded hover:bg-blue-200 dark:hover:bg-blue-800"
title="Reconnect server"
>
Reconnect
</button>
)}
</div>
{connection.tools.length > 0 && (

View file

@ -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<InferenceResponse> => {
if (!provider) {
throw new Error('No inference provider configured');

View file

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

View file

@ -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<Model[]> {
@ -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');
}
}