Fix schema parsing bug for Pydantic .model_json_schema() (browserbase#1873)

# why
JSON schema parsing for schemas generated from Pydantic
`.model_json_schema()` were dropping the nested references

# what changed
Before (from test script) with schema: 
```python
class RequiredInputType(StrEnum):
    USERNAME = "username"
    EMAIL = "email"
    PHONE = "phone"
    OTP = "otp"
    PASSWORD = "password"
    CARD_LAST_8 = "card_last_8"
    SOMETHING_ELSE = "something_else"

class PageState(StrEnum):
    LOGIN_FORM = "login_form"
    ERROR = "error"
    AUTHENTICATED = "authenticated"
    UNKNOWN = "unknown"

class PageAnalysis(BaseModel):
    page_state: PageState
    required_input: RequiredInputType | None = None
    contact_id_of_required_input: str | None = None
    error_message: str | None = None
    reasoning: str = ""
```
<img width="654" height="373" alt="Screenshot 2026-03-23 at 1 09 01 PM"
src="https://github.com/user-attachments/assets/29a24eea-cfda-40a8-8e43-10f959c71ad2"
/>

After: 
<img width="650" height="310" alt="Screenshot 2026-03-23 at 1 06 17 PM"
src="https://github.com/user-attachments/assets/9f62d1ad-04bb-41e4-9dbc-85d7b68c84aa"
/>

1. Added resolveRefs() function that inlines $ref/$defs before Zod
conversion (Pydantic
  v2's model_json_schema() generates these)
2. Changed z.string().refine() → z.enum() so enum constraints are
expressed in the
  schema sent to the LLM, not just validated post-hoc


# test plan

<!-- This is an auto-generated description by cubic. -->
---
## Summary by cubic
Fixes missing nested `$ref` when converting Pydantic v2
`.model_json_schema()` to Zod in `@browserbasehq/stagehand-server-v3`,
and emits proper `z.enum()` types so LLM-facing schemas are accurate.
Adds a `test:unit` script, moves tests to `tests/`, and updates the
runner to scan directories or single files.

- **Bug Fixes**
- Resolves `$ref` from `$defs` before conversion via `resolveRefs()`
with a cycle guard, handling refs in objects, arrays,
`anyOf`/`oneOf`/`allOf`, and root-level refs (unknown refs fall back
safely).
- Replaces `z.string().refine()` with `z.enum()` so enum values are
surfaced in generated schemas and enforced.

<sup>Written for commit ae9e2e5.
Summary will update on new commits. <a
href="https://cubic.dev/pr/browserbase/stagehand/pull/1873">Review in
cubic</a></sup>

<!-- End of auto-generated description by cubic. -->

Author: miguelg719

.changeset/tough-buttons-return.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand-server-v3": patch
+---
+
+Fix schema parsing bug for Pydantic `.model_json_schema()` on missing nested references

packages/server-v3/package.json

@@ -17,6 +17,7 @@
     "test:server": "tsx scripts/test-server.ts",
     "test:integration": "pnpm run test:server -- packages/server-v3/dist/tests/integration",
     "test:integration:local": "STAGEHAND_SERVER_TARGET=local pnpm run test:server -- packages/server-v3/dist/tests/integration",
+    "test:unit": "pnpm run test:server -- packages/server-v3/dist/tests/unit",
     "test:integration:sea": "STAGEHAND_SERVER_TARGET=sea pnpm run test:server -- packages/server-v3/dist/tests/integration",
     "gen:openapi": "tsx scripts/gen-openapi.ts"
   },

packages/server-v3/scripts/test-server.ts

@@ -93,7 +93,7 @@ const writeCtrfFromJunit = (junitPath: string, tool: string) => {
   }
 };
 
-const sourceTestsDir = `${repoRoot}/packages/server-v3/test`;
+const sourceTestsDir = `${repoRoot}/packages/server-v3/tests`;
 const sourceUnitDir = `${sourceTestsDir}/unit`;
 const sourceIntegrationDir = `${sourceTestsDir}/integration`;
 const unitDir = `${repoRoot}/packages/server-v3/dist/tests/unit`;
@@ -224,7 +224,12 @@ const nodeOptions = [process.env.NODE_OPTIONS, baseNodeOptions]
 
 const allPaths =
   paths.length > 0
-    ? paths.map(resolveRepoRelative)
+    ? paths.flatMap((p) => {
+        const abs = resolveRepoRelative(p);
+        return fs.existsSync(abs) && fs.statSync(abs).isDirectory()
+          ? collectFiles(abs, ".test.js")
+          : [abs];
+      })
     : [
         ...collectFiles(unitDir, ".test.js"),
         ...collectFiles(integrationDir, ".test.js"),

packages/server-v3/src/lib/utils.ts

@@ -18,6 +18,62 @@ interface JSONSchema {
   anyOf?: JSONSchema[];
   oneOf?: JSONSchema[];
   allOf?: JSONSchema[];
+  $defs?: Record<string, JSONSchema>;
+  $ref?: string;
+}
+
+/**
+ * Resolves all $ref pointers in a JSON Schema against its $defs, returning
+ * a self-contained schema with no $ref or $defs remaining.
+ */
+function resolveRefs(
+  schema: JSONSchema,
+  defs: Record<string, JSONSchema>,
+  seen: Set<string> = new Set(),
+): JSONSchema {
+  if (schema.$ref) {
+    const match = schema.$ref.match(/^#\/\$defs\/(.+)$/);
+    if (match && match[1] && defs[match[1]]) {
+      if (seen.has(match[1])) {
+        return {};
+      }
+      seen.add(match[1]);
+      return resolveRefs(defs[match[1]], defs, seen);
+    }
+    return {};
+  }
+
+  const resolved: JSONSchema = { ...schema };
+
+  if (resolved.properties) {
+    const props: Record<string, JSONSchema> = {};
+    for (const [key, val] of Object.entries(resolved.properties)) {
+      props[key] = resolveRefs(val, defs, new Set(seen));
+    }
+    resolved.properties = props;
+  }
+  if (resolved.items) {
+    resolved.items = resolveRefs(resolved.items, defs, new Set(seen));
+  }
+  if (resolved.anyOf) {
+    resolved.anyOf = resolved.anyOf.map((s) =>
+      resolveRefs(s, defs, new Set(seen)),
+    );
+  }
+  if (resolved.oneOf) {
+    resolved.oneOf = resolved.oneOf.map((s) =>
+      resolveRefs(s, defs, new Set(seen)),
+    );
+  }
+  if (resolved.allOf) {
+    resolved.allOf = resolved.allOf.map((s) =>
+      resolveRefs(s, defs, new Set(seen)),
+    );
+  }
+
+  delete resolved.$defs;
+  delete resolved.$ref;
+  return resolved;
 }
 
 /**
@@ -26,10 +82,19 @@ interface JSONSchema {
  * @returns A Zod schema equivalent to the input JSON Schema
  */
 export function jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
+  // Resolve $ref/$defs before converting so all types are inlined
+  const resolved =
+    schema.$defs || schema.$ref
+      ? resolveRefs(schema, schema.$defs ?? {})
+      : schema;
+  return _jsonSchemaToZod(resolved);
+}
+
+function _jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
   if (Array.isArray(schema.type)) {
     const subSchemas = schema.type.map((singleType) => {
       const sub = { ...schema, type: singleType };
-      return jsonSchemaToZod(sub);
+      return _jsonSchemaToZod(sub);
     });
 
     if (subSchemas.length === 0) {
@@ -45,7 +110,7 @@ export function jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
   }
 
   if (schema.anyOf && Array.isArray(schema.anyOf)) {
-    const subSchemas = schema.anyOf.map((sub) => jsonSchemaToZod(sub));
+    const subSchemas = schema.anyOf.map((sub) => _jsonSchemaToZod(sub));
     if (subSchemas.length === 0) {
       return z.any();
     } else if (subSchemas.length === 1) {
@@ -59,7 +124,7 @@ export function jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
   }
 
   if (schema.oneOf && Array.isArray(schema.oneOf)) {
-    const subSchemas = schema.oneOf.map((sub) => jsonSchemaToZod(sub));
+    const subSchemas = schema.oneOf.map((sub) => _jsonSchemaToZod(sub));
     if (subSchemas.length === 0) {
       return z.any();
     } else if (subSchemas.length === 1) {
@@ -84,7 +149,7 @@ export function jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
               StatusCodes.BAD_REQUEST,
             );
           }
-          shape[key] = jsonSchemaToZod(subSchema);
+          shape[key] = _jsonSchemaToZod(subSchema);
         }
         let zodObject = z.object(shape);
 
@@ -109,7 +174,7 @@ export function jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
 
     case "array":
       if (schema.items) {
-        let zodArray = z.array(jsonSchemaToZod(schema.items));
+        let zodArray = z.array(_jsonSchemaToZod(schema.items));
         if (schema.description) {
           zodArray = zodArray.describe(schema.description);
         }
@@ -118,8 +183,8 @@ export function jsonSchemaToZod(schema: JSONSchema): ZodTypeAny {
       return z.array(z.any());
 
     case "string": {
-      if (schema.enum) {
-        return z.string().refine((val) => schema.enum?.includes(val) ?? false);
+      if (schema.enum && schema.enum.length > 0) {
+        return z.enum(schema.enum as [string, ...string[]]);
       }
       let zodString = z.string();