fix(playground): chat turns show identical token counts across turns

web/packages/agenta-playground/src/state/execution/executionItems.ts

@@ -1466,9 +1466,14 @@ export const handleExecutionResultAtom = atom(
             // shared user message" is racy in comparison mode because the
             // first-to-finish variant auto-appends a blank user message
             // before the second variant's result arrives.
-            // rowId format in comparison mode: "turn-<entityUUID>-lt-<msgId>"
+            // rowId format in comparison mode: "turn-<entityUUID>-<logicalId>"
+            // logicalId is "msg-<uuid>" (current) or "lt-<id>" (legacy).
+            // Neither prefix can appear in a hex UUID, so searching for
+            // "-msg-" / "-lt-" unambiguously locates the logical boundary.
+            const msgIndex = rowId.indexOf("-msg-")
             const ltIndex = rowId.indexOf("-lt-")
-            const logicalRowId = ltIndex >= 0 ? rowId.slice(ltIndex + 1) : rowId
+            const sepIndex = msgIndex >= 0 ? msgIndex : ltIndex
+            const logicalRowId = sepIndex >= 0 ? rowId.slice(sepIndex + 1) : rowId
             const flatIds = get(messageIdsAtomFamily(loadableId))
             const flatById = get(messagesByIdAtomFamily(loadableId))
 

web/packages/agenta-playground/src/state/execution/selectors.ts

@@ -922,7 +922,8 @@ export const runStatusByRowEntityAtom = selectAtom(
             const sepIdx = key.indexOf(":sess:")
             if (sepIdx === -1) continue
             const stepId = key.slice(0, sepIdx)
-            const entityId = key.slice(sepIdx + 5)
+            // ":sess:" is 6 characters; +5 was off-by-one (left a leading ":").
+            const entityId = key.slice(sepIdx + 6)
             const status = result?.status
             const isRunning = status === "running" || status === "pending"
             mapped[`${stepId}:${entityId}`] = {

web/packages/agenta-playground/src/state/execution/webWorkerIntegration.ts

@@ -179,10 +179,10 @@ export const triggerExecutionAtom = atom(
         // Multi-entity fan-out: when no specific revision is requested and
         // multiple entities are shown side-by-side, trigger each one.
         if (!requestedRevisionId && Array.isArray(entityIds) && entityIds.length > 1) {
-            const sessionMatch = /^turn-([^-]+)-(lt-.+)$/.exec(String(rowId))
-            const logicalIdFromRow =
-                sessionMatch?.[2] || (String(rowId).startsWith("lt-") ? String(rowId) : "")
-            const lid = logicalIdFromRow || String(rowId)
+            // Extract the logical row ID from a possibly-compound turn-ID.
+            // Use the same sentinel search as extractLogicalRowId so that
+            // both "msg-<uuid>" (current) and "lt-<id>" (legacy) are handled.
+            const lid = extractLogicalRowId(String(rowId))
             for (const revId of entityIds) {
                 if (!revId) continue
                 const rid = `turn-${revId}-${lid}`
@@ -810,9 +810,15 @@ export const handleExecutionResultFromWorkerAtom = atom(
 /**
  * Extract the logical row ID from a turn-style row ID.
  * Turn IDs have format: `turn-<entityId>-<logicalId>`.
- * If it's already a logical ID (starts with "lt-"), return as-is.
+ * logicalId is "msg-<uuid>" (current) or "lt-<id>" (legacy).
+ * Neither prefix can appear inside a hex UUID, so "-msg-" / "-lt-"
+ * unambiguously marks the boundary between the entity segment and the
+ * logical row ID regardless of whether the entity UUID itself contains hyphens.
  */
-function extractLogicalRowId(rowId: string): string {
-    const match = /^turn-([^-]+)-(lt-.+)$/.exec(rowId)
-    return match?.[2] || rowId
+export function extractLogicalRowId(rowId: string): string {
+    const msgIdx = rowId.indexOf("-msg-")
+    if (msgIdx >= 0) return rowId.slice(msgIdx + 1)
+    const ltIdx = rowId.indexOf("-lt-")
+    if (ltIdx >= 0) return rowId.slice(ltIdx + 1)
+    return rowId
 }

web/packages/agenta-playground/tests/unit/chatTurnTokenKeys.test.ts

@@ -0,0 +1,147 @@
+/**
+ * Regression tests for issue #5789 -- token-usage numbers identical across
+ * every chat turn in a session.
+ *
+ * Three bugs caused the symptom:
+ *
+ * 1. `extractLogicalRowId` in webWorkerIntegration used a regex that only
+ *    matched "lt-" prefixed logical IDs; current message IDs start with "msg-",
+ *    so the function returned the full compound turn ID unchanged.  The wrong
+ *    ID was then used as the execution step key, making all turns collide on
+ *    the last shared user message.
+ *
+ * 2. `handleExecutionResultAtom` in executionItems only searched for "-lt-" to
+ *    locate the logical boundary in a compound "turn-<entityId>-<logicalId>"
+ *    rowId.  Same root cause; fallback always resolved to the last shared user
+ *    message, so every turn wrote its result at the same key and therefore
+ *    showed the same trace / same token counts.
+ *
+ * 3. `runStatusByRowEntityAtom` in selectors sliced result keys with
+ *    `sepIdx + 5` instead of `sepIdx + 6`, leaving a leading ":" on the
+ *    entityId segment and producing a malformed lookup key that could never
+ *    match the UI-side `"${rowId}:${entityId}"` key.
+ */
+
+import {describe, expect, it} from "vitest"
+
+import {extractLogicalRowId} from "../../src/state/execution/webWorkerIntegration"
+
+// ── extractLogicalRowId (Bug 1 fix) ─────────────────────────────────────────
+
+describe("extractLogicalRowId", () => {
+    it("returns a plain msg-<uuid> rowId unchanged (single-entity mode)", () => {
+        expect(extractLogicalRowId("msg-abc123")).toBe("msg-abc123")
+    })
+
+    it("returns a plain lt-<id> rowId unchanged (legacy single-entity)", () => {
+        expect(extractLogicalRowId("lt-old-id")).toBe("lt-old-id")
+    })
+
+    it("strips the turn-<entityId>- prefix from a current msg-<uuid> logical ID", () => {
+        expect(extractLogicalRowId("turn-rev1-msg-abc123")).toBe("msg-abc123")
+    })
+
+    it("handles a UUID entity ID with hyphens in the compound rowId", () => {
+        // Entity IDs are hex UUIDs; "msg-" cannot appear in a UUID segment,
+        // so the sentinel search is unambiguous even when revId contains hyphens.
+        expect(extractLogicalRowId("turn-a1b2c3d4-e5f6-msg-00000000-dead-beef")).toBe(
+            "msg-00000000-dead-beef",
+        )
+    })
+
+    it("strips the turn-<entityId>- prefix from a legacy lt-<id> logical ID", () => {
+        expect(extractLogicalRowId("turn-rev1-lt-old-id")).toBe("lt-old-id")
+    })
+
+    it("prefers msg- over lt- when both appear (current IDs take precedence)", () => {
+        // Contrived but guards against a mis-ordered search.
+        expect(extractLogicalRowId("turn-rev1-msg-foo-lt-bar")).toBe("msg-foo-lt-bar")
+    })
+
+    it("returns a non-turn rowId unchanged (no sentinel found)", () => {
+        expect(extractLogicalRowId("step-abc")).toBe("step-abc")
+        expect(extractLogicalRowId("")).toBe("")
+    })
+})
+
+// ── resultsByKey entity extraction (Bug 3 fix) ──────────────────────────────
+//
+// `runStatusByRowEntityAtom` parses keys of the form "stepId:sess:entityId".
+// The separator ":sess:" is 6 characters; the old code used `sepIdx + 5`,
+// leaving a leading ":" on the entityId and producing a malformed map key.
+// We test the parsing logic directly here as a pure string operation.
+
+function parseResultKey(key: string): {stepId: string; entityId: string} | null {
+    const sepIdx = key.indexOf(":sess:")
+    if (sepIdx === -1) return null
+    return {
+        stepId: key.slice(0, sepIdx),
+        // fixed: was sepIdx + 5 (off-by-one)
+        entityId: key.slice(sepIdx + 6),
+    }
+}
+
+describe("resultsByKey entity extraction (sepIdx + 6)", () => {
+    it("extracts stepId and entityId from a well-formed result key", () => {
+        expect(parseResultKey("msg-abc:sess:entity123")).toEqual({
+            stepId: "msg-abc",
+            entityId: "entity123",
+        })
+    })
+
+    it("handles a UUID entityId with hyphens", () => {
+        expect(parseResultKey("msg-abc:sess:a1b2-c3d4-e5f6")).toEqual({
+            stepId: "msg-abc",
+            entityId: "a1b2-c3d4-e5f6",
+        })
+    })
+
+    it("returns null when the key has no :sess: segment", () => {
+        expect(parseResultKey("msg-abc:entity123")).toBeNull()
+    })
+
+    it("produces a map key that matches the UI-side lookup", () => {
+        const key = "msg-abc:sess:entity123"
+        const parsed = parseResultKey(key)!
+        const mapKey = `${parsed.stepId}:${parsed.entityId}`
+        // The UI (TurnMessageAdapter) looks up by `"${rowId}:${entityId}"`.
+        expect(mapKey).toBe("msg-abc:entity123")
+    })
+})
+
+// ── logical rowId extraction in handleExecutionResultAtom (Bug 2 fix) ────────
+//
+// In comparison mode rowId = "turn-<entityUUID>-<logicalId>".  The old code
+// only searched for "-lt-"; with current message IDs ("msg-<uuid>") the
+// search always failed, leaving logicalRowId equal to the full compound string.
+// flatById[compoundId] is always undefined, so the fallback fired and returned
+// the LAST shared user message -- causing every turn to collide on that key.
+
+function extractLogicalFromRowId(rowId: string): string {
+    const msgIndex = rowId.indexOf("-msg-")
+    const ltIndex = rowId.indexOf("-lt-")
+    const sepIndex = msgIndex >= 0 ? msgIndex : ltIndex
+    return sepIndex >= 0 ? rowId.slice(sepIndex + 1) : rowId
+}
+
+describe("handleExecutionResultAtom logical rowId extraction", () => {
+    it("returns a plain msg-<uuid> rowId unchanged", () => {
+        expect(extractLogicalFromRowId("msg-abc123")).toBe("msg-abc123")
+    })
+
+    it("extracts the msg-<uuid> logical ID from a comparison-mode compound rowId", () => {
+        expect(extractLogicalFromRowId("turn-rev1-msg-abc123")).toBe("msg-abc123")
+    })
+
+    it("extracts the lt-<id> logical ID from a legacy comparison-mode compound rowId", () => {
+        expect(extractLogicalFromRowId("turn-rev1-lt-old-id")).toBe("lt-old-id")
+    })
+
+    it("handles a UUID entity ID with hyphens", () => {
+        expect(extractLogicalFromRowId("turn-a1b2-c3d4-msg-000-beef")).toBe("msg-000-beef")
+    })
+
+    it("returns the rowId unchanged when no sentinel is present", () => {
+        expect(extractLogicalFromRowId("step-abc")).toBe("step-abc")
+    })
+})