From 9a95d0217a7f4572d353a54742e5ef661ca0421d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matej=20=C5=A0poler?= Date: Fri, 23 Oct 2026 03:06:07 +0200 Subject: [PATCH] fix(step): fall back to the session's model when a step's start was missed `session.step.ended` and `session.step.failed` carry no model, so the model comes from the `stepMeta` entry `session.step.started` writes. When that entry is missing (the plugin subscribed after the step started, or the execution ended and its sweep cleared the entry before the step ended), `api_request`, `api_error`, the token, cost and message metrics all report model and provider `unknown`. The model now falls back to the in-flight model request of the session (`activeLlm`), then to the session's last known model, kept on the session identity from each step start and from `session.get` for a session the plugin looks up, and only then to `unknown`. Co-Authored-By: Claude Opus 5.5 --- src/handlers/step.ts | 24 +++++++++++++---- src/types.ts | 1 + src/util.ts | 9 ++++++- tests/handlers/step.test.ts | 51 ++++++++++++++++++++++++++++++++++++- tests/state.test.ts | 10 ++++++++ 5 files changed, 88 insertions(+), 7 deletions(-) diff --git a/src/handlers/step.ts b/src/handlers/step.ts index b87d11b..02809ec 100644 --- a/src/handlers/step.ts +++ b/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)) diff --git a/src/types.ts b/src/types.ts index 07468e0..dadd9bb 100644 --- a/src/types.ts +++ b/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 diff --git a/src/util.ts b/src/util.ts index 9a45d3b..a036a2d 100644 --- a/src/util.ts +++ b/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 { 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 { diff --git a/tests/handlers/step.test.ts b/tests/handlers/step.test.ts index ee36269..2428da4 100644 --- a/tests/handlers/step.test.ts +++ b/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["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() diff --git a/tests/state.test.ts b/tests/state.test.ts index b1e7514..2b7ad3f 100644 --- a/tests/state.test.ts +++ b/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[] = []