fix: improve model selection reliability, make gateway model source of truth

Author: jonathanlab

apps/twig/src/main/services/agent/schemas.ts

@@ -49,9 +49,20 @@ export const startSessionInput = z.object({
 
 export type StartSessionInput = z.infer<typeof startSessionInput>;
 
+export const modelOptionSchema = z.object({
+  modelId: z.string(),
+  name: z.string(),
+  description: z.string().nullish(),
+  provider: z.string().optional(),
+});
+
+export type ModelOption = z.infer<typeof modelOptionSchema>;
+
 export const sessionResponseSchema = z.object({
   sessionId: z.string(),
   channel: z.string(),
+  availableModels: z.array(modelOptionSchema).optional(),
+  currentModelId: z.string().optional(),
 });
 
 export type SessionResponse = z.infer<typeof sessionResponseSchema>;
@@ -220,3 +231,10 @@ export const sessionInfoSchema = z.object({
 });
 
 export const listSessionsOutput = z.array(sessionInfoSchema);
+
+export const getGatewayModelsInput = z.object({
+  apiHost: z.string(),
+  apiKey: z.string(),
+});
+
+export const getGatewayModelsOutput = z.array(modelOptionSchema);

apps/twig/src/main/services/agent/service.ts

@@ -10,7 +10,14 @@ import {
   type RequestPermissionRequest,
   type RequestPermissionResponse,
 } from "@agentclientprotocol/sdk";
-import { Agent, getLlmGatewayUrl, type OnLogCallback } from "@posthog/agent";
+import { Agent } from "@posthog/agent/agent";
+import {
+  fetchGatewayModels,
+  formatGatewayModelName,
+  getProviderName,
+} from "@posthog/agent/gateway-models";
+import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
+import type { OnLogCallback } from "@posthog/agent/types";
 import { app } from "electron";
 import { injectable } from "inversify";
 import type { AcpMessage } from "../../../shared/types/session-events.js";
@@ -152,6 +159,12 @@ interface ManagedSession {
   needsRecreation: boolean;
   promptPending: boolean;
   pendingContext?: string;
+  availableModels?: Array<{
+    modelId: string;
+    name: string;
+    description?: string | null;
+  }>;
+  currentModelId?: string;
 }
 
 function getClaudeCliPath(): string {
@@ -383,25 +396,43 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
 
       const mcpServers = this.buildMcpServers(credentials);
 
+      let availableModels:
+        | Array<{ modelId: string; name: string; description?: string | null }>
+        | undefined;
+      let currentModelId: string | undefined;
+
       if (isReconnect) {
-        await connection.extMethod("_posthog/session/resume", {
-          sessionId: taskRunId,
-          cwd: repoPath,
-          mcpServers,
-          _meta: {
-            ...(logUrl && {
-              persistence: { taskId, runId: taskRunId, logUrl },
-            }),
-            ...(sdkSessionId && { sdkSessionId }),
-            ...(additionalDirectories?.length && {
-              claudeCode: {
-                options: { additionalDirectories },
-              },
-            }),
+        const resumeResponse = await connection.extMethod(
+          "_posthog/session/resume",
+          {
+            sessionId: taskRunId,
+            cwd: repoPath,
+            mcpServers,
+            _meta: {
+              ...(logUrl && {
+                persistence: { taskId, runId: taskRunId, logUrl },
+              }),
+              ...(sdkSessionId && { sdkSessionId }),
+              ...(additionalDirectories?.length && {
+                claudeCode: {
+                  options: { additionalDirectories },
+                },
+              }),
+            },
           },
-        });
+        );
+        const resumeMeta = resumeResponse?._meta as
+          | {
+              models?: {
+                availableModels?: typeof availableModels;
+                currentModelId?: string;
+              };
+            }
+          | undefined;
+        availableModels = resumeMeta?.models?.availableModels;
+        currentModelId = resumeMeta?.models?.currentModelId;
       } else {
-        await connection.newSession({
+        const newSessionResponse = await connection.newSession({
           cwd: repoPath,
           mcpServers,
           _meta: {
@@ -415,6 +446,8 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
             }),
           },
         });
+        availableModels = newSessionResponse.models?.availableModels;
+        currentModelId = newSessionResponse.models?.currentModelId;
       }
 
       const session: ManagedSession = {
@@ -430,6 +463,8 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
         config,
         needsRecreation: false,
         promptPending: false,
+        availableModels,
+        currentModelId,
       };
 
       this.sessions.set(taskRunId, session);
@@ -999,7 +1034,12 @@ For git operations while detached:
   }
 
   private toSessionResponse(session: ManagedSession): SessionResponse {
-    return { sessionId: session.taskRunId, channel: session.channel };
+    return {
+      sessionId: session.taskRunId,
+      channel: session.channel,
+      availableModels: session.availableModels,
+      currentModelId: session.currentModelId,
+    };
   }
 
   /**
@@ -1109,4 +1149,16 @@ For git operations while detached:
       log.debug("Error in PR URL detection", { taskRunId, error: err });
     }
   }
+
+  async getGatewayModels(apiHost: string, _apiKey: string) {
+    const gatewayUrl = getLlmGatewayUrl(apiHost);
+    const models = await fetchGatewayModels({ gatewayUrl });
+
+    return models.map((model) => ({
+      modelId: model.id,
+      name: formatGatewayModelName(model),
+      description: `Context: ${model.context_window.toLocaleString()} tokens`,
+      provider: getProviderName(model.owned_by),
+    }));
+  }
 }

apps/twig/src/main/services/folders/service.ts

@@ -2,7 +2,7 @@ import { exec } from "node:child_process";
 import fs from "node:fs";
 import path from "node:path";
 import { promisify } from "node:util";
-import { WorktreeManager } from "@posthog/agent";
+import { WorktreeManager } from "@posthog/agent/worktree-manager";
 import { dialog } from "electron";
 import { injectable } from "inversify";
 import { generateId } from "../../../shared/utils/id.js";

apps/twig/src/main/services/workspace/service.ts

@@ -3,7 +3,7 @@ import * as fs from "node:fs";
 import * as fsPromises from "node:fs/promises";
 import path from "node:path";
 import { promisify } from "node:util";
-import { WorktreeManager } from "@posthog/agent";
+import { WorktreeManager } from "@posthog/agent/worktree-manager";
 import { inject, injectable } from "inversify";
 import type {
   TaskFolderAssociation,

apps/twig/src/main/trpc/routers/agent.ts

@@ -6,6 +6,8 @@ import {
   cancelPermissionInput,
   cancelPromptInput,
   cancelSessionInput,
+  getGatewayModelsInput,
+  getGatewayModelsOutput,
   listSessionsInput,
   listSessionsOutput,
   notifySessionContextInput,
@@ -140,4 +142,11 @@ export const agentRouter = router({
   markAllForRecreation: publicProcedure.mutation(() =>
     getService().markAllSessionsForRecreation(),
   ),
+
+  getGatewayModels: publicProcedure
+    .input(getGatewayModelsInput)
+    .output(getGatewayModelsOutput)
+    .query(({ input }) =>
+      getService().getGatewayModels(input.apiHost, input.apiKey),
+    ),
 });

apps/twig/src/main/utils/store.ts

@@ -1,4 +1,4 @@
-import { WorktreeManager } from "@posthog/agent";
+import { WorktreeManager } from "@posthog/agent/worktree-manager";
 import { app } from "electron";
 import Store from "electron-store";
 import type {

apps/twig/src/renderer/features/sessions/components/ModelSelector.tsx

@@ -1,11 +1,6 @@
-import { useSettingsStore } from "@features/settings/stores/settingsStore";
 import { Select, Text } from "@radix-ui/themes";
-import {
-  AVAILABLE_MODELS,
-  getModelsByProvider,
-  type ModelProvider,
-} from "@shared/types/models";
 import { Fragment } from "react";
+import { useModelsStore } from "../stores/modelsStore";
 import { useSessionActions, useSessionForTask } from "../stores/sessionStore";
 
 interface ModelSelectorProps {
@@ -19,31 +14,26 @@ export function ModelSelector({
   disabled,
   onModelChange,
 }: ModelSelectorProps) {
-  const defaultModel = useSettingsStore((state) => state.defaultModel);
-  const setDefaultModel = useSettingsStore((state) => state.setDefaultModel);
   const { setSessionModel } = useSessionActions();
   const session = useSessionForTask(taskId);
 
-  // Use session model if available, otherwise fall back to default
-  const activeModel = session?.model ?? defaultModel;
+  const groupedModels = useModelsStore((s) => s.groupedModels);
+  const models = useModelsStore((s) => s.models);
+  const selectedModel = useModelsStore((s) => s.selectedModel);
+  const setSelectedModel = useModelsStore((s) => s.setSelectedModel);
+
+  const activeModel = session?.model ?? selectedModel;
 
   const handleChange = (value: string) => {
-    // Always update the default
-    setDefaultModel(value);
+    setSelectedModel(value);
     onModelChange?.(value);
 
-    // If there's an active session, update the model mid-session
     if (taskId && session?.status === "connected" && !session.isCloud) {
       setSessionModel(taskId, value);
     }
   };
 
-  const modelsByProvider = getModelsByProvider();
-  const providers = (Object.keys(modelsByProvider) as ModelProvider[]).filter(
-    (provider) => modelsByProvider[provider].models.length > 0,
-  );
-
-  const currentModel = AVAILABLE_MODELS.find((m) => m.id === activeModel);
+  const currentModel = models.find((m) => m.modelId === activeModel);
   const displayName = currentModel?.name ?? activeModel;
 
   return (
@@ -69,13 +59,13 @@ export function ModelSelector({
         </Text>
       </Select.Trigger>
       <Select.Content position="popper" sideOffset={4}>
-        {providers.map((provider, index) => (
-          <Fragment key={provider}>
+        {groupedModels.map((group, index) => (
+          <Fragment key={group.provider}>
             {index > 0 && <Select.Separator />}
             <Select.Group>
-              <Select.Label>{modelsByProvider[provider].name}</Select.Label>
-              {modelsByProvider[provider].models.map((model) => (
-                <Select.Item key={model.id} value={model.id}>
+              <Select.Label>{group.provider}</Select.Label>
+              {group.models.map((model) => (
+                <Select.Item key={model.modelId} value={model.modelId}>
                   {model.name}
                 </Select.Item>
               ))}