fix(codex): preserve gateway model availability

Authenticate the gateway catalog so restricted models remain gated, and keep the catalog available when native model discovery fails.

Generated-By: PostHog Code
Task-Id: dfa9be76-95bf-4166-8adc-529483e3d525

Author: richardsolomou

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,7 +231,7 @@ function createCodexConnection(config: AcpConnectionConfig): AcpConnection {
       },
       model: codexOptions.model,
       reasoningEffort: codexOptions.reasoningEffort,
-      allowedModelIds: config.allowedModelIds,
+      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,7 +204,7 @@ export interface CodexAppServerAgentOptions {
   processOptions: CodexAppServerProcessOptions;
   model?: string;
   reasoningEffort?: string;
-  allowedModelIds?: ReadonlySet<string>;
+  gatewayModels?: ReadonlyArray<ModelInfo>;
   processCallbacks?: ProcessSpawnedCallback;
   logger?: Logger;
   onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
@@ -266,7 +267,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
     this.config = new SessionConfigState(
       options.model ?? DEFAULT_CODEX_MODEL,
       options.reasoningEffort,
-      options.allowedModelIds,
+      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,
@@ -136,22 +137,66 @@ describe("SessionConfigState", () => {
   });
 
   it("uses gateway models when the app-server model list is stale", () => {
-    const config = new SessionConfigState(
-      "gpt-5.5",
-      undefined,
-      new Set(["gpt-5.5", "gpt-5.6-sol"]),
-    );
+    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",
     );
-    expect(modelOption?.type === "select" ? modelOption.options : []).toEqual([
+    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" },
     ]);
   });
 });

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,19 +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,
-    allowedModelIds?: ReadonlySet<string>,
+    gatewayModels?: ReadonlyArray<ModelInfo>,
   ) {
     this._model = model;
     this._effort = effort;
-    this.allowedModelIds = allowedModelIds?.size ? allowedModelIds : undefined;
+    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();
   }
 
@@ -270,7 +300,7 @@ export class SessionConfigState {
     if (typeof value === "string") {
       if (
         configId === "model" &&
-        (!this.allowedModelIds || this.allowedModelIds.has(value))
+        (!this.gatewayModels || this.allowedModelIds?.has(value))
       ) {
         this._model = value;
       } else if (configId === "effort") this._effort = value;
@@ -296,14 +326,17 @@ export class SessionConfigState {
         id: (m.id ?? m.model) as string,
         name: (m.displayName ?? m.id ?? m.model) as string,
       }));
-    if (this.allowedModelIds) {
+    if (this.gatewayModels) {
       const liveModelsById = new Map(
         liveModels.map((model) => [model.id, model]),
       );
-      this.models = [...this.allowedModelIds].map(
-        (modelId) =>
-          liveModelsById.get(modelId) ?? { id: modelId, name: modelId },
-      );
+      this.models = this.gatewayModels.map((gatewayModel) => ({
+        ...(liveModelsById.get(gatewayModel.id) ?? {
+          id: gatewayModel.id,
+          name: gatewayModel.id,
+        }),
+        ...(gatewayModel.allowed ? {} : { _meta: restrictedModelMeta() }),
+      }));
     } else {
       this.models = liveModels;
     }
@@ -321,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,