Skip to content

Commit 3aa91fd

Browse files
dialupdisasterdialupdisaster
authored
Merge pull request #74 from DEVtheOPS/fix/run-scoped-traces
2 parents c9fee7f + 0e43ecc commit 3aa91fd

9 files changed

Lines changed: 372 additions & 119 deletions

File tree

‎.github/dependabot.yml‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,7 @@
55

66
version: 2
77
updates:
8-
- package-ecosystem: "" # See documentation for possible values
8+
- package-ecosystem: "bun" # See documentation for possible values
99
directory: "/" # Location of package manifests
1010
schedule:
1111
interval: "weekly"
12-

‎src/handlers/message.ts‎

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { SeverityNumber } from "@opentelemetry/api-logs"
2-
import { SpanStatusCode, SpanKind, trace } from "@opentelemetry/api"
2+
import { SpanStatusCode, SpanKind } from "@opentelemetry/api"
33
import type { AssistantMessage, EventMessageUpdated, EventMessagePartUpdated, ToolPart } from "@opencode-ai/sdk"
44
import {
55
AGENT_NAME,
@@ -27,7 +27,16 @@ import {
2727
TOOL_NAME,
2828
TOOL_PARAMETERS,
2929
} from "@arizeai/openinference-semantic-conventions"
30-
import { agentAttrs, errorSummary, setBoundedMap, accumulateSessionTotals, getSessionAgentMeta, isMetricEnabled, isTraceEnabled } from "../util.ts"
30+
import {
31+
agentAttrs,
32+
errorSummary,
33+
setBoundedMap,
34+
accumulateSessionTotals,
35+
getSessionAgentMeta,
36+
isMetricEnabled,
37+
isTraceEnabled,
38+
resolveSessionTraceContext,
39+
} from "../util.ts"
3140
import type { HandlerContext } from "../types.ts"
3241

3342
const OPENINFERENCE_SPAN_KIND = SemanticConventions.OPENINFERENCE_SPAN_KIND
@@ -52,6 +61,7 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext
5261
const msg = e.properties.info
5362
if (msg.role !== "assistant") return
5463
const assistant = msg as AssistantMessage
64+
setBoundedMap(ctx.assistantRuns, assistant.id, assistant.parentID)
5565
if (!assistant.time.completed) return
5666

5767
const { sessionID, modelID, providerID } = assistant
@@ -260,11 +270,6 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
260270
const { agentName, agentType } = getSessionAgentMeta(toolPart.sessionID, ctx)
261271
const toolSpan = isTraceEnabled("tool", ctx)
262272
? (() => {
263-
const sessionSpan = ctx.sessionSpans.get(toolPart.sessionID)
264-
const baseCtx = ctx.rootContext()
265-
const parentCtx = sessionSpan
266-
? trace.setSpan(baseCtx, sessionSpan)
267-
: baseCtx
268273
return ctx.tracer.startSpan(
269274
`${ctx.tracePrefix}tool.${toolPart.tool}`,
270275
{
@@ -283,7 +288,9 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
283288
...ctx.commonAttrs,
284289
},
285290
},
286-
parentCtx,
291+
resolveSessionTraceContext(toolPart.sessionID, ctx, {
292+
assistantMessageID: toolPart.messageID,
293+
}),
287294
)
288295
})()
289296
: undefined
@@ -319,11 +326,6 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
319326

320327
if (isTraceEnabled("tool", ctx)) {
321328
const toolSpan = pending?.span ?? (() => {
322-
const sessionSpan = ctx.sessionSpans.get(toolPart.sessionID)
323-
const baseCtx = ctx.rootContext()
324-
const parentCtx = sessionSpan
325-
? trace.setSpan(baseCtx, sessionSpan)
326-
: baseCtx
327329
return ctx.tracer.startSpan(
328330
`${ctx.tracePrefix}tool.${toolPart.tool}`,
329331
{
@@ -340,7 +342,9 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
340342
...ctx.commonAttrs,
341343
},
342344
},
343-
parentCtx,
345+
resolveSessionTraceContext(toolPart.sessionID, ctx, {
346+
assistantMessageID: toolPart.messageID,
347+
}),
344348
)
345349
})()
346350
toolSpan.setAttributes({ [AGENT_NAME]: agentName, "agent.type": agentType })
@@ -403,14 +407,16 @@ export function handleMessagePartUpdated(e: EventMessagePartUpdated, ctx: Handle
403407

404408
/**
405409
* Starts an LLM span for an assistant message when it first appears in `message.updated`.
406-
* The span is parented to the session span and carries `gen_ai.*` semantic attributes for
407-
* the model and provider. It is ended in `handleMessageUpdated` once the message completes.
410+
* The span is parented to the active run or subagent span and carries `gen_ai.*` semantic
411+
* attributes for the model and provider. It is ended in `handleMessageUpdated` once the
412+
* message completes.
408413
*
409414
* Only called for assistant messages that have not yet completed (`time.completed` absent).
410415
*/
411416
export function startMessageSpan(
412417
sessionID: string,
413418
messageID: string,
419+
parentID: string,
414420
modelID: string,
415421
providerID: string,
416422
startTime: number,
@@ -419,12 +425,9 @@ export function startMessageSpan(
419425
if (!isTraceEnabled("llm", ctx)) return
420426
const msgKey = `${sessionID}:${messageID}`
421427
if (ctx.messageSpans.has(msgKey)) return
428+
setBoundedMap(ctx.assistantRuns, messageID, parentID)
422429
const { agentName, agentType } = getSessionAgentMeta(sessionID, ctx)
423-
const sessionSpan = ctx.sessionSpans.get(sessionID)
424-
const baseCtx = ctx.rootContext()
425-
const parentCtx = sessionSpan
426-
? trace.setSpan(baseCtx, sessionSpan)
427-
: baseCtx
430+
const inputText = ctx.runInputs.get(parentID)
428431

429432
const msgSpan = ctx.tracer.startSpan(
430433
`${ctx.tracePrefix}llm`,
@@ -439,17 +442,17 @@ export function startMessageSpan(
439442
[LLM_SYSTEM]: providerID,
440443
[LLM_PROVIDER]: providerID,
441444
[LLM_MODEL_NAME]: modelID,
442-
...(ctx.sessionInputs.has(sessionID)
445+
...(inputText
443446
? {
444-
[INPUT_VALUE]: ctx.sessionInputs.get(sessionID)!,
447+
[INPUT_VALUE]: inputText,
445448
[INPUT_MIME_TYPE]: MimeType.TEXT,
446-
[LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: ctx.sessionInputs.get(sessionID)! }]),
449+
[LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: inputText }]),
447450
}
448451
: {}),
449452
...ctx.commonAttrs,
450453
},
451454
},
452-
parentCtx,
455+
resolveSessionTraceContext(sessionID, ctx, { runID: parentID, assistantMessageID: messageID }),
453456
)
454457
setBoundedMap(ctx.messageSpans, msgKey, msgSpan)
455458
}

‎src/handlers/session.ts‎

Lines changed: 109 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,86 @@
11
import { SeverityNumber } from "@opentelemetry/api-logs"
2-
import { SpanStatusCode, trace } from "@opentelemetry/api"
2+
import { SpanStatusCode } from "@opentelemetry/api"
33
import type { EventSessionCreated, EventSessionIdle, EventSessionError, EventSessionStatus } from "@opencode-ai/sdk"
4-
import { AGENT_NAME, OpenInferenceSpanKind, SemanticConventions, SESSION_ID } from "@arizeai/openinference-semantic-conventions"
5-
import { agentAttrs, errorSummary, getSessionAgentMeta, setBoundedMap, isMetricEnabled, isTraceEnabled } from "../util.ts"
4+
import {
5+
AGENT_NAME,
6+
INPUT_MIME_TYPE,
7+
INPUT_VALUE,
8+
LLM_INPUT_MESSAGES,
9+
MimeType,
10+
OpenInferenceSpanKind,
11+
SemanticConventions,
12+
SESSION_ID,
13+
} from "@arizeai/openinference-semantic-conventions"
14+
import {
15+
agentAttrs,
16+
errorSummary,
17+
getSessionAgentMeta,
18+
setBoundedMap,
19+
isMetricEnabled,
20+
isTraceEnabled,
21+
resolveSessionTraceContext,
22+
} from "../util.ts"
623
import type { HandlerContext, SessionAgentType } from "../types.ts"
724

825
const OPENINFERENCE_SPAN_KIND = SemanticConventions.OPENINFERENCE_SPAN_KIND
926

27+
/** Starts or refreshes the root run span for a single user turn, keyed by the user message ID. */
28+
export function handleRunStarted(
29+
runID: string,
30+
sessionID: string,
31+
agent: string,
32+
promptText: string,
33+
model: string,
34+
startTime: number,
35+
ctx: HandlerContext,
36+
) {
37+
ctx.activeRuns.set(sessionID, runID)
38+
ctx.pendingRuns.delete(sessionID)
39+
if (promptText) setBoundedMap(ctx.runInputs, runID, promptText)
40+
if (!isTraceEnabled("session", ctx)) return
41+
const existing = ctx.runSpans.get(runID)
42+
if (existing) {
43+
existing.setAttributes({
44+
[AGENT_NAME]: agent,
45+
...(promptText
46+
? {
47+
[INPUT_VALUE]: promptText,
48+
[INPUT_MIME_TYPE]: MimeType.TEXT,
49+
[LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: promptText }]),
50+
}
51+
: {}),
52+
model,
53+
})
54+
return
55+
}
56+
57+
const runSpan = ctx.tracer.startSpan(
58+
`${ctx.tracePrefix}session`,
59+
{
60+
startTime,
61+
attributes: {
62+
[OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.AGENT,
63+
[SESSION_ID]: sessionID,
64+
[AGENT_NAME]: agent,
65+
"agent.type": "primary",
66+
"session.is_subagent": false,
67+
...(promptText
68+
? {
69+
[INPUT_VALUE]: promptText,
70+
[INPUT_MIME_TYPE]: MimeType.TEXT,
71+
[LLM_INPUT_MESSAGES]: JSON.stringify([{ role: "user", content: promptText }]),
72+
}
73+
: {}),
74+
model,
75+
...ctx.commonAttrs,
76+
},
77+
},
78+
ctx.rootContext(),
79+
)
80+
ctx.runSpans.set(runID, runSpan)
81+
setBoundedMap(ctx.runSpanContexts, runID, runSpan.spanContext())
82+
}
83+
1084
/** Increments the session counter, records start time, starts the root session span, and emits a `session.created` log event. */
1185
export function handleSessionCreated(e: EventSessionCreated, ctx: HandlerContext) {
1286
const { id: sessionID, time, parentID } = e.properties.info
@@ -18,16 +92,7 @@ export function handleSessionCreated(e: EventSessionCreated, ctx: HandlerContext
1892
}
1993
setBoundedMap(ctx.sessionTotals, sessionID, { startMs: createdAt, tokens: 0, cost: 0, messages: 0, agent: "unknown", agentType })
2094

21-
// WARNING: disabling "session" traces while "llm" or "tool" traces remain enabled
22-
// leaves those child spans without a local session parent. If OPENCODE_TRACEPARENT
23-
// is set, they fall back to that remote parent; otherwise they become root spans.
24-
if (isTraceEnabled("session", ctx)) {
25-
const parentSpan = parentID ? ctx.sessionSpans.get(parentID) : undefined
26-
const baseCtx = ctx.rootContext()
27-
const spanCtx = parentSpan
28-
? trace.setSpan(baseCtx, parentSpan)
29-
: baseCtx
30-
95+
if (isTraceEnabled("session", ctx) && parentID) {
3196
const sessionSpan = ctx.tracer.startSpan(
3297
`${ctx.tracePrefix}session`,
3398
{
@@ -41,9 +106,10 @@ export function handleSessionCreated(e: EventSessionCreated, ctx: HandlerContext
41106
...ctx.commonAttrs,
42107
},
43108
},
44-
spanCtx,
109+
resolveSessionTraceContext(parentID, ctx),
45110
)
46-
setBoundedMap(ctx.sessionSpans, sessionID, sessionSpan)
111+
ctx.sessionSpans.set(sessionID, sessionSpan)
112+
setBoundedMap(ctx.sessionSpanContexts, sessionID, sessionSpan.spanContext())
47113
}
48114

49115
ctx.emitLog({
@@ -74,7 +140,7 @@ function sweepSession(sessionID: string, ctx: HandlerContext) {
74140
ctx.pendingToolSpans.delete(key)
75141
}
76142
}
77-
ctx.sessionInputs.delete(sessionID)
143+
ctx.pendingRuns.delete(sessionID)
78144
const msgPrefix = `${sessionID}:`
79145
for (const [key, span] of ctx.messageSpans) {
80146
if (key.startsWith(msgPrefix)) {
@@ -128,6 +194,23 @@ export function handleSessionIdle(e: EventSessionIdle, ctx: HandlerContext) {
128194
sessionSpan.end()
129195
ctx.sessionSpans.delete(sessionID)
130196
}
197+
const runID = ctx.activeRuns.get(sessionID)
198+
if (runID) ctx.activeRuns.delete(sessionID)
199+
const runSpan = runID ? ctx.runSpans.get(runID) : undefined
200+
if (runSpan) {
201+
if (totals) {
202+
runSpan.setAttributes({
203+
[AGENT_NAME]: totals.agent,
204+
"agent.type": totals.agentType,
205+
"session.total_tokens": totals.tokens,
206+
"session.total_cost_usd": totals.cost,
207+
"session.total_messages": totals.messages,
208+
})
209+
}
210+
runSpan.setStatus({ code: SpanStatusCode.OK })
211+
runSpan.end()
212+
ctx.runSpans.delete(runID!)
213+
}
131214

132215
ctx.emitLog({
133216
severityNumber: SeverityNumber.INFO,
@@ -173,6 +256,16 @@ export function handleSessionError(e: EventSessionError, ctx: HandlerContext) {
173256
sessionSpan.end()
174257
ctx.sessionSpans.delete(rawID)
175258
}
259+
const runID = ctx.activeRuns.get(rawID)
260+
if (runID) ctx.activeRuns.delete(rawID)
261+
const runSpan = runID ? ctx.runSpans.get(runID) : undefined
262+
if (runSpan) {
263+
if (totals) runSpan.setAttributes({ [AGENT_NAME]: totals.agent, "agent.type": totals.agentType })
264+
runSpan.setStatus({ code: SpanStatusCode.ERROR, message: error })
265+
runSpan.setAttribute("error", error)
266+
runSpan.end()
267+
ctx.runSpans.delete(runID!)
268+
}
176269
}
177270

178271
ctx.emitLog({

0 commit comments

Comments
 (0)