refactor: extract session and loop logic into reusable agent base classes

cli/src/agent/loopBase.ts

@@ -0,0 +1,40 @@
+import { logger } from '@/ui/logger';
+import type { AgentSessionBase } from './sessionBase';
+
+export type LoopLauncher<TSession> = (session: TSession) => Promise<'switch' | 'exit'>;
+
+export async function runLocalRemoteLoop<TSession extends AgentSessionBase<any>>(opts: {
+    session: TSession;
+    startingMode?: 'local' | 'remote';
+    logTag: string;
+    runLocal: LoopLauncher<TSession>;
+    runRemote: LoopLauncher<TSession>;
+}): Promise<void> {
+    let mode: 'local' | 'remote' = opts.startingMode ?? 'local';
+
+    while (true) {
+        logger.debug(`[${opts.logTag}] Iteration with mode: ${mode}`);
+
+        if (mode === 'local') {
+            const reason = await opts.runLocal(opts.session);
+            if (reason === 'exit') {
+                return;
+            }
+
+            mode = 'remote';
+            opts.session.onModeChange(mode);
+            continue;
+        }
+
+        if (mode === 'remote') {
+            const reason = await opts.runRemote(opts.session);
+            if (reason === 'exit') {
+                return;
+            }
+
+            mode = 'local';
+            opts.session.onModeChange(mode);
+            continue;
+        }
+    }
+}

cli/src/agent/sessionBase.ts

@@ -0,0 +1,79 @@
+import { ApiClient, ApiSessionClient } from '@/lib';
+import { MessageQueue2 } from '@/utils/MessageQueue2';
+import type { Metadata } from '@/api/types';
+import { logger } from '@/ui/logger';
+
+export type AgentSessionBaseOptions<Mode> = {
+    api: ApiClient;
+    client: ApiSessionClient;
+    path: string;
+    logPath: string;
+    sessionId: string | null;
+    messageQueue: MessageQueue2<Mode>;
+    onModeChange: (mode: 'local' | 'remote') => void;
+    mode?: 'local' | 'remote';
+    sessionLabel: string;
+    sessionIdLabel: string;
+    applySessionIdToMetadata: (metadata: Metadata, sessionId: string) => Metadata;
+};
+
+export class AgentSessionBase<Mode> {
+    readonly path: string;
+    readonly logPath: string;
+    readonly api: ApiClient;
+    readonly client: ApiSessionClient;
+    readonly queue: MessageQueue2<Mode>;
+    protected readonly _onModeChange: (mode: 'local' | 'remote') => void;
+
+    sessionId: string | null;
+    mode: 'local' | 'remote' = 'local';
+    thinking: boolean = false;
+
+    private readonly applySessionIdToMetadata: (metadata: Metadata, sessionId: string) => Metadata;
+    private readonly sessionLabel: string;
+    private readonly sessionIdLabel: string;
+    private keepAliveInterval: NodeJS.Timeout | null = null;
+
+    constructor(opts: AgentSessionBaseOptions<Mode>) {
+        this.path = opts.path;
+        this.api = opts.api;
+        this.client = opts.client;
+        this.logPath = opts.logPath;
+        this.sessionId = opts.sessionId;
+        this.queue = opts.messageQueue;
+        this._onModeChange = opts.onModeChange;
+        this.applySessionIdToMetadata = opts.applySessionIdToMetadata;
+        this.sessionLabel = opts.sessionLabel;
+        this.sessionIdLabel = opts.sessionIdLabel;
+        this.mode = opts.mode ?? 'local';
+
+        this.client.keepAlive(this.thinking, this.mode);
+        this.keepAliveInterval = setInterval(() => {
+            this.client.keepAlive(this.thinking, this.mode);
+        }, 2000);
+    }
+
+    onThinkingChange = (thinking: boolean) => {
+        this.thinking = thinking;
+        this.client.keepAlive(thinking, this.mode);
+    };
+
+    onModeChange = (mode: 'local' | 'remote') => {
+        this.mode = mode;
+        this.client.keepAlive(this.thinking, mode);
+        this._onModeChange(mode);
+    };
+
+    onSessionFound = (sessionId: string) => {
+        this.sessionId = sessionId;
+        this.client.updateMetadata((metadata) => this.applySessionIdToMetadata(metadata, sessionId));
+        logger.debug(`[${this.sessionLabel}] ${this.sessionIdLabel} session ID ${sessionId} added to metadata`);
+    };
+
+    stopKeepAlive = (): void => {
+        if (this.keepAliveInterval) {
+            clearInterval(this.keepAliveInterval);
+            this.keepAliveInterval = null;
+        }
+    };
+}

cli/src/claude/loop.ts

@@ -1,6 +1,7 @@
 import { ApiSessionClient } from "@/api/apiSession"
 import { MessageQueue2 } from "@/utils/MessageQueue2"
 import { logger } from "@/ui/logger"
+import { runLocalRemoteLoop } from "@/agent/loopBase"
 import { Session } from "./session"
 import { claudeLocalLauncher } from "./claudeLocalLauncher"
 import { claudeRemoteLauncher } from "./claudeRemoteLauncher"
@@ -49,7 +50,8 @@ export async function loop(opts: LoopOptions) {
         logPath: logPath,
         messageQueue: opts.messageQueue,
         allowedTools: opts.allowedTools,
-        onModeChange: opts.onModeChange
+        onModeChange: opts.onModeChange,
+        mode: opts.startingMode
     });
 
     // Notify that session is ready
@@ -57,38 +59,11 @@ export async function loop(opts: LoopOptions) {
         opts.onSessionReady(session);
     }
 
-    let mode: 'local' | 'remote' = opts.startingMode ?? 'local';
-    while (true) {
-        logger.debug(`[loop] Iteration with mode: ${mode}`);
-
-        // Run local mode if applicable
-        if (mode === 'local') {
-            let reason = await claudeLocalLauncher(session);
-            if (reason === 'exit') { // Normal exit - Exit loop
-                return;
-            }
-
-            // Non "exit" reason means we need to switch to remote mode
-            mode = 'remote';
-            if (opts.onModeChange) {
-                opts.onModeChange(mode);
-            }
-            continue;
-        }
-
-        // Start remote mode
-        if (mode === 'remote') {
-            let reason = await claudeRemoteLauncher(session);
-            if (reason === 'exit') { // Normal exit - Exit loop
-                return;
-            }
-
-            // Non "exit" reason means we need to switch to local mode
-            mode = 'local';
-            if (opts.onModeChange) {
-                opts.onModeChange(mode);
-            }
-            continue;
-        }
-    }
+    await runLocalRemoteLoop({
+        session,
+        startingMode: opts.startingMode,
+        logTag: 'loop',
+        runLocal: claudeLocalLauncher,
+        runRemote: claudeRemoteLauncher
+    });
 }

cli/src/claude/session.ts

@@ -1,76 +1,50 @@
-import { ApiClient, ApiSessionClient } from "@/lib";
-import { MessageQueue2 } from "@/utils/MessageQueue2";
-import { EnhancedMode } from "./loop";
-import { logger } from "@/ui/logger";
+import { ApiClient, ApiSessionClient } from '@/lib';
+import { MessageQueue2 } from '@/utils/MessageQueue2';
+import { logger } from '@/ui/logger';
+import { AgentSessionBase } from '@/agent/sessionBase';
+import type { EnhancedMode } from './loop';
 
-export class Session {
-    readonly path: string;
-    readonly logPath: string;
-    readonly api: ApiClient;
-    readonly client: ApiSessionClient;
-    readonly queue: MessageQueue2<EnhancedMode>;
+export class Session extends AgentSessionBase<EnhancedMode> {
     readonly claudeEnvVars?: Record<string, string>;
-    claudeArgs?: string[];  // Made mutable to allow filtering
+    claudeArgs?: string[];
     readonly mcpServers: Record<string, any>;
     readonly allowedTools?: string[];
-    readonly _onModeChange: (mode: 'local' | 'remote') => void;
-
-    sessionId: string | null;
-    mode: 'local' | 'remote' = 'local';
-    thinking: boolean = false;
 
     constructor(opts: {
-        api: ApiClient,
-        client: ApiSessionClient,
-        path: string,
-        logPath: string,
-        sessionId: string | null,
-        claudeEnvVars?: Record<string, string>,
-        claudeArgs?: string[],
-        mcpServers: Record<string, any>,
-        messageQueue: MessageQueue2<EnhancedMode>,
-        onModeChange: (mode: 'local' | 'remote') => void,
-        allowedTools?: string[],
+        api: ApiClient;
+        client: ApiSessionClient;
+        path: string;
+        logPath: string;
+        sessionId: string | null;
+        claudeEnvVars?: Record<string, string>;
+        claudeArgs?: string[];
+        mcpServers: Record<string, any>;
+        messageQueue: MessageQueue2<EnhancedMode>;
+        onModeChange: (mode: 'local' | 'remote') => void;
+        allowedTools?: string[];
+        mode?: 'local' | 'remote';
     }) {
-        this.path = opts.path;
-        this.api = opts.api;
-        this.client = opts.client;
-        this.logPath = opts.logPath;
-        this.sessionId = opts.sessionId;
-        this.queue = opts.messageQueue;
+        super({
+            api: opts.api,
+            client: opts.client,
+            path: opts.path,
+            logPath: opts.logPath,
+            sessionId: opts.sessionId,
+            messageQueue: opts.messageQueue,
+            onModeChange: opts.onModeChange,
+            mode: opts.mode,
+            sessionLabel: 'Session',
+            sessionIdLabel: 'Claude Code',
+            applySessionIdToMetadata: (metadata, sessionId) => ({
+                ...metadata,
+                claudeSessionId: sessionId
+            })
+        });
+
         this.claudeEnvVars = opts.claudeEnvVars;
         this.claudeArgs = opts.claudeArgs;
         this.mcpServers = opts.mcpServers;
         this.allowedTools = opts.allowedTools;
-        this._onModeChange = opts.onModeChange;
-
-        // Start keep alive
-        this.client.keepAlive(this.thinking, this.mode);
-        setInterval(() => {
-            this.client.keepAlive(this.thinking, this.mode);
-        }, 2000);
-    }
-
-    onThinkingChange = (thinking: boolean) => {
-        this.thinking = thinking;
-        this.client.keepAlive(thinking, this.mode);
-    }
-
-    onModeChange = (mode: 'local' | 'remote') => {
-        this.mode = mode;
-        this.client.keepAlive(this.thinking, mode);
-        this._onModeChange(mode);
-    }
-
-    onSessionFound = (sessionId: string) => {
-        this.sessionId = sessionId;
-        
-        // Update metadata with Claude Code session ID
-        this.client.updateMetadata((metadata) => ({
-            ...metadata,
-            claudeSessionId: sessionId
-        }));
-        logger.debug(`[Session] Claude Code session ID ${sessionId} added to metadata`);
     }
 
     /**
@@ -79,7 +53,7 @@ export class Session {
     clearSessionId = (): void => {
         this.sessionId = null;
         logger.debug('[Session] Session ID cleared');
-    }
+    };
 
     /**
      * Consume one-time Claude flags from claudeArgs after Claude spawn
@@ -87,7 +61,7 @@ export class Session {
      */
     consumeOneTimeFlags = (): void => {
         if (!this.claudeArgs) return;
-        
+
         const filteredArgs: string[] = [];
         for (let i = 0; i < this.claudeArgs.length; i++) {
             if (this.claudeArgs[i] === '--resume') {
@@ -111,8 +85,8 @@ export class Session {
                 filteredArgs.push(this.claudeArgs[i]);
             }
         }
-        
+
         this.claudeArgs = filteredArgs.length > 0 ? filteredArgs : undefined;
         logger.debug(`[Session] Consumed one-time flags, remaining args:`, this.claudeArgs);
-    }
-}
+    };
+}

cli/src/codex/loop.ts

@@ -1,5 +1,6 @@
 import { MessageQueue2 } from '@/utils/MessageQueue2';
 import { logger } from '@/ui/logger';
+import { runLocalRemoteLoop } from '@/agent/loopBase';
 import { CodexSession } from './session';
 import { codexLocalLauncher } from './codexLocalLauncher';
 import { codexRemoteLauncher } from './codexRemoteLauncher';
@@ -39,29 +40,11 @@ export async function loop(opts: LoopOptions): Promise<void> {
         opts.onSessionReady(session);
     }
 
-    let mode: 'local' | 'remote' = opts.startingMode ?? 'local';
-
-    while (true) {
-        logger.debug(`[codex-loop] Iteration with mode: ${mode}`);
-
-        if (mode === 'local') {
-            const reason = await codexLocalLauncher(session);
-            if (reason === 'exit') {
-                return;
-            }
-            mode = 'remote';
-            session.onModeChange(mode);
-            continue;
-        }
-
-        if (mode === 'remote') {
-            const reason = await codexRemoteLauncher(session);
-            if (reason === 'exit') {
-                return;
-            }
-            mode = 'local';
-            session.onModeChange(mode);
-            continue;
-        }
-    }
+    await runLocalRemoteLoop({
+        session,
+        startingMode: opts.startingMode,
+        logTag: 'codex-loop',
+        runLocal: codexLocalLauncher,
+        runRemote: codexRemoteLauncher
+    });
 }

cli/src/codex/session.ts

@@ -1,21 +1,9 @@
 import { ApiClient, ApiSessionClient } from '@/lib';
 import { MessageQueue2 } from '@/utils/MessageQueue2';
-import { logger } from '@/ui/logger';
+import { AgentSessionBase } from '@/agent/sessionBase';
 import type { EnhancedMode } from './loop';
 
-export class CodexSession {
-    readonly path: string;
-    readonly logPath: string;
-    readonly api: ApiClient;
-    readonly client: ApiSessionClient;
-    readonly queue: MessageQueue2<EnhancedMode>;
-    readonly _onModeChange: (mode: 'local' | 'remote') => void;
-
-    sessionId: string | null;
-    mode: 'local' | 'remote' = 'local';
-    thinking: boolean = false;
-    private keepAliveInterval: NodeJS.Timeout | null = null;
-
+export class CodexSession extends AgentSessionBase<EnhancedMode> {
     constructor(opts: {
         api: ApiClient;
         client: ApiSessionClient;
@@ -26,41 +14,24 @@ export class CodexSession {
         onModeChange: (mode: 'local' | 'remote') => void;
         mode?: 'local' | 'remote';
     }) {
-        this.path = opts.path;
-        this.api = opts.api;
-        this.client = opts.client;
-        this.logPath = opts.logPath;
-        this.sessionId = opts.sessionId;
-        this.queue = opts.messageQueue;
-        this._onModeChange = opts.onModeChange;
-        this.mode = opts.mode ?? 'local';
-
-        this.client.keepAlive(this.thinking, this.mode);
-        this.keepAliveInterval = setInterval(() => {
-            this.client.keepAlive(this.thinking, this.mode);
-        }, 2000);
+        super({
+            api: opts.api,
+            client: opts.client,
+            path: opts.path,
+            logPath: opts.logPath,
+            sessionId: opts.sessionId,
+            messageQueue: opts.messageQueue,
+            onModeChange: opts.onModeChange,
+            mode: opts.mode,
+            sessionLabel: 'CodexSession',
+            sessionIdLabel: 'Codex',
+            applySessionIdToMetadata: (metadata, sessionId) => ({
+                ...metadata,
+                codexSessionId: sessionId
+            })
+        });
     }
 
-    onThinkingChange = (thinking: boolean) => {
-        this.thinking = thinking;
-        this.client.keepAlive(thinking, this.mode);
-    };
-
-    onModeChange = (mode: 'local' | 'remote') => {
-        this.mode = mode;
-        this.client.keepAlive(this.thinking, mode);
-        this._onModeChange(mode);
-    };
-
-    onSessionFound = (sessionId: string) => {
-        this.sessionId = sessionId;
-        this.client.updateMetadata((metadata) => ({
-            ...metadata,
-            codexSessionId: sessionId
-        }));
-        logger.debug(`[CodexSession] Codex session ID ${sessionId} added to metadata`);
-    };
-
     sendCodexMessage = (message: unknown): void => {
         this.client.sendCodexMessage(message);
     };
@@ -68,11 +39,4 @@ export class CodexSession {
     sendSessionEvent = (event: Parameters<ApiSessionClient['sendSessionEvent']>[0]): void => {
         this.client.sendSessionEvent(event);
     };
-
-    stopKeepAlive = (): void => {
-        if (this.keepAliveInterval) {
-            clearInterval(this.keepAliveInterval);
-            this.keepAliveInterval = null;
-        }
-    };
 }