Skip to content

Commit 6eb5bed

Browse files
authored
fix(ai): send requests without auth relock (#65)
Reuse the initialized auth session for AI dispatch, preserve transport error mapping, and keep Sparring dedupe stable across token refreshes.
1 parent 9e1e016 commit 6eb5bed

15 files changed

Lines changed: 226 additions & 20 deletions

‎src/lib/ai/client.ts‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import {
2020
type SourceProvenance,
2121
type SourceSemanticMetadata,
2222
} from '../stores/session-store';
23-
import { supabase } from '@/lib/supabase/client';
23+
import { getActiveAccessToken, getActiveUserId, invokeAuthenticatedFunction, supabase } from '@/lib/supabase/client';
2424
import { toast } from '@/hooks/use-toast';
2525
import { captureException, captureProductEvent } from '@/lib/observability';
2626
import { validateUrl } from '@/lib/validation/url';
@@ -469,8 +469,9 @@ async function callSupabaseAI<T = unknown>(
469469
else options.signal?.addEventListener('abort', abortFromCaller, { once: true });
470470

471471
try {
472-
const { data: sessionData } = await waitForAbortable(() => supabase.auth.getSession(), controller.signal);
473-
const accessToken = sessionData.session?.access_token;
472+
const accessToken = getActiveAccessToken() ?? (
473+
await waitForAbortable(() => supabase.auth.getSession(), controller.signal)
474+
).data.session?.access_token;
474475

475476
if (DEBUG) {
476477
console.log(`[AI Request ${requestId}]`, {
@@ -498,7 +499,7 @@ async function callSupabaseAI<T = unknown>(
498499
workspaceContext: options.workspaceContext,
499500
};
500501

501-
const { data, error } = await supabase.functions.invoke('ai', {
502+
const { data, error } = await invokeAuthenticatedFunction<Record<string, unknown>>('ai', {
502503
headers: { Authorization: `Bearer ${accessToken}` },
503504
body: requestBody,
504505
signal: controller.signal,
@@ -1360,13 +1361,12 @@ function validateSparringScoreResponse(output: unknown): SparringScoreResponse {
13601361
const pendingSparringTurns = new Map<string, Promise<SparringTurnResponse>>();
13611362
const pendingSparringScores = new Map<string, Promise<SparringScoreResponse>>();
13621363

1363-
async function sparringRequestKey(
1364+
function sparringRequestKey(
13641365
input: SparringRespondInput | SparringScoreInput,
13651366
config: ProviderConfig,
13661367
workspaceContext?: WorkspaceToolContext,
1367-
): Promise<string> {
1368-
const { data } = await supabase.auth.getSession();
1369-
const authScope = data.session?.user?.id ?? data.session?.access_token ?? 'signed-out';
1368+
): string {
1369+
const authScope = getActiveUserId() ?? getActiveAccessToken() ?? 'signed-out';
13701370
return JSON.stringify([authScope, workspaceContext?.workspaceId ?? null, input, config.provider, config.model, config.temperature, config.reasoningEffort]);
13711371
}
13721372

@@ -1386,7 +1386,7 @@ export async function generateSparringTurn(
13861386
details: { issues: inputResult.success ? [{ message: 'Expected respond action.' }] : inputResult.error.issues },
13871387
});
13881388
}
1389-
const requestKey = await sparringRequestKey(inputResult.data, config, workspaceContext);
1389+
const requestKey = sparringRequestKey(inputResult.data, config, workspaceContext);
13901390
const pending = pendingSparringTurns.get(requestKey);
13911391
if (pending) return pending;
13921392

@@ -1437,7 +1437,7 @@ export async function judgeSparringSession(
14371437
details: { issues: inputResult.success ? [{ message: 'Expected score action.' }] : inputResult.error.issues },
14381438
});
14391439
}
1440-
const requestKey = await sparringRequestKey(inputResult.data, config, workspaceContext);
1440+
const requestKey = sparringRequestKey(inputResult.data, config, workspaceContext);
14411441
const pending = pendingSparringScores.get(requestKey);
14421442
if (pending) return pending;
14431443

‎src/lib/ai/streaming.ts‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { supabase, SUPABASE_URL } from '@/lib/supabase/client';
1+
import { getActiveAccessToken, supabase, SUPABASE_URL } from '@/lib/supabase/client';
22
import { captureProductEvent } from '@/lib/observability';
33
import { AIAuthError, AIResponseError, REQUEST_TIMEOUT_MS, TimeoutError, mapAIError, reportAIFailure, withOutputLanguage } from './client';
44
import type { ProviderConfig } from './providers/types';
@@ -130,14 +130,16 @@ export async function streamAskChat(
130130
throw error;
131131
}
132132

133-
let sessionData: Awaited<ReturnType<typeof supabase.auth.getSession>>['data'];
133+
let accessToken = getActiveAccessToken();
134134
try {
135-
({ data: sessionData } = await waitForCaller(supabase.auth.getSession(), handlers.signal));
135+
if (!accessToken) {
136+
const { data } = await waitForCaller(supabase.auth.getSession(), handlers.signal);
137+
accessToken = data.session?.access_token ?? null;
138+
}
136139
} catch (error) {
137140
captureStreamFailure(error, provider, config.model, startTime, requestId, Boolean(handlers.signal?.aborted));
138141
throw error;
139142
}
140-
const accessToken = sessionData.session?.access_token;
141143
if (!accessToken) {
142144
const error = new AIAuthError('Please sign in to use AI features', { requestId });
143145
captureStreamFailure(error, provider, config.model, startTime, requestId, Boolean(handlers.signal?.aborted));

‎src/lib/supabase/AuthProvider.tsx‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
/* eslint-disable react-refresh/only-export-components */
22
import { createContext, useContext, useEffect, useMemo, useRef, useState } from 'react';
33
import type { Session, User } from '@supabase/supabase-js';
4-
import { supabase } from './client';
4+
import { setActiveAccessToken, supabase } from './client';
55
import { useSessionStore } from '@/lib/stores/session-store';
66
import { useSettingsStore } from '@/lib/stores/settings-store';
77
import { useSparringStore } from '@/lib/stores/sparring-store';
@@ -163,6 +163,7 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
163163
if (!mounted) return;
164164

165165
if (oauthSession) {
166+
setActiveAccessToken(oauthSession.access_token, oauthSession.user.id);
166167
setSession(oauthSession);
167168
currentUserIdRef.current = oauthSession.user?.id;
168169
setJamUserId(oauthSession.user?.id);
@@ -183,6 +184,7 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
183184
if (!mounted) return;
184185
const { data, error } = getSessionResult;
185186
if (error) console.warn('[Auth] getSession error:', error.message);
187+
setActiveAccessToken(data.session?.access_token ?? null, data.session?.user.id ?? null);
186188
setSession(data.session ?? null);
187189
currentUserIdRef.current = data.session?.user?.id;
188190
setJamUserId(data.session?.user?.id);
@@ -192,6 +194,7 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
192194
} catch (err) {
193195
console.warn('[Auth] init failed:', err);
194196
if (!mounted) return;
197+
setActiveAccessToken(null);
195198
setSession(null);
196199
currentUserIdRef.current = undefined;
197200
setJamUserId(undefined);
@@ -210,6 +213,7 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
210213
clearCrossfireDraftOnAuthChange(prevUserId, nextUserId);
211214
currentUserIdRef.current = nextUserId;
212215
setJamUserId(nextUserId);
216+
setActiveAccessToken(newSession?.access_token ?? null, newSession?.user.id ?? null);
213217
setSession(newSession);
214218
setLoading(false);
215219

‎src/lib/supabase/client.ts‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
1-
import { createClient } from '@supabase/supabase-js';
1+
import {
2+
createClient,
3+
FunctionsFetchError,
4+
FunctionsHttpError,
5+
FunctionsRelayError,
6+
} from '@supabase/supabase-js';
27
import { getActiveEnv, SUPABASE_URL, SUPABASE_ANON_KEY } from './env';
38

49
export { SUPABASE_URL };
@@ -13,3 +18,53 @@ export const supabase = createClient(SUPABASE_URL, SUPABASE_ANON_KEY, {
1318
storageKey: `debate-auth-${getActiveEnv()}`,
1419
},
1520
});
21+
22+
let activeAccessToken: string | null = null;
23+
let activeUserId: string | null = null;
24+
25+
export function setActiveAccessToken(accessToken: string | null, userId: string | null = null): void {
26+
activeAccessToken = accessToken;
27+
activeUserId = userId;
28+
}
29+
30+
export function getActiveAccessToken(): string | null {
31+
return activeAccessToken;
32+
}
33+
34+
export function getActiveUserId(): string | null {
35+
return activeUserId;
36+
}
37+
38+
type AuthenticatedFunctionResult<T> =
39+
| { data: T; error: null }
40+
| { data: null; error: FunctionsFetchError | FunctionsHttpError | FunctionsRelayError };
41+
42+
export async function invokeAuthenticatedFunction<T = unknown>(
43+
functionName: string,
44+
options: { body: Record<string, unknown>; headers: Record<string, string>; signal?: AbortSignal },
45+
): Promise<AuthenticatedFunctionResult<T>> {
46+
let response: Response;
47+
try {
48+
response = await fetch(`${SUPABASE_URL}/functions/v1/${functionName}`, {
49+
method: 'POST',
50+
headers: {
51+
apikey: SUPABASE_ANON_KEY,
52+
'Content-Type': 'application/json',
53+
...options.headers,
54+
},
55+
body: JSON.stringify(options.body),
56+
signal: options.signal,
57+
});
58+
} catch (error) {
59+
if (error instanceof Error && error.name === 'AbortError') throw error;
60+
return { data: null, error: new FunctionsFetchError(error) };
61+
}
62+
63+
if (response.headers.get('x-relay-error') === 'true') {
64+
return { data: null, error: new FunctionsRelayError(response) };
65+
}
66+
if (!response.ok) return { data: null, error: new FunctionsHttpError(response) };
67+
68+
const data: T = await response.json();
69+
return { data, error: null };
70+
}

‎src/test/ai/client-analytics.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ vi.mock('@/lib/observability', () => ({
2020
}));
2121

2222
vi.mock('@/lib/supabase/client', () => ({
23+
getActiveAccessToken: (): string | null => null,
24+
invokeAuthenticatedFunction: mocks.invoke,
2325
supabase: {
2426
auth: { getSession: mocks.getSession },
2527
functions: { invoke: mocks.invoke },

‎src/test/ai/client-branches.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ const {
3030
});
3131

3232
vi.mock('@/lib/supabase/client', () => ({
33+
getActiveAccessToken: (): string | null => null,
34+
invokeAuthenticatedFunction: invokeMock,
3335
supabase: {
3436
auth: {
3537
getSession: getSessionMock,

‎src/test/ai/client.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ const clientMocks = vi.hoisted(() => ({
3232
// actually it's better to use doMock or inline return.
3333

3434
vi.mock('@/lib/supabase/client', () => ({
35+
getActiveAccessToken: (): string | null => null,
36+
invokeAuthenticatedFunction: supabaseMock.functions.invoke,
3537
supabase: supabaseMock
3638
}));
3739

‎src/test/ai/control-payload-matrix.test.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,11 @@
66
import { beforeEach, describe, expect, it, vi } from 'vitest';
77
import { supabaseMock, resetMockState, setMockSessionToken } from '../mocks/supabase';
88

9-
vi.mock('@/lib/supabase/client', () => ({ supabase: supabaseMock }));
9+
vi.mock('@/lib/supabase/client', () => ({
10+
getActiveAccessToken: (): string | null => null,
11+
invokeAuthenticatedFunction: supabaseMock.functions.invoke,
12+
supabase: supabaseMock,
13+
}));
1014

1115
import { generateCase, attackCase, brainstorm, factoids, askChat } from '@/lib/ai/client';
1216
import type { ProviderConfig } from '@/lib/ai/providers/types';

‎src/test/ai/crossfire-client.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ const mocks = vi.hoisted(() => ({
1818
}));
1919

2020
vi.mock('@/lib/supabase/client', () => ({
21+
getActiveAccessToken: (): string | null => null,
22+
invokeAuthenticatedFunction: mocks.invoke,
2123
supabase: {
2224
auth: { getSession: mocks.getSession },
2325
functions: { invoke: mocks.invoke },
Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import { createElement } from 'react';
2+
import { render, screen, waitFor } from '@testing-library/react';
3+
import { afterEach, expect, it, vi } from 'vitest';
4+
5+
const sessionState = {
6+
sourceLibrary: [] as unknown[],
7+
hydrateFromServer: vi.fn(),
8+
resetLocalData: vi.fn(),
9+
};
10+
11+
const settingsState = {
12+
customSystemPrompt: '',
13+
outputLanguage: 'en',
14+
hydrateFromServer: vi.fn(),
15+
resetLocalData: vi.fn(),
16+
};
17+
18+
vi.mock('@/lib/stores/settings-store', () => ({
19+
useSettingsStore: Object.assign(
20+
(selector: (state: typeof settingsState) => unknown) => selector(settingsState),
21+
{ getState: () => settingsState },
22+
),
23+
}));
24+
25+
vi.mock('@/lib/stores/session-store', () => ({
26+
useSessionStore: Object.assign(
27+
(selector: (state: typeof sessionState) => unknown) => selector(sessionState),
28+
{ getState: () => sessionState },
29+
),
30+
}));
31+
32+
vi.mock('@/lib/stores/sparring-store', () => ({
33+
useSparringStore: (selector: (state: Record<string, unknown>) => unknown) => selector({
34+
hydrateSavedSessions: vi.fn(),
35+
prepareForUser: vi.fn(),
36+
resetLocalData: vi.fn(),
37+
}),
38+
}));
39+
40+
vi.mock('@/lib/observability', () => ({
41+
captureException: vi.fn(),
42+
captureProductEvent: vi.fn(),
43+
}));
44+
45+
import { askChat } from '@/lib/ai/client';
46+
import { AuthProvider, useAuth } from '@/lib/supabase/AuthProvider';
47+
import { setActiveAccessToken, supabase } from '@/lib/supabase/client';
48+
49+
afterEach(() => {
50+
setActiveAccessToken(null);
51+
vi.restoreAllMocks();
52+
});
53+
54+
function AuthState() {
55+
const { loading, user } = useAuth();
56+
return createElement('span', { 'data-testid': 'auth-state' }, loading ? 'loading' : user?.id ?? 'signed-out');
57+
}
58+
59+
it('calls the AI fetch boundary without another session lookup', async () => {
60+
const getSession = vi.spyOn(supabase.auth, 'getSession')
61+
.mockResolvedValue({
62+
data: { session: { access_token: 'test-access-token', user: { id: 'user-1' } } },
63+
error: null,
64+
} as Awaited<ReturnType<typeof supabase.auth.getSession>>);
65+
vi.spyOn(supabase.auth, 'onAuthStateChange').mockReturnValue({
66+
data: { subscription: { id: 'test', callback: vi.fn(), unsubscribe: vi.fn() } },
67+
} as ReturnType<typeof supabase.auth.onAuthStateChange>);
68+
const fetchRequest = vi.spyOn(globalThis, 'fetch').mockResolvedValue(new Response(JSON.stringify({
69+
output: {
70+
response: 'Mock AI response.',
71+
followUpQuestions: [],
72+
uncertainties: [],
73+
},
74+
}), {
75+
status: 200,
76+
headers: { 'Content-Type': 'application/json' },
77+
}));
78+
79+
render(createElement(AuthProvider, null, createElement(AuthState)));
80+
await waitFor(() => expect(screen.getByTestId('auth-state')).toHaveTextContent('user-1'));
81+
const authInitializationLookups = getSession.mock.calls.length;
82+
83+
const result = await askChat('Hello', [], 'general', {
84+
provider: 'openrouter',
85+
model: 'openai/gpt-5.6-luna',
86+
temperature: 0.7,
87+
reasoningEffort: 'low',
88+
});
89+
90+
expect(fetchRequest).toHaveBeenCalledOnce();
91+
expect(result).toMatchObject({ response: 'Mock AI response.' });
92+
expect(getSession).toHaveBeenCalledTimes(authInitializationLookups);
93+
});
94+
95+
it('maps a rejected fetch through the Edge Function error contract', async () => {
96+
setActiveAccessToken('test-access-token');
97+
vi.spyOn(globalThis, 'fetch').mockRejectedValue(new TypeError('Failed to fetch'));
98+
99+
await expect(askChat('Hello', [], 'general', {
100+
provider: 'openrouter',
101+
model: 'openai/gpt-5.6-luna',
102+
temperature: 0.7,
103+
reasoningEffort: 'low',
104+
})).rejects.toMatchObject({
105+
name: 'AIError',
106+
message: 'AI service error: Failed to send a request to the Edge Function. Please try again.',
107+
statusCode: 500,
108+
});
109+
});

0 commit comments

Comments
 (0)