Skip to content

Commit 8faf2da

Browse files
murtaza-umurtaza-u
committed
fix: ensure single trace spans across a single run
1 parent a123a37 commit 8faf2da

6 files changed

Lines changed: 708 additions & 174 deletions

File tree

‎src/handlers/message.ts‎

Lines changed: 129 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
import { SeverityNumber } from "@opentelemetry/api-logs"
22
import { SpanStatusCode, SpanKind, context, trace } from "@opentelemetry/api"
3-
import type { AssistantMessage, EventMessageUpdated, EventMessagePartUpdated, ToolPart } from "@opencode-ai/sdk"
3+
import type {
4+
AssistantMessage,
5+
EventMessageUpdated,
6+
EventMessagePartUpdated,
7+
ToolPart,
8+
} from "@opencode-ai/sdk"
49
import {
510
AGENT_NAME,
611
INPUT_MIME_TYPE,
@@ -27,8 +32,15 @@ import {
2732
TOOL_NAME,
2833
TOOL_PARAMETERS,
2934
} from "@arizeai/openinference-semantic-conventions"
30-
import { errorSummary, setBoundedMap, accumulateSessionTotals, isMetricEnabled, isTraceEnabled } from "../util.ts"
35+
import {
36+
errorSummary,
37+
setBoundedMap,
38+
accumulateSessionTotals,
39+
isMetricEnabled,
40+
isTraceEnabled,
41+
} from "../util.ts"
3142
import type { HandlerContext } from "../types.ts"
43+
import { resolveSessionTraceContext } from "./session.ts"
3244

3345
const OPENINFERENCE_SPAN_KIND = SemanticConventions.OPENINFERENCE_SPAN_KIND
3446
const LLM_FINISH_REASON = "llm.finish_reason"
@@ -58,37 +70,99 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext
5870
const duration = assistant.time.completed - assistant.time.created
5971
const agent = ctx.sessionTotals.get(sessionID)?.agent ?? "unknown"
6072

61-
const totalTokens = assistant.tokens.input + assistant.tokens.output + assistant.tokens.reasoning
62-
+ assistant.tokens.cache.read + assistant.tokens.cache.write
73+
const totalTokens =
74+
assistant.tokens.input +
75+
assistant.tokens.output +
76+
assistant.tokens.reasoning +
77+
assistant.tokens.cache.read +
78+
assistant.tokens.cache.write
6379

6480
if (isMetricEnabled("token.usage", ctx)) {
6581
const { tokenCounter } = ctx.instruments
66-
tokenCounter.add(assistant.tokens.input, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "input" })
67-
tokenCounter.add(assistant.tokens.output, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "output" })
68-
tokenCounter.add(assistant.tokens.reasoning, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "reasoning" })
69-
tokenCounter.add(assistant.tokens.cache.read, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "cacheRead" })
70-
tokenCounter.add(assistant.tokens.cache.write, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "cacheCreation" })
82+
tokenCounter.add(assistant.tokens.input, {
83+
...ctx.commonAttrs,
84+
"session.id": sessionID,
85+
model: modelID,
86+
agent,
87+
type: "input",
88+
})
89+
tokenCounter.add(assistant.tokens.output, {
90+
...ctx.commonAttrs,
91+
"session.id": sessionID,
92+
model: modelID,
93+
agent,
94+
type: "output",
95+
})
96+
tokenCounter.add(assistant.tokens.reasoning, {
97+
...ctx.commonAttrs,
98+
"session.id": sessionID,
99+
model: modelID,
100+
agent,
101+
type: "reasoning",
102+
})
103+
tokenCounter.add(assistant.tokens.cache.read, {
104+
...ctx.commonAttrs,
105+
"session.id": sessionID,
106+
model: modelID,
107+
agent,
108+
type: "cacheRead",
109+
})
110+
tokenCounter.add(assistant.tokens.cache.write, {
111+
...ctx.commonAttrs,
112+
"session.id": sessionID,
113+
model: modelID,
114+
agent,
115+
type: "cacheCreation",
116+
})
71117
}
72118

73119
if (isMetricEnabled("cost.usage", ctx)) {
74-
ctx.instruments.costCounter.add(assistant.cost, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent })
120+
ctx.instruments.costCounter.add(assistant.cost, {
121+
...ctx.commonAttrs,
122+
"session.id": sessionID,
123+
model: modelID,
124+
agent,
125+
})
75126
}
76127

77128
if (isMetricEnabled("cache.count", ctx)) {
78129
if (assistant.tokens.cache.read > 0) {
79-
ctx.instruments.cacheCounter.add(1, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "cacheRead" })
130+
ctx.instruments.cacheCounter.add(1, {
131+
...ctx.commonAttrs,
132+
"session.id": sessionID,
133+
model: modelID,
134+
agent,
135+
type: "cacheRead",
136+
})
80137
}
81138
if (assistant.tokens.cache.write > 0) {
82-
ctx.instruments.cacheCounter.add(1, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent, type: "cacheCreation" })
139+
ctx.instruments.cacheCounter.add(1, {
140+
...ctx.commonAttrs,
141+
"session.id": sessionID,
142+
model: modelID,
143+
agent,
144+
type: "cacheCreation",
145+
})
83146
}
84147
}
85148

86149
if (isMetricEnabled("message.count", ctx)) {
87-
ctx.instruments.messageCounter.add(1, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, agent })
150+
ctx.instruments.messageCounter.add(1, {
151+
...ctx.commonAttrs,
152+
"session.id": sessionID,
153+
model: modelID,
154+
agent,
155+
})
88156
}
89157

90158
if (isMetricEnabled("model.usage", ctx)) {
91-
ctx.instruments.modelUsageCounter.add(1, { ...ctx.commonAttrs, "session.id": sessionID, model: modelID, provider: providerID, agent })
159+
ctx.instruments.modelUsageCounter.add(1, {
160+
...ctx.commonAttrs,
161+
"session.id": sessionID,
162+
model: modelID,
163+
provider: providerID,
164+
agent,
165+
})
92166
}
93167

94168
accumulateSessionTotals(sessionID, totalTokens, assistant.cost, ctx)
@@ -129,7 +203,10 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext
129203
duration_ms: duration,
130204
})
131205
if (assistant.error) {
132-
msgSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorSummary(assistant.error) })
206+
msgSpan.setStatus({
207+
code: SpanStatusCode.ERROR,
208+
message: errorSummary(assistant.error),
209+
})
133210
} else {
134211
msgSpan.setStatus({ code: SpanStatusCode.OK })
135212
}
@@ -256,9 +333,6 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
256333
const toolSpan = isTraceEnabled("tool", ctx)
257334
? (() => {
258335
const sessionSpan = ctx.sessionSpans.get(toolPart.sessionID)
259-
const parentCtx = sessionSpan
260-
? trace.setSpan(context.active(), sessionSpan)
261-
: context.active()
262336
return ctx.tracer.startSpan(
263337
`${ctx.tracePrefix}tool.${toolPart.tool}`,
264338
{
@@ -275,7 +349,7 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
275349
...ctx.commonAttrs,
276350
},
277351
},
278-
parentCtx,
352+
resolveSessionTraceContext(toolPart.sessionID, ctx) ?? context.active(),
279353
)
280354
})()
281355
: undefined
@@ -285,7 +359,11 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
285359
startMs: toolPart.state.time.start,
286360
span: toolSpan,
287361
})
288-
ctx.log("debug", "otel: tool span started", { sessionID: toolPart.sessionID, tool: toolPart.tool, key })
362+
ctx.log("debug", "otel: tool span started", {
363+
sessionID: toolPart.sessionID,
364+
tool: toolPart.tool,
365+
key,
366+
})
289367
return
290368
}
291369

@@ -309,30 +387,28 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
309387
}
310388

311389
if (isTraceEnabled("tool", ctx)) {
312-
const toolSpan = pending?.span ?? (() => {
313-
const sessionSpan = ctx.sessionSpans.get(toolPart.sessionID)
314-
const parentCtx = sessionSpan
315-
? trace.setSpan(context.active(), sessionSpan)
316-
: context.active()
317-
return ctx.tracer.startSpan(
318-
`${ctx.tracePrefix}tool.${toolPart.tool}`,
319-
{
320-
startTime: start,
321-
kind: SpanKind.INTERNAL,
322-
attributes: {
323-
[OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
324-
[SESSION_ID]: toolPart.sessionID,
325-
[TOOL_ID]: toolPart.callID,
326-
[TOOL_NAME]: toolPart.tool,
327-
[TOOL_PARAMETERS]: JSON.stringify(toolPart.state.input),
328-
[INPUT_VALUE]: JSON.stringify(toolPart.state.input),
329-
[INPUT_MIME_TYPE]: MimeType.JSON,
330-
...ctx.commonAttrs,
390+
const toolSpan =
391+
pending?.span ??
392+
(() => {
393+
return ctx.tracer.startSpan(
394+
`${ctx.tracePrefix}tool.${toolPart.tool}`,
395+
{
396+
startTime: start,
397+
kind: SpanKind.INTERNAL,
398+
attributes: {
399+
[OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
400+
[SESSION_ID]: toolPart.sessionID,
401+
[TOOL_ID]: toolPart.callID,
402+
[TOOL_NAME]: toolPart.tool,
403+
[TOOL_PARAMETERS]: JSON.stringify(toolPart.state.input),
404+
[INPUT_VALUE]: JSON.stringify(toolPart.state.input),
405+
[INPUT_MIME_TYPE]: MimeType.JSON,
406+
...ctx.commonAttrs,
407+
},
331408
},
332-
},
333-
parentCtx,
334-
)
335-
})()
409+
resolveSessionTraceContext(toolPart.sessionID, ctx) ?? context.active(),
410+
)
411+
})()
336412
toolSpan.setAttribute("tool.success", success)
337413
if (success) {
338414
const output = (toolPart.state as { output: string }).output
@@ -355,7 +431,12 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
355431
}
356432

357433
const sizeAttr = success
358-
? { tool_result_size_bytes: Buffer.byteLength((toolPart.state as { output: string }).output, "utf8") }
434+
? {
435+
tool_result_size_bytes: Buffer.byteLength(
436+
(toolPart.state as { output: string }).output,
437+
"utf8",
438+
),
439+
}
359440
: { error: (toolPart.state as { error: string }).error }
360441

361442
ctx.logger.emit({
@@ -407,11 +488,6 @@ export function startMessageSpan(
407488
if (!isTraceEnabled("llm", ctx)) return
408489
const msgKey = `${sessionID}:${messageID}`
409490
if (ctx.messageSpans.has(msgKey)) return
410-
const sessionSpan = ctx.sessionSpans.get(sessionID)
411-
const parentCtx = sessionSpan
412-
? trace.setSpan(context.active(), sessionSpan)
413-
: context.active()
414-
415491
const msgSpan = ctx.tracer.startSpan(
416492
`${ctx.tracePrefix}llm`,
417493
{
@@ -428,13 +504,15 @@ export function startMessageSpan(
428504
? {
429505
[INPUT_VALUE]: ctx.sessionInputs.get(sessionID)!,
430506
[INPUT_MIME_TYPE]: MimeType.TEXT,
431-
[LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: ctx.sessionInputs.get(sessionID)! }]),
507+
[LLM_INPUT_MESSAGES]: JSON.stringify([
508+
{ role: "user", content: ctx.sessionInputs.get(sessionID)! },
509+
]),
432510
}
433511
: {}),
434512
...ctx.commonAttrs,
435513
},
436514
},
437-
parentCtx,
515+
resolveSessionTraceContext(sessionID, ctx) ?? context.active(),
438516
)
439517
setBoundedMap(ctx.messageSpans, msgKey, msgSpan)
440518
}

0 commit comments

Comments
 (0)