diff --git a/.github/dependabot.yml b/.github/dependabot.yml index d698d8d..c7be077 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -5,8 +5,7 @@ version: 2 updates: - - package-ecosystem: "" # See documentation for possible values + - package-ecosystem: "bun" # See documentation for possible values directory: "/" # Location of package manifests schedule: interval: "weekly" - diff --git a/src/handlers/message.ts b/src/handlers/message.ts index f6dce25..568536e 100644 --- a/src/handlers/message.ts +++ b/src/handlers/message.ts @@ -1,5 +1,5 @@ import { SeverityNumber } from "@opentelemetry/api-logs" -import { SpanStatusCode, SpanKind, trace } from "@opentelemetry/api" +import { SpanStatusCode, SpanKind } from "@opentelemetry/api" import type { AssistantMessage, EventMessageUpdated, EventMessagePartUpdated, ToolPart } from "@opencode-ai/sdk" import { AGENT_NAME, @@ -27,7 +27,16 @@ import { TOOL_NAME, TOOL_PARAMETERS, } from "@arizeai/openinference-semantic-conventions" -import { agentAttrs, errorSummary, setBoundedMap, accumulateSessionTotals, getSessionAgentMeta, isMetricEnabled, isTraceEnabled } from "../util.ts" +import { + agentAttrs, + errorSummary, + setBoundedMap, + accumulateSessionTotals, + getSessionAgentMeta, + isMetricEnabled, + isTraceEnabled, + resolveSessionTraceContext, +} from "../util.ts" import type { HandlerContext } from "../types.ts" const OPENINFERENCE_SPAN_KIND = SemanticConventions.OPENINFERENCE_SPAN_KIND @@ -52,6 +61,7 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext const msg = e.properties.info if (msg.role !== "assistant") return const assistant = msg as AssistantMessage + setBoundedMap(ctx.assistantRuns, assistant.id, assistant.parentID) if (!assistant.time.completed) return const { sessionID, modelID, providerID } = assistant @@ -260,11 +270,6 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle const { agentName, agentType } = getSessionAgentMeta(toolPart.sessionID, ctx) const toolSpan = isTraceEnabled("tool", ctx) ? (() => { - const sessionSpan = ctx.sessionSpans.get(toolPart.sessionID) - const baseCtx = ctx.rootContext() - const parentCtx = sessionSpan - ? trace.setSpan(baseCtx, sessionSpan) - : baseCtx return ctx.tracer.startSpan( `${ctx.tracePrefix}tool.${toolPart.tool}`, { @@ -283,7 +288,9 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle ...ctx.commonAttrs, }, }, - parentCtx, + resolveSessionTraceContext(toolPart.sessionID, ctx, { + assistantMessageID: toolPart.messageID, + }), ) })() : undefined @@ -319,11 +326,6 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle if (isTraceEnabled("tool", ctx)) { const toolSpan = pending?.span ?? (() => { - const sessionSpan = ctx.sessionSpans.get(toolPart.sessionID) - const baseCtx = ctx.rootContext() - const parentCtx = sessionSpan - ? trace.setSpan(baseCtx, sessionSpan) - : baseCtx return ctx.tracer.startSpan( `${ctx.tracePrefix}tool.${toolPart.tool}`, { @@ -340,7 +342,9 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle ...ctx.commonAttrs, }, }, - parentCtx, + resolveSessionTraceContext(toolPart.sessionID, ctx, { + assistantMessageID: toolPart.messageID, + }), ) })() toolSpan.setAttributes({ [AGENT_NAME]: agentName, "agent.type": agentType }) @@ -403,14 +407,16 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle /** * Starts an LLM span for an assistant message when it first appears in `message.updated`. - * The span is parented to the session span and carries `gen_ai.*` semantic attributes for - * the model and provider. It is ended in `handleMessageUpdated` once the message completes. + * The span is parented to the active run or subagent span and carries `gen_ai.*` semantic + * attributes for the model and provider. It is ended in `handleMessageUpdated` once the + * message completes. * * Only called for assistant messages that have not yet completed (`time.completed` absent). */ export function startMessageSpan( sessionID: string, messageID: string, + parentID: string, modelID: string, providerID: string, startTime: number, @@ -419,12 +425,9 @@ export function startMessageSpan( if (!isTraceEnabled("llm", ctx)) return const msgKey = `${sessionID}:${messageID}` if (ctx.messageSpans.has(msgKey)) return + setBoundedMap(ctx.assistantRuns, messageID, parentID) const { agentName, agentType } = getSessionAgentMeta(sessionID, ctx) - const sessionSpan = ctx.sessionSpans.get(sessionID) - const baseCtx = ctx.rootContext() - const parentCtx = sessionSpan - ? trace.setSpan(baseCtx, sessionSpan) - : baseCtx + const inputText = ctx.runInputs.get(parentID) const msgSpan = ctx.tracer.startSpan( `${ctx.tracePrefix}llm`, @@ -439,17 +442,17 @@ export function startMessageSpan( [LLM_SYSTEM]: providerID, [LLM_PROVIDER]: providerID, [LLM_MODEL_NAME]: modelID, - ...(ctx.sessionInputs.has(sessionID) + ...(inputText ? { - [INPUT_VALUE]: ctx.sessionInputs.get(sessionID)!, + [INPUT_VALUE]: inputText, [INPUT_MIME_TYPE]: MimeType.TEXT, - [LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: ctx.sessionInputs.get(sessionID)! }]), + [LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: inputText }]), } : {}), ...ctx.commonAttrs, }, }, - parentCtx, + resolveSessionTraceContext(sessionID, ctx, { runID: parentID, assistantMessageID: messageID }), ) setBoundedMap(ctx.messageSpans, msgKey, msgSpan) } diff --git a/src/handlers/session.ts b/src/handlers/session.ts index f6bdb9b..b03a14c 100644 --- a/src/handlers/session.ts +++ b/src/handlers/session.ts @@ -1,12 +1,86 @@ import { SeverityNumber } from "@opentelemetry/api-logs" -import { SpanStatusCode, trace } from "@opentelemetry/api" +import { SpanStatusCode } from "@opentelemetry/api" import type { EventSessionCreated, EventSessionIdle, EventSessionError, EventSessionStatus } from "@opencode-ai/sdk" -import { AGENT_NAME, OpenInferenceSpanKind, SemanticConventions, SESSION_ID } from "@arizeai/openinference-semantic-conventions" -import { agentAttrs, errorSummary, getSessionAgentMeta, setBoundedMap, isMetricEnabled, isTraceEnabled } from "../util.ts" +import { + AGENT_NAME, + INPUT_MIME_TYPE, + INPUT_VALUE, + LLM_INPUT_MESSAGES, + MimeType, + OpenInferenceSpanKind, + SemanticConventions, + SESSION_ID, +} from "@arizeai/openinference-semantic-conventions" +import { + agentAttrs, + errorSummary, + getSessionAgentMeta, + setBoundedMap, + isMetricEnabled, + isTraceEnabled, + resolveSessionTraceContext, +} from "../util.ts" import type { HandlerContext, SessionAgentType } from "../types.ts" const OPENINFERENCE_SPAN_KIND = SemanticConventions.OPENINFERENCE_SPAN_KIND +/** Starts or refreshes the root run span for a single user turn, keyed by the user message ID. */ +export function handleRunStarted( + runID: string, + sessionID: string, + agent: string, + promptText: string, + model: string, + startTime: number, + ctx: HandlerContext, +) { + ctx.activeRuns.set(sessionID, runID) + ctx.pendingRuns.delete(sessionID) + if (promptText) setBoundedMap(ctx.runInputs, runID, promptText) + if (!isTraceEnabled("session", ctx)) return + const existing = ctx.runSpans.get(runID) + if (existing) { + existing.setAttributes({ + [AGENT_NAME]: agent, + ...(promptText + ? { + [INPUT_VALUE]: promptText, + [INPUT_MIME_TYPE]: MimeType.TEXT, + [LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: promptText }]), + } + : {}), + model, + }) + return + } + + const runSpan = ctx.tracer.startSpan( + `${ctx.tracePrefix}session`, + { + startTime, + attributes: { + [OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.AGENT, + [SESSION_ID]: sessionID, + [AGENT_NAME]: agent, + "agent.type": "primary", + "session.is_subagent": false, + ...(promptText + ? { + [INPUT_VALUE]: promptText, + [INPUT_MIME_TYPE]: MimeType.TEXT, + [LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: promptText }]), + } + : {}), + model, + ...ctx.commonAttrs, + }, + }, + ctx.rootContext(), + ) + ctx.runSpans.set(runID, runSpan) + setBoundedMap(ctx.runSpanContexts, runID, runSpan.spanContext()) +} + /** Increments the session counter, records start time, starts the root session span, and emits a `session.created` log event. */ export function handleSessionCreated(e: EventSessionCreated, ctx: HandlerContext) { const { id: sessionID, time, parentID } = e.properties.info @@ -18,16 +92,7 @@ export function handleSessionCreated(e: EventSessionCreated, ctx: HandlerContext } setBoundedMap(ctx.sessionTotals, sessionID, { startMs: createdAt, tokens: 0, cost: 0, messages: 0, agent: "unknown", agentType }) - // WARNING: disabling "session" traces while "llm" or "tool" traces remain enabled - // leaves those child spans without a local session parent. If OPENCODE_TRACEPARENT - // is set, they fall back to that remote parent; otherwise they become root spans. - if (isTraceEnabled("session", ctx)) { - const parentSpan = parentID ? ctx.sessionSpans.get(parentID) : undefined - const baseCtx = ctx.rootContext() - const spanCtx = parentSpan - ? trace.setSpan(baseCtx, parentSpan) - : baseCtx - + if (isTraceEnabled("session", ctx) && parentID) { const sessionSpan = ctx.tracer.startSpan( `${ctx.tracePrefix}session`, { @@ -41,9 +106,10 @@ export function handleSessionCreated(e: EventSessionCreated, ctx: HandlerContext ...ctx.commonAttrs, }, }, - spanCtx, + resolveSessionTraceContext(parentID, ctx), ) - setBoundedMap(ctx.sessionSpans, sessionID, sessionSpan) + ctx.sessionSpans.set(sessionID, sessionSpan) + setBoundedMap(ctx.sessionSpanContexts, sessionID, sessionSpan.spanContext()) } ctx.emitLog({ @@ -74,7 +140,7 @@ function sweepSession(sessionID: string, ctx: HandlerContext) { ctx.pendingToolSpans.delete(key) } } - ctx.sessionInputs.delete(sessionID) + ctx.pendingRuns.delete(sessionID) const msgPrefix = `${sessionID}:` for (const [key, span] of ctx.messageSpans) { if (key.startsWith(msgPrefix)) { @@ -128,6 +194,23 @@ export function handleSessionIdle(e: EventSessionIdle, ctx: HandlerContext) { sessionSpan.end() ctx.sessionSpans.delete(sessionID) } + const runID = ctx.activeRuns.get(sessionID) + if (runID) ctx.activeRuns.delete(sessionID) + const runSpan = runID ? ctx.runSpans.get(runID) : undefined + if (runSpan) { + if (totals) { + runSpan.setAttributes({ + [AGENT_NAME]: totals.agent, + "agent.type": totals.agentType, + "session.total_tokens": totals.tokens, + "session.total_cost_usd": totals.cost, + "session.total_messages": totals.messages, + }) + } + runSpan.setStatus({ code: SpanStatusCode.OK }) + runSpan.end() + ctx.runSpans.delete(runID!) + } ctx.emitLog({ severityNumber: SeverityNumber.INFO, @@ -173,6 +256,16 @@ export function handleSessionError(e: EventSessionError, ctx: HandlerContext) { sessionSpan.end() ctx.sessionSpans.delete(rawID) } + const runID = ctx.activeRuns.get(rawID) + if (runID) ctx.activeRuns.delete(rawID) + const runSpan = runID ? ctx.runSpans.get(runID) : undefined + if (runSpan) { + if (totals) runSpan.setAttributes({ [AGENT_NAME]: totals.agent, "agent.type": totals.agentType }) + runSpan.setStatus({ code: SpanStatusCode.ERROR, message: error }) + runSpan.setAttribute("error", error) + runSpan.end() + ctx.runSpans.delete(runID!) + } } ctx.emitLog({ diff --git a/src/index.ts b/src/index.ts index 87fafd7..1ca32cf 100644 --- a/src/index.ts +++ b/src/index.ts @@ -21,11 +21,12 @@ import { loadConfig, parseAttributePairs, resolveHelperPath, resolveLogLevel } f import { probeEndpoint } from "./probe.ts" import { setupOtel, createInstruments } from "./otel.ts" import { remoteParentContext } from "./trace-context.ts" -import { handleSessionCreated, handleSessionIdle, handleSessionError, handleSessionStatus } from "./handlers/session.ts" +import { handleSessionCreated, handleSessionIdle, handleSessionError, handleSessionStatus, handleRunStarted } from "./handlers/session.ts" import { handleMessageUpdated, handleMessagePartUpdated, startMessageSpan } from "./handlers/message.ts" import { handlePermissionUpdated, handlePermissionReplied } from "./handlers/permission.ts" import { handleSessionDiff, handleCommandExecuted } from "./handlers/activity.ts" -import { agentAttrs, getSessionAgentMeta } from "./util.ts" +import { agentAttrs, getSessionAgentMeta, setBoundedMap } from "./util.ts" +import type { SessionTotals } from "./types.ts" const PLUGIN_VERSION: string = (pkg as { version?: string }).version ?? "unknown" @@ -103,9 +104,15 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree const pendingPermissions = new Map() const sessionTotals = new Map() const sessionDiffTotals = new Map() + const runSpans = new Map() + const runSpanContexts = new Map() + const activeRuns = new Map() + const assistantRuns = new Map() + const pendingRuns = new Map() + const runInputs = new Map() const sessionSpans = new Map() + const sessionSpanContexts = new Map() const messageSpans = new Map() - const sessionInputs = new Map() const messageOutputs = new Map() const { disabledMetrics, disabledTraces } = config const commonAttrs = { @@ -139,9 +146,15 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree tracer, tracePrefix: config.metricPrefix, rootContext, + runSpans, + runSpanContexts, + activeRuns, + assistantRuns, + pendingRuns, + runInputs, sessionSpans, + sessionSpanContexts, messageSpans, - sessionInputs, messageOutputs, } @@ -183,9 +196,18 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree "chat.message": safe("chat.message", async (input, output) => { const agent = input.agent ?? "unknown" + const startTime = Date.now() + const existingTotals = sessionTotals.get(input.sessionID) + const nextTotals: SessionTotals = { + startMs: existingTotals?.startMs ?? startTime, + tokens: existingTotals?.tokens ?? 0, + cost: existingTotals?.cost ?? 0, + messages: existingTotals?.messages ?? 0, + agent, + agentType: existingTotals?.agentType ?? "primary", + } + setBoundedMap(sessionTotals, input.sessionID, nextTotals) const { agentType } = getSessionAgentMeta(input.sessionID, ctx) - const totals = sessionTotals.get(input.sessionID) - if (totals) totals.agent = agent const sessionSpan = sessionSpans.get(input.sessionID) if (sessionSpan) sessionSpan.setAttributes({ [AGENT_NAME]: agent, "agent.type": agentType }) const promptText = output.parts.map((part) => { @@ -202,13 +224,33 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree return "" } }).filter(Boolean).join("\n") - sessionInputs.set(input.sessionID, promptText) + if (!sessionSpan) { + const model = input.model ? `${input.model.providerID}/${input.model.modelID}` : "unknown" + if (input.messageID) { + handleRunStarted( + input.messageID, + input.sessionID, + agent, + promptText, + model, + startTime, + ctx, + ) + } else { + setBoundedMap(pendingRuns, input.sessionID, { + agent, + promptText, + model, + startTime, + }) + } + } const promptLength = promptText.length emitLog({ severityNumber: SeverityNumber.INFO, severityText: "INFO", - timestamp: Date.now(), - observedTimestamp: Date.now(), + timestamp: startTime, + observedTimestamp: startTime, body: "user_prompt", attributes: { "event.name": "user_prompt", @@ -252,10 +294,26 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree case "message.updated": { const msgEvt = event as EventMessageUpdated const info = msgEvt.properties.info + if (info.role === "user") { + const pendingRun = pendingRuns.get(info.sessionID) + if (!sessionSpans.has(info.sessionID) && (pendingRun || activeRuns.get(info.sessionID) !== info.id)) { + handleRunStarted( + info.id, + info.sessionID, + pendingRun?.agent ?? info.agent, + pendingRun?.promptText ?? "", + pendingRun?.model ?? `${info.model.providerID}/${info.model.modelID}`, + pendingRun?.startTime ?? info.time.created, + ctx, + ) + } + break + } if (info.role === "assistant" && !info.time?.completed) { startMessageSpan( info.sessionID, info.id, + info.parentID, info.modelID ?? "unknown", info.providerID ?? "unknown", info.time?.created ?? Date.now(), diff --git a/src/types.ts b/src/types.ts index df7d66c..2aa33cb 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,4 +1,4 @@ -import type { Context, Counter, Gauge, Histogram, Span, Tracer } from "@opentelemetry/api" +import type { Context, Counter, Gauge, Histogram, Span, SpanContext, Tracer } from "@opentelemetry/api" import type { LogRecord } from "@opentelemetry/api-logs" /** Numeric priority map for log levels; higher value = higher severity. */ @@ -67,6 +67,14 @@ export type SessionTotals = { agentType: SessionAgentType } +/** Pending root-run metadata captured from `chat.message` until the user message ID is known. */ +export type PendingRun = { + agent: string + promptText: string + model: string + startTime: number +} + /** Shared context threaded through every event handler. */ export type HandlerContext = { log: PluginLogger @@ -82,8 +90,14 @@ export type HandlerContext = { tracer: Tracer tracePrefix: string rootContext: () => Context + runSpans: Map + runSpanContexts: Map + activeRuns: Map + assistantRuns: Map + pendingRuns: Map + runInputs: Map sessionSpans: Map + sessionSpanContexts: Map messageSpans: Map - sessionInputs: Map messageOutputs: Map } diff --git a/src/util.ts b/src/util.ts index f79fe97..3dc1ce5 100644 --- a/src/util.ts +++ b/src/util.ts @@ -1,3 +1,4 @@ +import { trace } from "@opentelemetry/api" import { MAX_PENDING } from "./types.ts" import type { HandlerContext, SessionAgentType } from "./types.ts" @@ -22,6 +23,35 @@ export function setBoundedMap(map: Map, key: K, value: V) { map.set(key, value) } +/** Resolves a root-run context from the live span first, then from the retained ended span context. */ +export function resolveRunTraceContext(runID: string, ctx: Pick) { + const baseCtx = ctx.rootContext() + const runSpan = ctx.runSpans.get(runID) + if (runSpan) return trace.setSpan(baseCtx, runSpan) + const runSpanContext = ctx.runSpanContexts.get(runID) + return runSpanContext ? trace.setSpanContext(baseCtx, runSpanContext) : baseCtx +} + +/** Resolves the best available trace parent for a session event or message/tool child span. */ +export function resolveSessionTraceContext( + sessionID: string, + ctx: HandlerContext, + input?: { assistantMessageID?: string; runID?: string }, +) { + const baseCtx = ctx.rootContext() + const sessionSpan = ctx.sessionSpans.get(sessionID) + if (sessionSpan) return trace.setSpan(baseCtx, sessionSpan) + const sessionSpanContext = ctx.sessionSpanContexts.get(sessionID) + if (sessionSpanContext) return trace.setSpanContext(baseCtx, sessionSpanContext) + if (input?.runID) return resolveRunTraceContext(input.runID, ctx) + const assistantRunID = input?.assistantMessageID + ? ctx.assistantRuns.get(input.assistantMessageID) + : undefined + if (assistantRunID) return resolveRunTraceContext(assistantRunID, ctx) + const activeRunID = ctx.activeRuns.get(sessionID) + return activeRunID ? resolveRunTraceContext(activeRunID, ctx) : baseCtx +} + /** * Returns `true` if the metric name (without prefix) is not in the disabled set. * The `name` should be the suffix after the metric prefix, e.g. `"session.count"`. diff --git a/tests/handlers/disabled-signals.test.ts b/tests/handlers/disabled-signals.test.ts index e7b4852..0bd847e 100644 --- a/tests/handlers/disabled-signals.test.ts +++ b/tests/handlers/disabled-signals.test.ts @@ -106,7 +106,7 @@ describe("disabled traces", () => { test("disabling all three trace types suppresses every span", async () => { const { ctx, tracer } = makeCtx("proj_test", [], ["session", "llm", "tool"]) handleSessionCreated(makeSessionCreated("ses_1"), ctx) - startMessageSpan("ses_1", "msg_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) await handleMessagePartUpdated(makeToolPartUpdated("running"), ctx) expect(tracer.spans).toHaveLength(0) }) diff --git a/tests/handlers/spans.test.ts b/tests/handlers/spans.test.ts index c980bcb..f80a7fa 100644 --- a/tests/handlers/spans.test.ts +++ b/tests/handlers/spans.test.ts @@ -16,7 +16,7 @@ import { TOOL_NAME, } from "@arizeai/openinference-semantic-conventions" import type { Span } from "@opentelemetry/api" -import { handleSessionCreated, handleSessionIdle, handleSessionError } from "../../src/handlers/session.ts" +import { handleSessionCreated, handleSessionIdle, handleSessionError, handleRunStarted } from "../../src/handlers/session.ts" import { handleMessageUpdated, handleMessagePartUpdated, startMessageSpan } from "../../src/handlers/message.ts" import { remoteParentContext } from "../../src/trace-context.ts" import { makeCtx, makeTracer, type SpySpan } from "../helpers.ts" @@ -50,6 +50,7 @@ function makeSessionError(sessionID?: string, error?: { name: string }): EventSe function makeAssistantMessageUpdated(overrides: { id?: string + parentID?: string sessionID?: string modelID?: string providerID?: string @@ -63,6 +64,7 @@ function makeAssistantMessageUpdated(overrides: { properties: { info: { id: overrides.id ?? "msg_1", + parentID: overrides.parentID ?? "user_1", role: "assistant", sessionID: overrides.sessionID ?? "ses_1", modelID: overrides.modelID ?? "claude-3-5-sonnet", @@ -78,9 +80,10 @@ function makeAssistantMessageUpdated(overrides: { function makeToolPartUpdated( status: "running" | "completed" | "error", - overrides: { sessionID?: string; callID?: string; tool?: string; startMs?: number; endMs?: number; output?: string } = {}, + overrides: { sessionID?: string; messageID?: string; callID?: string; tool?: string; startMs?: number; endMs?: number; output?: string } = {}, ): EventMessagePartUpdated { const sessionID = overrides.sessionID ?? "ses_1" + const messageID = overrides.messageID ?? "msg_1" const callID = overrides.callID ?? "call_1" const start = overrides.startMs ?? 1000 const end = overrides.endMs ?? 2000 @@ -92,38 +95,37 @@ function makeToolPartUpdated( : { status: "error", time: { start, end }, error: "fail" } return { type: "message.part.updated", - properties: { part: { type: "tool", sessionID, callID, tool: overrides.tool ?? "bash", state } }, + properties: { part: { type: "tool", sessionID, messageID, callID, tool: overrides.tool ?? "bash", state } }, } as unknown as EventMessagePartUpdated } describe("session spans", () => { - test("starts a session span on session.created", () => { + test("does not start a root trace span on session.created for primary sessions", () => { const { ctx, tracer } = makeCtx() handleSessionCreated(makeSessionCreated("ses_1", 5000), ctx) - expect(tracer.spans).toHaveLength(1) - expect(tracer.spans[0]!.name).toBe("opencode.session") - expect(tracer.spans[0]!.startTime).toBe(5000) - expect(ctx.sessionSpans.has("ses_1")).toBe(true) + expect(tracer.spans).toHaveLength(0) + expect(ctx.sessionSpans.has("ses_1")).toBe(false) }) - test("session span carries session.id attribute", () => { + test("subagent session span carries session.id attribute", () => { const { ctx, tracer } = makeCtx("proj_test", [], [], true, { team: "platform" }) - handleSessionCreated(makeSessionCreated("ses_1"), ctx) - expect(tracer.spans[0]!.attributes["session.id"]).toBe("ses_1") - expect(tracer.spans[0]!.attributes[SESSION_ID]).toBe("ses_1") - expect(tracer.spans[0]!.attributes["team"]).toBe("platform") + handleRunStarted("user_parent", "ses_parent", "build", "prompt", "anthropic/claude", 900, ctx) + handleSessionCreated(makeSessionCreated("ses_1", 1000, "ses_parent"), ctx) + expect(tracer.spans[1]!.attributes["session.id"]).toBe("ses_1") + expect(tracer.spans[1]!.attributes[SESSION_ID]).toBe("ses_1") + expect(tracer.spans[1]!.attributes["team"]).toBe("platform") }) - test("session span is tagged as an OpenInference agent span", () => { + test("run span is tagged as an OpenInference agent span", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) expect(tracer.spans[0]!.attributes[OPENINFERENCE_SPAN_KIND]).toBe(OpenInferenceSpanKind.AGENT) - expect(tracer.spans[0]!.attributes[AGENT_NAME]).toBe("unknown") + expect(tracer.spans[0]!.attributes[AGENT_NAME]).toBe("build") }) - test("session span carries is_subagent=false for root session", () => { + test("run span carries is_subagent=false for root session", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_root"), ctx) + handleRunStarted("user_root", "ses_root", "build", "prompt", "anthropic/claude", 1000, ctx) expect(tracer.spans[0]!.attributes["session.is_subagent"]).toBe(false) }) @@ -133,17 +135,17 @@ describe("session spans", () => { expect(tracer.spans[0]!.attributes["session.is_subagent"]).toBe(true) }) - test("root session span is parented to injected remote context", () => { + test("run span is parented to injected remote context", () => { const { ctx, tracer } = makeCtx() const rootContext = remoteParentContext("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", undefined) expect(rootContext).toBeDefined() ctx.rootContext = () => rootContext! - handleSessionCreated(makeSessionCreated("ses_1"), ctx) - expect(tracer.spans[0]!.parentSpan?.spanContext().traceId).toBe("0af7651916cd43dd8448eb211c80319c") - expect(tracer.spans[0]!.parentSpan?.spanContext().spanId).toBe("b7ad6b7169203331") + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) + expect(tracer.spans[0]!.parentSpanContext?.traceId).toBe("0af7651916cd43dd8448eb211c80319c") + expect(tracer.spans[0]!.parentSpanContext?.spanId).toBe("b7ad6b7169203331") }) - test("root session span resolves root context at span creation", () => { + test("run span resolves root context at span creation", () => { const { ctx, tracer } = makeCtx() let rootContext = context.active() ctx.rootContext = () => rootContext @@ -152,24 +154,24 @@ describe("session spans", () => { spanId: "00f067aa0ba902b7", traceFlags: TraceFlags.SAMPLED, }) - handleSessionCreated(makeSessionCreated("ses_1"), ctx) - expect(tracer.spans[0]!.parentSpan?.spanContext().traceId).toBe("4bf92f3577b34da6a3ce929d0e0e4736") - expect(tracer.spans[0]!.parentSpan?.spanContext().spanId).toBe("00f067aa0ba902b7") + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) + expect(tracer.spans[0]!.parentSpanContext?.traceId).toBe("4bf92f3577b34da6a3ce929d0e0e4736") + expect(tracer.spans[0]!.parentSpanContext?.spanId).toBe("00f067aa0ba902b7") }) - test("ends session span with OK status on session.idle", () => { + test("ends run span with OK status on session.idle", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) handleSessionIdle(makeSessionIdle("ses_1"), ctx) const span = tracer.spans[0]! expect(span.ended).toBe(true) expect(span.status.code).toBe(SpanStatusCode.OK) - expect(ctx.sessionSpans.has("ses_1")).toBe(false) + expect(ctx.runSpans.has("user_1")).toBe(false) }) - test("sets session total attributes before ending on idle", () => { + test("sets session total attributes on the run span before ending on idle", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) ctx.sessionTotals.set("ses_1", { startMs: Date.now() - 100, tokens: 250, cost: 0.05, messages: 3, agent: "build", agentType: "primary" }) handleSessionIdle(makeSessionIdle("ses_1"), ctx) const span = tracer.spans[0]! @@ -180,19 +182,19 @@ describe("session spans", () => { expect(span.attributes["agent.type"]).toBe("primary") }) - test("ends session span with ERROR status on session.error", () => { + test("ends run span with ERROR status on session.error", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) handleSessionError(makeSessionError("ses_1", { name: "NetworkError" }), ctx) const span = tracer.spans[0]! expect(span.ended).toBe(true) expect(span.status.code).toBe(SpanStatusCode.ERROR) - expect(ctx.sessionSpans.has("ses_1")).toBe(false) + expect(ctx.runSpans.has("user_1")).toBe(false) }) - test("error message is propagated to session span status", () => { + test("error message is propagated to run span status", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) handleSessionError(makeSessionError("ses_1", { name: "TimeoutError" }), ctx) expect(tracer.spans[0]!.status.message).toBe("TimeoutError") }) @@ -205,25 +207,51 @@ describe("session spans", () => { test("session.error with undefined sessionID does not end any span", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) handleSessionError(makeSessionError(undefined, { name: "UnknownError" }), ctx) - expect(ctx.sessionSpans.has("ses_1")).toBe(true) + expect(ctx.runSpans.has("user_1")).toBe(true) expect(tracer.spans[0]!.ended).toBe(false) }) - test("subagent span — parent session span is in sessionSpans before child is created", () => { + test("subagent span is parented to the active run span", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_parent"), ctx) + handleRunStarted("user_parent", "ses_parent", "build", "prompt", "anthropic/claude", 1000, ctx) handleSessionCreated(makeSessionCreated("ses_child", 2000, "ses_parent"), ctx) expect(tracer.spans).toHaveLength(2) expect(tracer.spans[1]!.name).toBe("opencode.session") expect(tracer.spans[1]!.parentSpan).toBe(tracer.spans[0]) }) - test("subagent span — no error when parent session span is absent", () => { + test("subagent span falls back to a root trace when parent run is absent", () => { const { ctx, tracer } = makeCtx() expect(() => handleSessionCreated(makeSessionCreated("ses_child", 1000, "ses_missing_parent"), ctx)).not.toThrow() expect(tracer.spans).toHaveLength(1) + expect(tracer.spans[0]!.parentSpan).toBeUndefined() + expect(tracer.spans[0]!.parentSpanContext).toBeUndefined() + }) + + test("late child spans reuse the ended run trace context", () => { + const { ctx, tracer } = makeCtx() + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 1000, ctx) + handleSessionIdle(makeSessionIdle("ses_1"), ctx) + startMessageSpan("ses_1", "msg_late", "user_1", "claude", "anthropic", 2000, ctx) + expect(tracer.spans).toHaveLength(2) + expect(tracer.spans[1]!.parentSpanContext).toBeDefined() + expect(tracer.spans[1]!.parentSpanContext?.spanId).toBe(tracer.spans[0]!.spanContext().spanId) + expect(tracer.spans[1]!.parentSpanContext?.traceId).toBe(tracer.spans[0]!.spanContext().traceId) + }) + + test("late child events stay on the previous run after a new run starts", () => { + const { ctx, tracer } = makeCtx() + handleRunStarted("user_1", "ses_1", "build", "prompt one", "anthropic/claude", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1100, ctx) + handleSessionIdle(makeSessionIdle("ses_1"), ctx) + handleRunStarted("user_2", "ses_1", "build", "prompt two", "anthropic/claude", 2000, ctx) + handleMessagePartUpdated(makeToolPartUpdated("running", { sessionID: "ses_1", messageID: "msg_1", callID: "call_late" }), ctx) + expect(tracer.spans).toHaveLength(4) + expect(tracer.spans[3]!.name).toBe("opencode.tool.bash") + expect(tracer.spans[3]!.parentSpanContext?.spanId).toBe(tracer.spans[0]!.spanContext().spanId) + expect(tracer.spans[3]!.parentSpanContext?.spanId).not.toBe(tracer.spans[2]!.spanContext().spanId) }) }) @@ -300,7 +328,7 @@ describe("tool spans", () => { test("tool span is parented to session span when available", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 900, ctx) handleMessagePartUpdated(makeToolPartUpdated("running", { sessionID: "ses_1" }), ctx) expect(tracer.spans).toHaveLength(2) expect(tracer.spans[1]!.name).toBe("opencode.tool.bash") @@ -309,7 +337,7 @@ describe("tool spans", () => { test("out-of-order tool span is parented to session span when available", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 900, ctx) handleMessagePartUpdated(makeToolPartUpdated("completed", { sessionID: "ses_1", startMs: 500, endMs: 1500 }), ctx) expect(tracer.spans).toHaveLength(2) expect(tracer.spans[1]!.name).toBe("opencode.tool.bash") @@ -320,7 +348,7 @@ describe("tool spans", () => { describe("message (LLM) spans", () => { test("startMessageSpan creates an llm span", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) expect(tracer.spans).toHaveLength(1) expect(tracer.spans[0]!.name).toBe("opencode.llm") expect(ctx.messageSpans.has("ses_1:msg_1")).toBe(true) @@ -328,7 +356,7 @@ describe("message (LLM) spans", () => { test("startMessageSpan sets OpenInference LLM attributes", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "gpt-4o", "openai", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "gpt-4o", "openai", 1000, ctx) expect(tracer.spans[0]!.attributes[OPENINFERENCE_SPAN_KIND]).toBe(OpenInferenceSpanKind.LLM) expect(tracer.spans[0]!.attributes[LLM_SYSTEM]).toBe("openai") expect(tracer.spans[0]!.attributes[LLM_PROVIDER]).toBe("openai") @@ -337,14 +365,14 @@ describe("message (LLM) spans", () => { test("startMessageSpan is a no-op when span already exists for sessionID:messageID", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "claude", "anthropic", 1000, ctx) - startMessageSpan("ses_1", "msg_1", "claude", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1000, ctx) expect(tracer.spans).toHaveLength(1) }) test("handleMessageUpdated ends message span on completion", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) handleMessageUpdated(makeAssistantMessageUpdated({ id: "msg_1", time: { created: 1000, completed: 2000 } }), ctx) const span = tracer.spans[0]! expect(span.ended).toBe(true) @@ -354,14 +382,14 @@ describe("message (LLM) spans", () => { test("handleMessageUpdated sets OK status on success", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) handleMessageUpdated(makeAssistantMessageUpdated({ id: "msg_1" }), ctx) expect(tracer.spans[0]!.status.code).toBe(SpanStatusCode.OK) }) test("handleMessageUpdated sets ERROR status on api error", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) handleMessageUpdated(makeAssistantMessageUpdated({ id: "msg_1", error: { name: "RateLimitError" } }), ctx) expect(tracer.spans[0]!.status.code).toBe(SpanStatusCode.ERROR) expect(tracer.spans[0]!.status.message).toBe("RateLimitError") @@ -369,7 +397,7 @@ describe("message (LLM) spans", () => { test("handleMessageUpdated sets OpenInference token attributes on span", () => { const { ctx, tracer } = makeCtx() - startMessageSpan("ses_1", "msg_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) ctx.sessionTotals.set("ses_1", { startMs: 0, tokens: 0, cost: 0, messages: 0, agent: "review", agentType: "subagent" }) handleMessageUpdated( makeAssistantMessageUpdated({ @@ -399,8 +427,8 @@ describe("message (LLM) spans", () => { test("message span is parented to session span when available", () => { const { ctx, tracer } = makeCtx() - handleSessionCreated(makeSessionCreated("ses_1"), ctx) - startMessageSpan("ses_1", "msg_1", "claude", "anthropic", 1000, ctx) + handleRunStarted("user_1", "ses_1", "build", "prompt", "anthropic/claude", 900, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1000, ctx) expect(tracer.spans).toHaveLength(2) expect(tracer.spans[1]!.name).toBe("opencode.llm") expect(tracer.spans[1]!.parentSpan).toBe(tracer.spans[0]) @@ -444,7 +472,7 @@ describe("orphaned span cleanup", () => { test("pending message spans are ended with ERROR on session.idle", () => { const { ctx, tracer } = makeCtx() handleSessionCreated(makeSessionCreated("ses_1"), ctx) - startMessageSpan("ses_1", "msg_orphan", "claude", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_orphan", "user_1", "claude", "anthropic", 1000, ctx) handleSessionIdle(makeSessionIdle("ses_1"), ctx) expect(ctx.messageSpans.has("ses_1:msg_orphan")).toBe(false) const msgSpan = tracer.spans.find(s => s.name === "opencode.llm")! @@ -455,7 +483,7 @@ describe("orphaned span cleanup", () => { test("pending message spans are ended with ERROR on session.error", () => { const { ctx, tracer } = makeCtx() handleSessionCreated(makeSessionCreated("ses_1"), ctx) - startMessageSpan("ses_1", "msg_orphan", "claude", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_orphan", "user_1", "claude", "anthropic", 1000, ctx) handleSessionError(makeSessionError("ses_1"), ctx) expect(ctx.messageSpans.has("ses_1:msg_orphan")).toBe(false) const msgSpan = tracer.spans.find(s => s.name === "opencode.llm")! @@ -499,16 +527,17 @@ describe("OPENCODE_DISABLE_TRACES=session", () => { test("llm spans become root spans (no parent) when session traces disabled but llm enabled", () => { const { ctx, tracer } = makeCtx("proj_test", [], ["session"]) handleSessionCreated(makeSessionCreated("ses_1"), ctx) - startMessageSpan("ses_1", "msg_1", "claude", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1000, ctx) expect(tracer.spans).toHaveLength(1) expect(tracer.spans[0]!.name).toBe("opencode.llm") + expect(tracer.spans[0]!.parentSpan).toBeUndefined() }) }) describe("OPENCODE_DISABLE_TRACES=llm", () => { test("startMessageSpan is a no-op", () => { const { ctx, tracer } = makeCtx("proj_test", [], ["llm"]) - startMessageSpan("ses_1", "msg_1", "claude", "anthropic", 1000, ctx) + startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1000, ctx) expect(tracer.spans).toHaveLength(0) expect(ctx.messageSpans.has("msg_1")).toBe(false) }) @@ -539,8 +568,7 @@ describe("OPENCODE_DISABLE_TRACES=llm", () => { test("session spans still created when only llm disabled", () => { const { ctx, tracer } = makeCtx("proj_test", [], ["llm"]) handleSessionCreated(makeSessionCreated("ses_1"), ctx) - expect(tracer.spans).toHaveLength(1) - expect(tracer.spans[0]!.name).toBe("opencode.session") + expect(tracer.spans).toHaveLength(0) }) }) @@ -583,7 +611,6 @@ describe("OPENCODE_DISABLE_TRACES=tool", () => { test("session spans still created when only tool disabled", () => { const { ctx, tracer } = makeCtx("proj_test", [], ["tool"]) handleSessionCreated(makeSessionCreated("ses_1"), ctx) - expect(tracer.spans).toHaveLength(1) - expect(tracer.spans[0]!.name).toBe("opencode.session") + expect(tracer.spans).toHaveLength(0) }) }) diff --git a/tests/helpers.ts b/tests/helpers.ts index 0113660..14d2bff 100644 --- a/tests/helpers.ts +++ b/tests/helpers.ts @@ -36,6 +36,7 @@ export type SpySpan = { status: SpanStatus attributes: Record parentSpan: SpySpan | undefined + parentSpanContext: SpanContext | undefined setStatus(status: SpanStatus): SpySpan setAttribute(key: string, value: unknown): SpySpan setAttributes(attrs: Attributes): SpySpan @@ -80,7 +81,18 @@ function makePluginLog(): SpyPluginLog { return spy } -function makeSpan(name: string, startTime?: number, parentSpan?: SpySpan): SpySpan { +function makeSpan( + name: string, + startTime?: number, + parentSpan?: SpySpan, + parentSpanContext?: SpanContext, + ownSpanContext?: SpanContext, +): SpySpan { + const context = ownSpanContext ?? { + traceId: "00000000000000000000000000000001", + spanId: "0000000000000001", + traceFlags: 1, + } const span: SpySpan = { name, startTime, @@ -89,12 +101,13 @@ function makeSpan(name: string, startTime?: number, parentSpan?: SpySpan): SpySp status: { code: SpanStatusCode.UNSET }, attributes: {}, parentSpan, + parentSpanContext, setStatus(s) { span.status = s; return span }, setAttribute(k, v) { span.attributes[k] = v; return span }, setAttributes(attrs) { Object.assign(span.attributes, attrs); return span }, end(t) { span.ended = true; span.endTime = t }, isRecording() { return !span.ended }, - spanContext() { return { traceId: "00000000000000000000000000000001", spanId: "0000000000000001", traceFlags: 1 } }, + spanContext() { return context }, addEvent() { return span }, recordException() { return span }, updateName(n) { span.name = n; return span }, @@ -103,14 +116,24 @@ function makeSpan(name: string, startTime?: number, parentSpan?: SpySpan): SpySp } export function makeTracer(): SpyTracer { + let nextSpanID = 1 const tracer: SpyTracer = { spans: [], startSpan(name, options, ctx) { const parentFromCtx = ctx ? trace.getSpan(ctx) as SpySpan | undefined : undefined + const parentSpanContext = ctx ? trace.getSpanContext(ctx) ?? undefined : undefined + const ownSpanContext: SpanContext = { + traceId: parentSpanContext?.traceId ?? "00000000000000000000000000000001", + spanId: (nextSpanID++).toString(16).padStart(16, "0"), + traceFlags: parentSpanContext?.traceFlags ?? 1, + ...(parentSpanContext?.traceState ? { traceState: parentSpanContext.traceState } : {}), + } const span = makeSpan( name, typeof options?.startTime === "number" ? options.startTime : undefined, parentFromCtx, + parentSpanContext, + ownSpanContext, ) if (options?.attributes) Object.assign(span.attributes, options.attributes) tracer.spans.push(span) @@ -209,9 +232,15 @@ export function makeCtx( tracer: tracer as unknown as Tracer, tracePrefix: "opencode.", rootContext: () => ROOT_CONTEXT, + runSpans: new Map(), + runSpanContexts: new Map(), + activeRuns: new Map(), + assistantRuns: new Map(), + pendingRuns: new Map(), + runInputs: new Map(), sessionSpans: new Map(), + sessionSpanContexts: new Map(), messageSpans: new Map(), - sessionInputs: new Map(), messageOutputs: new Map(), }