Harden plan and goal session state

Restore collaboration mode when sessions are loaded and publish authoritative goal snapshots without allowing stale reads to overwrite live updates. Keep goal control responses typed and preserve stable goal identity for UI replay.

Author: nikita-ashihmin

src/CodexAcpClient.ts

@@ -336,6 +336,7 @@ export class CodexAcpClient {
             sessionId: request.sessionId,
             currentModelId: currentModelId,
             models: codexModels,
+            collaborationMode: this.getCollaborationMode(response.thread.id),
             modelProvider: response.modelProvider,
             currentServiceTier: response.serviceTier as ServiceTier ?? null,
             additionalDirectories,
@@ -363,6 +364,7 @@ export class CodexAcpClient {
             sessionId: request.sessionId,
             currentModelId: currentModelId,
             models: codexModels,
+            collaborationMode: this.getCollaborationMode(response.thread.id),
             modelProvider: response.modelProvider,
             currentServiceTier: response.serviceTier as ServiceTier ?? null,
             thread: historyResponse.thread,
@@ -389,6 +391,7 @@ export class CodexAcpClient {
             sessionId: response.thread.id,
             currentModelId: currentModelId,
             models: codexModels,
+            collaborationMode: this.getCollaborationMode(response.thread.id),
             modelProvider: response.modelProvider,
             currentServiceTier: response.serviceTier as ServiceTier ?? null,
             additionalDirectories,
@@ -440,11 +443,18 @@ export class CodexAcpClient {
         }, onTurnStarted);
     }
 
-    async setGoalStatus(sessionId: string, status: ThreadGoalStatus): Promise<void> {
+    async setGoalStatus(sessionId: string, status: ThreadGoalStatus): Promise<ThreadGoal> {
+        let updatedGoal: ThreadGoal | null = null;
         await this.codexClient.runGoalSet({
             threadId: sessionId,
             status,
+        }, undefined, undefined, (goal) => {
+            updatedGoal = goal;
         });
+        if (updatedGoal === null) {
+            throw new Error(`Goal update for session ${sessionId} returned no goal`);
+        }
+        return updatedGoal;
     }
 
     async resumeGoal(
@@ -697,6 +707,10 @@ export class CodexAcpClient {
         });
     }
 
+    private getCollaborationMode(sessionId: string): ModeKind {
+        return this.codexClient.getThreadSettings(sessionId)?.collaborationMode.mode ?? "default";
+    }
+
     resolveTurnInterrupted(params: { threadId: string, turnId: string }): void {
         this.codexClient.resolveTurnInterrupted(params.threadId, params.turnId);
     }
@@ -863,6 +877,7 @@ export type SessionMetadata = {
     sessionId: string,
     currentModelId: string,
     models: Model[],
+    collaborationMode: ModeKind,
     modelProvider?: string | null,
     currentServiceTier?: ServiceTier | null,
     additionalDirectories: string[],

src/CodexAcpServer.ts

@@ -13,7 +13,6 @@ import type {
     Model,
     ReasoningEffortOption,
     Thread,
-    ThreadGoalStatus,
     ThreadItem,
     UserInput
 } from "./app-server/v2";
@@ -23,7 +22,6 @@ import {AgentMode, MODE_CONFIG_ID} from "./AgentMode";
 import {
     COLLABORATION_MODE_CONFIG_ID,
     createCollaborationModeConfigOption,
-    DEFAULT_COLLABORATION_MODE,
     parseCollaborationMode,
 } from "./CollaborationModeConfig";
 import type {ModeKind} from "./app-server/ModeKind";
@@ -82,14 +80,11 @@ import {
     createAgentTextThoughtChunk,
     createUserMessageChunk,
 } from "./ContentChunks";
-
-export interface ThreadGoalSnapshot {
-    objective: string;
-    status: ThreadGoalStatus;
-    tokenBudget: number | null;
-    timeUsedSeconds: number;
-    controlMethod: typeof GOAL_CONTROL_METHOD;
-}
+import {
+    sameThreadGoalSnapshot,
+    type ThreadGoalSnapshot,
+    toThreadGoalSnapshot,
+} from "./ThreadGoalSnapshot";
 
 export interface SessionState {
     sessionId: string,
@@ -114,6 +109,7 @@ export interface SessionState {
     sessionMcpServers?: Array<string>;
     terminalOutputMode: TerminalOutputMode;
     currentGoal?: ThreadGoalSnapshot | null;
+    goalRevision: number;
     sessionTitle: string | null;
     sessionTitleSource: "unset" | "fallback" | "explicit" | "unknown";
 }
@@ -263,10 +259,17 @@ export class CodexAcpServer {
                 if (!sessionState) {
                     throw RequestError.invalidParams(undefined, `Unknown session: ${methodRequest.params.sessionId}`);
                 }
+                const sessionGeneration = this.getSessionGeneration(sessionState.sessionId);
                 if (methodRequest.params.action === "pause") {
-                    await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionState.sessionId, "paused"));
-                } else {
+                    const goal = await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionState.sessionId, "paused"));
+                    if (this.goalPublishIsCurrent(sessionState, sessionGeneration)) {
+                        await this.publishGoalSnapshot(sessionState, toThreadGoalSnapshot(goal), false);
+                    }
+                } else if (methodRequest.params.action === "clear") {
                     await this.runWithProcessCheck(() => this.codexAcpClient.clearGoal(sessionState.sessionId));
+                    if (this.goalPublishIsCurrent(sessionState, sessionGeneration)) {
+                        await this.publishGoalSnapshot(sessionState, null, false);
+                    }
                 }
                 return {};
             }
@@ -428,7 +431,7 @@ export class CodexAcpServer {
             supportedReasoningEfforts: currentModel?.supportedReasoningEfforts ?? [],
             supportedInputModalities: currentModel?.inputModalities ?? ["text", "image"],
             agentMode: AgentMode.getInitialAgentMode(),
-            collaborationMode: DEFAULT_COLLABORATION_MODE,
+            collaborationMode: sessionMetadata.collaborationMode,
             currentTurnId: null,
             lastTokenUsage: null,
             totalTokenUsage: null,
@@ -443,6 +446,7 @@ export class CodexAcpServer {
             currentModelSupportsFast: currentModelSupportsFast,
             sessionMcpServers: sessionMcpServers,
             terminalOutputMode: this.terminalOutputMode,
+            goalRevision: 0,
             sessionTitle: null,
             sessionTitleSource: "sessionId" in request ? "unknown" : "unset",
         };
@@ -459,7 +463,7 @@ export class CodexAcpServer {
 
         this.publishAvailableCommandsAsync(sessionState);
         if ("sessionId" in request) {
-            this.publishCurrentGoalAsync(sessionState);
+            this.publishCurrentGoalAsync(sessionState, sessionGeneration);
         }
         const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId);
         const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
@@ -897,21 +901,55 @@ export class CodexAcpServer {
         void this.availableCommands.publish(sessionState);
     }
 
-    private publishCurrentGoalAsync(sessionState: SessionState): void {
-        void this.publishCurrentGoal(sessionState).catch((err) => {
+    private publishCurrentGoalAsync(sessionState: SessionState, sessionGeneration: number): void {
+        void this.publishCurrentGoalBestEffort(sessionState, sessionGeneration, true);
+    }
+
+    private async publishCurrentGoalBestEffort(
+        sessionState: SessionState,
+        sessionGeneration: number,
+        force: boolean,
+    ): Promise<void> {
+        try {
+            await this.publishCurrentGoal(sessionState, sessionGeneration, force);
+        } catch (err) {
             logger.error(`Failed to publish current goal for session ${sessionState.sessionId}`, err);
-        });
+        }
     }
 
-    private async publishCurrentGoal(sessionState: SessionState): Promise<void> {
+    private async publishCurrentGoal(
+        sessionState: SessionState,
+        sessionGeneration: number,
+        force: boolean,
+    ): Promise<void> {
+        const requestRevision = ++sessionState.goalRevision;
         const goal = await this.runWithProcessCheck(() => this.codexAcpClient.getGoal(sessionState.sessionId));
-        const snapshot: ThreadGoalSnapshot | null = goal === null ? null : {
-            objective: goal.objective.trim(),
-            status: goal.status,
-            tokenBudget: goal.tokenBudget,
-            timeUsedSeconds: goal.timeUsedSeconds,
-            controlMethod: GOAL_CONTROL_METHOD,
-        };
+        const snapshot = goal === null ? null : toThreadGoalSnapshot(goal);
+        if (!this.goalPublishIsCurrent(sessionState, sessionGeneration)
+            || sessionState.goalRevision !== requestRevision) {
+            return;
+        }
+        await this.publishGoalSnapshot(sessionState, snapshot, force, false);
+    }
+
+    private goalPublishIsCurrent(sessionState: SessionState, sessionGeneration: number): boolean {
+        return this.sessions.get(sessionState.sessionId) === sessionState
+            && this.getSessionGeneration(sessionState.sessionId) === sessionGeneration
+            && !this.sessionIsClosing(sessionState.sessionId);
+    }
+
+    private async publishGoalSnapshot(
+        sessionState: SessionState,
+        snapshot: ThreadGoalSnapshot | null,
+        force: boolean,
+        incrementRevision = true,
+    ): Promise<void> {
+        if (incrementRevision) {
+            sessionState.goalRevision += 1;
+        }
+        if (!force && sameThreadGoalSnapshot(sessionState.currentGoal, snapshot)) {
+            return;
+        }
         sessionState.currentGoal = snapshot;
         const session = new ACPSessionConnection(this.connection, sessionState.sessionId);
         await session.update({
@@ -1007,7 +1045,7 @@ export class CodexAcpServer {
             supportedReasoningEfforts: currentModel?.supportedReasoningEfforts ?? [],
             supportedInputModalities: currentModel?.inputModalities ?? ["text", "image"],
             agentMode: AgentMode.getInitialAgentMode(),
-            collaborationMode: DEFAULT_COLLABORATION_MODE,
+            collaborationMode: sessionMetadata.collaborationMode,
             currentTurnId: null,
             lastTokenUsage: null,
             totalTokenUsage: null,
@@ -1022,6 +1060,7 @@ export class CodexAcpServer {
             currentModelSupportsFast: currentModelSupportsFast,
             sessionMcpServers: sessionMcpServers,
             terminalOutputMode: this.terminalOutputMode,
+            goalRevision: 0,
             sessionTitle: null,
             sessionTitleSource: "unset",
         };
@@ -1037,7 +1076,7 @@ export class CodexAcpServer {
         }
 
         await this.availableCommands.publish(sessionState);
-        await this.publishCurrentGoal(sessionState);
+        await this.publishCurrentGoalBestEffort(sessionState, requestedSessionGeneration, true);
         const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId);
         const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
 

src/CodexAppServerClient.ts

@@ -49,6 +49,7 @@ import type {
     ThreadReadResponse,
     ThreadResumeParams,
     ThreadResumeResponse,
+    ThreadSettings,
     ThreadStartParams,
     ThreadStartResponse,
     ThreadUnsubscribeParams,
@@ -141,6 +142,7 @@ export class CodexAppServerClient {
     private readonly threadStatusCaptures = new Map<string, Set<(status: ThreadStatus) => void>>();
     private readonly threadGoalUpdateCaptures = new Map<string, Set<(event: ThreadGoalUpdatedNotification) => void>>();
     private readonly threadGoalClearedCaptures = new Map<string, Set<() => void>>();
+    private readonly threadSettings = new Map<string, ThreadSettings>();
     private readonly staleTurnIds = new Map<string, Set<string>>();
 
     constructor(connection: MessageConnection) {
@@ -171,6 +173,9 @@ export class CodexAppServerClient {
             if (isThreadGoalClearedNotification(serverNotification)) {
                 this.recordThreadGoalCleared(serverNotification.params);
             }
+            if (serverNotification.method === "thread/settings/updated") {
+                this.threadSettings.set(serverNotification.params.threadId, serverNotification.params.threadSettings);
+            }
             const routing = extractTurnRouting(serverNotification);
             if (this.handleStaleTurnNotification(serverNotification, routing)) {
                 return;
@@ -312,6 +317,7 @@ export class CodexAppServerClient {
         params: ThreadGoalSetParams,
         onTurnStarted?: (turnId: string) => void,
         runtimeEffectsGraceMs = GOAL_RUNTIME_EFFECTS_GRACE_MS,
+        onGoalSet?: (goal: ThreadGoal) => void,
     ): Promise<TurnCompletedNotification | null> {
         let goalTurnId: string | null = null;
         const capturedCompletions: Array<TurnCompletedNotification> = [];
@@ -363,6 +369,7 @@ export class CodexAppServerClient {
         try {
             const goalSetResponse = await this.threadGoalSet(params);
             expectedGoal = goalSetResponse.goal;
+            onGoalSet?.(expectedGoal);
             if (capturedGoalUpdates.some(event => goalsMatch(event.goal, expectedGoal!))) {
                 goalUpdateHandled = true;
                 resolveGoalUpdateHandled();
@@ -515,6 +522,10 @@ export class CodexAppServerClient {
         return await this.sendRequest({ method: "thread/resume", params: params });
     }
 
+    getThreadSettings(threadId: string): ThreadSettings | undefined {
+        return this.threadSettings.get(threadId);
+    }
+
     async threadSettingsUpdate(params: ExperimentalThreadSettingsUpdateParams): Promise<void> {
         await this.connection.sendRequest("thread/settings/update", params);
     }

src/CodexCommands.ts

@@ -204,7 +204,10 @@ export class CodexCommands {
                     await this.sendCommandUsageMessage(commandName, "no arguments", sessionId);
                     return { handled: true };
                 }
-                await options.setConfigOption?.(COLLABORATION_MODE_CONFIG_ID, PLAN_COLLABORATION_MODE);
+                const mode = sessionState.collaborationMode === PLAN_COLLABORATION_MODE
+                    ? DEFAULT_COLLABORATION_MODE
+                    : PLAN_COLLABORATION_MODE;
+                await options.setConfigOption?.(COLLABORATION_MODE_CONFIG_ID, mode);
                 return { handled: options.setConfigOption !== undefined };
             }
             case "compact": {

src/CodexEventHandler.ts

@@ -3,8 +3,7 @@ import type {
     FuzzyFileSearchSessionUpdatedNotification,
     ServerNotification
 } from "./app-server";
-import type {SessionState, ThreadGoalSnapshot} from "./CodexAcpServer";
-import {GOAL_CONTROL_METHOD} from "./AcpExtensions";
+import type {SessionState} from "./CodexAcpServer";
 import {type PlanEntry, RequestError} from "@agentclientprotocol/sdk";
 import {ACPSessionConnection, type AcpClientConnection, type UpdateSessionEvent} from "./ACPSessionConnection";
 import type {
@@ -63,6 +62,7 @@ import {
     createAgentTextMessageChunk,
     createAgentTextThoughtChunk,
 } from "./ContentChunks";
+import {sameThreadGoalSnapshot, toThreadGoalSnapshot} from "./ThreadGoalSnapshot";
 
 export { stripShellPrefix };
 
@@ -257,8 +257,9 @@ export class CodexEventHandler {
     }
 
     private createThreadGoalUpdatedEvent(event: ThreadGoalUpdatedNotification): UpdateSessionEvent | null {
-        const goalSnapshot = this.createThreadGoalSnapshot(event);
-        if (this.sameThreadGoalSnapshot(this.sessionState.currentGoal, goalSnapshot)) {
+        this.sessionState.goalRevision += 1;
+        const goalSnapshot = toThreadGoalSnapshot(event.goal);
+        if (sameThreadGoalSnapshot(this.sessionState.currentGoal, goalSnapshot)) {
             return null;
         }
         this.sessionState.currentGoal = goalSnapshot;
@@ -269,6 +270,7 @@ export class CodexEventHandler {
     }
 
     private createThreadGoalClearedEvent(_event: ThreadGoalClearedNotification): UpdateSessionEvent | null {
+        this.sessionState.goalRevision += 1;
         if (this.sessionState.currentGoal === null) {
             return null;
         }
@@ -279,27 +281,6 @@ export class CodexEventHandler {
         });
     }
 
-    private createThreadGoalSnapshot(event: ThreadGoalUpdatedNotification): ThreadGoalSnapshot {
-        return {
-            objective: event.goal.objective.trim(),
-            status: event.goal.status,
-            tokenBudget: event.goal.tokenBudget,
-            timeUsedSeconds: event.goal.timeUsedSeconds,
-            controlMethod: GOAL_CONTROL_METHOD,
-        };
-    }
-
-    private sameThreadGoalSnapshot(
-        left: ThreadGoalSnapshot | null | undefined,
-        right: ThreadGoalSnapshot
-    ): boolean {
-        return left !== null
-            && left !== undefined
-            && left.objective === right.objective
-            && left.status === right.status
-            && left.tokenBudget === right.tokenBudget;
-    }
-
     private createReasoningDeltaEvent(
         event: ReasoningSummaryTextDeltaNotification | ReasoningTextDeltaNotification
     ): UpdateSessionEvent {