feat: persist session mode (#471)

Co-authored-by: Charles Vien <charles.v@posthog.com>

Author: jonathanlab

apps/array/src/renderer/features/sessions/stores/sessionModeStore.ts

@@ -0,0 +1,53 @@
+import { create } from "zustand";
+import { persist } from "zustand/middleware";
+import type { ExecutionMode } from "./sessionStore";
+
+interface SessionModeState {
+  /** Map of taskId -> last used execution mode */
+  taskModes: Record<string, ExecutionMode>;
+}
+
+interface SessionModeActions {
+  /** Save the mode for a task */
+  setTaskMode: (taskId: string, mode: ExecutionMode) => void;
+  /** Get the saved mode for a task */
+  getTaskMode: (taskId: string) => ExecutionMode | undefined;
+}
+
+type SessionModeStore = SessionModeState & SessionModeActions;
+
+export const useSessionModeStore = create<SessionModeStore>()(
+  persist(
+    (set, get) => ({
+      taskModes: {},
+
+      setTaskMode: (taskId, mode) => {
+        set((state) => ({
+          taskModes: { ...state.taskModes, [taskId]: mode },
+        }));
+      },
+
+      getTaskMode: (taskId) => {
+        return get().taskModes[taskId];
+      },
+    }),
+    {
+      name: "session-mode-storage",
+    },
+  ),
+);
+
+/** Non-hook accessor for getting task mode */
+export function getPersistedTaskMode(
+  taskId: string,
+): ExecutionMode | undefined {
+  return useSessionModeStore.getState().getTaskMode(taskId);
+}
+
+/** Non-hook accessor for setting task mode */
+export function setPersistedTaskMode(
+  taskId: string,
+  mode: ExecutionMode,
+): void {
+  useSessionModeStore.getState().setTaskMode(taskId, mode);
+}

apps/array/src/renderer/features/sessions/stores/sessionStore.ts

@@ -27,6 +27,7 @@ import {
   findPendingPermissions,
   type PermissionRequest,
 } from "../utils/parseSessionLogs";
+import { getPersistedTaskMode, setPersistedTaskMode } from "./sessionModeStore";
 
 const log = logger.scope("session-store");
 const CLOUD_POLLING_INTERVAL_MS = 500;
@@ -149,6 +150,7 @@ function subscribeToChannel(taskRunId: string) {
                   newMode === "acceptEdits"
                 ) {
                   session.currentMode = newMode;
+                  setPersistedTaskMode(session.taskId, newMode);
                   log.info("Session mode updated", { taskRunId, newMode });
                 }
               }
@@ -595,10 +597,14 @@ const useStore = create<SessionStore>()(
         });
       }
 
+      const persistedMode = getPersistedTaskMode(taskId);
       const session = createBaseSession(taskRunId, taskId, false);
       session.events = events;
       session.logUrl = logUrl;
       session.pendingPermissions = pendingPermissions;
+      if (persistedMode) {
+        session.currentMode = persistedMode;
+      }
 
       addSession(session);
       subscribeToChannel(taskRunId);
@@ -616,6 +622,24 @@ const useStore = create<SessionStore>()(
 
       if (result) {
         updateSession(taskRunId, { status: "connected" });
+        if (persistedMode) {
+          try {
+            await trpcVanilla.agent.setMode.mutate({
+              sessionId: taskRunId,
+              modeId: persistedMode,
+            });
+            log.info("Restored persisted mode after reconnect", {
+              taskId,
+              taskRunId,
+              mode: persistedMode,
+            });
+          } catch (error) {
+            log.warn("Failed to restore persisted mode after reconnect", {
+              taskId,
+              error,
+            });
+          }
+        }
       } else {
         unsubscribeFromChannel(taskRunId);
         removeSession(taskRunId);
@@ -640,6 +664,9 @@ const useStore = create<SessionStore>()(
         return;
       }
 
+      const persistedMode = getPersistedTaskMode(taskId);
+      const effectiveMode = executionMode ?? persistedMode;
+
       const { defaultModel, defaultFramework } = useSettingsStore.getState();
       const result = await trpcVanilla.agent.start.mutate({
         taskId,
@@ -650,19 +677,22 @@ const useStore = create<SessionStore>()(
         projectId: auth.projectId,
         model: defaultModel,
         framework: defaultFramework,
-        executionMode,
+        executionMode: effectiveMode,
       });
 
       const session = createBaseSession(
         taskRun.id,
         taskId,
         false,
         defaultFramework,
-        executionMode,
+        effectiveMode === "default" ? undefined : effectiveMode,
       );
       session.channel = result.channel;
       session.status = "connected";
       session.model = defaultModel;
+      if (persistedMode && !executionMode) {
+        session.currentMode = persistedMode;
+      }
 
       addSession(session);
       subscribeToChannel(taskRun.id);
@@ -903,6 +933,7 @@ const useStore = create<SessionStore>()(
               modeId,
             });
             updateSession(session.taskRunId, { currentMode: modeId });
+            setPersistedTaskMode(taskId, modeId);
           } catch (error) {
             log.error("Failed to change session mode", {
               taskId,