Skip to content

Commit 4f46438

Browse files
charlesviencharlesvien
authored
fix: Fix agent using stale oauth tokens (#579)
1 parent b212408 commit 4f46438

File tree

10 files changed

+370
-47
lines changed

10 files changed

+370
-47
lines changed

‎apps/twig/src/main/index.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import { setMainWindowGetter } from "./trpc/context.js";
3333
import { trpcRouter } from "./trpc/index.js";
3434

3535
import "./services/index.js";
36+
import type { AgentService } from "./services/agent/service.js";
3637
import type { AppLifecycleService } from "./services/app-lifecycle/service.js";
3738
import type { DeepLinkService } from "./services/deep-link/service.js";
3839
import type { ExternalAppsService } from "./services/external-apps/service.js";
@@ -257,6 +258,19 @@ function createWindow(): void {
257258
container.get<UIService>(MAIN_TOKENS.UIService).clearStorage();
258259
},
259260
},
261+
{
262+
label: "Mark all agent sessions for recreation",
263+
click: () => {
264+
const count = container
265+
.get<AgentService>(MAIN_TOKENS.AgentService)
266+
.markAllSessionsForRecreation();
267+
dialog.showMessageBox({
268+
type: "info",
269+
title: "Sessions Marked",
270+
message: `Marked ${count} session(s) for recreation.\n\nThey will be recreated on the next prompt.`,
271+
});
272+
},
273+
},
260274
],
261275
},
262276
],

‎apps/twig/src/main/services/agent/schemas.ts‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -117,10 +117,9 @@ export const reconnectSessionInput = z.object({
117117

118118
export type ReconnectSessionInput = z.infer<typeof reconnectSessionInput>;
119119

120-
// Token refresh input
121-
export const tokenRefreshInput = z.object({
122-
taskRunId: z.string(),
123-
newToken: z.string(),
120+
// Token update input - updates the global token for all agent operations
121+
export const tokenUpdateInput = z.object({
122+
token: z.string(),
124123
});
125124

126125
// Set model input

‎apps/twig/src/main/services/agent/service.ts‎

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,22 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
189189
});
190190
}
191191

192+
/**
193+
* Mark all sessions for recreation (developer tool for testing token refresh).
194+
* Sessions will be recreated before their next prompt.
195+
*/
196+
public markAllSessionsForRecreation(): number {
197+
let count = 0;
198+
for (const session of this.sessions.values()) {
199+
session.needsRecreation = true;
200+
count++;
201+
}
202+
log.info("Marked all sessions for recreation (dev tool)", {
203+
sessionCount: count,
204+
});
205+
return count;
206+
}
207+
192208
/**
193209
* Respond to a pending permission request from the UI.
194210
* This resolves the promise that the agent is waiting on.
@@ -450,7 +466,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
450466
const config = existing.config;
451467
const pendingContext = existing.pendingContext;
452468

453-
this.cleanupSession(taskRunId);
469+
await this.cleanupSession(taskRunId);
454470

455471
const newSession = await this.getOrCreateSession(config, true);
456472
if (!newSession) {
@@ -532,8 +548,12 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
532548
const session = this.sessions.get(sessionId);
533549
if (!session) return false;
534550

535-
this.cleanupSession(sessionId);
536-
return true;
551+
try {
552+
await this.cleanupSession(sessionId);
553+
return true;
554+
} catch (_err) {
555+
return false;
556+
}
537557
}
538558

539559
async cancelPrompt(
@@ -783,9 +803,25 @@ For git operations while detached:
783803
}
784804
}
785805

786-
private cleanupSession(taskRunId: string): void {
806+
private async cleanupSession(taskRunId: string): Promise<void> {
787807
const session = this.sessions.get(taskRunId);
788808
if (session) {
809+
// Cancel any ongoing operations
810+
try {
811+
if (!session.connection.signal.aborted) {
812+
await session.connection.cancel({ sessionId: taskRunId });
813+
}
814+
} catch {
815+
// Ignore cancel errors
816+
}
817+
818+
// Cleanup agent (closes streams, aborts subprocess)
819+
try {
820+
await session.agent.cleanup();
821+
} catch {
822+
// Ignore cleanup errors
823+
}
824+
789825
this.cleanupMockNodeEnvironment(session.mockNodeDir);
790826
this.sessions.delete(taskRunId);
791827
}

‎apps/twig/src/main/services/oauth/schemas.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,13 @@ export type CloudRegion = z.infer<typeof cloudRegion>;
66
/**
77
* Error codes for OAuth operations.
88
* - network_error: Transient network issue, should retry
9+
* - server_error: Server error (5xx), should retry
910
* - auth_error: Authentication failed (invalid token, 401/403), should logout
1011
* - unknown_error: Other errors
1112
*/
1213
export const oAuthErrorCode = z.enum([
1314
"network_error",
15+
"server_error",
1416
"auth_error",
1517
"unknown_error",
1618
]);

‎apps/twig/src/main/services/oauth/service.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,19 @@ export class OAuthService {
180180
if (!response.ok) {
181181
// 401/403 are auth errors - the token is invalid
182182
const isAuthError = response.status === 401 || response.status === 403;
183+
// 5xx are server errors - should be retried
184+
const isServerError = response.status >= 500;
185+
log.warn(
186+
`Token refresh failed: ${response.status} ${response.statusText}`,
187+
);
183188
return {
184189
success: false,
185-
error: `Token refresh failed: ${response.statusText}`,
186-
errorCode: isAuthError ? "auth_error" : "unknown_error",
190+
error: `Token refresh failed: ${response.status} ${response.statusText}`,
191+
errorCode: isAuthError
192+
? "auth_error"
193+
: isServerError
194+
? "server_error"
195+
: "unknown_error",
187196
};
188197
}
189198

‎apps/twig/src/main/trpc/routers/agent.ts‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import {
1818
setModelInput,
1919
startSessionInput,
2020
subscribeSessionInput,
21-
tokenRefreshInput,
21+
tokenUpdateInput,
2222
} from "../../services/agent/schemas.js";
2323
import type { AgentService } from "../../services/agent/service.js";
2424
import { publicProcedure, router } from "../trpc.js";
@@ -53,11 +53,9 @@ export const agentRouter = router({
5353
.output(sessionResponseSchema.nullable())
5454
.mutation(({ input }) => getService().reconnectSession(input)),
5555

56-
refreshToken: publicProcedure
57-
.input(tokenRefreshInput)
58-
.mutation(({ input }) => {
59-
getService().updateToken(input.newToken);
60-
}),
56+
updateToken: publicProcedure.input(tokenUpdateInput).mutation(({ input }) => {
57+
getService().updateToken(input.token);
58+
}),
6159

6260
setModel: publicProcedure
6361
.input(setModelInput)
@@ -138,4 +136,8 @@ export const agentRouter = router({
138136
.mutation(({ input }) =>
139137
getService().notifySessionContext(input.sessionId, input.context),
140138
),
139+
140+
markAllForRecreation: publicProcedure.mutation(() =>
141+
getService().markAllSessionsForRecreation(),
142+
),
141143
});

‎apps/twig/src/renderer/features/auth/stores/authStore.ts‎

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ interface AuthState {
3838
tokenExpiry: number | null; // Unix timestamp in milliseconds
3939
cloudRegion: CloudRegion | null;
4040
storedTokens: StoredTokens | null;
41+
staleTokens: StoredTokens | null;
4142

4243
// PostHog client
4344
isAuthenticated: boolean;
@@ -66,6 +67,7 @@ export const useAuthStore = create<AuthState>()(
6667
tokenExpiry: null,
6768
cloudRegion: null,
6869
storedTokens: null,
70+
staleTokens: null,
6971

7072
// PostHog client
7173
isAuthenticated: false,
@@ -126,6 +128,10 @@ export const useAuthStore = create<AuthState>()(
126128
projectId,
127129
});
128130

131+
trpcVanilla.agent.updateToken
132+
.mutate({ token: tokenResponse.access_token })
133+
.catch((err) => log.warn("Failed to update agent token", err));
134+
129135
// Clear any cached data from previous sessions AFTER setting new auth
130136
queryClient.clear();
131137

@@ -191,15 +197,24 @@ export const useAuthStore = create<AuthState>()(
191197
});
192198

193199
if (!result.success || !result.data) {
194-
// Network errors should retry, auth errors should logout immediately
195-
if (result.errorCode === "network_error") {
200+
// Network/server errors should retry, auth errors should logout immediately
201+
if (
202+
result.errorCode === "network_error" ||
203+
result.errorCode === "server_error"
204+
) {
196205
log.warn(
197-
`Token refresh network error (attempt ${attempt + 1}): ${result.error}`,
206+
`Token refresh ${result.errorCode} (attempt ${attempt + 1}/${REFRESH_MAX_RETRIES}): ${result.error}`,
207+
);
208+
lastError = new Error(
209+
result.error || "Token refresh failed",
198210
);
199211
continue; // Retry
200212
}
201213

202214
// Auth error or unknown - logout
215+
log.error(
216+
`Token refresh failed with ${result.errorCode}: ${result.error}`,
217+
);
203218
get().logout();
204219
throw new Error(result.error || "Token refresh failed");
205220
}
@@ -244,6 +259,12 @@ export const useAuthStore = create<AuthState>()(
244259
...(projectId && { projectId }),
245260
});
246261

262+
trpcVanilla.agent.updateToken
263+
.mutate({ token: tokenResponse.access_token })
264+
.catch((err) =>
265+
log.warn("Failed to update agent token", err),
266+
);
267+
247268
get().scheduleTokenRefresh();
248269
return; // Success
249270
} catch (error) {
@@ -263,7 +284,9 @@ export const useAuthStore = create<AuthState>()(
263284
}
264285

265286
// All retries exhausted
266-
log.error("Token refresh failed after all retries");
287+
log.error(
288+
`Token refresh failed after all retries: ${lastError?.message || "Unknown error"}`,
289+
);
267290
get().logout();
268291
throw lastError || new Error("Token refresh failed");
269292
};
@@ -384,6 +407,12 @@ export const useAuthStore = create<AuthState>()(
384407
projectId,
385408
});
386409

410+
trpcVanilla.agent.updateToken
411+
.mutate({ token: currentTokens.accessToken })
412+
.catch((err) =>
413+
log.warn("Failed to update agent token", err),
414+
);
415+
387416
get().scheduleTokenRefresh();
388417

389418
// Use distinct_id to match web sessions (same as PostHog web app)
@@ -457,12 +486,15 @@ export const useAuthStore = create<AuthState>()(
457486

458487
useNavigationStore.getState().navigateToTaskInput();
459488

489+
const currentTokens = get().storedTokens;
490+
460491
set({
461492
oauthAccessToken: null,
462493
oauthRefreshToken: null,
463494
tokenExpiry: null,
464495
cloudRegion: null,
465496
storedTokens: null,
497+
staleTokens: currentTokens,
466498
isAuthenticated: false,
467499
client: null,
468500
projectId: null,
@@ -475,6 +507,7 @@ export const useAuthStore = create<AuthState>()(
475507
partialize: (state) => ({
476508
cloudRegion: state.cloudRegion,
477509
storedTokens: state.storedTokens,
510+
staleTokens: state.staleTokens,
478511
projectId: state.projectId,
479512
}),
480513
},

‎apps/twig/src/renderer/features/sessions/stores/sessionStore.ts‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1360,30 +1360,3 @@ export const useCurrentModeForTask = (
13601360
return session?.currentMode;
13611361
});
13621362
};
1363-
1364-
// Token refresh subscription
1365-
let lastKnownToken: string | null = null;
1366-
useAuthStore.subscribe(
1367-
(state) => state.oauthAccessToken,
1368-
(newToken) => {
1369-
if (!newToken || newToken === lastKnownToken) return;
1370-
lastKnownToken = newToken;
1371-
1372-
const sessions = useStore.getState().sessions;
1373-
for (const session of Object.values(sessions)) {
1374-
if (session.status === "connected" && !session.isCloud) {
1375-
trpcVanilla.agent.refreshToken
1376-
.mutate({
1377-
taskRunId: session.taskRunId,
1378-
newToken,
1379-
})
1380-
.catch((err) => {
1381-
log.warn("Failed to update session token", {
1382-
taskRunId: session.taskRunId,
1383-
error: err,
1384-
});
1385-
});
1386-
}
1387-
}
1388-
},
1389-
);

0 commit comments

Comments
 (0)