fix: missing URLs for extract() with array schema

[seanmcguire12]

why

• before this change, when we convert z.string().url() to an ID, if it was inside a z.array(), it was not getting converted back into a URL
• this meant that if you defined a schema like this:

schema: z.object({
  records: z.array(z.string().url()),
})

you would receive an array like this:

{
  records: [
    '0-302', '0-309',
    '0-316', '0-323',
    '0-330', '0-337',
    '0-344', '0-351',
    '0-358', '0-365'
  ]
}

• with this change, you will now receive the actual URLs, ie:

{
  records: [
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10003-10041.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10004-10143%20(C06932208).pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10004-10143.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10004-10156.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10004-10213.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10005-10321.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10006-10247.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10007-10345.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10009-10021.pdf',
    'https://www.archives.gov/files/research/jfk/releases/2025/0318/104-10009-10222.pdf'
  ]
}

what changed

• updated the injectUrls function so that when it hits an array and there is not deeper path, it loops through the array and injects the URLs

test plan

• evals

[changeset-bot]

🦋 Changeset detected

Latest commit: 6d6c7aa

The changes in this PR will be included in the next version bump.

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

[greptile-apps]

Greptile Overview

Summary

This PR fixes a critical bug in Stagehand's URL extraction system for array schemas. The issue occurred when using `z.array(z.string().url())` in extraction schemas - instead of returning actual URLs, users received numeric identifiers like '0-302', '0-309', etc.

The bug was in the injectUrls function in lib/utils.ts, which handles converting numeric IDs back to URLs after LLM extraction. Stagehand's extraction pipeline works by first transforming URL fields into numeric IDs (making them easier for LLMs to process), then converting those IDs back to actual URLs using a mapping. While the schema transformation correctly identified URL fields in arrays and created appropriate path segments like ['*'], the injectUrls function wasn't handling the terminal case where the path ends at an array wildcard.

The fix adds logic to handle this scenario: when the path terminates at an array (rest.length === 0), the function now loops through each array element and converts any valid IDs back to their corresponding URLs. Additionally, a toId helper function was extracted to standardize the ID detection logic that was previously duplicated.

This change integrates seamlessly with Stagehand's existing URL extraction architecture and ensures that array-based URL schemas work as expected, maintaining consistency with non-array URL field behavior.

Important Files Changed

Changed Files

Filename	Score	Overview
lib/utils.ts	4/5	Fixed URL injection for array schemas by adding array element processing logic and extracting a toId helper function

Confidence score: 4/5

• This PR is safe to merge with minimal risk as it fixes a clear bug without affecting existing functionality
• Score reflects focused bug fix with good code organization and clear understanding of the existing system
• No files require special attention as the change is contained and well-implemented

Sequence Diagram

sequenceDiagram
    participant User
    participant ExtractHandler as "StagehandExtractHandler"
    participant Utils as "utils.ts"
    participant Schema as "Zod Schema"
    participant LLM as "LLM Client"

    User->>ExtractHandler: "extract() with array schema containing z.string().url()"
    ExtractHandler->>ExtractHandler: "domExtract()"
    ExtractHandler->>Utils: "transformSchema(schema, [])"
    
    Note over Utils: Transform z.string().url() to z.number() in arrays
    Utils->>Utils: "isKind(schema, Kind.ZodArray)"
    Utils->>Utils: "transformSchema(itemType, [...currentPath, '*'])"
    Utils->>Utils: "Replace z.string().url() with z.number()"
    Utils-->>ExtractHandler: "[transformedSchema, urlPaths]"
    
    ExtractHandler->>LLM: "extract data using transformed schema"
    LLM-->>ExtractHandler: "extracted data with numeric IDs"
    
    Note over ExtractHandler: Data contains numeric IDs like ['0-302', '0-309']
    
    ExtractHandler->>Utils: "injectUrls(extractedData, urlPaths, idToUrlMapping)"
    Utils->>Utils: "Check if path[0] === '*' and rest.length === 0"
    
    Note over Utils: NEW: Handle arrays directly when no deeper path
    Utils->>Utils: "Loop through array items"
    Utils->>Utils: "Replace numeric IDs with actual URLs"
    
    Utils-->>ExtractHandler: "Data with URLs injected"
    ExtractHandler-->>User: "Array with actual URLs instead of numeric IDs"

Loading

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}