diff --git a/src/handlers/message.ts b/src/handlers/message.ts index 248228a..182558c 100644 --- a/src/handlers/message.ts +++ b/src/handlers/message.ts @@ -1,5 +1,5 @@ import { SeverityNumber } from "@opentelemetry/api-logs" -import { SpanStatusCode, SpanKind } from "@opentelemetry/api" +import { SpanStatusCode, SpanKind, type Span } from "@opentelemetry/api" import type { AssistantMessage, EventMessageUpdated, EventMessagePartUpdated, ToolPart } from "@opencode-ai/sdk" import { AGENT_NAME, @@ -29,6 +29,7 @@ import { } from "@arizeai/openinference-semantic-conventions" import { agentAttrs, + contextForSpanContext, errorSummary, genAiProviderName, setBoundedMap, @@ -119,6 +120,7 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext const msgKey = `${sessionID}:${assistant.id}` const msgSpan = ctx.messageSpans.get(msgKey) + const msgSpanContext = msgSpan?.spanContext() if (msgSpan) { const outputText = ctx.messageOutputs.get(msgKey) msgSpan.setAttributes({ @@ -166,6 +168,7 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext timestamp: assistant.time.created, observedTimestamp: Date.now(), body: "api_error", + context: contextForSpanContext(msgSpanContext, ctx), attributes: { "event.name": "api_error", "session.id": sessionID, @@ -193,6 +196,7 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext timestamp: assistant.time.created, observedTimestamp: Date.now(), body: "api_request", + context: contextForSpanContext(msgSpanContext, ctx), attributes: { "event.name": "api_request", "session.id": sessionID, @@ -334,29 +338,28 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle }) } + let toolSpan: Span | undefined if (isTraceEnabled("tool", ctx)) { - const toolSpan = pending?.span ?? (() => { - return ctx.tracer.startSpan( - `${ctx.tracePrefix}tool.${toolPart.tool}`, - { - startTime: start, - kind: SpanKind.INTERNAL, - attributes: { - [OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL, - [SESSION_ID]: toolPart.sessionID, - [TOOL_ID]: toolPart.callID, - [TOOL_NAME]: toolPart.tool, - [TOOL_PARAMETERS]: JSON.stringify(toolPart.state.input), - [INPUT_VALUE]: JSON.stringify(toolPart.state.input), - [INPUT_MIME_TYPE]: MimeType.JSON, - ...ctx.commonAttrs, - }, + toolSpan = pending?.span ?? ctx.tracer.startSpan( + `${ctx.tracePrefix}tool.${toolPart.tool}`, + { + startTime: start, + kind: SpanKind.INTERNAL, + attributes: { + [OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL, + [SESSION_ID]: toolPart.sessionID, + [TOOL_ID]: toolPart.callID, + [TOOL_NAME]: toolPart.tool, + [TOOL_PARAMETERS]: JSON.stringify(toolPart.state.input), + [INPUT_VALUE]: JSON.stringify(toolPart.state.input), + [INPUT_MIME_TYPE]: MimeType.JSON, + ...ctx.commonAttrs, }, - resolveSessionTraceContext(toolPart.sessionID, ctx, { - assistantMessageID: toolPart.messageID, - }), - ) - })() + }, + resolveSessionTraceContext(toolPart.sessionID, ctx, { + assistantMessageID: toolPart.messageID, + }), + ) toolSpan.setAttributes({ [AGENT_NAME]: agentName, "agent.type": agentType }) toolSpan.setAttribute("tool.success", success) if (success) { @@ -389,6 +392,7 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle timestamp: start, observedTimestamp: Date.now(), body: "tool_result", + context: contextForSpanContext(toolSpan?.spanContext(), ctx), attributes: { "event.name": "tool_result", "session.id": toolPart.sessionID, diff --git a/src/index.ts b/src/index.ts index 194270d..49ebd22 100644 --- a/src/index.ts +++ b/src/index.ts @@ -26,7 +26,7 @@ import { handleMessageUpdated, handleMessagePartUpdated, startMessageSpan } from import { handlePermissionUpdated, handlePermissionReplied } from "./handlers/permission.ts" import { handleSessionDiff, handleCommandExecuted } from "./handlers/activity.ts" import { handleChatHeaders } from "./handlers/chat-headers.ts" -import { agentAttrs, getSessionAgentMeta, setBoundedMap } from "./util.ts" +import { agentAttrs, getSessionAgentMeta, resolveLogContext, setBoundedMap } from "./util.ts" import type { SessionTotals } from "./types.ts" const PLUGIN_VERSION: string = (pkg as { version?: string }).version ?? "unknown" @@ -94,7 +94,8 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree const logger = logs.getLogger("com.opencode") const emitLog: HandlerContext["emitLog"] = (record) => { if (!config.logsEnabled) return - logger.emit(record) + const context = resolveLogContext(record, ctx) + logger.emit(context ? { ...record, context } : record) } const tracer = trace.getTracer("com.opencode") const remoteContext = remoteParentContext(config.traceparent, config.tracestate) diff --git a/src/util.ts b/src/util.ts index 27958a0..cb74dc0 100644 --- a/src/util.ts +++ b/src/util.ts @@ -1,4 +1,5 @@ -import { trace } from "@opentelemetry/api" +import { trace, type Context, type SpanContext } from "@opentelemetry/api" +import type { LogRecord } from "@opentelemetry/api-logs" import { MAX_PENDING } from "./types.ts" import type { HandlerContext, SessionAgentType } from "./types.ts" @@ -68,6 +69,33 @@ export function resolveSessionTraceContext( return activeRunID ? resolveRunTraceContext(activeRunID, ctx) : baseCtx } +/** + * Wraps a span context in an OTel `Context` anchored at the plugin root. + * Returns `undefined` when there is no span context, so callers can leave the + * log record's `context` unset and fall back to `resolveLogContext`. + */ +export function contextForSpanContext( + spanContext: SpanContext | undefined, + ctx: Pick, +): Context | undefined { + return spanContext ? trace.setSpanContext(ctx.rootContext(), spanContext) : undefined +} + +/** + * Resolves the OTel `Context` a log record should be correlated with. + * + * Log records emitted without a context inherit `context.active()`, which is + * always empty under opencode's event dispatch — so each record landed as its + * own root trace instead of nesting under the session, run, tool, or LLM span + * that produced it. An explicit `record.context` wins; otherwise the record is + * parented to the session/run context named by its `session.id` attribute. + */ +export function resolveLogContext(record: LogRecord, ctx: HandlerContext): Context | undefined { + if (record.context) return record.context + const sessionID = record.attributes?.["session.id"] + return typeof sessionID === "string" ? resolveSessionTraceContext(sessionID, ctx) : undefined +} + /** * 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/message.test.ts b/tests/handlers/message.test.ts index 0fed30a..a9661cf 100644 --- a/tests/handlers/message.test.ts +++ b/tests/handlers/message.test.ts @@ -1,5 +1,6 @@ import { describe, test, expect } from "bun:test" -import { handleMessageUpdated, handleMessagePartUpdated } from "../../src/handlers/message.ts" +import { trace } from "@opentelemetry/api" +import { handleMessageUpdated, handleMessagePartUpdated, startMessageSpan } from "../../src/handlers/message.ts" import { makeCtx } from "../helpers.ts" import type { EventMessageUpdated, EventMessagePartUpdated } from "@opencode-ai/sdk" @@ -233,6 +234,17 @@ describe("handleMessageUpdated", () => { expect(pluginLog.calls.find(c => c.level === "error")?.level).toBe("error") }) + test("correlates the api_request log record with the LLM span", async () => { + const { ctx, tracer, logger } = makeCtx() + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx, "build") + await handleMessageUpdated(makeAssistantMessageUpdated({}), ctx) + const record = logger.records.find((r) => r.body === "api_request")! + const llmSpan = tracer.spans.find((s) => s.name === "opencode.llm")! + expect(record.context).toBeDefined() + expect(trace.getSpanContext(record.context!)?.spanId).toBe(llmSpan.spanContext().spanId) + expect(trace.getSpanContext(record.context!)?.traceId).toBe(llmSpan.spanContext().traceId) + }) + test("uses assistant.time.created as log timestamp", async () => { const { ctx, logger } = makeCtx() await handleMessageUpdated( @@ -302,6 +314,17 @@ describe("handleMessagePartUpdated", () => { expect(pluginLog.calls.find(c => c.level === "error")?.level).toBe("error") }) + test("correlates the tool_result log record with the tool span", async () => { + const { ctx, tracer, logger } = makeCtx() + await handleMessagePartUpdated(makeToolPartUpdated("running"), ctx) + await handleMessagePartUpdated(makeToolPartUpdated("completed"), ctx) + const record = logger.records.find((r) => r.body === "tool_result")! + const toolSpan = tracer.spans.find((s) => s.name === "opencode.tool.bash")! + expect(record.context).toBeDefined() + expect(trace.getSpanContext(record.context!)?.spanId).toBe(toolSpan.spanContext().spanId) + expect(trace.getSpanContext(record.context!)?.traceId).toBe(toolSpan.spanContext().traceId) + }) + test("removes entry from pendingToolSpans after completion", async () => { const { ctx } = makeCtx() await handleMessagePartUpdated(makeToolPartUpdated("running"), ctx) diff --git a/tests/util.test.ts b/tests/util.test.ts index 04fada1..da3d687 100644 --- a/tests/util.test.ts +++ b/tests/util.test.ts @@ -1,6 +1,16 @@ import { describe, test, expect } from "bun:test" -import { errorSummary, genAiProviderName, setBoundedMap, isMetricEnabled, isTraceEnabled } from "../src/util.ts" +import { ROOT_CONTEXT, trace, type Span } from "@opentelemetry/api" +import { + contextForSpanContext, + errorSummary, + genAiProviderName, + resolveLogContext, + setBoundedMap, + isMetricEnabled, + isTraceEnabled, +} from "../src/util.ts" import { MAX_PENDING } from "../src/types.ts" +import { makeCtx } from "./helpers.ts" describe("errorSummary", () => { test("returns 'unknown' for undefined", () => { @@ -119,6 +129,72 @@ describe("isMetricEnabled", () => { }) }) +describe("contextForSpanContext", () => { + test("returns undefined when there is no span context", () => { + const { ctx } = makeCtx() + expect(contextForSpanContext(undefined, ctx)).toBeUndefined() + }) + + test("activates the span context on the plugin root context", () => { + const { ctx } = makeCtx() + const spanContext = { + traceId: "0000000000000000000000000000000a", + spanId: "000000000000000b", + traceFlags: 1, + } + const context = contextForSpanContext(spanContext, ctx)! + expect(trace.getSpanContext(context)).toMatchObject(spanContext) + }) +}) + +describe("resolveLogContext", () => { + test("returns undefined when the record has no session id", () => { + const { ctx } = makeCtx() + expect(resolveLogContext({ body: "orphan" }, ctx)).toBeUndefined() + }) + + test("parents a session-scoped record to the active run span", () => { + const { ctx, tracer } = makeCtx() + const runSpan = tracer.startSpan("opencode.session", {}, ROOT_CONTEXT) + ctx.runSpans.set("user_1", runSpan as unknown as Span) + ctx.activeRuns.set("ses_1", "user_1") + + const context = resolveLogContext( + { body: "user_prompt", attributes: { "session.id": "ses_1" } }, + ctx, + )! + + expect(trace.getSpanContext(context)?.spanId).toBe(runSpan.spanContext().spanId) + expect(trace.getSpanContext(context)?.traceId).toBe(runSpan.spanContext().traceId) + }) + + test("prefers an explicit record context over session resolution", () => { + const { ctx } = makeCtx() + const explicit = { + traceId: "0000000000000000000000000000000c", + spanId: "000000000000000d", + traceFlags: 1, + } + const record = { + body: "tool_result", + attributes: { "session.id": "ses_1" }, + context: trace.setSpanContext(ROOT_CONTEXT, explicit), + } + + expect(trace.getSpanContext(resolveLogContext(record, ctx)!)).toMatchObject(explicit) + }) + + test("falls back to the root context for an unknown session", () => { + const { ctx } = makeCtx() + const context = resolveLogContext( + { body: "session.idle", attributes: { "session.id": "ses_unknown" } }, + ctx, + )! + expect(context).toBeDefined() + expect(trace.getSpanContext(context)).toBeUndefined() + }) +}) + describe("isTraceEnabled", () => { test("returns true when disabled set is empty", () => { expect(isTraceEnabled("session", { disabledTraces: new Set() })).toBe(true)