From 1f13bab58adb1265d5174cbcb95629f70cc1a982 Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Fri, 17 Jul 2026 23:58:56 +0200 Subject: [PATCH] fix(realtime): isolate notification failures --- src/server/realtime/sessionEventHub.test.ts | 42 ++++++++++ src/server/realtime/sessionEventHub.ts | 18 ++-- src/server/sessiond.ts | 2 +- src/server/sessions/authService.test.ts | 84 +++++++++++++++++-- src/server/sessions/authService.ts | 39 ++++++--- .../sessions/oauthLoginFlowService.test.ts | 19 +++++ src/server/sessions/oauthLoginFlowService.ts | 7 +- 7 files changed, 187 insertions(+), 24 deletions(-) diff --git a/src/server/realtime/sessionEventHub.test.ts b/src/server/realtime/sessionEventHub.test.ts index bf190e5..08641b9 100644 --- a/src/server/realtime/sessionEventHub.test.ts +++ b/src/server/realtime/sessionEventHub.test.ts @@ -54,6 +54,28 @@ describe("SessionEventHub", () => { expect(removed.send).not.toHaveBeenCalled(); }); + it("continues publishing session events when one socket send fails", () => { + const hub = new SessionEventHub(); + const failed = new FakeSocket(); + const healthy = new FakeSocket(); + failed.send.mockImplementation(() => { throw new Error("socket closed"); }); + hub.add("s1", failed); + hub.add("s1", healthy); + + hub.publish("s1", { type: "assistant.delta", text: "hello" }); + + expect(failed.send).toHaveBeenCalledOnce(); + expect(healthy.send).toHaveBeenCalledWith(JSON.stringify({ type: "assistant.delta", text: "hello", seq: 1 })); + expect(hub.currentSeq("s1")).toBe(1); + + failed.send.mockClear(); + hub.publish("s1", { type: "assistant.delta", text: "again" }); + + expect(failed.send).not.toHaveBeenCalled(); + expect(healthy.send).toHaveBeenLastCalledWith(JSON.stringify({ type: "assistant.delta", text: "again", seq: 2 })); + expect(hub.currentSeq("s1")).toBe(2); + }); + it("publishes global events only to global sockets", () => { const hub = new SessionEventHub(); const globalSocket = new FakeSocket(); @@ -78,6 +100,26 @@ describe("SessionEventHub", () => { expect(sessionSocket.send).not.toHaveBeenCalled(); }); + it("continues publishing unstamped global events when one socket send fails", () => { + const hub = new SessionEventHub(); + const failed = new FakeSocket(); + const healthy = new FakeSocket(); + failed.send.mockImplementation(() => { throw new Error("socket closed"); }); + hub.addGlobal(failed); + hub.addGlobal(healthy); + + hub.publishGlobal({ type: "session.name", sessionId: "s1", name: "Renamed" }); + + expect(failed.send).toHaveBeenCalledOnce(); + expect(healthy.send).toHaveBeenCalledWith(JSON.stringify({ type: "session.name", sessionId: "s1", name: "Renamed" })); + + failed.send.mockClear(); + hub.publishGlobal({ type: "session.name", sessionId: "s1", name: "Renamed again" }); + + expect(failed.send).not.toHaveBeenCalled(); + expect(healthy.send).toHaveBeenLastCalledWith(JSON.stringify({ type: "session.name", sessionId: "s1", name: "Renamed again" })); + }); + it("stamps a monotonically increasing per-session seq on published events", () => { const hub = new SessionEventHub(); const socket = new FakeSocket(); diff --git a/src/server/realtime/sessionEventHub.ts b/src/server/realtime/sessionEventHub.ts index b57e6d1..9df8a1a 100644 --- a/src/server/realtime/sessionEventHub.ts +++ b/src/server/realtime/sessionEventHub.ts @@ -34,9 +34,7 @@ export class SessionEventHub { const seq = (this.seqBySession.get(sessionId) ?? 0) + 1; this.seqBySession.set(sessionId, seq); const payload = JSON.stringify({ ...projectBrowserSessionEvent(event), seq }); - for (const socket of this.socketsBySession.get(sessionId) ?? []) { - if (socket.readyState === socket.OPEN) socket.send(payload); - } + this.sendToSockets(this.socketsBySession.get(sessionId), payload); } /** @@ -55,8 +53,18 @@ export class SessionEventHub { publishRealtime(event: RealtimeEvent): void { const payload = JSON.stringify(event); - for (const socket of this.globalSockets) { - if (socket.readyState === socket.OPEN) socket.send(payload); + this.sendToSockets(this.globalSockets, payload); + } + + private sendToSockets(sockets: Set | undefined, payload: string): void { + if (sockets === undefined) return; + for (const socket of sockets) { + if (socket.readyState !== socket.OPEN) continue; + try { + socket.send(payload); + } catch { + sockets.delete(socket); + } } } } diff --git a/src/server/sessiond.ts b/src/server/sessiond.ts index 0e31a9d..e35fde1 100644 --- a/src/server/sessiond.ts +++ b/src/server/sessiond.ts @@ -39,7 +39,7 @@ await runSessionDaemonStartup({ async createRuntime() { const eventHub = new SessionEventHub(); const workspaceActivity = new WorkspaceActivityService(eventHub); - const auth = await AuthService.create({ agentDir: activeAgentProfile.dir }); + const auth = await AuthService.create({ agentDir: activeAgentProfile.dir, logger: app.log }); const spawnTargets = config.spawnSessions ? new ProjectScopedSpawnTargetResolver({ projects: new ProjectService(new ProjectStore()), workspaces: new WorkspaceService() }) : undefined; diff --git a/src/server/sessions/authService.test.ts b/src/server/sessions/authService.test.ts index 1751677..599faae 100644 --- a/src/server/sessions/authService.test.ts +++ b/src/server/sessions/authService.test.ts @@ -5,7 +5,7 @@ import { ModelRuntime } from "@earendil-works/pi-coding-agent"; import { InMemoryCredentialStore, type AuthPrompt, type Credential } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { OAuthFlowState } from "../../shared/apiTypes.js"; -import { AuthService, type AuthChange } from "./authService.js"; +import { AuthService, type AuthChange, type AuthServiceLogger } from "./authService.js"; import { OAuthLoginFlowService } from "./oauthLoginFlowService.js"; const tempDirs: string[] = []; @@ -41,6 +41,54 @@ describe("AuthService", () => { auth.dispose(); }); + it("persists an API key and attempts every listener when propagation fails", async () => { + const error = vi.fn(); + const logger: AuthServiceLogger = { error }; + const { auth, credentials, changes } = await createAuthService({}, logger); + const failure = new Error("session auth refresh failed"); + const attempts: string[] = []; + auth.subscribe(() => { + attempts.push("throwing"); + throw failure; + }); + auth.subscribe(async () => { + await Promise.resolve(); + attempts.push("healthy"); + }); + + await expect(auth.saveApiKey("anthropic", "sk-test")).resolves.toEqual({ accepted: true }); + + await expect(credentials.read("anthropic")).resolves.toEqual({ type: "api_key", key: "sk-test" }); + expect(changes).toEqual([{}]); + expect(attempts).toEqual(["throwing", "healthy"]); + expect(error).toHaveBeenCalledWith( + { err: failure, operation: "login", providerId: "anthropic", authType: "api_key" }, + "auth-change listener failed", + ); + auth.dispose(); + }); + + it("removes a credential when auth-change propagation rejects", async () => { + const error = vi.fn(); + const logger: AuthServiceLogger = { error }; + const { auth, credentials, changes } = await createAuthService( + { anthropic: { type: "api_key", key: "sk-test" } }, + logger, + ); + const failure = new Error("session logout refresh failed"); + auth.subscribe(() => Promise.reject(failure)); + + await expect(auth.logoutProvider("anthropic")).resolves.toEqual({ accepted: true }); + + await expect(credentials.read("anthropic")).resolves.toBeUndefined(); + expect(changes).toEqual([{ removedProviderId: "anthropic" }]); + expect(error).toHaveBeenCalledWith( + { err: failure, operation: "logout", providerId: "anthropic" }, + "auth-change listener failed", + ); + auth.dispose(); + }); + it("rejects blank API keys", async () => { const { auth, changes } = await createAuthService(); @@ -236,22 +284,48 @@ describe("AuthService", () => { refresh.mockClear(); if (startOptions.onComplete === undefined) throw new Error("Expected OAuth completion callback"); - startOptions.onComplete(); - await vi.waitFor(() => { expect(changes).toEqual([{}]); }); + await startOptions.onComplete(); + expect(changes).toEqual([{}]); expect(refresh).not.toHaveBeenCalled(); auth.dispose(); expect(authFlows.disposed).toBe(true); }); + + it("completes OAuth when an auth-change listener rejects", async () => { + const error = vi.fn(); + const logger: AuthServiceLogger = { error }; + const { auth, runtime, changes } = await createAuthService({}, logger); + const provider = runtime.getProviders().find((option) => option.id === "anthropic" && option.auth.oauth !== undefined); + if (provider === undefined) throw new Error("Expected built-in OAuth provider"); + vi.spyOn(runtime, "login").mockResolvedValue({ + type: "oauth", + refresh: "refresh-token", + access: "access-token", + expires: Date.now() + 60_000, + }); + const failure = new Error("session OAuth refresh failed"); + auth.subscribe(() => Promise.reject(failure)); + + const state = await auth.startOAuthLogin(provider.id); + await vi.waitFor(() => { expect(auth.oauthFlow(state.flowId).status).toBe("complete"); }); + + expect(changes).toEqual([{}]); + expect(error).toHaveBeenCalledWith( + { err: failure, operation: "login", providerId: provider.id, authType: "oauth" }, + "auth-change listener failed", + ); + auth.dispose(); + }); }); -async function createAuthService(seed: Record = {}) { +async function createAuthService(seed: Record = {}, logger?: AuthServiceLogger) { const credentials = new InMemoryCredentialStore(); for (const [providerId, credential] of Object.entries(seed)) { await credentials.modify(providerId, () => Promise.resolve(credential)); } const runtime = await ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false }); - const auth = await AuthService.create({ runtime }); + const auth = await AuthService.create({ runtime, ...(logger === undefined ? {} : { logger }) }); const changes: AuthChange[] = []; auth.subscribe((change) => { changes.push(change); }); return { auth, runtime, credentials, changes }; diff --git a/src/server/sessions/authService.ts b/src/server/sessions/authService.ts index ad7d650..6f91d85 100644 --- a/src/server/sessions/authService.ts +++ b/src/server/sessions/authService.ts @@ -9,14 +9,28 @@ export interface AuthChange { removedProviderId?: string; } -type AuthChangeListener = (change: AuthChange) => void; +type AuthChangeListener = (change: AuthChange) => void | Promise; export interface AuthServiceDependencies { agentDir?: string; runtime?: ModelRuntime; authFlows?: OAuthLoginFlowService; + logger?: AuthServiceLogger; } +/** Minimal structured-logging seam for non-fatal auth propagation failures. */ +export interface AuthServiceLogger { + error(details: Record, message: string): void; +} + +interface AuthChangeContext { + operation: "login" | "logout"; + providerId: string; + authType?: AuthType; +} + +const noopLogger: AuthServiceLogger = { error() { /* no-op */ } }; + export function createModelRuntimeForAgentDir(agentDir: string): Promise { return ModelRuntime.create({ authPath: join(agentDir, "auth.json"), modelsPath: join(agentDir, "models.json") }); } @@ -24,17 +38,19 @@ export function createModelRuntimeForAgentDir(agentDir: string): Promise(); - private constructor(runtime: ModelRuntime, authFlows: OAuthLoginFlowService) { + private constructor(runtime: ModelRuntime, authFlows: OAuthLoginFlowService, logger: AuthServiceLogger) { this.runtime = runtime; this.authFlows = authFlows; + this.logger = logger; } static async create(deps: AuthServiceDependencies = {}): Promise { const runtime = deps.runtime ?? (deps.agentDir === undefined ? await ModelRuntime.create({}) : await createModelRuntimeForAgentDir(deps.agentDir)); const authFlows = deps.authFlows ?? new OAuthLoginFlowService(); - return new AuthService(runtime, authFlows); + return new AuthService(runtime, authFlows, deps.logger ?? noopLogger); } subscribe(listener: AuthChangeListener): () => void { @@ -74,13 +90,13 @@ export class AuthService { notify: () => undefined, }; await this.runtime.login(providerId, "api_key", interaction); - this.emit({}); + await this.emit({}, { operation: "login", providerId, authType: "api_key" }); return { accepted: true }; } async logoutProvider(providerId: string): Promise<{ accepted: true }> { await this.runtime.logout(providerId); - this.emit({ removedProviderId: providerId }); + await this.emit({ removedProviderId: providerId }, { operation: "logout", providerId }); return { accepted: true }; } @@ -90,9 +106,7 @@ export class AuthService { providerId, providerName: provider.name, runtime: this.runtime, - onComplete: () => { - this.emit({}); - }, + onComplete: () => this.emit({}, { operation: "login", providerId, authType: "oauth" }), }); } @@ -108,8 +122,13 @@ export class AuthService { return this.authFlows.cancel(flowId); } - private emit(change: AuthChange): void { - for (const listener of this.listeners) listener(change); + private async emit(change: AuthChange, context: AuthChangeContext): Promise { + const results = await Promise.allSettled([...this.listeners].map(async (listener) => listener(change))); + for (const result of results) { + if (result.status === "rejected") { + this.logger.error({ err: result.reason, ...context }, "auth-change listener failed"); + } + } } private async requireApiKeyLoginProvider(providerId: string) { diff --git a/src/server/sessions/oauthLoginFlowService.test.ts b/src/server/sessions/oauthLoginFlowService.test.ts index 6349b38..d42fbe8 100644 --- a/src/server/sessions/oauthLoginFlowService.test.ts +++ b/src/server/sessions/oauthLoginFlowService.test.ts @@ -41,6 +41,25 @@ describe("OAuthLoginFlowService", () => { service.dispose(); }); + it("awaits async completion propagation before marking the flow complete", async () => { + const completion = deferred(); + const service = new OAuthLoginFlowService(); + const state = service.start({ + providerId: "test-provider", + providerName: "Test Provider", + runtime: fakeRuntime(() => Promise.resolve()), + onComplete: () => completion.promise, + }); + + await flushAsyncLogin(); + expect(service.get(state.flowId).status).toBe("running"); + + completion.resolve(undefined); + await flushAsyncLogin(); + expect(service.get(state.flowId).status).toBe("complete"); + service.dispose(); + }); + it("allows blank text responses for providers that use blank as a default", async () => { let domain: string | undefined; const service = new OAuthLoginFlowService(); diff --git a/src/server/sessions/oauthLoginFlowService.ts b/src/server/sessions/oauthLoginFlowService.ts index e5d9786..214f694 100644 --- a/src/server/sessions/oauthLoginFlowService.ts +++ b/src/server/sessions/oauthLoginFlowService.ts @@ -52,7 +52,7 @@ export class OAuthLoginFlowService { providerId: string; providerName: string; runtime: OAuthLoginRuntime; - onComplete?: () => void; + onComplete?: () => void | Promise; }): OAuthFlowState { const flowId = crypto.randomUUID(); const abort = new AbortController(); @@ -81,11 +81,12 @@ export class OAuthLoginFlowService { }; void options.runtime.login(options.providerId, "oauth", interaction) - .then(() => { + .then(async () => { if (!this.isCurrentRunning(record)) return; this.clearPending(record); + await options.onComplete?.(); + if (!this.isCurrentRunning(record)) return; this.markTerminal(record, { ...withoutInteraction(record.state), status: "complete", progress: [...record.state.progress, "Login complete"] }); - options.onComplete?.(); }) .catch((error: unknown) => { if (this.flows.get(record.flowId) !== record) return;