fix(step): fall back to the session's model when a step's start was missed

src/handlers/step.ts

@@ -23,7 +23,7 @@ import {
   SemanticConventions,
   SESSION_ID,
 } from "@arizeai/openinference-semantic-conventions"
-import type { EventOf, HandlerContext, SessionTotals } from "../types.ts"
+import type { EventOf, HandlerContext, SessionTotals, StepMeta } from "../types.ts"
 import { ensureSession } from "./session.ts"
 import { applyModelContext } from "./chat-headers.ts"
 import {
@@ -83,6 +83,13 @@ export function handleStepStarted(e: EventOf<"session.step.started">, ctx: Handl
   if (totals && totals.agent !== d.agent) {
     setBoundedMap(ctx.tracing.sessionTotals, d.sessionID, { ...totals, agent: d.agent })
   }
+  const identity = ctx.tracing.sessionIdentity.get(d.sessionID)
+  if (identity) {
+    setBoundedMap(ctx.tracing.sessionIdentity, d.sessionID, {
+      ...identity,
+      model: { id: d.model.id, providerID: d.model.providerID },
+    })
+  }
 
   if (!isTraceEnabled("llm", ctx)) return
   const promptText = ctx.tracing.activePrompts.get(d.sessionID)?.text
@@ -141,6 +148,15 @@ export function handleTextEnded(e: EventOf<"session.text.ended">, ctx: HandlerCo
   setBoundedMap(ctx.tracing.stepOutputs, assistantMessageID, output)
 }
 
+function stepModel(meta: StepMeta | undefined, sessionID: string, ctx: HandlerContext): { modelID: string; providerID: string } {
+  if (meta) return { modelID: meta.modelID, providerID: meta.providerID }
+  const active = ctx.tracing.activeLlm.get(sessionID)
+  if (active) return { modelID: active.modelID, providerID: active.providerID }
+  const model = ctx.tracing.sessionIdentity.get(sessionID)?.model
+  if (model) return { modelID: model.id, providerID: model.providerID }
+  return { modelID: "unknown", providerID: "unknown" }
+}
+
 function countMessage(messageID: string, sessionID: string, modelID: string, providerID: string, agent: string, ctx: HandlerContext): boolean {
   if (ctx.tracing.countedMessages.has(messageID)) return false
   markSeen(ctx.tracing.countedMessages, messageID)
@@ -166,8 +182,7 @@ export function handleStepEnded(e: EventOf<"session.step.ended">, ctx: HandlerCo
   const sessionTotals = ctx.tracing.sessionTotals.get(d.sessionID)
   const agent = meta?.agent ?? sessionTotals?.agent ?? "unknown"
   const agentType = meta?.agentType ?? sessionTotals?.agentType ?? "unknown"
-  const modelID = meta?.modelID ?? "unknown"
-  const providerID = meta?.providerID ?? "unknown"
+  const { modelID, providerID } = stepModel(meta, d.sessionID, ctx)
   const durationMs = Math.max(0, e.created - (meta?.startMs ?? e.created))
 
   recordUsageMetrics(d.sessionID, modelID, agent, d.tokens, d.cost, ctx)
@@ -233,8 +248,7 @@ export function handleStepFailed(e: EventOf<"session.step.failed">, ctx: Handler
   const sessionTotals = ctx.tracing.sessionTotals.get(d.sessionID)
   const agent = meta?.agent ?? sessionTotals?.agent ?? "unknown"
   const agentType = meta?.agentType ?? sessionTotals?.agentType ?? "unknown"
-  const modelID = meta?.modelID ?? "unknown"
-  const providerID = meta?.providerID ?? "unknown"
+  const { modelID, providerID } = stepModel(meta, d.sessionID, ctx)
   const error = errorSummary(d.error)
   const durationMs = Math.max(0, e.created - (meta?.startMs ?? e.created))
 

src/types.ts

@@ -81,6 +81,7 @@ export type SessionIdentity = {
   agent: string
   agentType: SessionAgentType
   parentID?: string
+  model?: { id: string; providerID: string }
   startMs: number
   tokens?: number
   cost?: number

src/util.ts

@@ -178,7 +178,13 @@ export function getSessionAgentMeta(
 export async function contextForSession(
   sessionID: string,
   ctx: HandlerContext,
-  getSession: (sessionID: string) => Promise<{ projectID: string; agent?: string; parentID?: string; time: { created: number } }>,
+  getSession: (sessionID: string) => Promise<{
+    projectID: string
+    agent?: string
+    parentID?: string
+    model?: { id: string; providerID: string }
+    time: { created: number }
+  }>,
 ): Promise<HandlerContext> {
   let projectID = ctx.tracing.sessionProjects.get(sessionID)
   if (!projectID || !ctx.tracing.sessionIdentity.has(sessionID)) {
@@ -190,6 +196,7 @@ export async function contextForSession(
         agent: session.agent ?? "unknown",
         agentType: session.parentID ? "subagent" : "primary",
         ...(session.parentID ? { parentID: session.parentID } : {}),
+        ...(session.model ? { model: { id: session.model.id, providerID: session.model.providerID } } : {}),
         startMs: session.time.created,
       })
     } catch {

tests/handlers/step.test.ts

@@ -1,7 +1,7 @@
 import { describe, test, expect } from "bun:test"
 import { SpanStatusCode } from "@opentelemetry/api"
 import { handleStepEnded, handleStepFailed, handleStepStarted, handleTextEnded } from "../../src/handlers/step.ts"
-import { handleExecutionStarted, handleSessionCreated } from "../../src/handlers/session.ts"
+import { finalizeSession, handleExecutionEnded, handleExecutionStarted, handleSessionCreated } from "../../src/handlers/session.ts"
 import { makeCtx, evt, tokens } from "../helpers.ts"
 
 function seedSession(ctx: ReturnType<typeof makeCtx>["ctx"], sessionID = "ses_1") {
@@ -70,6 +70,55 @@ describe("handleStepEnded", () => {
   })
 })
 
+describe("step model fallback", () => {
+  function ended(assistantMessageID = "msg_2") {
+    return evt("session.step.ended", { sessionID: "ses_1", assistantMessageID, finish: "stop", cost: 0, tokens: tokens() })
+  }
+
+  test("uses the session's last model when the step start was missed", () => {
+    const { ctx, logger, counters } = makeCtx()
+    seedSession(ctx)
+    handleStepStarted(stepStarted("ses_1", "msg_1"), ctx)
+    handleStepEnded(ended("msg_1"), ctx)
+    handleStepEnded(ended("msg_2"), ctx)
+    const request = logger.records.at(-1)!
+    expect(request.attributes?.["model"]).toBe("claude-sonnet")
+    expect(request.attributes?.["provider"]).toBe("anthropic")
+    expect(counters.token.calls.at(-1)!.attrs["model"]).toBe("claude-sonnet")
+  })
+
+  test("keeps the model of a step that ends after its execution was swept", () => {
+    const { ctx, logger } = makeCtx()
+    seedSession(ctx)
+    handleStepStarted(stepStarted("ses_1", "msg_1"), ctx)
+    handleExecutionEnded(evt("session.execution.interrupted", { sessionID: "ses_1", reason: "user" }), ctx, { type: "interrupted", reason: "user" })
+    finalizeSession("ses_1", ctx)
+    handleStepEnded(ended("msg_1"), ctx)
+    expect(logger.records.find((record) => record.body === "api_request")?.attributes?.["model"]).toBe("claude-sonnet")
+  })
+
+  test("uses the in-flight model request when the step start was missed", () => {
+    const { ctx, logger } = makeCtx()
+    seedSession(ctx)
+    ctx.tracing.activeLlm.set("ses_1", {
+      agent: "build",
+      modelID: "gpt-5",
+      providerID: "openai",
+      spanContext: { traceId: "0".repeat(32), spanId: "0".repeat(16), traceFlags: 1 },
+    })
+    handleStepFailed(evt("session.step.failed", { sessionID: "ses_1", assistantMessageID: "msg_2", error: { type: "x", message: "y" } }), ctx)
+    expect(logger.records.at(-1)!.attributes?.["model"]).toBe("gpt-5")
+    expect(logger.records.at(-1)!.attributes?.["provider"]).toBe("openai")
+  })
+
+  test("reports unknown when nothing names the model", () => {
+    const { ctx, logger } = makeCtx()
+    seedSession(ctx)
+    handleStepEnded(ended(), ctx)
+    expect(logger.records.at(-1)!.attributes?.["model"]).toBe("unknown")
+  })
+})
+
 describe("handleStepFailed", () => {
   test("records usage, ends the span with error, and emits api_error", () => {
     const { ctx, tracer, counters, logger } = makeCtx()

tests/state.test.ts

@@ -54,6 +54,16 @@ describe("multi-location telemetry", () => {
     expect(scoped.commonAttrs["project.id"]).toBe("other-project")
   })
 
+  test("keeps the model of a session it looks up", async () => {
+    const { ctx } = makeCtx()
+    await contextForSession("resumed", ctx, async () => ({
+      projectID: "project",
+      model: { id: "claude-opus", providerID: "anthropic" },
+      time: { created: 500 },
+    }))
+    expect(ctx.tracing.sessionIdentity.get("resumed")?.model).toEqual({ id: "claude-opus", providerID: "anthropic" })
+  })
+
   test("serializes duplicate subscribers across an asynchronous lookup", async () => {
     const { ctx } = makeCtx()
     const order: string[] = []