Skip to content

Commit cc28a89

Browse files
cliffhallcliffhall
authored
Merge branch 'main' into 450-template-queries
2 parents 9508cc4 + 3772ff3 commit cc28a89

File tree

10 files changed

+500
-162
lines changed

10 files changed

+500
-162
lines changed

‎client/src/App.tsx‎

Lines changed: 65 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ import {
1919
} from "@modelcontextprotocol/sdk/types.js";
2020
import { OAuthTokensSchema } from "@modelcontextprotocol/sdk/shared/auth.js";
2121
import { SESSION_KEYS, getServerSpecificKey } from "./lib/constants";
22-
import { AuthDebuggerState } from "./lib/auth-types";
22+
import { AuthDebuggerState, EMPTY_DEBUGGER_STATE } from "./lib/auth-types";
23+
import { OAuthStateMachine } from "./lib/oauth-state-machine";
2324
import { cacheToolOutputSchemas } from "./utils/schemaUtils";
2425
import React, {
2526
Suspense,
@@ -121,19 +122,8 @@ const App = () => {
121122
const [isAuthDebuggerVisible, setIsAuthDebuggerVisible] = useState(false);
122123

123124
// Auth debugger state
124-
const [authState, setAuthState] = useState<AuthDebuggerState>({
125-
isInitiatingAuth: false,
126-
oauthTokens: null,
127-
loading: true,
128-
oauthStep: "metadata_discovery",
129-
oauthMetadata: null,
130-
oauthClientInfo: null,
131-
authorizationUrl: null,
132-
authorizationCode: "",
133-
latestError: null,
134-
statusMessage: null,
135-
validationError: null,
136-
});
125+
const [authState, setAuthState] =
126+
useState<AuthDebuggerState>(EMPTY_DEBUGGER_STATE);
137127

138128
// Helper function to update specific auth state properties
139129
const updateAuthState = (updates: Partial<AuthDebuggerState>) => {
@@ -243,27 +233,81 @@ const App = () => {
243233

244234
// Update OAuth debug state during debug callback
245235
const onOAuthDebugConnect = useCallback(
246-
({
236+
async ({
247237
authorizationCode,
248238
errorMsg,
239+
restoredState,
249240
}: {
250241
authorizationCode?: string;
251242
errorMsg?: string;
243+
restoredState?: AuthDebuggerState;
252244
}) => {
253245
setIsAuthDebuggerVisible(true);
254-
if (authorizationCode) {
246+
247+
if (errorMsg) {
255248
updateAuthState({
256-
authorizationCode,
257-
oauthStep: "token_request",
249+
latestError: new Error(errorMsg),
258250
});
251+
return;
259252
}
260-
if (errorMsg) {
253+
254+
if (restoredState && authorizationCode) {
255+
// Restore the previous auth state and continue the OAuth flow
256+
let currentState: AuthDebuggerState = {
257+
...restoredState,
258+
authorizationCode,
259+
oauthStep: "token_request",
260+
isInitiatingAuth: true,
261+
statusMessage: null,
262+
latestError: null,
263+
};
264+
265+
try {
266+
// Create a new state machine instance to continue the flow
267+
const stateMachine = new OAuthStateMachine(sseUrl, (updates) => {
268+
currentState = { ...currentState, ...updates };
269+
});
270+
271+
// Continue stepping through the OAuth flow from where we left off
272+
while (
273+
currentState.oauthStep !== "complete" &&
274+
currentState.oauthStep !== "authorization_code"
275+
) {
276+
await stateMachine.executeStep(currentState);
277+
}
278+
279+
if (currentState.oauthStep === "complete") {
280+
// After the flow completes or reaches a user-input step, update the app state
281+
updateAuthState({
282+
...currentState,
283+
statusMessage: {
284+
type: "success",
285+
message: "Authentication completed successfully",
286+
},
287+
isInitiatingAuth: false,
288+
});
289+
}
290+
} catch (error) {
291+
console.error("OAuth continuation error:", error);
292+
updateAuthState({
293+
latestError:
294+
error instanceof Error ? error : new Error(String(error)),
295+
statusMessage: {
296+
type: "error",
297+
message: `Failed to complete OAuth flow: ${error instanceof Error ? error.message : String(error)}`,
298+
},
299+
isInitiatingAuth: false,
300+
});
301+
}
302+
} else if (authorizationCode) {
303+
// Fallback to the original behavior if no state was restored
261304
updateAuthState({
262-
latestError: new Error(errorMsg),
305+
authorizationCode,
306+
oauthStep: "token_request",
263307
});
264308
}
265309
},
266-
[],
310+
[sseUrl],
267311
);
268312

269313
// Load OAuth tokens when sseUrl changes
@@ -285,8 +329,6 @@ const App = () => {
285329
}
286330
} catch (error) {
287331
console.error("Error loading OAuth tokens:", error);
288-
} finally {
289-
updateAuthState({ loading: false });
290332
}
291333
};
292334

‎client/src/components/AuthDebugger.tsx‎

Lines changed: 67 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1-
import { useCallback, useMemo } from "react";
1+
import { useCallback, useMemo, useEffect } from "react";
22
import { Button } from "@/components/ui/button";
33
import { DebugInspectorOAuthClientProvider } from "../lib/auth";
44
import { AlertCircle } from "lucide-react";
5-
import { AuthDebuggerState } from "../lib/auth-types";
5+
import { AuthDebuggerState, EMPTY_DEBUGGER_STATE } from "../lib/auth-types";
66
import { OAuthFlowProgress } from "./OAuthFlowProgress";
77
import { OAuthStateMachine } from "../lib/oauth-state-machine";
8+
import { SESSION_KEYS } from "../lib/constants";
89

910
export interface AuthDebuggerProps {
1011
serverUrl: string;
@@ -59,6 +60,27 @@ const AuthDebugger = ({
5960
authState,
6061
updateAuthState,
6162
}: AuthDebuggerProps) => {
63+
// Check for existing tokens on mount
64+
useEffect(() => {
65+
if (serverUrl && !authState.oauthTokens) {
66+
const checkTokens = async () => {
67+
try {
68+
const provider = new DebugInspectorOAuthClientProvider(serverUrl);
69+
const existingTokens = await provider.tokens();
70+
if (existingTokens) {
71+
updateAuthState({
72+
oauthTokens: existingTokens,
73+
oauthStep: "complete",
74+
});
75+
}
76+
} catch (error) {
77+
console.error("Failed to load existing OAuth tokens:", error);
78+
}
79+
};
80+
checkTokens();
81+
}
82+
}, [serverUrl, updateAuthState, authState.oauthTokens]);
83+
6284
const startOAuthFlow = useCallback(() => {
6385
if (!serverUrl) {
6486
updateAuthState({
@@ -141,6 +163,11 @@ const AuthDebugger = ({
141163
currentState.oauthStep === "authorization_code" &&
142164
currentState.authorizationUrl
143165
) {
166+
// Store the current auth state before redirecting
167+
sessionStorage.setItem(
168+
SESSION_KEYS.AUTH_DEBUGGER_STATE,
169+
JSON.stringify(currentState),
170+
);
144171
// Open the authorization URL automatically
145172
window.location.href = currentState.authorizationUrl;
146173
break;
@@ -178,13 +205,7 @@ const AuthDebugger = ({
178205
);
179206
serverAuthProvider.clear();
180207
updateAuthState({
181-
oauthTokens: null,
182-
oauthStep: "metadata_discovery",
183-
latestError: null,
184-
oauthClientInfo: null,
185-
authorizationCode: "",
186-
validationError: null,
187-
oauthMetadata: null,
208+
...EMPTY_DEBUGGER_STATE,
188209
statusMessage: {
189210
type: "success",
190211
message: "OAuth tokens cleared successfully",
@@ -224,52 +245,48 @@ const AuthDebugger = ({
224245
<StatusMessage message={authState.statusMessage} />
225246
)}
226247

227-
{authState.loading ? (
228-
<p>Loading authentication status...</p>
229-
) : (
230-
<div className="space-y-4">
231-
{authState.oauthTokens && (
232-
<div className="space-y-2">
233-
<p className="text-sm font-medium">Access Token:</p>
234-
<div className="bg-muted p-2 rounded-md text-xs overflow-x-auto">
235-
{authState.oauthTokens.access_token.substring(0, 25)}...
236-
</div>
248+
<div className="space-y-4">
249+
{authState.oauthTokens && (
250+
<div className="space-y-2">
251+
<p className="text-sm font-medium">Access Token:</p>
252+
<div className="bg-muted p-2 rounded-md text-xs overflow-x-auto">
253+
{authState.oauthTokens.access_token.substring(0, 25)}...
237254
</div>
238-
)}
239-
240-
<div className="flex gap-4">
241-
<Button
242-
variant="outline"
243-
onClick={startOAuthFlow}
244-
disabled={authState.isInitiatingAuth}
245-
>
246-
{authState.oauthTokens
247-
? "Guided Token Refresh"
248-
: "Guided OAuth Flow"}
249-
</Button>
255+
</div>
256+
)}
250257

251-
<Button
252-
onClick={handleQuickOAuth}
253-
disabled={authState.isInitiatingAuth}
254-
>
255-
{authState.isInitiatingAuth
256-
? "Initiating..."
257-
: authState.oauthTokens
258-
? "Quick Refresh"
259-
: "Quick OAuth Flow"}
260-
</Button>
258+
<div className="flex gap-4">
259+
<Button
260+
variant="outline"
261+
onClick={startOAuthFlow}
262+
disabled={authState.isInitiatingAuth}
263+
>
264+
{authState.oauthTokens
265+
? "Guided Token Refresh"
266+
: "Guided OAuth Flow"}
267+
</Button>
261268

262-
<Button variant="outline" onClick={handleClearOAuth}>
263-
Clear OAuth State
264-
</Button>
265-
</div>
269+
<Button
270+
onClick={handleQuickOAuth}
271+
disabled={authState.isInitiatingAuth}
272+
>
273+
{authState.isInitiatingAuth
274+
? "Initiating..."
275+
: authState.oauthTokens
276+
? "Quick Refresh"
277+
: "Quick OAuth Flow"}
278+
</Button>
266279

267-
<p className="text-xs text-muted-foreground">
268-
Choose "Guided" for step-by-step instructions or "Quick" for
269-
the standard automatic flow.
270-
</p>
280+
<Button variant="outline" onClick={handleClearOAuth}>
281+
Clear OAuth State
282+
</Button>
271283
</div>
272-
)}
284+
285+
<p className="text-xs text-muted-foreground">
286+
Choose "Guided" for step-by-step instructions or "Quick" for
287+
the standard automatic flow.
288+
</p>
289+
</div>
273290
</div>
274291

275292
<OAuthFlowProgress

‎client/src/components/OAuthDebugCallback.tsx‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,17 @@ import {
44
generateOAuthErrorDescription,
55
parseOAuthCallbackParams,
66
} from "@/utils/oauthUtils.ts";
7+
import { AuthDebuggerState } from "@/lib/auth-types";
78

89
interface OAuthCallbackProps {
910
onConnect: ({
1011
authorizationCode,
1112
errorMsg,
13+
restoredState,
1214
}: {
1315
authorizationCode?: string;
1416
errorMsg?: string;
17+
restoredState?: AuthDebuggerState;
1518
}) => void;
1619
}
1720

@@ -35,6 +38,21 @@ const OAuthDebugCallback = ({ onConnect }: OAuthCallbackProps) => {
3538

3639
const serverUrl = sessionStorage.getItem(SESSION_KEYS.SERVER_URL);
3740

41+
// Try to restore the auth state
42+
const storedState = sessionStorage.getItem(
43+
SESSION_KEYS.AUTH_DEBUGGER_STATE,
44+
);
45+
let restoredState = null;
46+
if (storedState) {
47+
try {
48+
restoredState = JSON.parse(storedState);
49+
// Clean up the stored state
50+
sessionStorage.removeItem(SESSION_KEYS.AUTH_DEBUGGER_STATE);
51+
} catch (e) {
52+
console.error("Failed to parse stored auth state:", e);
53+
}
54+
}
55+
3856
// ServerURL isn't set, this can happen if we've opened the
3957
// authentication request in a new tab, so we don't have the same
4058
// session storage
@@ -50,8 +68,8 @@ const OAuthDebugCallback = ({ onConnect }: OAuthCallbackProps) => {
5068
}
5169

5270
// Instead of storing in sessionStorage, pass the code directly
53-
// to the auth state manager through onConnect
54-
onConnect({ authorizationCode: params.code });
71+
// to the auth state manager through onConnect, along with restored state
72+
onConnect({ authorizationCode: params.code, restoredState });
5573
};
5674

5775
handleCallback().finally(() => {

0 commit comments

Comments
 (0)