2/2 Implement explicit model gateway

packages/sdk-python/src/stagehand/client_models.py

@@ -8,6 +8,7 @@
 from ._generated import models as _models
 from ._generated.models import (
     BrowserbaseBrowserSettings,
+    BrowserbaseModelConfig,
     BrowserbaseRegion,
     CustomModelConfig,
     DirectModelConfig,
@@ -155,6 +156,13 @@ def _model_config(
         if value is not None
     }
     if base_url is None:
+        if model.startswith("browserbase/"):
+            return ModelConfig(
+                root=BrowserbaseModelConfig.model_validate({
+                    "model_name": model,
+                    **connection,
+                })
+            )
         return ModelConfig(
             root=DirectModelConfig.model_validate({"model_name": model, **connection})
         )

packages/sdk-python/tests/test_stagehand.py

@@ -13,6 +13,7 @@
     Action,
     ActResult,
     ActResultData,
+    BrowserbaseModelConfig,
     BrowserGetVersionResult,
     ClientModelReference,
     DirectModelConfig,
@@ -92,6 +93,20 @@ def test_stagehand_constructor_builds_private_browser_and_model_models() -> None
         _ = cdp.browser
 
 
+@pytest.mark.parametrize(
+    "model_name",
+    ["browserbase/auto", "browserbase/openai/future-model"],
+)
+def test_stagehand_constructor_routes_browserbase_models_to_the_gateway(
+    model_name: str,
+) -> None:
+    stagehand = Stagehand(api_key="browserbase-key", model=model_name)
+
+    assert isinstance(stagehand.init_params.model, ModelConfig)
+    assert isinstance(stagehand.init_params.model.root, BrowserbaseModelConfig)
+    assert stagehand.init_params.model.root.model_name.model_dump() == model_name
+
+
 def test_stagehand_constructor_rejects_incomplete_flattened_options() -> None:
     with pytest.raises(TypeError, match="viewport_width and viewport_height"):
         Stagehand(browser="local", viewport_width=1280)

packages/server/clients/cacheClient.ts

@@ -1,6 +1,5 @@
 import type { Protocol } from "devtools-protocol";
 import { z } from "zod/v4";
-import type { BrowserbaseRegion } from "../../protocol/types.js";
 
 /**
  * HTTP client for the Stagehand API's stateless cache routes (POST /cache/get
@@ -12,22 +11,6 @@ import type { BrowserbaseRegion } from "../../protocol/types.js";
  * server-side contract.
  */
 
-/** Stagehand API base URL per Browserbase region; sessions must hit the API
- * deployment in the region their browser runs in. Mirrors the v3 SDK's
- * REGION_API_URLS. */
-const REGION_API_URLS: Record<BrowserbaseRegion, string> = {
-  "us-west-2": "https://api.stagehand.browserbase.com",
-  "us-east-1": "https://api.use1.stagehand.browserbase.com",
-  "eu-central-1": "https://api.euc1.stagehand.browserbase.com",
-  "ap-southeast-1": "https://api.apse1.stagehand.browserbase.com",
-};
-
-const DEFAULT_REGION: BrowserbaseRegion = "us-west-2";
-
-export function apiUrlForRegion(region: BrowserbaseRegion | undefined): string {
-  return `${REGION_API_URLS[region ?? DEFAULT_REGION]}/v1`;
-}
-
 /** Ceiling for a single cache HTTP round-trip so a slow cache can only ever
  * add a bounded delay to the action it fronts. */
 const CACHE_REQUEST_TIMEOUT_MS = 5_000;

packages/server/clients/stagehandApi.ts

@@ -0,0 +1,16 @@
+import type { BrowserbaseRegion } from "../../protocol/types.js";
+
+/** Stagehand API base URL per Browserbase region. */
+const REGION_API_URLS: Record<BrowserbaseRegion, string> = {
+  "us-west-2": "https://api.stagehand.browserbase.com",
+  "us-east-1": "https://api.use1.stagehand.browserbase.com",
+  "eu-central-1": "https://api.euc1.stagehand.browserbase.com",
+  "ap-southeast-1": "https://api.apse1.stagehand.browserbase.com",
+};
+
+const DEFAULT_REGION: BrowserbaseRegion = "us-west-2";
+
+/** Resolves the regional Stagehand API endpoint for a Browserbase session. */
+export function apiUrlForRegion(region: BrowserbaseRegion | undefined): string {
+  return `${REGION_API_URLS[region ?? DEFAULT_REGION]}/v1`;
+}

packages/server/controllers/stagehandController.ts

@@ -10,6 +10,7 @@ import type { StagehandRuntime } from "../runtime.js";
 import * as actService from "../services/actService.js";
 import * as cacheService from "../services/cacheService.js";
 import * as extractService from "../services/extractService.js";
+import { buildGatewayContext } from "../llm/gatewayClient.js";
 import * as observeService from "../services/observeService.js";
 
 export function createStagehandController(runtime: StagehandRuntime) {
@@ -46,6 +47,7 @@ export function createStagehandController(runtime: StagehandRuntime) {
       selfHeal: state.initParams.selfHeal,
       domSettleTimeoutMs: state.initParams.domSettleTimeoutMs,
       cache: cacheService.buildCacheContext(state.initParams),
+      gateway: buildGatewayContext(state.initParams),
     });
   }
 
@@ -69,6 +71,7 @@ export function createStagehandController(runtime: StagehandRuntime) {
       logger,
       systemPrompt: state.initParams.systemPrompt,
       cache: cacheService.buildCacheContext(state.initParams),
+      gateway: buildGatewayContext(state.initParams),
     });
   }
 
@@ -92,6 +95,7 @@ export function createStagehandController(runtime: StagehandRuntime) {
       logger,
       systemPrompt: state.initParams.systemPrompt,
       cache: cacheService.buildCacheContext(state.initParams),
+      gateway: buildGatewayContext(state.initParams),
     });
   }
 

packages/server/llm/aiSdkClient.ts

@@ -13,11 +13,10 @@ import {
   LLMGenerateResultSchema,
   ModelProviderSchema,
 } from "../../protocol/schemas.js";
-import type {
-  DirectModelConfig,
-  LLMGenerateParams,
-  LLMGenerateResult,
-} from "../../protocol/types.js";
+import type { LLMGenerateParams, LLMGenerateResult } from "../../protocol/types.js";
+import type { DirectModelTarget } from "./modelTarget.js";
+
+type DirectModelConnection = Omit<DirectModelTarget, "type">;
 
 const AiSdkMessagesSchema = z.array(LLMMessageSchema).transform((messages): ModelMessage[] => {
   const toolNames = new Map<string, string>();
@@ -123,7 +122,7 @@ const AiSdkGenerationSchema = z
   .strip();
 
 /** Creates a direct AI SDK model from a validated Stagehand model configuration. */
-export function createAiSdkLanguageModel(config: DirectModelConfig): LanguageModel {
+export function createAiSdkLanguageModel(config: DirectModelConnection): LanguageModel {
   const separator = config.modelName.indexOf("/");
   const provider = ModelProviderSchema.parse(config.modelName.slice(0, separator));
   const modelId = config.modelName.slice(separator + 1);
@@ -134,7 +133,7 @@ export function createAiSdkLanguageModel(config: DirectModelConfig): LanguageMod
 
   switch (provider) {
     case "openai":
-      return createOpenAI(connection).responses(modelId);
+      return createOpenAI(connection).chat(modelId);
     case "anthropic":
       return createAnthropic(connection)(modelId);
     case "google":

packages/server/llm/gatewayClient.ts

@@ -0,0 +1,37 @@
+import { createOpenAI } from "@ai-sdk/openai";
+import type { LanguageModel } from "ai";
+import type { StagehandInitParams } from "../../protocol/types.js";
+import { apiUrlForRegion } from "../clients/stagehandApi.js";
+import type { BrowserbaseModelTarget } from "./modelTarget.js";
+
+export type GatewayContext = {
+  apiUrl: string;
+  apiKey: string;
+  sessionId: string;
+};
+
+/** Builds the Browserbase gateway context for a Browserbase-backed session. */
+export function buildGatewayContext(initParams: StagehandInitParams): GatewayContext | undefined {
+  const sessionId = initParams.browser?.sessionId;
+  if (!initParams.apiKey || !sessionId) return undefined;
+  return {
+    apiUrl: apiUrlForRegion(initParams.browser?.region),
+    apiKey: initParams.apiKey,
+    sessionId,
+  };
+}
+
+/** Creates an AI SDK Responses model backed by Browserbase Model Gateway. */
+export function createGatewayLanguageModel(
+  target: BrowserbaseModelTarget,
+  gateway: GatewayContext,
+): LanguageModel {
+  return createOpenAI({
+    apiKey: gateway.apiKey,
+    baseURL: `${gateway.apiUrl}/llm`,
+    headers: {
+      "x-bb-api-key": gateway.apiKey,
+      "x-bb-session-id": gateway.sessionId,
+    },
+  }).responses(target.modelName);
+}

packages/server/llm/modelTarget.ts

@@ -0,0 +1,71 @@
+import type {
+  BrowserbaseModelConfig,
+  CustomModelConfig,
+  DirectModelConfig,
+  ModelConfig,
+} from "../../protocol/types.js";
+
+export type DirectModelTarget = {
+  type: "direct";
+  modelName: string;
+  apiKey: string;
+  headers?: Record<string, string>;
+};
+
+export type BrowserbaseModelTarget = {
+  type: "browserbase";
+  modelName: string;
+};
+
+export type OpenAICompatibleModelTarget = {
+  type: "openai-compatible";
+  modelName: string;
+  baseURL: string;
+  apiKey?: string;
+  headers?: Record<string, string>;
+};
+
+export type ModelTarget = DirectModelTarget | BrowserbaseModelTarget | OpenAICompatibleModelTarget;
+
+function isBrowserbaseModel(config: ModelConfig): config is BrowserbaseModelConfig {
+  return config.modelName.startsWith("browserbase/");
+}
+
+function isOpenAICompatibleModel(config: ModelConfig): config is CustomModelConfig {
+  return "baseURL" in config;
+}
+
+function isDirectModel(config: ModelConfig): config is DirectModelConfig {
+  return "apiKey" in config && !isOpenAICompatibleModel(config);
+}
+
+/** Resolves the public model-name namespace into the server's execution target. */
+export function resolveModelTarget(config: ModelConfig): ModelTarget {
+  if (isOpenAICompatibleModel(config)) {
+    return {
+      type: "openai-compatible",
+      modelName: config.modelName,
+      baseURL: config.baseURL,
+      apiKey: config.apiKey,
+      headers: config.headers,
+    };
+  }
+
+  if (isBrowserbaseModel(config)) {
+    return {
+      type: "browserbase",
+      modelName: config.modelName.slice("browserbase/".length),
+    };
+  }
+
+  if (isDirectModel(config)) {
+    return {
+      type: "direct",
+      modelName: config.modelName,
+      apiKey: config.apiKey,
+      headers: config.headers,
+    };
+  }
+
+  throw new Error("Model configuration did not resolve to a supported inference target");
+}

packages/server/llm/openAiCompatibleClient.ts

@@ -0,0 +1,15 @@
+import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
+import type { LanguageModel } from "ai";
+import type { OpenAICompatibleModelTarget } from "./modelTarget.js";
+
+/** Creates an OpenAI Chat Completions-compatible model for a caller-controlled endpoint. */
+export function createOpenAICompatibleLanguageModel(
+  target: OpenAICompatibleModelTarget,
+): LanguageModel {
+  return createOpenAICompatible({
+    name: "openai-compatible",
+    apiKey: target.apiKey,
+    baseURL: target.baseURL,
+    headers: target.headers,
+  }).chatModel(target.modelName);
+}

packages/server/package.json

@@ -13,6 +13,7 @@
     "@ai-sdk/google": "catalog:",
     "@ai-sdk/groq": "catalog:",
     "@ai-sdk/openai": "catalog:",
+    "@ai-sdk/openai-compatible": "catalog:",
     "@ai-sdk/provider": "catalog:",
     "@opentelemetry/api": "catalog:",
     "@opentelemetry/core": "catalog:",

packages/server/runtime.ts

@@ -97,6 +97,7 @@ import { bytesToBase64 } from "./understudy/fileUploadUtils.js";
 import { createStore } from "zustand/vanilla";
 import type { StagehandLogEmitter } from "./logger.js";
 import { StagehandLogger } from "./logger.js";
+import { buildGatewayContext } from "./llm/gatewayClient.js";
 import * as llmService from "./services/llmService.js";
 import { StagehandRuntimeStateSchema, type StagehandRuntimeState } from "./runtimeState.js";
 import { createStagehandTracing, type StagehandTracing } from "./tracing.js";
@@ -317,12 +318,21 @@ export class StagehandRuntime {
 
   async generateLlm(input: LLMGenerateParams): Promise<LLMGenerateResult> {
     const state = this.state.getState();
-    const model = state.status === "initialized" ? state.initParams.model : undefined;
+    if (state.status !== "initialized") {
+      throw new Error("Stagehand must be initialized before generating LLM output");
+    }
+
+    const model = state.initParams.model;
     if (!model) {
       throw new Error("An LLM was not configured during Stagehand initialization");
     }
 
-    return await llmService.generate(model, input, this.adapters.clientLLMGenerate);
+    return await llmService.generate(
+      model,
+      input,
+      this.adapters.clientLLMGenerate,
+      buildGatewayContext(state.initParams),
+    );
   }
 
   async contextPages(): Promise<ContextPagesResult> {

packages/server/services/actService.ts

@@ -16,6 +16,7 @@ import { createTimeoutGuard } from "../handlers/handlerUtils/timeoutGuard.js";
 import { resolveVariableValue } from "../handlers/handlerUtils/variables.js";
 import * as inference from "../inference.js";
 import type { ClientLlmRequest } from "../llm/clientLlmClient.js";
+import type { GatewayContext } from "../llm/gatewayClient.js";
 import type { StagehandLogger } from "../logger.js";
 import { buildActPrompt, buildStepTwoPrompt } from "../prompt.js";
 import type { EncodedId } from "../types/private/internal.js";
@@ -38,6 +39,7 @@ type ActContext = {
   selfHeal: boolean;
   domSettleTimeoutMs?: number;
   ensureTimeRemaining: () => void;
+  gateway?: GatewayContext;
 };
 
 export async function act({
@@ -50,6 +52,7 @@ export async function act({
   selfHeal = false,
   domSettleTimeoutMs,
   cache,
+  gateway,
 }: {
   params: StagehandActParams;
   page: Page;
@@ -60,6 +63,7 @@ export async function act({
   selfHeal?: boolean;
   domSettleTimeoutMs?: number;
   cache?: cacheService.CacheContext;
+  gateway?: GatewayContext;
 }): Promise<ActResult> {
   const { input, options } = params;
   const variables = options?.variables;
@@ -74,6 +78,7 @@ export async function act({
     selfHeal,
     domSettleTimeoutMs,
     ensureTimeRemaining,
+    gateway,
   };
 
   ensureTimeRemaining();
@@ -234,7 +239,8 @@ async function getActionFromLLM({
   const response = await inference.act({
     instruction,
     domElements,
-    generate: (input) => llmService.generate(context.model, input, context.clientLLMGenerate),
+    generate: (input) =>
+      llmService.generate(context.model, input, context.clientLLMGenerate, context.gateway),
     userProvidedInstructions: context.systemPrompt,
   });
 

packages/server/services/cacheService.ts

@@ -9,12 +9,12 @@ import type {
 } from "../../protocol/types.js";
 import {
   CacheClient,
-  apiUrlForRegion,
   type CacheGetResponse,
   type CacheLlmUsage,
   type CacheMethod,
   type CdpTree,
 } from "../clients/cacheClient.js";
+import { apiUrlForRegion } from "../clients/stagehandApi.js";
 import type { StagehandLogger } from "../logger.js";
 import type { Frame } from "../understudy/frame.js";
 import { FrameSelectorResolver } from "../understudy/selectorResolver.js";