diff --git a/src/daemon.ts b/src/daemon.ts index da2cf1f..56f8326 100644 --- a/src/daemon.ts +++ b/src/daemon.ts @@ -1727,14 +1727,39 @@ export class GlobalDaemon { `SessionEnd: session=${sessionId} reason=${(payload['reason'] as string | undefined) ?? 'unknown'} transcript_path=${session.transcript.resolvedPath} turns=${session.turnNumber} total_tools=${session.totalToolCalls} pending_tools=${session.pendingToolCalls.size} open_subagents=${session.subagents.size()}`, ); + this.finalizeSession(session, 'session_ended'); + + this.log('INFO', `Finished session ${sessionId}`); + + this.sessions.delete(sessionId); + this.sessionQueues.delete(sessionId); + session.transcript.close(); + } + + /** + * End every span still open on a session — pending tool calls, the active + * chat span, the current turn (root) span, and any tracked subagent + * `invoke_agent` spans — stamping `weave.claude_code.orphan_reason` so the + * trace records why each closed outside its normal path. The active chat + * span is finalized from the transcript (recovering its text + usage) like + * Stop does; only a failed parse falls back to a bare orphan close. + * + * Called from SessionEnd and from `drain` (daemon shutdown). Finalizing at + * shutdown is what keeps a turn's root span exported: without it, a turn + * interrupted by an inactivity/signal/restart shutdown leaks its still-open + * root, leaving its already-exported tool/chat children rootless. Idempotent + * per span — each builder ends at most once. + */ + private finalizeSession(session: SessionState, orphanReason: string): void { // Close any pending tool calls that were never completed for (const [toolUseId, pending] of session.pendingToolCalls) { resolvePermissionIfPending(pending, false); - pending.span.setAttribute(ATTR.WEAVE_ORPHAN_REASON, 'session_ended'); - pending.span.setStatus({ code: SpanStatusCode.ERROR, message: 'session ended before tool completed' }); + pending.span.setAttribute(ATTR.WEAVE_ORPHAN_REASON, orphanReason); + pending.span.setStatus({ code: SpanStatusCode.ERROR, message: 'tool did not complete before shutdown' }); pending.span.end(); this.log('DEBUG', `Closed orphaned tool span: ${toolUseId} (${pending.toolName})`); } + session.pendingToolCalls.clear(); // Finalize a chat span left open mid-turn (Stop never fired) from the // now-flushed transcript, like Stop does, so its text + usage aren't lost. @@ -1755,17 +1780,18 @@ export class GlobalDaemon { } } if (session.activeChatSpan) { - session.activeChatSpan.span.setAttribute(ATTR.WEAVE_ORPHAN_REASON, 'session_ended'); + session.activeChatSpan.span.setAttribute(ATTR.WEAVE_ORPHAN_REASON, orphanReason); session.activeChatSpan.span.end(); session.activeChatSpan = undefined; } - this.log('DEBUG', finalized ? `Finalized active chat span at SessionEnd` : `Closed orphaned chat span`); + this.log('DEBUG', finalized ? `Finalized active chat span` : `Closed orphaned chat span`); } - // Close the current turn if still open + // Close the current turn (root) span if still open if (session.currentTurnSpan) { - session.currentTurnSpan.setAttribute(ATTR.WEAVE_ORPHAN_REASON, 'session_ended'); + session.currentTurnSpan.setAttribute(ATTR.WEAVE_ORPHAN_REASON, orphanReason); session.currentTurnSpan.end(); + session.currentTurnSpan = undefined; this.log('DEBUG', `Closed orphaned turn span`); } @@ -1773,19 +1799,13 @@ export class GlobalDaemon { // or SubagentStop. Without this they'd leak open and never export. for (const tracker of session.subagents.all()) { if (tracker.invokeAgentSpan && !tracker.ended) { - tracker.invokeAgentSpan.setAttribute(ATTR.WEAVE_ORPHAN_REASON, 'session_ended'); - tracker.invokeAgentSpan.setStatus({ code: SpanStatusCode.ERROR, message: 'session ended before subagent completed' }); + tracker.invokeAgentSpan.setAttribute(ATTR.WEAVE_ORPHAN_REASON, orphanReason); + tracker.invokeAgentSpan.setStatus({ code: SpanStatusCode.ERROR, message: 'subagent did not complete before shutdown' }); tracker.invokeAgentSpan.end(); tracker.ended = true; } this.log('DEBUG', `Subagent tracker not stopped: ${tracker.agentId ?? '(unmatched)'} type=${tracker.subagentType}`); } - - this.log('INFO', `Finished session ${sessionId}`); - - this.sessions.delete(sessionId); - this.sessionQueues.delete(sessionId); - session.transcript.close(); } // ── lifecycle ───────────────────────────────────────────────────────────── @@ -1820,6 +1840,21 @@ export class GlobalDaemon { private async shutdown(reason: string): Promise { if (!this.running) return; this.running = false; + await this.drain(reason); + process.exit(0); + } + + /** + * Everything a shutdown does except the final `process.exit`: end in-flight + * spans, flush the exporter, and release the socket. Split out from + * `shutdown` so it can be exercised in tests without terminating the process. + * + * Order matters: open sessions are finalized (their root turn spans ended) + * BEFORE `provider.shutdown()` flushes, so those roots make the final export + * batch instead of being dropped — the fix for rootless traces left behind + * when the daemon exits mid-turn. + */ + private async drain(reason: string): Promise { this.log('INFO', `Shutdown: ${reason}`); this.server?.close(); // Backstop: close any queued team-member invoke_agent spans whose teammate @@ -1831,6 +1866,11 @@ export class GlobalDaemon { } } this.teamMembers.clear(); + // Finalize every live session's still-open spans (turn root, active chat, + // pending tools, subagents) so an interrupted turn keeps its exported root. + for (const session of this.sessions.values()) { + this.finalizeSession(session, 'daemon_shutdown'); + } if (this.provider) { try { await this.provider.shutdown(); @@ -1844,7 +1884,6 @@ export class GlobalDaemon { if (fs.existsSync(this.socketPath)) { fs.unlinkSync(this.socketPath); } - process.exit(0); } // ── helpers ─────────────────────────────────────────────────────────────── diff --git a/tests/daemon-shutdown-finalizes-turn.test.ts b/tests/daemon-shutdown-finalizes-turn.test.ts new file mode 100644 index 0000000..bd2032f --- /dev/null +++ b/tests/daemon-shutdown-finalizes-turn.test.ts @@ -0,0 +1,160 @@ +// SPDX-FileCopyrightText: 2026 CoreWeave, Inc. +// SPDX-License-Identifier: MIT +// SPDX-PackageName: weave-claude-code + +// A turn's root span (`invoke_agent claude-code`) is created at +// UserPromptSubmit and only ended at Stop or SessionEnd. When the daemon exits +// for any other reason — inactivity timeout, SIGTERM/SIGINT/SIGHUP, or a +// restart control message — its already-ended children (completed tool spans, +// finalized chat spans, closed subagent spans) have been exported, but the +// still-open root had not. The result was a rootless trace: tool spans with no +// user turn to attribute them to. +// +// The fix finalizes every live session (ending its turn root) inside the +// shutdown drain, before the exporter is flushed. These tests drive a turn to a +// mid-flight state, run the drain, and assert the root is exported. + +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import * as fs from 'node:fs'; +import * as os from 'node:os'; +import * as path from 'node:path'; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from '@opentelemetry/sdk-trace-base'; +import { GlobalDaemon } from '../src/daemon.ts'; +import { ATTR, OP } from '../src/genaiSpans.ts'; + +function setupTracer() { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] }); + return { tracer: provider.getTracer('test'), exporter, provider }; +} + +const USAGE = { input_tokens: 100, output_tokens: 50, cache_read_input_tokens: 0 }; + +function aLine(id: string, ts: string, block: Record, stop?: string) { + return { + type: 'assistant', + timestamp: ts, + message: { role: 'assistant', id, model: 'claude-opus-4-8', content: [block], usage: USAGE, ...(stop ? { stop_reason: stop } : {}) }, + }; +} +function userText(ts: string, text: string) { + return { type: 'user', timestamp: ts, message: { role: 'user', content: [{ type: 'text', text }] } }; +} + +function makeTranscript(sessionId: string): { file: string; append: (line: unknown) => void; dir: string } { + const dir = fs.mkdtempSync(path.join(os.homedir(), '.weave-shutdown-itest-')); + const file = path.join(dir, `${sessionId}.jsonl`); + fs.writeFileSync(file, ''); + return { file, dir, append: (line: unknown) => fs.appendFileSync(file, JSON.stringify(line) + '\n') }; +} + +interface Harness { + handleSessionStart(s: string, p: Record): Promise; + handleUserPromptSubmit(s: string, p: Record): Promise; + handlePreToolUse(s: string, a: string | undefined, p: Record): Promise; + handlePostToolUse(s: string, p: Record): Promise; + handleSessionEnd(s: string, p: Record): Promise; + drain(reason: string): Promise; + tracer: unknown; +} + +function makeDaemon(tracer: ReturnType['tracer']): Harness { + const logFile = path.join(os.tmpdir(), `wcp-shutdown-itest-${process.pid}.log`); + const d = new GlobalDaemon('/tmp/unused-shutdown.sock', logFile, 'e/p', 'k', 'https://x', false, 'claude-code'); + (d as unknown as { tracer: unknown }).tracer = tracer; + return d as unknown as Harness; +} + +/** Drive a session to a mid-turn state: turn open, one tool completed. */ +async function openTurnWithOneCompletedTool(d: Harness, sid: string, append: (l: unknown) => void, file: string) { + append(userText('2026-01-01T00:00:00.000Z', 'do the thing')); + await d.handleSessionStart(sid, { transcript_path: file, source: 'startup', cwd: '/x' }); + await d.handleUserPromptSubmit(sid, { prompt: 'do the thing' }); + append(aLine('msgA', '2026-01-01T00:00:02.000Z', { type: 'text', text: 'reading' })); + append(aLine('msgA', '2026-01-01T00:00:03.000Z', { type: 'tool_use', id: 'tool_1', name: 'Read', input: {} }, 'tool_use')); + await d.handlePreToolUse(sid, undefined, { tool_use_id: 'tool_1', tool_name: 'Read', tool_input: { file_path: '/foo' } }); + await d.handlePostToolUse(sid, { tool_use_id: 'tool_1', tool_response: 'ok' }); +} + +test('daemon shutdown mid-turn exports the turn root span (children are not left rootless)', async () => { + const sid = 'sess-shutdown'; + const { file, append, dir } = makeTranscript(sid); + const { tracer, exporter, provider } = setupTracer(); + const d = makeDaemon(tracer); + try { + await openTurnWithOneCompletedTool(d, sid, append, file); + + // Neither Stop nor SessionEnd fired: the daemon exits (inactivity / signal + // / restart). The drain must finalize the open turn before flushing. + await d.drain('inactivity'); + await provider.forceFlush(); + + const spans = exporter.getFinishedSpans(); + const tool = spans.find(s => s.attributes[ATTR.OPERATION_NAME] === OP.EXECUTE_TOOL); + assert.ok(tool, 'the completed tool span exported as a child'); + + const root = spans.find(s => s.name === `${OP.INVOKE_AGENT} claude-code`); + assert.ok(root, 'the turn root span must be exported on shutdown, not leaked'); + assert.equal(root!.attributes[ATTR.AGENT_NAME], 'claude-code'); + assert.equal(root!.attributes[ATTR.CONVERSATION_ID], sid); + assert.equal(root!.attributes[ATTR.WEAVE_ORPHAN_REASON], 'daemon_shutdown'); + + // The trace is well-formed: the child shares the exported root's trace id. + assert.equal(tool!.spanContext().traceId, root!.spanContext().traceId, 'child and root share one trace'); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } +}); + +test('daemon shutdown ends an open subagent invoke_agent span under the same trace', async () => { + const sid = 'sess-shutdown-subagent'; + const { file, append, dir } = makeTranscript(sid); + const { tracer, exporter, provider } = setupTracer(); + const d = makeDaemon(tracer); + try { + append(userText('2026-01-01T00:00:00.000Z', 'spawn a reviewer')); + await d.handleSessionStart(sid, { transcript_path: file, source: 'startup', cwd: '/x' }); + await d.handleUserPromptSubmit(sid, { prompt: 'spawn a reviewer' }); + + // Agent tool with subagent_type opens a nested invoke_agent span that a + // mid-flight shutdown would otherwise leave open. + append(aLine('msgA', '2026-01-01T00:00:02.000Z', { type: 'tool_use', id: 'agent_1', name: 'Agent', input: { subagent_type: 'code-reviewer', prompt: 'review' } }, 'tool_use')); + await d.handlePreToolUse(sid, undefined, { tool_use_id: 'agent_1', tool_name: 'Agent', tool_input: { subagent_type: 'code-reviewer', prompt: 'review' } }); + + await d.drain('SIGTERM'); + await provider.forceFlush(); + + const spans = exporter.getFinishedSpans(); + const root = spans.find(s => s.name === `${OP.INVOKE_AGENT} claude-code`); + const sub = spans.find(s => s.name === `${OP.INVOKE_AGENT} code-reviewer`); + assert.ok(root, 'turn root exported'); + assert.ok(sub, 'open subagent invoke_agent span exported on shutdown'); + assert.equal(sub!.spanContext().traceId, root!.spanContext().traceId, 'subagent nests under the same trace as the root'); + assert.equal(sub!.attributes[ATTR.WEAVE_ORPHAN_REASON], 'daemon_shutdown'); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } +}); + +test('SessionEnd still exports the turn root span after the finalize refactor', async () => { + const sid = 'sess-sessionend'; + const { file, append, dir } = makeTranscript(sid); + const { tracer, exporter, provider } = setupTracer(); + const d = makeDaemon(tracer); + try { + await openTurnWithOneCompletedTool(d, sid, append, file); + await d.handleSessionEnd(sid, { reason: 'clear' }); + await provider.forceFlush(); + + const root = exporter.getFinishedSpans().find(s => s.name === `${OP.INVOKE_AGENT} claude-code`); + assert.ok(root, 'SessionEnd exports the turn root'); + assert.equal(root!.attributes[ATTR.WEAVE_ORPHAN_REASON], 'session_ended'); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } +});