Prevent fallback failure source misattribution

Author: TheGreatAxios

docs/TELEMETRY.md

@@ -167,9 +167,13 @@ completed message run discards it. `inference.done` settles success and only
 then applies successful-generation sampling. A failed `message.run.ended`
 settles an unresolved turn once as an unsampled terminal failure. Attribution
 uses the latest `inference.usage` source, the first lifecycle payload carrying
-the runtime-resolved provider/model pair for an attempt; it does not infer
-fallback from the externally selected source. Therefore a parent turn emits at
-most one terminal `$ai_generation`, including retry and failover paths.
+the runtime-resolved provider/model pair for an attempt. Each `inference.start`
+clears that authoritative source and records the newly attempted model. If a
+fallback fails before usage exposes its source, telemetry retains that actual
+model but uses the fixed `unknown` provider/source bucket rather than attributing
+it to the previously selected provider. When the attempted model still matches
+the selected source, that full source remains valid. Therefore a parent turn
+emits at most one terminal `$ai_generation`, including retry and failover paths.
 
 ## What's never collected
 

src/session/run-sink.test.ts

@@ -25,6 +25,42 @@ const enabledHook: LifecycleHookStatus = {
   enabled: true,
 };
 
+function attributionHarness(selectedSource = { provider: "provider-a", model: "model-a" }) {
+  const captured: { event: string; properties: Record<string, unknown> }[] = [];
+  const telemetry: Telemetry = {
+    enabled: true,
+    installationId: "test",
+    capture: (capturedEvent, properties = {}) => {
+      captured.push({ event: capturedEvent, properties });
+    },
+    captureIntentional: () => false,
+    flush: async () => {},
+    discard: () => {},
+  };
+  const observer = createTurnObserver({
+    telemetry: () => telemetry,
+    getSessionId: () => "session-1",
+    getSource: () => selectedSource,
+  });
+  const runSink = createRunSink({
+    emitter: new EventEmitter(),
+    hookManager: stubHookManager([]),
+    ...observer,
+  });
+  return { captured, runSink };
+}
+
+function failMessageRun(runSink: ReturnType<typeof createRunSink>): void {
+  runSink.sink(event("inference.error", { error: { message: "attempt failed" } }));
+  runSink.sink(
+    event("message.run.ended", {
+      messageRunId: "run-1",
+      messageId: "message-1",
+      status: "failed",
+    }),
+  );
+}
+
 describe("createRunSink", () => {
   test("allocates no turn collector when no lifecycle hooks are configured", () => {
     const runSink = createRunSink({
@@ -138,52 +174,72 @@ describe("createRunSink", () => {
     expect(failures).toEqual([{ turnIndex: 0, error: "429 rate limit" }]);
   });
 
-  test("attributes terminal retry failure to the latest attempted source", () => {
-    const captured: { event: string; properties: Record<string, unknown> }[] = [];
-    const telemetry: Telemetry = {
-      enabled: true,
-      installationId: "test",
-      capture: (capturedEvent, properties = {}) => {
-        captured.push({ event: capturedEvent, properties });
-      },
-      captureIntentional: () => false,
-      flush: async () => {},
-      discard: () => {},
-    };
-    const selectedSource = { provider: "provider-a", model: "model-a" };
-    const observer = createTurnObserver({
-      telemetry: () => telemetry,
-      getSessionId: () => "session-1",
-      getSource: () => selectedSource,
-    });
-    const runSink = createRunSink({
-      emitter: new EventEmitter(),
-      hookManager: stubHookManager([]),
-      ...observer,
+  test("uses unknown attribution when a fallback model fails before usage", () => {
+    const { captured, runSink } = attributionHarness();
+
+    runSink.sink(event("inference.start", { model: "model-b" }));
+    failMessageRun(runSink);
+
+    expect(captured).toHaveLength(1);
+    expect(captured[0]?.event).toBe("$ai_generation");
+    expect(captured[0]?.properties).toMatchObject({
+      $ai_provider: "unknown",
+      $ai_model: "model-b",
+      $ai_is_error: true,
     });
+  });
+
+  test("uses the selected source when its model fails before usage", () => {
+    const { captured, runSink } = attributionHarness();
 
     runSink.sink(event("inference.start", { model: "model-a" }));
-    runSink.sink(event("inference.error", { error: { message: "attempt a failed" } }));
+    failMessageRun(runSink);
+
+    expect(captured).toHaveLength(1);
+    expect(captured[0]?.properties).toMatchObject({
+      $ai_provider: "provider-a",
+      $ai_model: "model-a",
+      $ai_is_error: true,
+    });
+  });
+
+  test("uses authoritative usage attribution for a failed fallback", () => {
+    const { captured, runSink } = attributionHarness();
+
     runSink.sink(event("inference.start", { model: "model-b" }));
     runSink.sink(
       event("inference.usage", {
         usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, thinking: 0 },
         source: { sourceId: "fallback", provider: "provider-b", model: "model-b" },
       }),
     );
-    runSink.sink(event("inference.error", { error: { message: "attempt b failed" } }));
+    failMessageRun(runSink);
+
+    expect(captured).toHaveLength(1);
+    expect(captured[0]?.properties).toMatchObject({
+      $ai_provider: "provider-b",
+      $ai_model: "model-b",
+      $ai_is_error: true,
+    });
+  });
+
+  test("does not leak authoritative source attribution across retry attempts", () => {
+    const { captured, runSink } = attributionHarness();
+
+    runSink.sink(event("inference.start", { model: "model-a" }));
     runSink.sink(
-      event("message.run.ended", {
-        messageRunId: "run-1",
-        messageId: "message-1",
-        status: "failed",
+      event("inference.usage", {
+        usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, thinking: 0 },
+        source: { sourceId: "selected", provider: "provider-a", model: "model-a" },
       }),
     );
+    runSink.sink(event("inference.error", { error: { message: "retry" } }));
+    runSink.sink(event("inference.start", { model: "model-b" }));
+    failMessageRun(runSink);
 
     expect(captured).toHaveLength(1);
-    expect(captured[0]?.event).toBe("$ai_generation");
     expect(captured[0]?.properties).toMatchObject({
-      $ai_provider: "provider-b",
+      $ai_provider: "unknown",
       $ai_model: "model-b",
       $ai_is_error: true,
     });

src/session/run-sink.ts

@@ -22,10 +22,9 @@ export interface RunSinkArgs {
   // run goes wrong. The turn index is the collector's current count: the
   // in-flight turn is the one that would have been recorded next.
   onTurnFailed?: (info: { turnIndex: number; error: string }) => void;
-  // Fired when inference for a turn begins. The turn index is the collector's
-  // current count (the in-flight turn that has not completed yet) — used to
-  // stamp parent_trace_id on subagent_end while tools still run.
-  onTurnStarted?: (info: { turnIndex: number }) => void;
+  // Fired for every inference attempt. The model comes from inference.start,
+  // while the turn index is the collector's current in-flight turn count.
+  onTurnStarted?: (info: { turnIndex: number; model: string }) => void;
   // inference.usage is the first attempt event carrying the runtime-resolved
   // provider/model pair. It remains authoritative even when the selected source
   // outside the reactor has not changed during fallback.
@@ -150,7 +149,7 @@ export function createRunSink(args: RunSinkArgs): RunSink {
     perfObserver.observe(event);
     if (event.type === "inference.start") {
       turnInFlight = true;
-      onTurnStarted?.({ turnIndex: turnCollector.getTurnCount() });
+      onTurnStarted?.({ turnIndex: turnCollector.getTurnCount(), model: event.data.model });
     }
     if (event.type === "inference.usage") {
       onTurnSourceObserved?.({

src/telemetry/ai-observability.test.ts

@@ -249,7 +249,7 @@ describe("createTurnObserver", () => {
       getSource: () => ({ provider: "openai-compatible", model: "model-x" }),
     });
 
-    observer.onTurnStarted({ turnIndex: 2 });
+    observer.onTurnStarted({ turnIndex: 2, model: "model-x" });
     expect(getCurrentTurnTraceId()).toBe(`${SESSION_ID}:turn:2`);
     observer.onTurnComplete(fakeTurnContext({ turnIndex: 2, toolCalls: [], toolResults: [] }));
     expect(getCurrentTurnTraceId()).toBeUndefined();

src/telemetry/ai-observability.ts

@@ -219,39 +219,65 @@ export interface CreateTurnObserverOptions {
   // session id in this same process, and a trace id built from a captured
   // one would file the new session's turns under the old session's traces.
   getSessionId: () => string;
-  // The source the next inference will run against, which is the best
-  // available attribution for a turn that failed before producing one.
+  // The currently selected source. It is safe failure attribution only when
+  // its model matches the model named by the latest inference.start.
   getSource: () => TurnSource;
 }
 
-// Binds the emitters to the live session and source, giving the run sink two
-// plain callbacks and keeping the "read it now, do not capture it" rule in
-// one place instead of at each call site.
+const UNKNOWN_INFERENCE_PROVIDER = "unknown";
+
+function failedAttemptSource(
+  attemptedModel: string | undefined,
+  observedSource: TurnSource | undefined,
+  selectedSource: TurnSource,
+): TurnSource {
+  if (observedSource !== undefined) return observedSource;
+  if (attemptedModel === undefined || attemptedModel === selectedSource.model) {
+    return selectedSource;
+  }
+  return { provider: UNKNOWN_INFERENCE_PROVIDER, model: attemptedModel };
+}
+
+// Binds the emitters to the live session and source, keeping the "read it now,
+// do not capture it" rule in one place instead of at each call site.
 export function createTurnObserver(options: CreateTurnObserverOptions): {
-  onTurnStarted: (info: { turnIndex: number }) => void;
+  onTurnStarted: (info: { turnIndex: number; model: string }) => void;
   onTurnSourceObserved: (info: { turnIndex: number; source: TurnSource }) => void;
   onTurnComplete: (ctx: TurnContext) => void;
   onTurnFailed: (info: { turnIndex: number; error: string }) => void;
 } {
+  let latestAttemptModel: string | undefined;
   let latestAttemptSource: TurnSource | undefined;
+
+  function clearAttempt(): void {
+    latestAttemptModel = undefined;
+    latestAttemptSource = undefined;
+  }
+
   return {
     onTurnStarted: (info) => {
+      latestAttemptModel = info.model;
+      latestAttemptSource = undefined;
       noteCurrentTurnTraceId(turnTraceId(options.getSessionId(), info.turnIndex));
     },
     onTurnSourceObserved: (info) => {
       latestAttemptSource = { ...info.source };
     },
     onTurnComplete: (ctx) => {
-      latestAttemptSource = undefined;
+      clearAttempt();
       clearCurrentTurnTraceId();
       emitAiObservability(options.telemetry(), ctx, {
         sessionId: options.getSessionId(),
       });
     },
     onTurnFailed: (info) => {
       clearCurrentTurnTraceId();
-      const source = latestAttemptSource ?? options.getSource();
-      latestAttemptSource = undefined;
+      const source = failedAttemptSource(
+        latestAttemptModel,
+        latestAttemptSource,
+        options.getSource(),
+      );
+      clearAttempt();
       emitAiTurnFailure(options.telemetry(), {
         sessionId: options.getSessionId(),
         source,

src/tui/runner.ts

@@ -1590,6 +1590,7 @@ export async function runTUI(initialConfig: Config): Promise<number> {
       hookManager,
       initialTurnCount: resumeSeed.turnsUsed,
       onTurnStarted: turnObserver.onTurnStarted,
+      onTurnSourceObserved: turnObserver.onTurnSourceObserved,
       onTurnComplete: turnObserver.onTurnComplete,
       onTurnFailed: turnObserver.onTurnFailed,
       // persistRunSnapshot is defined below but not invoked until the stream