[fix]: make `count()` respect `nth()` and `first()`

AB
1d ago
+52-4

Hey Stagehand team, hope you're doing well!

Asking a narrowed locator like locator("li").nth(3) how many elements it matches returned the count for the whole list, and this PR makes it return 1 when that element exists and 0 when it does not.

What this code is for

Locator.count() in packages/extension/understudy/locator.ts answers locator.count calls from the SDKs. The TypeScript and Python SDKs send the selector plus an optional nth over the protocol. resolveLocator in runtime.ts turns that into page.deepLocator(selector).nth(nth), and DeepLocatorDelegate.count() builds the real Locator with that index and calls its count(). So every SDK call, and every deep or iframe locator, ends up in this one method.

Before this PR

A script loads a page with five <li> items and calls page.locator("li").nth(3).count().

  1. The SDK sends locator.count with selector: "li" and nth: 3.
  2. resolveLocator returns a deep locator narrowed to index 3, which builds a Locator with nthIndex set to 3.
  3. count() calls this.selectorResolver.count(this.selectorQuery), which counts every li on the page: 5.
  4. count() returns that 5 directly and never looks at nthIndex.
  5. The script gets 5. nth(10).count() and first().count() also return 5, so a check like if (await item.count() > 0) passes for an index that does not exist.

After this PR

Same script, same call.

  1. Same as above through step 3: the selector still matches 5 elements.
  2. count() now checks nthIndex. It is 3, and 3 is below 5, so it returns 1.
  3. nth(5).count() and nth(10).count() return 0, since that index has no element.
  4. first().count() returns 1 when anything matches and 0 on an empty page.
  5. page.locator("li").count() with no nth() still returns 5.

The diagram

flowchart LR

classDef box fill:#eef1ff,stroke:#4c5bd4,color:#111827
classDef bad fill:#ffd6d6,stroke:#cc0000,color:#111827,stroke-width:3px
classDef good fill:#c8f0c8,stroke:#008800,color:#111827,stroke-width:3px

subgraph BEFORE["BEFORE: nth(3).count() returns 5"]
direction TB
  B1["locator('li').nth(3).count()"] --> B2["resolveLocator applies nth(3)"]
  B2 --> B3["selectorResolver.count gives 5"]
  B3 --> B4{"nthIndex is never read"}
  B4 -->|"nth(3), nth(10), first()"| B5["returns 5"]
  B4 -->|"no nth()"| B6["locator('li').count() returns 5"]
end

subgraph AFTER["AFTER: nth(3).count() returns 1"]
direction TB
  A1["locator('li').nth(3).count()"] --> A2["resolveLocator applies nth(3)"]
  A2 --> A3["selectorResolver.count gives 5"]
  A3 --> A4{"THE CHANGE: nthIndex < total"}
  A4 -->|"nth(3) gives 1, nth(10) gives 0"| A5["returns 1 or 0"]
  A4 -->|"no nth()"| A6["locator('li').count() returns 5"]
end

B1 ~~~ A1

class B1,B2,B3,B6 box
class A1,A2,A3,A6 box
class B4,B5 bad
class A4,A5 good

linkStyle 3 stroke:#cc0000,stroke-width:5px
linkStyle 8 stroke:#008800,stroke-width:5px

style BEFORE fill:#fff5f5,stroke:#cc0000,color:#111827
style AFTER fill:#f3fff3,stroke:#008800,color:#111827

The change

Locator.count() keeps the existing selectorResolver.count call and then reads nthIndex: when it is -1 (no nth() was called) it returns the total as before, otherwise it returns nthIndex < total ? 1 : 0. That is the whole fix, 3 lines. A changeset marks it as a patch for the extension and both SDKs that bundle it.

What changes for a user

  • locator(...).count() without nth() returns exactly what it did before, for CSS, XPath and text= selectors.
  • nth(i).count() and first().count() now return 1 or 0, which is what Playwright returns for a narrowed locator and what the docs describe as "the elements matched by this locator".
  • Every SDK gets the fix with no SDK code change, since the TypeScript and Python SDKs both send locator.count to the extension and never count on their own.
  • Actions on a narrowed locator (textContent, isVisible, click and the rest) already used nthIndex through resolveNode, so count() now agrees with them.

why

count() was the only Locator method that ignored nth(), so it disagreed with textContent() and isVisible() on the same narrowed locator.

what changed

packages/extension/understudy/locator.ts: count() returns 1 or 0 for a locator narrowed by nth() or first(). Plus one unit test table, one integration test, and a changeset.

test plan

  • New unit cases in locatorActions.test.ts for no nth() (5 and 0 matches), nth(0), nth(3), nth(4) (last), nth(5) (equal to the total), nth(10), and nth(0) on 0 matches. Against unpatched main, the 5 narrowed cases fail with the total returned; with the fix all 159 tests in the file pass.
  • New integration test in locatorCount.test.ts against a real headless Chrome with five <li> items, covering CSS, XPath and text= selectors plus first() on no matches. Before the fix it fails with expected 5 to be 1; after, locatorCount and locatorNth pass (16 of 16).
  • pnpm --filter ./packages/extension test:unit: 1198 passed. pnpm typecheck, extension and SDK typecheck, pnpm lint (no new findings), oxfmt --check, and check-changesets.ts all pass.

I'd love your view on one thing: I followed Playwright and return 1 or 0 for a narrowed locator. Would you rather count() keep reporting the total, with the docs updated instead? Happy to switch if that fits Stagehand better.

This touches the same code as two open PRs, and I want to make sure it plays nicely with both:

  • #3417 by @brennanbutler01 adds negative nth() (thank you, separating hasExplicitNth from the -1 sentinel is a really nice cleanup). The two merge without a git conflict, but they do interact: with #3417, nth(-1) keeps nthIndex at -1, so the nthIndex < 0 check here would return the total for it. Whichever lands second should check !this.hasExplicitNth instead, and count a negative index as 1 when its absolute value is at most the total and 0 otherwise. I'm glad to do that rebase and add the negative cases in whichever order you prefer.
  • #3420 by arunjmoorthy changes the two enable lines right above this one in count() and adds a test at the same spot in locatorCount.test.ts, so whichever lands second needs a small, mechanical rebase. Happy to take that on.

Any feedback or tips are very welcome, thank you for taking a look!


Summary by cubic

Fixes Locator.count() so a locator narrowed with nth() or first() returns 1 when the index exists and 0 when it does not, instead of the count for the whole selector match. count() was the only Locator method that ignored nth(), so it disagreed with textContent() and isVisible() on the same narrowed locator.

Behavior

  • count() without nth() still returns the total match count.
  • The TypeScript and Python SDKs get the fix without changes, since they delegate locator.count to the extension.
  • The count() docs in packages/docs/v4/reference/locator.mdx now describe the narrowed behavior for the TypeScript, Python, and Go SDKs.

Testing

  • Unit table in locatorActions.test.ts covers no nth(), nth(0), nth(3), nth(4), nth(5), nth(10), and 0 matches.
  • Integration test in locatorCount.test.ts covers CSS, XPath, and text= selectors plus first() on no matches.