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:
parent
d028846a14
commit
2038ea5957
4 changed files with 96 additions and 10 deletions
|
|
@ -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 && (
|
||||
|
|
|
|||
|
|
@ -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');
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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');
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue