fix(codex): restore gateway models after reconnect (#3566)

Co-authored-by: posthog[bot] <206114724+posthog[bot]@users.noreply.github.com>
Co-authored-by: richardsolomou <2622273+richardsolomou@users.noreply.github.com>

Author: richardsolomou

apps/code/snapshots.yml

@@ -27,7 +27,7 @@ snapshots:
     archive-archivedtasksview--many-tasks--dark:
         hash: v1.k4693efd2.40d24f07a24c6400c32ae68b39fd908a056beb5dc6df8bcb237ec2ab2b494f4a.l2vjfSQDalCPRtDNV7pbfTKGlMsQoI81jI-UdJfLHWE
     archive-archivedtasksview--many-tasks--light:
-        hash: v1.k4693efd2.52ec9963bfe6f248ffcdaa4c26945d2b7e303115825b65d978aa5aefd4af1007.ZLtOfTZUrznhYeB-N2SIHzQWRzF6TDSaipNkfd8uzjM
+        hash: v1.k4693efd2.74c25b303262f2d4c0f22890d351e76f482827548548f7c023bc309327c927b8.DP9VVrvBz1naIJMwHORkfqm69BlO785ulOt4o6zFs94
     archive-archivedtasksview--mixed-modes--dark:
         hash: v1.k4693efd2.d94039b8cc17a4ad1b720364f41fee58a4843aa9a901443997ba62602f698794.441LWZhT-2WQZJHni4LoOgQsOSr2O103NZOk1pF8ME4
     archive-archivedtasksview--mixed-modes--light:

packages/agent/src/adapters/acp-connection.ts

@@ -1,5 +1,6 @@
 import { AgentSideConnection, ndJsonStream } from "@agentclientprotocol/sdk";
 import type { Adapter } from "@posthog/shared";
+import type { ModelInfo } from "../gateway-models";
 import type { SessionLogWriter } from "../session-log-writer";
 import type { PostHogAPIConfig, ProcessSpawnedCallback } from "../types";
 import { Logger } from "../utils/logger";
@@ -24,7 +25,7 @@ export type AcpConnectionConfig = {
   logger?: Logger;
   processCallbacks?: ProcessSpawnedCallback;
   codexOptions?: CodexOptions;
-  allowedModelIds?: Set<string>;
+  codexModels?: ReadonlyArray<ModelInfo>;
   /** Callback invoked when the agent calls the create_output tool for structured output */
   onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
   /** PostHog API config; when set, enables file-read enrichment unless disabled. */
@@ -230,6 +231,7 @@ function createCodexConnection(config: AcpConnectionConfig): AcpConnection {
       },
       model: codexOptions.model,
       reasoningEffort: codexOptions.reasoningEffort,
+      gatewayModels: config.codexModels,
       processCallbacks: config.processCallbacks,
       onStructuredOutput: config.onStructuredOutput,
       logger: config.logger?.child("CodexAppServerAgent"),

packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts

@@ -29,6 +29,7 @@ import {
   type NativeGoalState,
   POSTHOG_NOTIFICATIONS,
 } from "../../acp-extensions";
+import type { ModelInfo } from "../../gateway-models";
 import { DEFAULT_CODEX_MODEL } from "../../gateway-models";
 import type { ProcessSpawnedCallback } from "../../types";
 import { ALLOW_BYPASS } from "../../utils/common";
@@ -203,6 +204,7 @@ export interface CodexAppServerAgentOptions {
   processOptions: CodexAppServerProcessOptions;
   model?: string;
   reasoningEffort?: string;
+  gatewayModels?: ReadonlyArray<ModelInfo>;
   processCallbacks?: ProcessSpawnedCallback;
   logger?: Logger;
   onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
@@ -265,6 +267,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
     this.config = new SessionConfigState(
       options.model ?? DEFAULT_CODEX_MODEL,
       options.reasoningEffort,
+      options.gatewayModels,
     );
     this.onStructuredOutput = options.onStructuredOutput;
     this.developerInstructions = options.processOptions.developerInstructions;

packages/agent/src/adapters/codex-app-server/session-config.test.ts

@@ -1,3 +1,4 @@
+import { isRestrictedModelOption } from "@posthog/shared";
 import { describe, expect, it } from "vitest";
 import {
   buildCodexModes,
@@ -134,6 +135,70 @@ describe("SessionConfigState", () => {
       config.options.find((option) => option.category === "mode")?.currentValue,
     ).toBe("full-access");
   });
+
+  it("uses gateway models when the app-server model list is stale", () => {
+    const config = new SessionConfigState("gpt-5.5", undefined, [
+      { id: "gpt-5.5", allowed: true },
+      { id: "gpt-5.6-sol", allowed: true },
+      { id: "gpt-5.6-terra", allowed: false },
+    ]);
+
+    config.loadModels([
+      { id: "gpt-5.5", model: "gpt-5.5", displayName: "GPT-5.5" },
+      {
+        id: "gpt-5.6-terra",
+        model: "gpt-5.6-terra",
+        displayName: "GPT-5.6 Terra",
+      },
+    ]);
+
+    const modelOption = config.options.find(
+      (option) => option.category === "model",
+    );
+    const modelOptions =
+      modelOption?.type === "select"
+        ? (modelOption.options as Array<{
+            name: string;
+            value: string;
+            _meta?: Record<string, unknown>;
+          }>)
+        : [];
+    expect(modelOptions).toEqual([
+      { name: "GPT-5.5", value: "gpt-5.5" },
+      { name: "gpt-5.6-sol", value: "gpt-5.6-sol" },
+      {
+        name: "GPT-5.6 Terra",
+        value: "gpt-5.6-terra",
+        _meta: { "posthog.code/restrictedModel": true },
+      },
+    ]);
+
+    config.setOption("model", "gpt-5.6-terra");
+
+    expect(config.model).toBe("gpt-5.5");
+    expect(
+      isRestrictedModelOption(
+        modelOptions.find((option) => option.value === "gpt-5.6-terra")?._meta,
+      ),
+    ).toBe(true);
+  });
+
+  it("keeps gateway models when the app-server model list fails", () => {
+    const config = new SessionConfigState("gpt-5.5", undefined, [
+      { id: "gpt-5.5", allowed: true },
+      { id: "gpt-5.6-sol", allowed: true },
+    ]);
+
+    config.clearModels();
+
+    const modelOption = config.options.find(
+      (option) => option.category === "model",
+    );
+    expect(modelOption?.type === "select" ? modelOption.options : []).toEqual([
+      { name: "gpt-5.5", value: "gpt-5.5" },
+      { name: "gpt-5.6-sol", value: "gpt-5.6-sol" },
+    ]);
+  });
 });
 
 describe("buildConfigOptions", () => {

packages/agent/src/adapters/codex-app-server/session-config.ts

@@ -1,11 +1,19 @@
-import type { SessionConfigOption } from "@agentclientprotocol/sdk";
+import type {
+  SessionConfigOption,
+  SessionConfigSelectOption,
+} from "@agentclientprotocol/sdk";
 import {
   CODEX_MODE_PRESETS,
   type CodexModePreset,
   type ExecutionMode,
   resolveCloudInitialPermissionMode,
+  restrictedModelMeta,
 } from "@posthog/shared";
-import { type GatewayModel, isOpenAIModel } from "../../gateway-models";
+import {
+  type GatewayModel,
+  isOpenAIModel,
+  type ModelInfo,
+} from "../../gateway-models";
 import { getReasoningEffortOptions } from "./models";
 
 /**
@@ -158,7 +166,11 @@ export interface ConfigSelectors {
   model: string;
   effort?: string;
   /** From model/list; falls back to the single current model when empty. */
-  models: Array<{ id: string; name: string }>;
+  models: Array<{
+    id: string;
+    name: string;
+    _meta?: Record<string, unknown>;
+  }>;
   efforts: string[];
 }
 
@@ -195,7 +207,13 @@ export function buildConfigOptions(s: ConfigSelectors): SessionConfigOption[] {
       name: "Model",
       category: "model",
       currentValue: s.model,
-      options: models.map((m) => ({ name: m.name, value: m.id })),
+      options: models.map(
+        (m): SessionConfigSelectOption => ({
+          name: m.name,
+          value: m.id,
+          ...(m._meta ? { _meta: m._meta } : {}),
+        }),
+      ),
     } as unknown as SessionConfigOption,
     {
       type: "select",
@@ -226,13 +244,31 @@ export class SessionConfigState {
   private _model: string;
   private _effort?: string;
   private _mode = DEFAULT_MODE;
-  private models: Array<{ id: string; name: string }> = [];
+  private models: Array<{
+    id: string;
+    name: string;
+    _meta?: Record<string, unknown>;
+  }> = [];
   private efforts: string[] = [];
   private _options: SessionConfigOption[] = [];
+  private readonly gatewayModels?: ReadonlyArray<ModelInfo>;
+  private readonly allowedModelIds?: ReadonlySet<string>;
 
-  constructor(model: string, effort?: string) {
+  constructor(
+    model: string,
+    effort?: string,
+    gatewayModels?: ReadonlyArray<ModelInfo>,
+  ) {
     this._model = model;
     this._effort = effort;
+    this.gatewayModels = gatewayModels?.length ? gatewayModels : undefined;
+    this.allowedModelIds = this.gatewayModels
+      ? new Set(
+          this.gatewayModels
+            .filter((gatewayModel) => gatewayModel.allowed)
+            .map((gatewayModel) => gatewayModel.id),
+        )
+      : undefined;
     this.rebuild();
   }
 
@@ -262,8 +298,12 @@ export class SessionConfigState {
   ): { modeChanged: boolean } {
     let modeChanged = false;
     if (typeof value === "string") {
-      if (configId === "model") this._model = value;
-      else if (configId === "effort") this._effort = value;
+      if (
+        configId === "model" &&
+        (!this.gatewayModels || this.allowedModelIds?.has(value))
+      ) {
+        this._model = value;
+      } else if (configId === "effort") this._effort = value;
       else if (configId === "mode") {
         this._mode = resolveCodexMode(value);
         modeChanged = true;
@@ -279,13 +319,27 @@ export class SessionConfigState {
    * populate efforts, so fall back to the shared codex model→effort map.
    */
   loadModels(rawModels: RawModel[]): void {
-    this.models = rawModels
+    const liveModels = rawModels
       .filter((m) => !m?.hidden)
       .filter((m) => isOpenAIModel(m as unknown as GatewayModel))
       .map((m) => ({
         id: (m.id ?? m.model) as string,
         name: (m.displayName ?? m.id ?? m.model) as string,
       }));
+    if (this.gatewayModels) {
+      const liveModelsById = new Map(
+        liveModels.map((model) => [model.id, model]),
+      );
+      this.models = this.gatewayModels.map((gatewayModel) => ({
+        ...(liveModelsById.get(gatewayModel.id) ?? {
+          id: gatewayModel.id,
+          name: gatewayModel.id,
+        }),
+        ...(gatewayModel.allowed ? {} : { _meta: restrictedModelMeta() }),
+      }));
+    } else {
+      this.models = liveModels;
+    }
     const current = rawModels.find(
       (m) => m.id === this._model || m.model === this._model,
     );
@@ -300,7 +354,12 @@ export class SessionConfigState {
 
   /** Reset the model/effort lists (model/list failed); keeps the current model. */
   clearModels(): void {
-    this.models = [];
+    this.models =
+      this.gatewayModels?.map((gatewayModel) => ({
+        id: gatewayModel.id,
+        name: gatewayModel.id,
+        ...(gatewayModel.allowed ? {} : { _meta: restrictedModelMeta() }),
+      })) ?? [];
     this.efforts = [];
     this.rebuild();
   }

packages/agent/src/agent.test.ts

@@ -54,5 +54,14 @@ describe("Agent", () => {
         reasoningEffort: "xhigh",
       }),
     );
+    expect(config.codexModels).toEqual([
+      expect.objectContaining({ id: "gpt-5.5", allowed: true }),
+    ]);
+    expect(fetchMock).toHaveBeenCalledWith(
+      expect.any(String),
+      expect.objectContaining({
+        headers: { Authorization: "Bearer token" },
+      }),
+    );
   });
 });

packages/agent/src/agent.ts

@@ -9,6 +9,7 @@ import {
   DEFAULT_GATEWAY_MODEL,
   fetchModelsList,
   isBlockedModelId,
+  type ModelInfo,
 } from "./gateway-models";
 import { PostHogAPIClient, type TaskRunUpdate } from "./posthog-api";
 import { SessionLogWriter } from "./session-log-writer";
@@ -78,33 +79,35 @@ export class Agent {
     const gatewayConfig = await this._resolveGatewayConfig(options.gatewayUrl);
     this.taskRunId = taskRunId;
 
-    let allowedModelIds: Set<string> | undefined;
+    let codexModels: ModelInfo[] | undefined;
     let sanitizedModel =
       options.model && !isBlockedModelId(options.model)
         ? options.model
         : undefined;
     if (options.adapter === "codex" && gatewayConfig) {
       const models = await fetchModelsList({
         gatewayUrl: gatewayConfig.gatewayUrl,
+        authToken: gatewayConfig.apiKey,
       });
-      const codexModelIds = models
-        .filter((model) => {
-          if (isBlockedModelId(model.id)) return false;
-          if (model.owned_by) {
-            return model.owned_by === "openai";
-          }
-          return model.id.startsWith("gpt-") || model.id.startsWith("openai/");
-        })
+      const gatewayCodexModels = models.filter((model) => {
+        if (isBlockedModelId(model.id)) return false;
+        if (model.owned_by) {
+          return model.owned_by === "openai";
+        }
+        return model.id.startsWith("gpt-") || model.id.startsWith("openai/");
+      });
+      const allowedModelIds = gatewayCodexModels
+        .filter((model) => model.allowed)
         .map((model) => model.id);
 
-      if (codexModelIds.length > 0) {
-        allowedModelIds = new Set(codexModelIds);
+      if (gatewayCodexModels.length > 0) {
+        codexModels = gatewayCodexModels;
       }
 
-      if (!sanitizedModel || !allowedModelIds?.has(sanitizedModel)) {
-        sanitizedModel = codexModelIds.includes(DEFAULT_CODEX_MODEL)
+      if (!sanitizedModel || !allowedModelIds.includes(sanitizedModel)) {
+        sanitizedModel = allowedModelIds.includes(DEFAULT_CODEX_MODEL)
           ? DEFAULT_CODEX_MODEL
-          : codexModelIds[0];
+          : (allowedModelIds[0] ?? sanitizedModel);
       }
     }
     if (!sanitizedModel && options.adapter !== "codex") {
@@ -130,7 +133,7 @@ export class Agent {
       logger: this.logger,
       processCallbacks: options.processCallbacks,
       onStructuredOutput: options.onStructuredOutput,
-      allowedModelIds,
+      codexModels,
       posthogApiConfig: this.posthogApiConfig,
       enricherEnabled: this.enricherEnabled,
       claudeGatewayEnv,