Skip to content

Commit c38c5a6

Browse files
jonathanlabjonathanlab
authored
fix: handle tokens properly and recreate session on token expiry (#261)
1 parent 3dba9da commit c38c5a6

File tree

7 files changed

+121
-47
lines changed

7 files changed

+121
-47
lines changed

‎apps/array/src/main/services/session-manager.ts‎

Lines changed: 100 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,13 @@ import { logger } from "../lib/logger";
2323

2424
const log = logger.scope("session-manager");
2525

26+
function isAuthError(error: unknown): boolean {
27+
return (
28+
error instanceof Error &&
29+
error.message.startsWith("Authentication required")
30+
);
31+
}
32+
2633
type MessageCallback = (message: unknown) => void;
2734

2835
class NdJsonTap {
@@ -135,6 +142,7 @@ export interface ManagedSession {
135142
createdAt: number;
136143
lastActivityAt: number;
137144
mockNodeDir: string;
145+
config: SessionConfig;
138146
}
139147

140148
function getClaudeCliPath(): string {
@@ -146,7 +154,7 @@ function getClaudeCliPath(): string {
146154

147155
export class SessionManager {
148156
private sessions = new Map<string, ManagedSession>();
149-
private sessionTokens = new Map<string, string>();
157+
private currentToken: string | null = null;
150158
private getMainWindow: () => BrowserWindow | null;
151159
private onLog: OnLogCallback;
152160

@@ -155,23 +163,20 @@ export class SessionManager {
155163
this.onLog = onLog;
156164
}
157165

158-
public updateSessionToken(taskRunId: string, newToken: string): void {
159-
this.sessionTokens.set(taskRunId, newToken);
160-
log.info("Session token updated", { taskRunId });
166+
public updateToken(newToken: string): void {
167+
this.currentToken = newToken;
168+
log.info("Session token updated");
161169
}
162170

163-
private getSessionToken(taskRunId: string, fallback: string): string {
164-
return this.sessionTokens.get(taskRunId) || fallback;
171+
private getToken(fallback: string): string {
172+
return this.currentToken || fallback;
165173
}
166174

167-
private buildMcpServers(
168-
credentials: PostHogCredentials,
169-
taskRunId: string,
170-
): AcpMcpServer[] {
175+
private buildMcpServers(credentials: PostHogCredentials): AcpMcpServer[] {
171176
const servers: AcpMcpServer[] = [];
172177

173178
const mcpUrl = this.getPostHogMcpUrl(credentials.apiHost);
174-
const token = this.getSessionToken(taskRunId, credentials.apiKey);
179+
const token = this.getToken(credentials.apiKey);
175180

176181
servers.push({
177182
name: "posthog",
@@ -211,6 +216,7 @@ export class SessionManager {
211216
private async getOrCreateSession(
212217
config: SessionConfig,
213218
isReconnect: boolean,
219+
isRetry = false,
214220
): Promise<ManagedSession | null> {
215221
const {
216222
taskId,
@@ -222,9 +228,11 @@ export class SessionManager {
222228
model,
223229
} = config;
224230

225-
const existing = this.sessions.get(taskRunId);
226-
if (existing) {
227-
return existing;
231+
if (!isRetry) {
232+
const existing = this.sessions.get(taskRunId);
233+
if (existing) {
234+
return existing;
235+
}
228236
}
229237

230238
const channel = `agent-event:${taskRunId}`;
@@ -234,7 +242,7 @@ export class SessionManager {
234242
const agent = new Agent({
235243
workingDirectory: repoPath,
236244
posthogApiUrl: credentials.apiHost,
237-
posthogApiKey: credentials.apiKey,
245+
getPosthogApiKey: () => this.getToken(credentials.apiKey),
238246
posthogProjectId: credentials.projectId,
239247
debug: !app.isPackaged,
240248
onLog: this.onLog,
@@ -256,11 +264,9 @@ export class SessionManager {
256264
clientCapabilities: {},
257265
});
258266

259-
const mcpServers = this.buildMcpServers(credentials, taskRunId);
267+
const mcpServers = this.buildMcpServers(credentials);
260268

261269
if (isReconnect) {
262-
// Use our custom extension method to resume without replaying history.
263-
// Client fetches history from S3 directly.
264270
await connection.extMethod("_posthog/session/resume", {
265271
sessionId: taskRunId,
266272
cwd: repoPath,
@@ -290,38 +296,82 @@ export class SessionManager {
290296
createdAt: Date.now(),
291297
lastActivityAt: Date.now(),
292298
mockNodeDir,
299+
config,
293300
};
294301

295302
this.sessions.set(taskRunId, session);
303+
if (isRetry) {
304+
log.info("Session created after auth retry", { taskRunId });
305+
}
296306
return session;
297307
} catch (err) {
298308
this.cleanupMockNodeEnvironment(mockNodeDir);
309+
if (!isRetry && isAuthError(err)) {
310+
log.warn(
311+
`Auth error during ${isReconnect ? "reconnect" : "create"}, retrying`,
312+
{ taskRunId },
313+
);
314+
return this.getOrCreateSession(config, isReconnect, true);
315+
}
299316
log.error(
300-
`Failed to ${isReconnect ? "reconnect" : "create"} session`,
317+
`Failed to ${isReconnect ? "reconnect" : "create"} session${isRetry ? " after retry" : ""}`,
301318
err,
302319
);
303320
if (isReconnect) return null;
304321
throw err;
305322
}
306323
}
307324

325+
private async recreateSession(taskRunId: string): Promise<ManagedSession> {
326+
const existing = this.sessions.get(taskRunId);
327+
if (!existing) {
328+
throw new Error(`Session not found for recreation: ${taskRunId}`);
329+
}
330+
331+
log.info("Recreating session due to auth error", { taskRunId });
332+
333+
// Store config and cleanup old session
334+
const config = existing.config;
335+
this.cleanupSession(taskRunId);
336+
337+
// Reconnect to preserve Claude context via sdkSessionId
338+
const newSession = await this.getOrCreateSession(config, true);
339+
if (!newSession) {
340+
throw new Error(`Failed to recreate session: ${taskRunId}`);
341+
}
342+
343+
return newSession;
344+
}
345+
308346
async prompt(
309347
taskRunId: string,
310348
prompt: ContentBlock[],
311349
): Promise<{ stopReason: string }> {
312-
const session = this.sessions.get(taskRunId);
350+
let session = this.sessions.get(taskRunId);
313351
if (!session) {
314352
throw new Error(`Session not found: ${taskRunId}`);
315353
}
316354

317355
session.lastActivityAt = Date.now();
318356

319-
const result = await session.connection.prompt({
320-
sessionId: taskRunId, // Use taskRunId as ACP sessionId
321-
prompt,
322-
});
323-
324-
return { stopReason: result.stopReason };
357+
try {
358+
const result = await session.connection.prompt({
359+
sessionId: taskRunId,
360+
prompt,
361+
});
362+
return { stopReason: result.stopReason };
363+
} catch (err) {
364+
if (isAuthError(err)) {
365+
log.warn("Auth error during prompt, recreating session", { taskRunId });
366+
session = await this.recreateSession(taskRunId);
367+
const result = await session.connection.prompt({
368+
sessionId: taskRunId,
369+
prompt,
370+
});
371+
return { stopReason: result.stopReason };
372+
}
373+
throw err;
374+
}
325375
}
326376

327377
async cancelSession(taskRunId: string): Promise<boolean> {
@@ -398,11 +448,12 @@ export class SessionManager {
398448
credentials: PostHogCredentials,
399449
mockNodeDir: string,
400450
): void {
451+
const token = this.getToken(credentials.apiKey);
401452
const newPath = `${mockNodeDir}:${process.env.PATH || ""}`;
402453
process.env.PATH = newPath;
403-
process.env.POSTHOG_AUTH_HEADER = `Bearer ${credentials.apiKey}`;
404-
process.env.ANTHROPIC_API_KEY = credentials.apiKey;
405-
process.env.ANTHROPIC_AUTH_TOKEN = credentials.apiKey;
454+
process.env.POSTHOG_AUTH_HEADER = `Bearer ${token}`;
455+
process.env.ANTHROPIC_API_KEY = token;
456+
process.env.ANTHROPIC_AUTH_TOKEN = token;
406457

407458
const llmGatewayUrl =
408459
process.env.LLM_GATEWAY_URL ||
@@ -508,6 +559,24 @@ export class SessionManager {
508559
// No-op: session/update notifications are captured by the stream tap
509560
// and forwarded as acp_message events to avoid duplication
510561
},
562+
563+
extNotification: async (
564+
method: string,
565+
params: Record<string, unknown>,
566+
): Promise<void> => {
567+
if (method === "_posthog/sdk_session") {
568+
const { sessionId, sdkSessionId } = params as {
569+
sessionId: string;
570+
sdkSessionId: string;
571+
};
572+
// Store sdkSessionId in session config for recreation/reconnection
573+
const session = this.sessions.get(sessionId);
574+
if (session) {
575+
session.config.sdkSessionId = sdkSessionId;
576+
log.info("SDK session ID captured", { sessionId, sdkSessionId });
577+
}
578+
}
579+
},
511580
};
512581

513582
// Create client-side connection with tapped streams (bidirectional)
@@ -665,10 +734,10 @@ export function registerAgentIpc(
665734
"agent-token-refresh",
666735
async (
667736
_event: IpcMainInvokeEvent,
668-
taskRunId: string,
737+
_taskRunId: string,
669738
newToken: string,
670739
): Promise<void> => {
671-
sessionManager.updateSessionToken(taskRunId, newToken);
740+
sessionManager.updateToken(newToken);
672741
},
673742
);
674743

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,13 @@ export const useAuthStore = create<AuthState>()(
202202
...(projectId && { projectId }),
203203
});
204204

205+
// Notify main process of token refresh for active agent sessions
206+
window.electronAPI
207+
.agentTokenRefresh("", tokenResponse.access_token)
208+
.catch((err) => {
209+
log.warn("Failed to update agent token:", err);
210+
});
211+
205212
get().scheduleTokenRefresh();
206213
},
207214

‎packages/agent/example-client.ts‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,10 @@ import type { SessionPersistenceConfig } from "./src/session-store.js";
2525
import { Logger } from "./src/utils/logger.js";
2626

2727
// PostHog configuration - set via env vars
28+
const POSTHOG_API_KEY = process.env.POSTHOG_API_KEY || "";
2829
const POSTHOG_CONFIG = {
2930
apiUrl: process.env.POSTHOG_API_URL || "",
30-
apiKey: process.env.POSTHOG_API_KEY || "",
31+
getApiKey: () => POSTHOG_API_KEY,
3132
projectId: parseInt(process.env.POSTHOG_PROJECT_ID || "0", 10),
3233
};
3334

@@ -250,11 +251,7 @@ async function main() {
250251
}
251252
} else if (!existingSessionId) {
252253
// Create new Task/TaskRun for new sessions (only if PostHog is configured)
253-
if (
254-
POSTHOG_CONFIG.apiUrl &&
255-
POSTHOG_CONFIG.apiKey &&
256-
POSTHOG_CONFIG.projectId
257-
) {
254+
if (POSTHOG_CONFIG.apiUrl && POSTHOG_API_KEY && POSTHOG_CONFIG.projectId) {
258255
console.log("🔗 Connecting to PostHog...");
259256
const posthogClient = new PostHogAPIClient(POSTHOG_CONFIG);
260257

@@ -299,7 +296,7 @@ async function main() {
299296
logger.log(level, message, data, scope);
300297
},
301298
...(POSTHOG_CONFIG.apiUrl && { posthogApiUrl: POSTHOG_CONFIG.apiUrl }),
302-
...(POSTHOG_CONFIG.apiKey && { posthogApiKey: POSTHOG_CONFIG.apiKey }),
299+
...(POSTHOG_API_KEY && { getPosthogApiKey: POSTHOG_CONFIG.getApiKey }),
303300
...(POSTHOG_CONFIG.projectId && {
304301
posthogProjectId: POSTHOG_CONFIG.projectId,
305302
}),

‎packages/agent/example.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,10 +83,11 @@ async function testAgent() {
8383
});
8484
}
8585

86+
const apiKey = process.env.POSTHOG_API_KEY || "";
8687
const agent = new Agent({
8788
workingDirectory: REPO_PATH,
8889
posthogApiUrl: process.env.POSTHOG_API_URL || "http://localhost:8010",
89-
posthogApiKey: process.env.POSTHOG_API_KEY,
90+
getPosthogApiKey: () => apiKey,
9091
posthogProjectId: process.env.POSTHOG_PROJECT_ID
9192
? parseInt(process.env.POSTHOG_PROJECT_ID, 10)
9293
: 1,

‎packages/agent/src/agent.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,8 @@ export class Agent {
5757

5858
// Add auth if API key provided
5959
const headers: Record<string, string> = {};
60-
if (config.posthogApiKey) {
61-
headers.Authorization = `Bearer ${config.posthogApiKey}`;
60+
if (config.getPosthogApiKey) {
61+
headers.Authorization = `Bearer ${config.getPosthogApiKey()}`;
6262
}
6363

6464
const defaultMcpServers = {
@@ -93,12 +93,12 @@ export class Agent {
9393

9494
if (
9595
config.posthogApiUrl &&
96-
config.posthogApiKey &&
96+
config.getPosthogApiKey &&
9797
config.posthogProjectId
9898
) {
9999
this.posthogAPI = new PostHogAPIClient({
100100
apiUrl: config.posthogApiUrl,
101-
apiKey: config.posthogApiKey,
101+
getApiKey: config.getPosthogApiKey,
102102
projectId: config.posthogProjectId,
103103
});
104104

‎packages/agent/src/posthog-api.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ export class PostHogAPIClient {
4242

4343
private get headers(): Record<string, string> {
4444
return {
45-
Authorization: `Bearer ${this.config.apiKey}`,
45+
Authorization: `Bearer ${this.config.getApiKey()}`,
4646
"Content-Type": "application/json",
4747
};
4848
}
@@ -84,7 +84,7 @@ export class PostHogAPIClient {
8484
}
8585

8686
getApiKey(): string {
87-
return this.config.apiKey;
87+
return this.config.getApiKey();
8888
}
8989

9090
getLlmGatewayUrl(): string {

‎packages/agent/src/types.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ export interface AgentConfig {
197197

198198
// PostHog API configuration (optional - enables PostHog integration when provided)
199199
posthogApiUrl?: string;
200-
posthogApiKey?: string;
200+
getPosthogApiKey?: () => string;
201201
posthogProjectId?: number;
202202

203203
// PostHog MCP configuration
@@ -219,7 +219,7 @@ export interface AgentConfig {
219219

220220
export interface PostHogAPIConfig {
221221
apiUrl: string;
222-
apiKey: string;
222+
getApiKey: () => string;
223223
projectId: number;
224224
}
225225

0 commit comments

Comments
 (0)