Skip to content

Commit d3ed25f

Browse files
committed
fix: keep root session span open across session.idle
The opencode.session span and accumulated session totals were ended and deleted on the first session.idle event. Since a session keeps receiving turns after going idle, every subsequent turn's LLM/tool spans had no local session span to nest under and became orphaned root traces, and their token/cost contributions were silently dropped from session.total_*. Keep the session span and totals alive across session.idle (only snapshotting attributes), and finalize them in a new session.deleted handler, on session.error, or on process shutdown.
1 parent 7153352 commit d3ed25f

4 files changed

Lines changed: 136 additions & 27 deletions

File tree

src/handlers/session.ts

Lines changed: 50 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { SeverityNumber } from "@opentelemetry/api-logs"
22
import { SpanStatusCode, trace } from "@opentelemetry/api"
3-
import type { EventSessionCreated, EventSessionIdle, EventSessionError, EventSessionStatus } from "@opencode-ai/sdk"
3+
import type { EventSessionCreated, EventSessionIdle, EventSessionDeleted, EventSessionError, EventSessionStatus } from "@opencode-ai/sdk"
44
import { AGENT_NAME, OpenInferenceSpanKind, SemanticConventions, SESSION_ID } from "@arizeai/openinference-semantic-conventions"
55
import { errorSummary, setBoundedMap, isMetricEnabled, isTraceEnabled } from "../util.ts"
66
import type { HandlerContext } from "../types.ts"
@@ -80,12 +80,20 @@ function sweepSession(sessionID: string, ctx: HandlerContext) {
8080
}
8181
}
8282

83-
/** Emits a `session.idle` log event, records duration and session total histograms, ends the session span, and clears pending state. */
83+
/**
84+
* Emits a `session.idle` log event and records duration/total histograms for the turn
85+
* that just completed, then sweeps per-turn pending state (tool/message spans, pending
86+
* permissions, the cached user prompt).
87+
*
88+
* Unlike a one-shot turn, an opencode session stays alive and may receive further user
89+
* messages, so `session.total_*` totals and the root `opencode.session` span are kept
90+
* open across `session.idle` events: ending the span here would otherwise orphan every
91+
* subsequent turn's LLM/tool spans as new root traces. The span and totals are only
92+
* finalized in `handleSessionDeleted` (or `handleSessionError`/shutdown).
93+
*/
8494
export function handleSessionIdle(e: EventSessionIdle, ctx: HandlerContext) {
8595
const sessionID = e.properties.sessionID
8696
const totals = ctx.sessionTotals.get(sessionID)
87-
ctx.sessionTotals.delete(sessionID)
88-
ctx.sessionDiffTotals.delete(sessionID)
8997
sweepSession(sessionID, ctx)
9098

9199
const attrs = { ...ctx.commonAttrs, "session.id": sessionID }
@@ -105,18 +113,13 @@ export function handleSessionIdle(e: EventSessionIdle, ctx: HandlerContext) {
105113
}
106114

107115
const sessionSpan = ctx.sessionSpans.get(sessionID)
108-
if (sessionSpan) {
109-
if (totals) {
110-
sessionSpan.setAttributes({
111-
[AGENT_NAME]: totals.agent,
112-
"session.total_tokens": totals.tokens,
113-
"session.total_cost_usd": totals.cost,
114-
"session.total_messages": totals.messages,
115-
})
116-
}
117-
sessionSpan.setStatus({ code: SpanStatusCode.OK })
118-
sessionSpan.end()
119-
ctx.sessionSpans.delete(sessionID)
116+
if (sessionSpan && totals) {
117+
sessionSpan.setAttributes({
118+
[AGENT_NAME]: totals.agent,
119+
"session.total_tokens": totals.tokens,
120+
"session.total_cost_usd": totals.cost,
121+
"session.total_messages": totals.messages,
122+
})
120123
}
121124

122125
ctx.emitLog({
@@ -140,6 +143,37 @@ export function handleSessionIdle(e: EventSessionIdle, ctx: HandlerContext) {
140143
})
141144
}
142145

146+
/**
147+
* Final cleanup when a session is removed: ends the root `opencode.session` span with the
148+
* last known totals, and clears the session's accumulated totals and pending state. This is
149+
* the counterpart to the "keep the span open across `session.idle`" behavior above — it's
150+
* where a long-lived session's span and totals actually get torn down.
151+
*/
152+
export function handleSessionDeleted(e: EventSessionDeleted, ctx: HandlerContext) {
153+
const sessionID = e.properties.info.id
154+
const totals = ctx.sessionTotals.get(sessionID)
155+
ctx.sessionTotals.delete(sessionID)
156+
ctx.sessionDiffTotals.delete(sessionID)
157+
sweepSession(sessionID, ctx)
158+
159+
const sessionSpan = ctx.sessionSpans.get(sessionID)
160+
if (sessionSpan) {
161+
if (totals) {
162+
sessionSpan.setAttributes({
163+
[AGENT_NAME]: totals.agent,
164+
"session.total_tokens": totals.tokens,
165+
"session.total_cost_usd": totals.cost,
166+
"session.total_messages": totals.messages,
167+
})
168+
}
169+
sessionSpan.setStatus({ code: SpanStatusCode.OK })
170+
sessionSpan.end()
171+
ctx.sessionSpans.delete(sessionID)
172+
}
173+
174+
ctx.log("debug", "otel: session.deleted", { sessionID })
175+
}
176+
143177
/** Emits a `session.error` log event, ends the session span with error status, and clears any pending state for the session. */
144178
export function handleSessionError(e: EventSessionError, ctx: HandlerContext) {
145179
const rawID = e.properties.sessionID

src/index.ts

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
import type { Plugin } from "@opencode-ai/plugin"
22
import { SeverityNumber } from "@opentelemetry/api-logs"
33
import { logs } from "@opentelemetry/api-logs"
4-
import { ROOT_CONTEXT, trace } from "@opentelemetry/api"
4+
import { ROOT_CONTEXT, SpanStatusCode, trace } from "@opentelemetry/api"
55
import { AGENT_NAME } from "@arizeai/openinference-semantic-conventions"
66
import pkg from "../package.json" with { type: "json" }
77
import type {
88
EventSessionCreated,
99
EventSessionIdle,
10+
EventSessionDeleted,
1011
EventSessionError,
1112
EventSessionStatus,
1213
EventMessageUpdated,
@@ -21,7 +22,7 @@ import { loadConfig, resolveHelperPath, resolveLogLevel } from "./config.ts"
2122
import { probeEndpoint } from "./probe.ts"
2223
import { setupOtel, createInstruments } from "./otel.ts"
2324
import { remoteParentContext } from "./trace-context.ts"
24-
import { handleSessionCreated, handleSessionIdle, handleSessionError, handleSessionStatus } from "./handlers/session.ts"
25+
import { handleSessionCreated, handleSessionIdle, handleSessionDeleted, handleSessionError, handleSessionStatus } from "./handlers/session.ts"
2526
import { handleMessageUpdated, handleMessagePartUpdated, startMessageSpan } from "./handlers/message.ts"
2627
import { handlePermissionUpdated, handlePermissionReplied } from "./handlers/permission.ts"
2728
import { handleSessionDiff, handleCommandExecuted } from "./handlers/activity.ts"
@@ -141,6 +142,23 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree
141142
}
142143

143144
async function shutdown() {
145+
// Session spans are kept open across `session.idle` so later turns nest under them
146+
// (see handleSessionIdle). On process exit, end any still-open session spans so
147+
// they're flushed rather than dropped.
148+
for (const [sessionID, sessionSpan] of sessionSpans) {
149+
const totals = sessionTotals.get(sessionID)
150+
if (totals) {
151+
sessionSpan.setAttributes({
152+
[AGENT_NAME]: totals.agent,
153+
"session.total_tokens": totals.tokens,
154+
"session.total_cost_usd": totals.cost,
155+
"session.total_messages": totals.messages,
156+
})
157+
}
158+
sessionSpan.setStatus({ code: SpanStatusCode.OK })
159+
sessionSpan.end()
160+
}
161+
sessionSpans.clear()
144162
await Promise.allSettled([meterProvider.shutdown(), loggerProvider.shutdown(), tracerProvider.shutdown()])
145163
}
146164

@@ -225,6 +243,9 @@ export const OtelPlugin: Plugin = async ({ project, client, directory, worktree
225243
case "session.idle":
226244
handleSessionIdle(event as EventSessionIdle, ctx)
227245
break
246+
case "session.deleted":
247+
handleSessionDeleted(event as EventSessionDeleted, ctx)
248+
break
228249
case "session.error":
229250
handleSessionError(event as EventSessionError, ctx)
230251
break

tests/handlers/session.test.ts

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { describe, test, expect } from "bun:test"
2-
import { handleSessionCreated, handleSessionIdle, handleSessionError, handleSessionStatus } from "../../src/handlers/session.ts"
2+
import { handleSessionCreated, handleSessionIdle, handleSessionDeleted, handleSessionError, handleSessionStatus } from "../../src/handlers/session.ts"
33
import { makeCtx, makeTracer } from "../helpers.ts"
4-
import type { EventSessionCreated, EventSessionIdle, EventSessionError, EventSessionStatus } from "@opencode-ai/sdk"
4+
import type { EventSessionCreated, EventSessionIdle, EventSessionDeleted, EventSessionError, EventSessionStatus } from "@opencode-ai/sdk"
55
import type { Span } from "@opentelemetry/api"
66

77
function makeSessionCreated(sessionID: string, createdAt = 1000, parentID?: string): EventSessionCreated {
@@ -151,12 +151,36 @@ describe("handleSessionIdle", () => {
151151
expect(gauges.sessionCost.calls).toHaveLength(0)
152152
})
153153

154-
test("removes sessionTotals entry on idle", async () => {
154+
test("keeps sessionTotals entry across idle so later turns keep accumulating", async () => {
155155
const { ctx } = makeCtx()
156156
await handleSessionCreated(makeSessionCreated("ses_1"), ctx)
157157
expect(ctx.sessionTotals.has("ses_1")).toBe(true)
158158
handleSessionIdle(makeSessionIdle("ses_1"), ctx)
159+
expect(ctx.sessionTotals.has("ses_1")).toBe(true)
160+
})
161+
})
162+
163+
describe("handleSessionDeleted", () => {
164+
function makeSessionDeleted(sessionID: string): EventSessionDeleted {
165+
return { type: "session.deleted", properties: { info: { id: sessionID } } } as unknown as EventSessionDeleted
166+
}
167+
168+
test("removes sessionTotals and sessionDiffTotals entries", async () => {
169+
const { ctx } = makeCtx()
170+
await handleSessionCreated(makeSessionCreated("ses_1"), ctx)
171+
ctx.sessionDiffTotals.set("ses_1", { additions: 1, deletions: 2 })
172+
handleSessionDeleted(makeSessionDeleted("ses_1"), ctx)
159173
expect(ctx.sessionTotals.has("ses_1")).toBe(false)
174+
expect(ctx.sessionDiffTotals.has("ses_1")).toBe(false)
175+
})
176+
177+
test("ends the session span with OK status", async () => {
178+
const { ctx, tracer } = makeCtx()
179+
await handleSessionCreated(makeSessionCreated("ses_1"), ctx)
180+
handleSessionDeleted(makeSessionDeleted("ses_1"), ctx)
181+
const span = tracer.spans[0]!
182+
expect(span.ended).toBe(true)
183+
expect(ctx.sessionSpans.has("ses_1")).toBe(false)
160184
})
161185
})
162186

tests/handlers/spans.test.ts

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,14 @@ import {
1616
TOOL_NAME,
1717
} from "@arizeai/openinference-semantic-conventions"
1818
import type { Span } from "@opentelemetry/api"
19-
import { handleSessionCreated, handleSessionIdle, handleSessionError } from "../../src/handlers/session.ts"
19+
import { handleSessionCreated, handleSessionIdle, handleSessionDeleted, handleSessionError } from "../../src/handlers/session.ts"
2020
import { handleMessageUpdated, handleMessagePartUpdated, startMessageSpan } from "../../src/handlers/message.ts"
2121
import { remoteParentContext } from "../../src/trace-context.ts"
2222
import { makeCtx, makeTracer, type SpySpan } from "../helpers.ts"
2323
import type {
2424
EventSessionCreated,
2525
EventSessionIdle,
26+
EventSessionDeleted,
2627
EventSessionError,
2728
EventMessageUpdated,
2829
EventMessagePartUpdated,
@@ -41,6 +42,10 @@ function makeSessionIdle(sessionID: string): EventSessionIdle {
4142
return { type: "session.idle", properties: { sessionID } } as EventSessionIdle
4243
}
4344

45+
function makeSessionDeleted(sessionID: string): EventSessionDeleted {
46+
return { type: "session.deleted", properties: { info: { id: sessionID } } } as unknown as EventSessionDeleted
47+
}
48+
4449
function makeSessionError(sessionID?: string, error?: { name: string }): EventSessionError {
4550
return {
4651
type: "session.error",
@@ -156,17 +161,17 @@ describe("session spans", () => {
156161
expect(tracer.spans[0]!.parentSpan?.spanContext().spanId).toBe("00f067aa0ba902b7")
157162
})
158163

159-
test("ends session span with OK status on session.idle", () => {
164+
test("keeps session span open across session.idle so later turns can nest under it", () => {
160165
const { ctx, tracer } = makeCtx()
161166
handleSessionCreated(makeSessionCreated("ses_1"), ctx)
162167
handleSessionIdle(makeSessionIdle("ses_1"), ctx)
163168
const span = tracer.spans[0]!
164-
expect(span.ended).toBe(true)
165-
expect(span.status.code).toBe(SpanStatusCode.OK)
166-
expect(ctx.sessionSpans.has("ses_1")).toBe(false)
169+
expect(span.ended).toBe(false)
170+
expect(ctx.sessionSpans.has("ses_1")).toBe(true)
171+
expect(ctx.sessionTotals.has("ses_1")).toBe(true)
167172
})
168173

169-
test("sets session total attributes before ending on idle", () => {
174+
test("sets session total attributes on session.idle without ending the span", () => {
170175
const { ctx, tracer } = makeCtx()
171176
handleSessionCreated(makeSessionCreated("ses_1"), ctx)
172177
ctx.sessionTotals.set("ses_1", { startMs: Date.now() - 100, tokens: 250, cost: 0.05, messages: 3, agent: "build" })
@@ -175,6 +180,31 @@ describe("session spans", () => {
175180
expect(span.attributes["session.total_tokens"]).toBe(250)
176181
expect(span.attributes["session.total_cost_usd"]).toBe(0.05)
177182
expect(span.attributes["session.total_messages"]).toBe(3)
183+
expect(span.ended).toBe(false)
184+
})
185+
186+
test("ends session span with OK status and clears totals on session.deleted", () => {
187+
const { ctx, tracer } = makeCtx()
188+
handleSessionCreated(makeSessionCreated("ses_1"), ctx)
189+
ctx.sessionTotals.set("ses_1", { startMs: Date.now() - 100, tokens: 250, cost: 0.05, messages: 3, agent: "build" })
190+
handleSessionIdle(makeSessionIdle("ses_1"), ctx)
191+
handleSessionDeleted(makeSessionDeleted("ses_1"), ctx)
192+
const span = tracer.spans[0]!
193+
expect(span.ended).toBe(true)
194+
expect(span.status.code).toBe(SpanStatusCode.OK)
195+
expect(span.attributes["session.total_tokens"]).toBe(250)
196+
expect(ctx.sessionSpans.has("ses_1")).toBe(false)
197+
expect(ctx.sessionTotals.has("ses_1")).toBe(false)
198+
})
199+
200+
test("a second turn's LLM span nests under the still-open session span after idle", () => {
201+
const { ctx, tracer } = makeCtx()
202+
handleSessionCreated(makeSessionCreated("ses_1"), ctx)
203+
handleSessionIdle(makeSessionIdle("ses_1"), ctx)
204+
startMessageSpan("ses_1", "msg_2", "claude-3-5-sonnet", "anthropic", 5000, ctx)
205+
const sessionSpan = tracer.spans[0]!
206+
const llmSpan = tracer.spans[1]!
207+
expect(llmSpan.parentSpan).toBe(sessionSpan)
178208
})
179209

180210
test("ends session span with ERROR status on session.error", () => {

0 commit comments

Comments
 (0)