fix(cli): re-attach session handlers registered while offline after reconnecting (#1858)

Backends register their user-message and RPC handlers once, on the offline
stub when the server is unreachable at startup. The real session that replaces
it never got them, so app messages, abort and kill were ignored.

Author: rophy

packages/happy-cli/src/utils/offlineSessionStub.ts

@@ -31,8 +31,18 @@ import type { ApiSessionClient } from '@/api/apiSession';
  * session = api.sessionSyncClient(response);
  * ```
  */
+type OfflineRegistrations = {
+    userMessageHandler: Parameters<ApiSessionClient['onUserMessage']>[0] | null;
+    rpcHandlers: Map<string, (...args: any[]) => any>;
+};
+
+const offlineRegistrations = new WeakMap<object, OfflineRegistrations>();
+
 export function createOfflineSessionStub(sessionTag: string): ApiSessionClient {
-    return {
+    // Backends register their handlers once, on whatever session they hold at
+    // startup. Record them so they can be attached to the real session later.
+    const registrations: OfflineRegistrations = { userMessageHandler: null, rpcHandlers: new Map() };
+    const stub = {
         sessionId: `offline-${sessionTag}`,
         sendCodexMessage: () => {},
         sendAgentMessage: () => {},
@@ -46,9 +56,31 @@ export function createOfflineSessionStub(sessionTag: string): ApiSessionClient {
         close: async () => {},
         updateMetadata: () => {},
         updateAgentState: () => {},
-        onUserMessage: () => {},
+        onUserMessage: (handler: OfflineRegistrations['userMessageHandler']) => {
+            registrations.userMessageHandler = handler;
+        },
         rpcHandlerManager: {
-            registerHandler: () => {}
+            registerHandler: (method: string, handler: (...args: any[]) => any) => {
+                registrations.rpcHandlers.set(method, handler);
+            }
         }
-    } as unknown as ApiSessionClient;
+    };
+    offlineRegistrations.set(stub, registrations);
+    return stub as unknown as ApiSessionClient;
+}
+
+/**
+ * Attaches the handlers registered on an offline stub to the real session that
+ * replaces it after reconnection. Without this, the real session would have no
+ * user-message listener and no RPC handlers (abort, kill, ...).
+ */
+export function transferOfflineRegistrations(stub: ApiSessionClient, session: ApiSessionClient): void {
+    const registrations = offlineRegistrations.get(stub);
+    if (!registrations) return;
+    for (const [method, handler] of registrations.rpcHandlers) {
+        session.rpcHandlerManager.registerHandler(method, handler);
+    }
+    if (registrations.userMessageHandler) {
+        session.onUserMessage(registrations.userMessageHandler);
+    }
 }

packages/happy-cli/src/utils/setupOfflineReconnection.test.ts

@@ -0,0 +1,97 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+let capturedOnReconnected: (() => Promise<unknown>) | null = null;
+
+vi.mock('@/configuration', () => ({
+    configuration: { serverUrl: 'http://server.test' },
+}));
+
+vi.mock('@/utils/serverConnectionErrors', () => ({
+    startOfflineReconnection: (config: { onReconnected: () => Promise<unknown> }) => {
+        capturedOnReconnected = config.onReconnected;
+        return { cancel: () => {}, getSession: () => null, isReconnected: () => false };
+    },
+}));
+
+import { setupOfflineReconnection } from './setupOfflineReconnection';
+
+function makeRealSession() {
+    return {
+        sessionId: 'real-session',
+        onUserMessage: vi.fn(),
+        rpcHandlerManager: { registerHandler: vi.fn() },
+    };
+}
+
+function setupOffline(realSession: ReturnType<typeof makeRealSession>, onSessionSwap = vi.fn()) {
+    const api = {
+        getOrCreateSession: vi.fn().mockResolvedValue({ id: 'real-session' }),
+        sessionSyncClient: vi.fn().mockReturnValue(realSession),
+    };
+    const result = setupOfflineReconnection({
+        api: api as any,
+        sessionTag: 'tag-1',
+        metadata: {} as any,
+        state: {} as any,
+        response: null,
+        onSessionSwap,
+    });
+    return { result, onSessionSwap };
+}
+
+describe('setupOfflineReconnection', () => {
+    beforeEach(() => {
+        capturedOnReconnected = null;
+    });
+
+    it('re-attaches the user message handler registered while offline', async () => {
+        const realSession = makeRealSession();
+        const { result } = setupOffline(realSession);
+        const handler = vi.fn();
+
+        result.session.onUserMessage(handler);
+        await capturedOnReconnected!();
+
+        expect(realSession.onUserMessage).toHaveBeenCalledWith(handler);
+    });
+
+    it('re-registers RPC handlers registered while offline', async () => {
+        const realSession = makeRealSession();
+        const { result } = setupOffline(realSession);
+        const abort = vi.fn();
+        const kill = vi.fn();
+
+        result.session.rpcHandlerManager.registerHandler('abort', abort);
+        result.session.rpcHandlerManager.registerHandler('killSession', kill);
+        await capturedOnReconnected!();
+
+        expect(realSession.rpcHandlerManager.registerHandler).toHaveBeenCalledWith('abort', abort);
+        expect(realSession.rpcHandlerManager.registerHandler).toHaveBeenCalledWith('killSession', kill);
+    });
+
+    it('attaches handlers before handing the real session to the caller', async () => {
+        const realSession = makeRealSession();
+        const order: string[] = [];
+        realSession.onUserMessage.mockImplementation(() => order.push('attach'));
+        const { result } = setupOffline(realSession, vi.fn(() => order.push('swap')));
+
+        result.session.onUserMessage(vi.fn());
+        await capturedOnReconnected!();
+
+        expect(order).toEqual(['attach', 'swap']);
+    });
+
+    it('re-attaches only the latest user message handler', async () => {
+        const realSession = makeRealSession();
+        const { result } = setupOffline(realSession);
+        const first = vi.fn();
+        const second = vi.fn();
+
+        result.session.onUserMessage(first);
+        result.session.onUserMessage(second);
+        await capturedOnReconnected!();
+
+        expect(realSession.onUserMessage).toHaveBeenCalledTimes(1);
+        expect(realSession.onUserMessage).toHaveBeenCalledWith(second);
+    });
+});

packages/happy-cli/src/utils/setupOfflineReconnection.ts

@@ -11,7 +11,7 @@ import type { ApiClient } from '@/api/api';
 import type { ApiSessionClient } from '@/api/apiSession';
 import type { AgentState, Metadata, Session } from '@/api/types';
 import { configuration } from '@/configuration';
-import { createOfflineSessionStub } from '@/utils/offlineSessionStub';
+import { createOfflineSessionStub, transferOfflineRegistrations } from '@/utils/offlineSessionStub';
 import { startOfflineReconnection } from '@/utils/serverConnectionErrors';
 
 /**
@@ -92,6 +92,7 @@ export function setupOfflineReconnection(opts: SetupOfflineReconnectionOptions):
                 const resp = await api.getOrCreateSession({ tag: sessionTag, metadata, state });
                 if (!resp) throw new Error('Server unavailable');
                 const realSession = api.sessionSyncClient(resp);
+                transferOfflineRegistrations(session, realSession);
                 // Notify caller to swap the session reference
                 onSessionSwap(realSession);
                 return realSession;