address feedback

src/snapshot.ts

@@ -168,14 +168,21 @@ export async function getEnhancedSnapshot(
  */
 interface RoleNameTracker {
   counts: Map<string, number>;
+  /** Maps role+name key to array of ref IDs that use it */
+  refsByKey: Map<string, string[]>;
   getKey(role: string, name?: string): string;
   getNextIndex(role: string, name?: string): number;
+  trackRef(role: string, name: string | undefined, ref: string): void;
+  /** Get all role+name keys that have duplicates */
+  getDuplicateKeys(): Set<string>;
 }
 
 function createRoleNameTracker(): RoleNameTracker {
   const counts = new Map<string, number>();
+  const refsByKey = new Map<string, string[]>();
   return {
     counts,
+    refsByKey,
     getKey(role: string, name?: string): string {
       return `${role}:${name ?? ''}`;
     },
@@ -185,6 +192,21 @@ function createRoleNameTracker(): RoleNameTracker {
       counts.set(key, current + 1);
       return current;
     },
+    trackRef(role: string, name: string | undefined, ref: string): void {
+      const key = this.getKey(role, name);
+      const refs = refsByKey.get(key) ?? [];
+      refs.push(ref);
+      refsByKey.set(key, refs);
+    },
+    getDuplicateKeys(): Set<string> {
+      const duplicates = new Set<string>();
+      for (const [key, refs] of refsByKey) {
+        if (refs.length > 1) {
+          duplicates.add(key);
+        }
+      }
+      return duplicates;
+    },
   };
 }
 
@@ -208,23 +230,28 @@ function processAriaTree(ariaTree: string, refs: RefMap, options: SnapshotOption
       if (INTERACTIVE_ROLES.has(roleLower)) {
         const ref = nextRef();
         const nth = tracker.getNextIndex(roleLower, name);
+        tracker.trackRef(roleLower, name, ref);
         refs[ref] = {
           selector: buildSelector(roleLower, name),
           role: roleLower,
           name,
-          // Only store nth if this is a duplicate (nth > 0)
-          ...(nth > 0 ? { nth } : {}),
+          nth, // Always store nth, we'll use it for duplicates
         };
 
         let enhanced = `- ${role}`;
         if (name) enhanced += ` "${name}"`;
         enhanced += ` [ref=${ref}]`;
+        // Only show nth in output if it's > 0 (for readability)
         if (nth > 0) enhanced += ` [nth=${nth}]`;
         if (suffix && suffix.includes('[')) enhanced += suffix;
 
         result.push(enhanced);
       }
     }
+
+    // Post-process: remove nth from refs that don't have duplicates
+    removeNthFromNonDuplicates(refs, tracker);
+
     return result.join('\n') || '(no interactive elements)';
   }
 
@@ -236,6 +263,9 @@ function processAriaTree(ariaTree: string, refs: RefMap, options: SnapshotOption
     }
   }
 
+  // Post-process: remove nth from refs that don't have duplicates
+  removeNthFromNonDuplicates(refs, tracker);
+
   // If compact mode, remove empty structural elements
   if (options.compact) {
     return compactTree(result.join('\n'));
@@ -244,6 +274,22 @@ function processAriaTree(ariaTree: string, refs: RefMap, options: SnapshotOption
   return result.join('\n');
 }
 
+/**
+ * Remove nth from refs that ended up not having duplicates
+ * This keeps single-element locators simple (no unnecessary .nth(0))
+ */
+function removeNthFromNonDuplicates(refs: RefMap, tracker: RoleNameTracker): void {
+  const duplicateKeys = tracker.getDuplicateKeys();
+
+  for (const [ref, data] of Object.entries(refs)) {
+    const key = tracker.getKey(data.role, data.name);
+    if (!duplicateKeys.has(key)) {
+      // Not a duplicate, remove nth to keep locator simple
+      delete refs[ref].nth;
+    }
+  }
+}
+
 /**
  * Get indentation level (number of spaces / 2)
  */
@@ -311,19 +357,20 @@ function processLine(
   if (shouldHaveRef) {
     const ref = nextRef();
     const nth = tracker.getNextIndex(roleLower, name);
+    tracker.trackRef(roleLower, name, ref);
 
     refs[ref] = {
       selector: buildSelector(roleLower, name),
       role: roleLower,
       name,
-      // Only store nth if this is a duplicate (nth > 0)
-      ...(nth > 0 ? { nth } : {}),
+      nth, // Always store nth, we'll clean up non-duplicates later
     };
 
     // Build enhanced line with ref
     let enhanced = `${prefix}${role}`;
     if (name) enhanced += ` "${name}"`;
     enhanced += ` [ref=${ref}]`;
+    // Only show nth in output if it's > 0 (for readability)
     if (nth > 0) enhanced += ` [nth=${nth}]`;
     if (suffix) enhanced += suffix;