refactor: extract permission handler base class for code reuse

cli/src/claude/utils/permissionHandler.ts

@@ -15,6 +15,11 @@ import { getToolName } from "./getToolName";
 import { EnhancedMode, PermissionMode } from "../loop";
 import { getToolDescriptor } from "./getToolDescriptor";
 import { delay } from "@/utils/time";
+import {
+    BasePermissionHandler,
+    type PendingPermissionRequest,
+    type PermissionCompletion
+} from "@/modules/common/permission/BasePermissionHandler";
 
 interface PermissionResponse {
     id: string;
@@ -91,17 +96,9 @@ function buildAskUserQuestionUpdatedInput(input: unknown, answers: Record<string
     };
 }
 
-interface PendingRequest {
-    resolve: (value: PermissionResult) => void;
-    reject: (error: Error) => void;
-    toolName: string;
-    input: unknown;
-}
-
-export class PermissionHandler {
+export class PermissionHandler extends BasePermissionHandler<PermissionResponse, PermissionResult> {
     private toolCalls: { id: string, name: string, input: any, used: boolean }[] = [];
     private responses = new Map<string, PermissionResponse>();
-    private pendingRequests = new Map<string, PendingRequest>();
     private session: Session;
     private allowedTools = new Set<string>();
     private allowedBashLiterals = new Set<string>();
@@ -110,8 +107,8 @@ export class PermissionHandler {
     private onPermissionRequestCallback?: (toolCallId: string) => void;
 
     constructor(session: Session) {
+        super(session.client);
         this.session = session;
-        this.setupClientHandler();
     }
     
     /**
@@ -129,10 +126,17 @@ export class PermissionHandler {
     /**
      * Handler response
      */
-    private handlePermissionResponse(
+    protected handlePermissionResponse(
         response: PermissionResponse,
-        pending: PendingRequest
-    ): void {
+        pending: PendingPermissionRequest<PermissionResult>
+    ): PermissionCompletion {
+        const completion: PermissionCompletion = {
+            status: response.approved ? 'approved' : 'denied',
+            reason: response.reason,
+            mode: response.mode,
+            allowTools: response.allowTools,
+            answers: response.answers
+        };
 
         // Update allowed tools
         if (response.allowTools && response.allowTools.length > 0) {
@@ -154,19 +158,20 @@ export class PermissionHandler {
             this.session.setPermissionMode(response.mode);
         }
 
-        // Handle 
+        // Handle ask_user_question
         if (isAskUserQuestionToolName(pending.toolName)) {
             const answers = response.answers ?? {};
             if (Object.keys(answers).length === 0) {
                 pending.resolve({ behavior: 'deny', message: 'No answers were provided.' });
-                return;
+                completion.status = 'denied';
+                completion.reason = completion.reason ?? 'No answers were provided.';
+            } else {
+                pending.resolve({
+                    behavior: 'allow',
+                    updatedInput: buildAskUserQuestionUpdatedInput(pending.input, answers)
+                });
             }
-
-            pending.resolve({
-                behavior: 'allow',
-                updatedInput: buildAskUserQuestionUpdatedInput(pending.input, answers)
-            });
-            return;
+            return completion;
         }
 
         if (pending.toolName === 'exit_plan_mode' || pending.toolName === 'ExitPlanMode') {
@@ -184,14 +189,16 @@ export class PermissionHandler {
             } else {
                 pending.resolve({ behavior: 'deny', message: response.reason || 'Plan rejected' });
             }
-        } else {
-            // Handle default case for all other tools
-            const result: PermissionResult = response.approved
-                ? { behavior: 'allow', updatedInput: (pending.input as Record<string, unknown>) || {} }
-                : { behavior: 'deny', message: response.reason || `The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). STOP what you are doing and wait for the user to tell you how to proceed.` };
-
-            pending.resolve(result);
+            return completion;
         }
+
+        // Handle default case for all other tools
+        const result: PermissionResult = response.approved
+            ? { behavior: 'allow', updatedInput: (pending.input as Record<string, unknown>) || {} }
+            : { behavior: 'deny', message: response.reason || `The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). STOP what you are doing and wait for the user to tell you how to proceed.` };
+
+        pending.resolve(result);
+        return completion;
     }
 
     /**
@@ -267,7 +274,7 @@ export class PermissionHandler {
             signal.addEventListener('abort', abortHandler, { once: true });
 
             // Store the pending request
-            this.pendingRequests.set(id, {
+            this.addPendingRequest(id, toolName, input, {
                 resolve: (result: PermissionResult) => {
                     signal.removeEventListener('abort', abortHandler);
                     resolve(result);
@@ -275,28 +282,8 @@ export class PermissionHandler {
                 reject: (error: Error) => {
                     signal.removeEventListener('abort', abortHandler);
                     reject(error);
-                },
-                toolName,
-                input
-            });
-
-            // Trigger callback to send delayed messages immediately
-            if (this.onPermissionRequestCallback) {
-                this.onPermissionRequestCallback(id);
-            }
-            
-            // Update agent state
-            this.session.client.updateAgentState((currentState) => ({
-                ...currentState,
-                requests: {
-                    ...currentState.requests,
-                    [id]: {
-                        tool: toolName,
-                        arguments: input,
-                        createdAt: Date.now()
-                    }
                 }
-            }));
+            });
 
             logger.debug(`Permission request sent for tool call ${id}: ${toolName}`);
         });
@@ -416,80 +403,9 @@ export class PermissionHandler {
         this.allowedBashLiterals.clear();
         this.allowedBashPrefixes.clear();
 
-        // Cancel all pending requests
-        for (const [, pending] of this.pendingRequests.entries()) {
-            pending.reject(new Error('Session reset'));
-        }
-        this.pendingRequests.clear();
-
-        // Move all pending requests to completedRequests with canceled status
-        this.session.client.updateAgentState((currentState) => {
-            const pendingRequests = currentState.requests || {};
-            const completedRequests = { ...currentState.completedRequests };
-
-            // Move each pending request to completed with canceled status
-            for (const [id, request] of Object.entries(pendingRequests)) {
-                completedRequests[id] = {
-                    ...request,
-                    completedAt: Date.now(),
-                    status: 'canceled',
-                    reason: 'Session switched to local mode'
-                };
-            }
-
-            return {
-                ...currentState,
-                requests: {}, // Clear all pending requests
-                completedRequests
-            };
-        });
-    }
-
-    /**
-     * Sets up the client handler for permission responses
-     */
-    private setupClientHandler(): void {
-        this.session.client.rpcHandlerManager.registerHandler<PermissionResponse, void>('permission', async (message) => {
-            logger.debug(`Permission response: ${JSON.stringify(message)}`);
-
-            const id = message.id;
-            const pending = this.pendingRequests.get(id);
-
-            if (!pending) {
-                logger.debug('Permission request not found or already resolved');
-                return;
-            }
-
-            // Store the response with timestamp
-            this.responses.set(id, { ...message, receivedAt: Date.now() });
-            this.pendingRequests.delete(id);
-
-            // Handle the permission response based on tool type
-            this.handlePermissionResponse(message, pending);
-
-            // Move processed request to completedRequests
-            this.session.client.updateAgentState((currentState) => {
-                const request = currentState.requests?.[id];
-                if (!request) return currentState;
-                let r = { ...currentState.requests };
-                delete r[id];
-                return {
-                    ...currentState,
-                    requests: r,
-                    completedRequests: {
-                        ...currentState.completedRequests,
-                        [id]: {
-                            ...request,
-                            completedAt: Date.now(),
-                            status: message.approved ? 'approved' : 'denied',
-                            reason: message.reason,
-                            mode: message.mode,
-                            allowTools: message.allowTools,
-                            answers: message.answers
-                        }
-                    }
-                };
-            });
+        this.cancelPendingRequests({
+            completedReason: 'Session switched to local mode',
+            rejectMessage: 'Session reset'
         });
     }
 
@@ -499,4 +415,19 @@ export class PermissionHandler {
     getResponses(): Map<string, PermissionResponse> {
         return this.responses;
     }
+
+    protected handleMissingPendingResponse(_response: PermissionResponse): void {
+        logger.debug('Permission request not found or already resolved');
+    }
+
+    protected onResponseReceived(response: PermissionResponse): void {
+        logger.debug(`Permission response: ${JSON.stringify(response)}`);
+        this.responses.set(response.id, { ...response, receivedAt: Date.now() });
+    }
+
+    protected onRequestRegistered(toolCallId: string): void {
+        if (this.onPermissionRequestCallback) {
+            this.onPermissionRequestCallback(toolCallId);
+        }
+    }
 }

cli/src/codex/utils/permissionHandler.ts

@@ -7,7 +7,11 @@
 
 import { logger } from "@/ui/logger";
 import { ApiSessionClient } from "@/api/apiSession";
-import { AgentState } from "@/api/types";
+import {
+    BasePermissionHandler,
+    type PendingPermissionRequest,
+    type PermissionCompletion
+} from "@/modules/common/permission/BasePermissionHandler";
 
 interface PermissionResponse {
     id: string;
@@ -16,25 +20,15 @@ interface PermissionResponse {
     reason?: string;
 }
 
-interface PendingRequest {
-    resolve: (value: PermissionResult) => void;
-    reject: (error: Error) => void;
-    toolName: string;
-    input: unknown;
-}
-
 interface PermissionResult {
     decision: 'approved' | 'approved_for_session' | 'denied' | 'abort';
     reason?: string;
 }
 
-export class CodexPermissionHandler {
-    private pendingRequests = new Map<string, PendingRequest>();
-    private session: ApiSessionClient;
+export class CodexPermissionHandler extends BasePermissionHandler<PermissionResponse, PermissionResult> {
 
     constructor(session: ApiSessionClient) {
-        this.session = session;
-        this.setupRpcHandler();
+        super(session);
     }
 
     /**
@@ -51,12 +45,7 @@ export class CodexPermissionHandler {
     ): Promise<PermissionResult> {
         return new Promise<PermissionResult>((resolve, reject) => {
             // Store the pending request
-            this.pendingRequests.set(toolCallId, {
-                resolve,
-                reject,
-                toolName,
-                input
-            });
+            this.addPendingRequest(toolCallId, toolName, input, { resolve, reject });
 
             // Send push notification
             // this.session.api.push().sendToAllDevices(
@@ -70,117 +59,49 @@ export class CodexPermissionHandler {
             //     }
             // );
 
-            // Update agent state with pending request
-            this.session.updateAgentState((currentState) => ({
-                ...currentState,
-                requests: {
-                    ...currentState.requests,
-                    [toolCallId]: {
-                        tool: toolName,
-                        arguments: input,
-                        createdAt: Date.now()
-                    }
-                }
-            }));
-
             logger.debug(`[Codex] Permission request sent for tool: ${toolName} (${toolCallId})`);
         });
     }
 
     /**
-     * Setup RPC handler for permission responses
+     * Handle permission responses
      */
-    private setupRpcHandler(): void {
-        this.session.rpcHandlerManager.registerHandler<PermissionResponse, void>(
-            'permission',
-            async (response) => {
-                // console.log(`[Codex] Permission response received:`, response);
-
-                const pending = this.pendingRequests.get(response.id);
-                if (!pending) {
-                    logger.debug('[Codex] Permission request not found or already resolved');
-                    return;
-                }
-
-                // Remove from pending
-                this.pendingRequests.delete(response.id);
-
-                // Resolve the permission request
-                const reason = typeof response.reason === 'string' ? response.reason : undefined;
-                const result: PermissionResult = response.approved
-                    ? {
-                        decision: response.decision === 'approved_for_session' ? 'approved_for_session' : 'approved',
-                        reason
-                    }
-                    : {
-                        decision: response.decision === 'denied' ? 'denied' : 'abort',
-                        reason
-                    };
-
-                pending.resolve(result);
-
-                // Move request to completed in agent state
-                this.session.updateAgentState((currentState) => {
-                    const request = currentState.requests?.[response.id];
-                    if (!request) return currentState;
-
-                    // console.log(`[Codex] Permission ${response.approved ? 'approved' : 'denied'} for ${pending.toolName}`);
-
-                    const { [response.id]: _, ...remainingRequests } = currentState.requests || {};
-
-                    let res = {
-                        ...currentState,
-                        requests: remainingRequests,
-                        completedRequests: {
-                            ...currentState.completedRequests,
-                            [response.id]: {
-                                ...request,
-                                completedAt: Date.now(),
-                                status: response.approved ? 'approved' : 'denied',
-                                decision: result.decision,
-                                reason: result.reason
-                            }
-                        }
-                    } satisfies AgentState;
-                    // console.log(`[Codex] Updated agent state:`, res);
-                    return res;
-                });
-
-                logger.debug(`[Codex] Permission ${response.approved ? 'approved' : 'denied'} for ${pending.toolName}`);
+    protected handlePermissionResponse(
+        response: PermissionResponse,
+        pending: PendingPermissionRequest<PermissionResult>
+    ): PermissionCompletion {
+        const reason = typeof response.reason === 'string' ? response.reason : undefined;
+        const result: PermissionResult = response.approved
+            ? {
+                decision: response.decision === 'approved_for_session' ? 'approved_for_session' : 'approved',
+                reason
             }
-        );
+            : {
+                decision: response.decision === 'denied' ? 'denied' : 'abort',
+                reason
+            };
+
+        pending.resolve(result);
+        logger.debug(`[Codex] Permission ${response.approved ? 'approved' : 'denied'} for ${pending.toolName}`);
+
+        return {
+            status: response.approved ? 'approved' : 'denied',
+            decision: result.decision,
+            reason: result.reason
+        };
+    }
+
+    protected handleMissingPendingResponse(_response: PermissionResponse): void {
+        logger.debug('[Codex] Permission request not found or already resolved');
     }
 
     /**
      * Reset state for new sessions
      */
     reset(): void {
-        // Reject all pending requests
-        for (const [id, pending] of this.pendingRequests.entries()) {
-            pending.reject(new Error('Session reset'));
-        }
-        this.pendingRequests.clear();
-
-        // Clear requests in agent state
-        this.session.updateAgentState((currentState) => {
-            const pendingRequests = currentState.requests || {};
-            const completedRequests = { ...currentState.completedRequests };
-
-            // Move all pending to completed as canceled
-            for (const [id, request] of Object.entries(pendingRequests)) {
-                completedRequests[id] = {
-                    ...request,
-                    completedAt: Date.now(),
-                    status: 'canceled',
-                    reason: 'Session reset'
-                };
-            }
-
-            return {
-                ...currentState,
-                requests: {},
-                completedRequests
-            };
+        this.cancelPendingRequests({
+            completedReason: 'Session reset',
+            rejectMessage: 'Session reset'
         });
 
         logger.debug('[Codex] Permission handler reset');

cli/src/modules/common/permission/BasePermissionHandler.ts

@@ -0,0 +1,150 @@
+import type { AgentState } from "@/api/types";
+
+type RpcHandlerManagerLike = {
+    registerHandler<TRequest = unknown, TResponse = unknown>(
+        method: string,
+        handler: (params: TRequest) => Promise<TResponse> | TResponse
+    ): void;
+};
+
+export type PermissionHandlerClient = {
+    rpcHandlerManager: RpcHandlerManagerLike;
+    updateAgentState: (handler: (state: AgentState) => AgentState) => void;
+};
+
+export type PendingPermissionRequest<TResult> = {
+    resolve: (value: TResult) => void;
+    reject: (error: Error) => void;
+    toolName: string;
+    input: unknown;
+};
+
+export type PermissionCompletion = {
+    status: 'approved' | 'denied' | 'canceled';
+    reason?: string;
+    mode?: string;
+    decision?: 'approved' | 'approved_for_session' | 'denied' | 'abort';
+    allowTools?: string[];
+    answers?: Record<string, string[]>;
+};
+
+export type CancelPendingRequestOptions = {
+    completedReason: string;
+    rejectMessage: string;
+};
+
+export abstract class BasePermissionHandler<TResponse extends { id: string }, TResult> {
+    protected readonly pendingRequests = new Map<string, PendingPermissionRequest<TResult>>();
+    protected readonly client: PermissionHandlerClient;
+
+    protected constructor(client: PermissionHandlerClient) {
+        this.client = client;
+        this.setupRpcHandler();
+    }
+
+    protected abstract handlePermissionResponse(
+        response: TResponse,
+        pending: PendingPermissionRequest<TResult>
+    ): PermissionCompletion;
+
+    protected abstract handleMissingPendingResponse(response: TResponse): void;
+
+    protected onRequestRegistered(_id: string, _toolName: string, _input: unknown): void {
+    }
+
+    protected onResponseReceived(_response: TResponse): void {
+    }
+
+    protected addPendingRequest(
+        id: string,
+        toolName: string,
+        input: unknown,
+        handlers: { resolve: (value: TResult) => void; reject: (error: Error) => void }
+    ): void {
+        this.pendingRequests.set(id, { ...handlers, toolName, input });
+        this.onRequestRegistered(id, toolName, input);
+        this.client.updateAgentState((currentState) => ({
+            ...currentState,
+            requests: {
+                ...currentState.requests,
+                [id]: {
+                    tool: toolName,
+                    arguments: input,
+                    createdAt: Date.now()
+                }
+            }
+        }));
+    }
+
+    protected finalizeRequest(id: string, completion: PermissionCompletion): void {
+        this.client.updateAgentState((currentState) => {
+            const request = currentState.requests?.[id];
+            if (!request) return currentState;
+
+            const nextRequests = { ...currentState.requests };
+            delete nextRequests[id];
+
+            return {
+                ...currentState,
+                requests: nextRequests,
+                completedRequests: {
+                    ...currentState.completedRequests,
+                    [id]: {
+                        ...request,
+                        completedAt: Date.now(),
+                        status: completion.status,
+                        reason: completion.reason,
+                        mode: completion.mode,
+                        decision: completion.decision,
+                        allowTools: completion.allowTools,
+                        answers: completion.answers
+                    }
+                }
+            };
+        });
+    }
+
+    protected cancelPendingRequests(options: CancelPendingRequestOptions): void {
+        for (const [, pending] of this.pendingRequests.entries()) {
+            pending.reject(new Error(options.rejectMessage));
+        }
+        this.pendingRequests.clear();
+
+        this.client.updateAgentState((currentState) => {
+            const pendingRequests = currentState.requests || {};
+            const completedRequests = { ...currentState.completedRequests };
+
+            for (const [id, request] of Object.entries(pendingRequests)) {
+                completedRequests[id] = {
+                    ...request,
+                    completedAt: Date.now(),
+                    status: 'canceled',
+                    reason: options.completedReason
+                };
+            }
+
+            return {
+                ...currentState,
+                requests: {},
+                completedRequests
+            };
+        });
+    }
+
+    private setupRpcHandler(): void {
+        this.client.rpcHandlerManager.registerHandler<TResponse, void>('permission', async (response) => {
+            const pending = this.pendingRequests.get(response.id);
+
+            if (!pending) {
+                this.handleMissingPendingResponse(response);
+                return;
+            }
+
+            this.onResponseReceived(response);
+            this.pendingRequests.delete(response.id);
+
+            const completion = this.handlePermissionResponse(response, pending);
+            this.finalizeRequest(response.id, completion);
+        });
+    }
+}