fix: correlate log records with their span context

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,

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)

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<HandlerContext, "rootContext">,
+): 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"`.

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)

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)