💵 fix: Bill Cross-Endpoint Labels at Their Own Rates

recordCollectedUsage never accepted an endpointTokenConfig, so the value
the activity-label caller passed was dropped and the balance transaction
was written at the primary agent's rates. Only the UI cost honored the
label endpoint, so a custom primary pointing activityEndpoint at another
endpoint showed one price and charged another. The parameter is now
accepted, and an explicit config wins outright over per-agent resolution:
that map is keyed by AGENT, so it cannot describe usage that ran on a
different endpoint.

Group identity is stable for id-less tool calls too. The previous pass
anchored the key to the first tool call ID; where a supported tool call
carries no id the fallback still used the block's first part index, which
shifts when a filled label absorbs the leading THINK part. The fallback now
anchors to the first TOOL entry's index, so only a block containing no
tool call at all keys off parts[0].

markActivityLabels is retried rather than fire-and-forget. It gates resume
gap reconciliation and is a SEPARATE write from the durable label append,
so a single lost write silently drops a label the content itself recorded.
The earlier "shared fate with content writes" reasoning was wrong. One
retry at run setup costs nothing and removes the only realistic way the
gate goes stale, without billing a content read to every resume.

api/server/controllers/agents/client.js

@@ -604,9 +604,21 @@ class AgentClient extends BaseClient {
       return undefined;
     }
     this.activityLabelPrompt = activityConfig.prompt;
-    /** Mark the job so a resume can reconcile label gaps without probing
-     *  content; fire-and-forget, the flag is only an optimization hint. */
-    void GenerationJobManager.markActivityLabels(streamId);
+    /**
+     * Mark the job so a resume can reconcile label gaps without probing
+     * content. Retried rather than fire-and-forget: this flag GATES that
+     * reconciliation, and it is a separate write from the durable label
+     * append — so a single lost write silently drops a label that the label
+     * content itself recorded perfectly well. One retry costs nothing at run
+     * setup and removes the only realistic way the gate goes stale.
+     */
+    void GenerationJobManager.markActivityLabels(streamId).catch(() =>
+      GenerationJobManager.markActivityLabels(streamId).catch(() => {
+        logger.warn(
+          `[AgentClient] Could not flag activity labels for ${streamId}; a label resolving during a resume gap may not be reconciled.`,
+        );
+      }),
+    );
     /** SDK support probe (steering-style): the Run method and the formatter
      *  replay skip ship together, so method presence is the capability. */
     const sdkCapable = typeof Run?.prototype?.generateActivityLabel === 'function';
@@ -1619,7 +1631,19 @@ class AgentClient extends BaseClient {
     transactions,
     context = 'message',
     collectedUsage = this.collectedUsage,
+    /**
+     * Rates for usage that did NOT run on the agent's endpoint — currently
+     * activity labels pointed at a different `activityEndpoint`. Without it
+     * the caller's config was dropped here and the balance transaction was
+     * written at the primary agent's rates while the UI cost was computed at
+     * the label's, so the two disagreed. `undefined` keeps the agent default.
+     */
+    endpointTokenConfig,
   }) {
+    /** Per-agent resolution keys off the AGENT's config map, which cannot
+     *  describe a label running on a different endpoint — so an explicit
+     *  config wins outright rather than being second-guessed per usage row. */
+    const overrideTokenConfig = endpointTokenConfig !== undefined;
     const result = await recordCollectedUsage(
       {
         spendTokens: db.spendTokens,
@@ -1636,8 +1660,12 @@ class AgentClient extends BaseClient {
         messageId: this.responseMessageId,
         balance,
         transactions,
-        endpointTokenConfig: this.options.endpointTokenConfig,
-        resolveEndpointTokenConfig: (usage) => this.resolveAgentEndpointTokenConfig(usage),
+        endpointTokenConfig: overrideTokenConfig
+          ? endpointTokenConfig
+          : this.options.endpointTokenConfig,
+        ...(overrideTokenConfig
+          ? {}
+          : { resolveEndpointTokenConfig: (usage) => this.resolveAgentEndpointTokenConfig(usage) }),
       },
     );
 

client/src/components/Chat/Messages/Content/ContentParts.tsx

@@ -34,13 +34,20 @@ const getToolGroupId = (parts: PartWithIndex[], fallbackScope: number): string =
    *  absorbs the block's leading THINK part when its text lands, so keying on
    *  `parts[0]` would flip the key mid-run — remounting the group and losing
    *  whatever the user had expanded. The tool calls themselves do not move. */
-  for (const { part } of parts) {
+  let firstToolIdx: number | undefined;
+  for (const { part, idx } of parts) {
     const toolCallId = getToolCallId(part);
     if (toolCallId) {
       return `tool:${toolCallId}`;
     }
+    if (firstToolIdx === undefined && part?.type === ContentTypes.TOOL_CALL) {
+      firstToolIdx = idx;
+    }
   }
-  return `fallback:${fallbackScope}:${firstPart.idx}`;
+  /** Same reasoning for id-less tool calls: anchor to the first TOOL entry's
+   *  index rather than the block's first part, which shifts when reasoning is
+   *  absorbed. Only a block with no tool call at all falls back to `parts[0]`. */
+  return `fallback:${fallbackScope}:${firstToolIdx ?? firstPart.idx}`;
 };
 
 type PartWithContextProps = {

packages/api/src/stream/GenerationJobManager.ts

@@ -2320,11 +2320,11 @@ class GenerationJobManagerClass {
        * queue skips the content re-read").
        *
        * The residual: if `markActivityLabels` lost its write AND the first
-       * label is claimed inside the gap, this is skipped. That flag write
-       * shares fate with the content writes the labels themselves live in, so
-       * the case implies a store already dropping data — not worth a read on
-       * every resume. When the steer pass above already fetched content, this
-       * check is free.
+       * label is claimed inside the gap, this is skipped. The flag is a
+       * SEPARATE write from the durable label append, so that is genuinely
+       * possible rather than implying a broken store — which is why the mark
+       * is retried at run setup instead of being fire-and-forget. When the
+       * steer pass above already fetched content, this check is free.
        */
       const snapshotHasActivityLabels =
         resumeState?.aggregatedContent?.some(