From ef720ee13a3f701e66a29846e0d3268b1f300b6d Mon Sep 17 00:00:00 2001 From: viettranx Date: Tue, 10 Mar 2026 13:52:24 +0700 Subject: [PATCH] =?UTF-8?q?feat(tracing):=20two-phase=20spans=20=E2=80=94?= =?UTF-8?q?=20show=20running=20state=20before=20completion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spans are now emitted BEFORE execution starts with status "running" and input visible, then updated when the step completes. This lets users see what the system is doing during long-running LLM calls and tool executions. - Add EmitSpanUpdate() to collector with separate update buffer - Flush ordering: batch INSERT new spans first, then process updates - Split LLM/tool/agent spans into start/end pairs (agent loop + subagent) - Emit tool span end inside goroutines for parallel calls (no orphans) - EmitSpanUpdate is a channel send — works after ctx cancellation --- internal/agent/loop.go | 14 +- internal/agent/loop_run.go | 25 ++- internal/agent/loop_tracing.go | 253 +++++++++++++++++++---------- internal/tools/subagent_exec.go | 20 ++- internal/tools/subagent_tracing.go | 220 ++++++++++++++++--------- internal/tracing/collector.go | 63 ++++++- 6 files changed, 404 insertions(+), 191 deletions(-) diff --git a/internal/agent/loop.go b/internal/agent/loop.go index f7611fa0..d127aafc 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -424,6 +424,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) var err error llmSpanStart := time.Now().UTC() + llmSpanID := l.emitLLMSpanStart(ctx, llmSpanStart, iteration, messages) if req.Stream { resp, err = l.provider.ChatStream(ctx, chatReq, func(chunk providers.StreamChunk) { @@ -449,11 +450,11 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } if err != nil { - l.emitLLMSpan(ctx, llmSpanStart, iteration, messages, nil, err) + l.emitLLMSpanEnd(ctx, llmSpanID, llmSpanStart, nil, err) return nil, fmt.Errorf("LLM call failed (iteration %d): %w", iteration, err) } - l.emitLLMSpan(ctx, llmSpanStart, iteration, messages, resp, nil) + l.emitLLMSpanEnd(ctx, llmSpanID, llmSpanStart, resp, nil) // For non-streaming responses, emit thinking and content as single events if !req.Stream { @@ -643,6 +644,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) argsHash := loopDetector.record(tc.Name, tc.Arguments) toolSpanStart := time.Now().UTC() + toolSpanID := l.emitToolSpanStart(ctx, toolSpanStart, tc.Name, tc.ID, string(argsJSON)) var result *tools.Result if allowedTools != nil && !allowedTools[tc.Name] { slog.Warn("security.tool_policy_blocked", "agent", l.id, "tool", tc.Name) @@ -651,7 +653,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) result = l.tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, req.Channel, req.ChatID, req.PeerKind, req.SessionKey, nil) } - l.emitToolSpan(ctx, toolSpanStart, tc.Name, tc.ID, string(argsJSON), result) + l.emitToolSpanEnd(ctx, toolSpanID, toolSpanStart, result) // Record result for loop detection. loopDetector.recordResult(argsHash, result.ForLLM) @@ -763,6 +765,9 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) argsJSON, _ := json.Marshal(tc.Arguments) slog.Info("tool call", "agent", l.id, "tool", tc.Name, "args_len", len(argsJSON), "parallel", true) spanStart := time.Now().UTC() + // Emit running span inside goroutine — goroutine-safe (channel send only). + // End is also emitted here to prevent orphans on ctx cancellation. + spanID := l.emitToolSpanStart(ctx, spanStart, tc.Name, tc.ID, string(argsJSON)) var result *tools.Result if allowedTools != nil && !allowedTools[tc.Name] { slog.Warn("security.tool_policy_blocked", "agent", l.id, "tool", tc.Name) @@ -770,6 +775,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } else { result = l.tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, req.Channel, req.ChatID, req.PeerKind, req.SessionKey, nil) } + l.emitToolSpanEnd(ctx, spanID, spanStart, result) resultCh <- indexedResult{idx: idx, tc: tc, result: result, argsJSON: string(argsJSON), spanStart: spanStart} }(i, tc) } @@ -789,9 +795,9 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) }) // 5. Process results sequentially: emit events, append messages, save to session + // Note: tool span start/end already emitted inside goroutines above. var loopStuck bool for _, r := range collected { - l.emitToolSpan(ctx, r.spanStart, r.tc.Name, r.tc.ID, r.argsJSON, r.result) // Record for loop detection. argsHash := loopDetector.record(r.tc.Name, r.tc.Arguments) diff --git a/internal/agent/loop_run.go b/internal/agent/loop_run.go index daa5c7f0..0607c877 100644 --- a/internal/agent/loop_run.go +++ b/internal/agent/loop_run.go @@ -43,13 +43,18 @@ func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) { var traceID uuid.UUID isChildTrace := req.ParentTraceID != uuid.Nil && l.traceCollector != nil + // agentSpanID holds the pre-generated root agent span ID. + // Used by emitAgentSpanEnd in the deferred finalizer below. + var agentSpanID uuid.UUID + if isChildTrace { // Announce run: reuse parent trace, don't create new trace record. // Spans will be added to the parent trace with proper nesting. traceID = req.ParentTraceID ctx = tracing.WithTraceID(ctx, traceID) ctx = tracing.WithCollector(ctx, l.traceCollector) - ctx = tracing.WithParentSpanID(ctx, store.GenNewID()) + agentSpanID = store.GenNewID() + ctx = tracing.WithParentSpanID(ctx, agentSpanID) if req.ParentRootSpanID != uuid.Nil { ctx = tracing.WithAnnounceParentSpanID(ctx, req.ParentRootSpanID) } @@ -87,8 +92,8 @@ func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) { ctx = tracing.WithCollector(ctx, l.traceCollector) // Pre-generate root "agent" span ID so LLM/tool spans can reference it as parent. - // The span itself is emitted after runLoop completes (with full timing data). - ctx = tracing.WithParentSpanID(ctx, store.GenNewID()) + agentSpanID = store.GenNewID() + ctx = tracing.WithParentSpanID(ctx, agentSpanID) } } @@ -99,11 +104,19 @@ func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) { } runStart := time.Now().UTC() + + // Emit running agent span immediately so it's visible in the trace UI. + if agentSpanID != uuid.Nil { + l.emitAgentSpanStart(ctx, agentSpanID, runStart, req.Message) + } + result, err := l.runLoop(ctx, req) - // Emit root "agent" span with full timing (parent for all LLM/tool spans). - if l.traceCollector != nil && traceID != uuid.Nil { - l.emitAgentSpan(ctx, runStart, result, err) + // Finalize the root agent span. Uses EmitSpanUpdate (channel send) so it + // succeeds even if ctx is cancelled. Must run before FinishTrace so + // aggregates include this span. + if agentSpanID != uuid.Nil { + l.emitAgentSpanEnd(ctx, agentSpanID, runStart, result, err) } if err != nil { diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index f9ad1036..029acb08 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -31,29 +31,32 @@ func (l *Loop) Model() string { return l.model } // IsRunning returns whether the agent is currently processing. func (l *Loop) IsRunning() bool { return l.activeRuns.Load() > 0 } -// emitLLMSpan records an LLM call span if tracing is active. -// When GOCLAW_TRACE_VERBOSE is set, messages are serialized as InputPreview. -func (l *Loop) emitLLMSpan(ctx context.Context, start time.Time, iteration int, messages []providers.Message, resp *providers.ChatResponse, callErr error) { - traceID := tracing.TraceIDFromContext(ctx) +// --------------------------------------------------------------------------- +// Two-phase LLM span: start (running) + end (completed/error) +// --------------------------------------------------------------------------- + +// emitLLMSpanStart emits a "running" LLM span before the LLM call begins. +// Returns the span ID so the caller can later call emitLLMSpanEnd to finalize it. +// Goroutine-safe: only reads immutable Loop fields and does a channel send. +func (l *Loop) emitLLMSpanStart(ctx context.Context, start time.Time, iteration int, messages []providers.Message) uuid.UUID { collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) if collector == nil || traceID == uuid.Nil { - return + return uuid.Nil } - now := time.Now().UTC() - dur := int(now.Sub(start).Milliseconds()) + spanID := store.GenNewID() span := store.SpanData{ - TraceID: traceID, - SpanType: store.SpanTypeLLMCall, - Name: fmt.Sprintf("%s/%s #%d", l.provider.Name(), l.model, iteration), - StartTime: start, - EndTime: &now, - DurationMS: dur, - Model: l.model, - Provider: l.provider.Name(), - Status: store.SpanStatusCompleted, - Level: store.SpanLevelDefault, - CreatedAt: now, + ID: spanID, + TraceID: traceID, + SpanType: store.SpanTypeLLMCall, + Name: fmt.Sprintf("%s/%s #%d", l.provider.Name(), l.model, iteration), + StartTime: start, + Status: store.SpanStatusRunning, + Level: store.SpanLevelDefault, + Model: l.model, + Provider: l.provider.Name(), + CreatedAt: start, } if parentID := tracing.ParentSpanIDFromContext(ctx); parentID != uuid.Nil { span.ParentSpanID = &parentID @@ -62,10 +65,8 @@ func (l *Loop) emitLLMSpan(ctx context.Context, start time.Time, iteration int, span.AgentID = &l.agentUUID } - // Verbose mode: serialize full messages and output. - // Strip base64 image data to avoid bloating traces and PostgreSQL encoding issues. - verbose := collector.Verbose() - if verbose && len(messages) > 0 { + // Verbose mode: include input messages (same stripping as emitLLMSpan). + if collector.Verbose() && len(messages) > 0 { stripped := make([]providers.Message, len(messages)) copy(stripped, messages) for i := range stripped { @@ -82,13 +83,37 @@ func (l *Loop) emitLLMSpan(ctx context.Context, start time.Time, iteration int, } } + collector.EmitSpan(span) + return spanID +} + +// emitLLMSpanEnd finalizes a running LLM span with results. +// Uses EmitSpanUpdate (channel send) — does NOT depend on ctx being alive, +// so it works correctly even after ctx cancellation or deadline exceeded. +func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.Time, resp *providers.ChatResponse, callErr error) { + if spanID == uuid.Nil { + return // tracing disabled — no running span was emitted + } + collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) + if collector == nil || traceID == uuid.Nil { + return + } + + now := time.Now().UTC() + updates := map[string]any{ + "end_time": now, + "duration_ms": int(now.Sub(start).Milliseconds()), + "status": store.SpanStatusCompleted, + } + if callErr != nil { - span.Status = store.SpanStatusError - span.Error = callErr.Error() + updates["status"] = store.SpanStatusError + updates["error"] = callErr.Error() } else if resp != nil { if resp.Usage != nil { - span.InputTokens = resp.Usage.PromptTokens - span.OutputTokens = resp.Usage.CompletionTokens + updates["input_tokens"] = resp.Usage.PromptTokens + updates["output_tokens"] = resp.Usage.CompletionTokens hasMeta := resp.Usage.CacheCreationTokens > 0 || resp.Usage.CacheReadTokens > 0 || resp.Usage.ThinkingTokens > 0 if hasMeta { meta := map[string]int{} @@ -102,54 +127,58 @@ func (l *Loop) emitLLMSpan(ctx context.Context, start time.Time, iteration int, meta["thinking_tokens"] = resp.Usage.ThinkingTokens } if b, err := json.Marshal(meta); err == nil { - span.Metadata = b + updates["metadata"] = b } } } - span.FinishReason = resp.FinishReason + updates["finish_reason"] = resp.FinishReason + verbose := collector.Verbose() if verbose { preview := resp.Content if resp.Thinking != "" { preview = "\n" + resp.Thinking + "\n\n" + resp.Content } - span.OutputPreview = truncateStr(preview, 100000) + updates["output_preview"] = truncateStr(preview, 100000) } else { - span.OutputPreview = truncateStr(resp.Content, 500) + updates["output_preview"] = truncateStr(resp.Content, 500) } } - collector.EmitSpan(span) + collector.EmitSpanUpdate(spanID, traceID, updates) } -// emitToolSpan records a tool call span if tracing is active. -// result is the full tool execution result, which may contain Usage from inner LLM calls. -func (l *Loop) emitToolSpan(ctx context.Context, start time.Time, toolName, toolCallID, input string, result *tools.Result) { - traceID := tracing.TraceIDFromContext(ctx) +// --------------------------------------------------------------------------- +// Two-phase tool span: start (running) + end (completed/error) +// --------------------------------------------------------------------------- + +// emitToolSpanStart emits a "running" tool span before tool execution begins. +// Returns the span ID so the caller can later call emitToolSpanEnd to finalize it. +// Goroutine-safe: only reads immutable Loop fields and does a channel send. +func (l *Loop) emitToolSpanStart(ctx context.Context, start time.Time, toolName, toolCallID, input string) uuid.UUID { collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) if collector == nil || traceID == uuid.Nil { - return + return uuid.Nil } - now := time.Now().UTC() - dur := int(now.Sub(start).Milliseconds()) previewLimit := 500 if collector.Verbose() { previewLimit = 100000 } + + spanID := store.GenNewID() span := store.SpanData{ - TraceID: traceID, - SpanType: store.SpanTypeToolCall, - Name: toolName, - StartTime: start, - EndTime: &now, - DurationMS: dur, - ToolName: toolName, - ToolCallID: toolCallID, - InputPreview: truncateStr(input, previewLimit), - OutputPreview: truncateStr(result.ForLLM, previewLimit), - Status: store.SpanStatusCompleted, - Level: "DEFAULT", - CreatedAt: now, + ID: spanID, + TraceID: traceID, + SpanType: store.SpanTypeToolCall, + Name: toolName, + StartTime: start, + ToolName: toolName, + ToolCallID: toolCallID, + InputPreview: truncateStr(input, previewLimit), + Status: store.SpanStatusRunning, + Level: store.SpanLevelDefault, + CreatedAt: start, } if parentID := tracing.ParentSpanIDFromContext(ctx); parentID != uuid.Nil { span.ParentSpanID = &parentID @@ -157,62 +186,91 @@ func (l *Loop) emitToolSpan(ctx context.Context, start time.Time, toolName, tool if l.agentUUID != uuid.Nil { span.AgentID = &l.agentUUID } + + collector.EmitSpan(span) + return spanID +} + +// emitToolSpanEnd finalizes a running tool span with execution results. +// Uses EmitSpanUpdate (channel send) — safe after ctx cancellation. +// Goroutine-safe: only does a channel send via EmitSpanUpdate. +func (l *Loop) emitToolSpanEnd(ctx context.Context, spanID uuid.UUID, start time.Time, result *tools.Result) { + if spanID == uuid.Nil { + return // tracing disabled + } + collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) + if collector == nil || traceID == uuid.Nil { + return + } + + now := time.Now().UTC() + previewLimit := 500 + if collector.Verbose() { + previewLimit = 100000 + } + + updates := map[string]any{ + "end_time": now, + "duration_ms": int(now.Sub(start).Milliseconds()), + "status": store.SpanStatusCompleted, + "output_preview": truncateStr(result.ForLLM, previewLimit), + } + if result.IsError { - span.Status = store.SpanStatusError - span.Error = truncateStr(result.ForLLM, 200) + updates["status"] = store.SpanStatusError + updates["error"] = truncateStr(result.ForLLM, 200) } // Record token usage from tools that make internal LLM calls (e.g. read_image). if result.Usage != nil { - span.InputTokens = result.Usage.PromptTokens - span.OutputTokens = result.Usage.CompletionTokens - span.Provider = result.Provider - span.Model = result.Model + updates["input_tokens"] = result.Usage.PromptTokens + updates["output_tokens"] = result.Usage.CompletionTokens + updates["provider"] = result.Provider + updates["model"] = result.Model if result.Usage.CacheCreationTokens > 0 || result.Usage.CacheReadTokens > 0 { meta := map[string]int{ "cache_creation_tokens": result.Usage.CacheCreationTokens, "cache_read_tokens": result.Usage.CacheReadTokens, } if b, err := json.Marshal(meta); err == nil { - span.Metadata = b + updates["metadata"] = b } } } - collector.EmitSpan(span) + collector.EmitSpanUpdate(spanID, traceID, updates) } -// emitAgentSpan records the root "agent" span that parents all LLM/tool spans in this request. -func (l *Loop) emitAgentSpan(ctx context.Context, start time.Time, result *RunResult, runErr error) { - traceID := tracing.TraceIDFromContext(ctx) +// --------------------------------------------------------------------------- +// Two-phase agent span: start (running) + end (completed/error) +// --------------------------------------------------------------------------- + +// emitAgentSpanStart emits a "running" root agent span at the beginning of a run. +// The span is identified by agentSpanID (pre-generated, same ID used as ParentSpanID +// for child LLM/tool spans). +func (l *Loop) emitAgentSpanStart(ctx context.Context, agentSpanID uuid.UUID, start time.Time, inputPreview string) { collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) if collector == nil || traceID == uuid.Nil { return } - agentSpanID := tracing.ParentSpanIDFromContext(ctx) - if agentSpanID == uuid.Nil { - return - } - - now := time.Now().UTC() - dur := int(now.Sub(start).Milliseconds()) spanName := l.id span := store.SpanData{ - ID: agentSpanID, - TraceID: traceID, - SpanType: store.SpanTypeAgent, - Name: spanName, - StartTime: start, - EndTime: &now, - DurationMS: dur, - Model: l.model, - Provider: l.provider.Name(), - Status: store.SpanStatusCompleted, - Level: store.SpanLevelDefault, - CreatedAt: now, + ID: agentSpanID, + TraceID: traceID, + SpanType: store.SpanTypeAgent, + Name: spanName, + StartTime: start, + Status: store.SpanStatusRunning, + Level: store.SpanLevelDefault, + Model: l.model, + Provider: l.provider.Name(), + InputPreview: truncateStr(inputPreview, 500), + CreatedAt: start, } - // Nest under parent root span if this is an announce run + // Nest under parent root span if this is an announce run. if announceParent := tracing.AnnounceParentSpanIDFromContext(ctx); announceParent != uuid.Nil { span.ParentSpanID = &announceParent span.Name = "announce:" + spanName @@ -220,20 +278,43 @@ func (l *Loop) emitAgentSpan(ctx context.Context, start time.Time, result *RunRe if l.agentUUID != uuid.Nil { span.AgentID = &l.agentUUID } + + collector.EmitSpan(span) +} + +// emitAgentSpanEnd finalizes the running root agent span with results. +// Uses EmitSpanUpdate (channel send) — safe after ctx cancellation. +func (l *Loop) emitAgentSpanEnd(ctx context.Context, agentSpanID uuid.UUID, start time.Time, result *RunResult, runErr error) { + if agentSpanID == uuid.Nil { + return + } + collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) + if collector == nil || traceID == uuid.Nil { + return + } + + now := time.Now().UTC() + updates := map[string]any{ + "end_time": now, + "duration_ms": int(now.Sub(start).Milliseconds()), + "status": store.SpanStatusCompleted, + } + if runErr != nil { - span.Status = store.SpanStatusError - span.Error = runErr.Error() + updates["status"] = store.SpanStatusError + updates["error"] = runErr.Error() } else if result != nil { limit := 500 if collector.Verbose() { limit = 100000 } - span.OutputPreview = truncateStr(result.Content, limit) + updates["output_preview"] = truncateStr(result.Content, limit) // Note: token counts are NOT set on agent spans to avoid double-counting // with child llm_call spans. Trace aggregation sums only llm_call spans. } - collector.EmitSpan(span) + collector.EmitSpanUpdate(agentSpanID, traceID, updates) } func truncateStr(s string, maxLen int) string { diff --git a/internal/tools/subagent_exec.go b/internal/tools/subagent_exec.go index 7bbea394..cf223382 100644 --- a/internal/tools/subagent_exec.go +++ b/internal/tools/subagent_exec.go @@ -120,9 +120,9 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) task.CompletedAt = time.Now().UnixMilli() sm.mu.Unlock() - // Always emit root subagent span on exit (uses traceCtx which is never cancelled). - sm.emitSubagentSpan(traceCtx, subRootSpanID, taskStart, task, model, finalContent) - slog.Debug("subagent tracing: root span emitted", + // Finalize root subagent span on exit (uses traceCtx which is never cancelled). + sm.emitSubagentSpanEnd(traceCtx, subRootSpanID, taskStart, task, finalContent) + slog.Debug("subagent tracing: root span finalized", "id", task.ID, "span_id", subRootSpanID, "trace_id", tracing.TraceIDFromContext(traceCtx), "status", task.Status, "iterations", iteration) @@ -157,6 +157,9 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) model = task.Model } + // Emit running subagent root span (after model resolution so span has correct model). + sm.emitSubagentSpanStart(traceCtx, subRootSpanID, taskStart, task, model) + // Build subagent system prompt (matching TS buildSubagentSystemPrompt pattern). systemPrompt := sm.buildSubagentSystemPrompt(task) @@ -191,8 +194,9 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) } llmStart := time.Now().UTC() + llmSpanID := sm.emitLLMSpanStart(subTraceCtx, llmStart, iteration, model, messages) resp, err := sm.provider.Chat(ctx, chatReq) - sm.emitLLMSpan(subTraceCtx, llmStart, iteration, model, messages, resp, err) + sm.emitLLMSpanEnd(subTraceCtx, llmSpanID, llmStart, resp, err) if err != nil { sm.mu.Lock() @@ -221,11 +225,11 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) for _, tc := range resp.ToolCalls { slog.Debug("subagent tool call", "id", task.ID, "tool", tc.Name) - toolStart := time.Now().UTC() - result := toolsReg.Execute(ctx, tc.Name, tc.Arguments) - argsJSON, _ := json.Marshal(tc.Arguments) - sm.emitToolSpan(subTraceCtx, toolStart, tc.Name, tc.ID, string(argsJSON), result.ForLLM, result.IsError) + toolStart := time.Now().UTC() + toolSpanID := sm.emitToolSpanStart(subTraceCtx, toolStart, tc.Name, tc.ID, string(argsJSON)) + result := toolsReg.Execute(ctx, tc.Name, tc.Arguments) + sm.emitToolSpanEnd(subTraceCtx, toolSpanID, toolStart, result.ForLLM, result.IsError) // Capture media file paths from tool results (e.g. image generation). if len(result.Media) > 0 { diff --git a/internal/tools/subagent_tracing.go b/internal/tools/subagent_tracing.go index 14cc9149..eca6323e 100644 --- a/internal/tools/subagent_tracing.go +++ b/internal/tools/subagent_tracing.go @@ -13,9 +13,48 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/tracing" ) -// emitLLMSpan records an LLM call span for a subagent iteration. -// When GOCLAW_TRACE_VERBOSE is set, messages are serialized as InputPreview. -func (sm *SubagentManager) emitLLMSpan(ctx context.Context, start time.Time, iteration int, model string, messages []providers.Message, resp *providers.ChatResponse, callErr error) { +// --------------------------------------------------------------------------- +// Two-phase subagent LLM span: start (running) + end (completed/error) +// --------------------------------------------------------------------------- + +// emitLLMSpanStart emits a "running" LLM span before the subagent LLM call. +func (sm *SubagentManager) emitLLMSpanStart(ctx context.Context, start time.Time, iteration int, model string, messages []providers.Message) uuid.UUID { + collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) + if collector == nil || traceID == uuid.Nil { + return uuid.Nil + } + + spanID := store.GenNewID() + span := store.SpanData{ + ID: spanID, + TraceID: traceID, + SpanType: store.SpanTypeLLMCall, + Name: fmt.Sprintf("%s/%s #%d", sm.provider.Name(), model, iteration), + StartTime: start, + Status: store.SpanStatusRunning, + Level: store.SpanLevelDefault, + Model: model, + Provider: sm.provider.Name(), + CreatedAt: start, + } + if parentID := tracing.ParentSpanIDFromContext(ctx); parentID != uuid.Nil { + span.ParentSpanID = &parentID + } + if collector.Verbose() && len(messages) > 0 { + if b, err := json.Marshal(messages); err == nil { + span.InputPreview = truncate(string(b), 50000) + } + } + collector.EmitSpan(span) + return spanID +} + +// emitLLMSpanEnd finalizes a running subagent LLM span. +func (sm *SubagentManager) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.Time, resp *providers.ChatResponse, callErr error) { + if spanID == uuid.Nil { + return + } collector := tracing.CollectorFromContext(ctx) traceID := tracing.TraceIDFromContext(ctx) if collector == nil || traceID == uuid.Nil { @@ -23,121 +62,144 @@ func (sm *SubagentManager) emitLLMSpan(ctx context.Context, start time.Time, ite } now := time.Now().UTC() - span := store.SpanData{ - TraceID: traceID, - SpanType: store.SpanTypeLLMCall, - Name: fmt.Sprintf("%s/%s #%d", sm.provider.Name(), model, iteration), - StartTime: start, - EndTime: &now, - DurationMS: int(now.Sub(start).Milliseconds()), - Model: model, - Provider: sm.provider.Name(), - Status: store.SpanStatusCompleted, - Level: store.SpanLevelDefault, - CreatedAt: now, + updates := map[string]any{ + "end_time": now, + "duration_ms": int(now.Sub(start).Milliseconds()), + "status": store.SpanStatusCompleted, } - if parentID := tracing.ParentSpanIDFromContext(ctx); parentID != uuid.Nil { - span.ParentSpanID = &parentID - } - - // Verbose mode: serialize full messages as InputPreview - if collector.Verbose() && len(messages) > 0 { - if b, err := json.Marshal(messages); err == nil { - span.InputPreview = truncate(string(b), 50000) - } - } - if callErr != nil { - span.Status = store.SpanStatusError - span.Error = callErr.Error() + updates["status"] = store.SpanStatusError + updates["error"] = callErr.Error() } else if resp != nil { if resp.Usage != nil { - span.InputTokens = resp.Usage.PromptTokens - span.OutputTokens = resp.Usage.CompletionTokens + updates["input_tokens"] = resp.Usage.PromptTokens + updates["output_tokens"] = resp.Usage.CompletionTokens if resp.Usage.CacheCreationTokens > 0 || resp.Usage.CacheReadTokens > 0 { meta := map[string]int{ "cache_creation_tokens": resp.Usage.CacheCreationTokens, "cache_read_tokens": resp.Usage.CacheReadTokens, } if b, err := json.Marshal(meta); err == nil { - span.Metadata = b + updates["metadata"] = b } } } - span.FinishReason = resp.FinishReason - span.OutputPreview = truncate(resp.Content, 500) + updates["finish_reason"] = resp.FinishReason + updates["output_preview"] = truncate(resp.Content, 500) } - collector.EmitSpan(span) + collector.EmitSpanUpdate(spanID, traceID, updates) } -// emitToolSpan records a tool call span for a subagent tool execution. -func (sm *SubagentManager) emitToolSpan(ctx context.Context, start time.Time, toolName, toolCallID, input, output string, isError bool) { +// --------------------------------------------------------------------------- +// Two-phase subagent tool span: start (running) + end (completed/error) +// --------------------------------------------------------------------------- + +// emitToolSpanStart emits a "running" tool span before subagent tool execution. +func (sm *SubagentManager) emitToolSpanStart(ctx context.Context, start time.Time, toolName, toolCallID, input string) uuid.UUID { collector := tracing.CollectorFromContext(ctx) traceID := tracing.TraceIDFromContext(ctx) if collector == nil || traceID == uuid.Nil { - return + return uuid.Nil } - now := time.Now().UTC() + spanID := store.GenNewID() span := store.SpanData{ - TraceID: traceID, - SpanType: store.SpanTypeToolCall, - Name: toolName, - StartTime: start, - EndTime: &now, - DurationMS: int(now.Sub(start).Milliseconds()), - ToolName: toolName, - ToolCallID: toolCallID, - InputPreview: truncate(input, 500), - OutputPreview: truncate(output, 500), - Status: store.SpanStatusCompleted, - Level: store.SpanLevelDefault, - CreatedAt: now, + ID: spanID, + TraceID: traceID, + SpanType: store.SpanTypeToolCall, + Name: toolName, + StartTime: start, + ToolName: toolName, + ToolCallID: toolCallID, + InputPreview: truncate(input, 500), + Status: store.SpanStatusRunning, + Level: store.SpanLevelDefault, + CreatedAt: start, } if parentID := tracing.ParentSpanIDFromContext(ctx); parentID != uuid.Nil { span.ParentSpanID = &parentID } - if isError { - span.Status = store.SpanStatusError - span.Error = truncate(output, 200) - } collector.EmitSpan(span) + return spanID } -// emitSubagentSpan records the root "agent" span for the subagent execution. -// This span parents all LLM/tool spans within this subagent run. -func (sm *SubagentManager) emitSubagentSpan(ctx context.Context, spanID uuid.UUID, start time.Time, task *SubagentTask, model string, output string) { +// emitToolSpanEnd finalizes a running subagent tool span. +func (sm *SubagentManager) emitToolSpanEnd(ctx context.Context, spanID uuid.UUID, start time.Time, output string, isError bool) { + if spanID == uuid.Nil { + return + } collector := tracing.CollectorFromContext(ctx) traceID := tracing.TraceIDFromContext(ctx) if collector == nil || traceID == uuid.Nil { return } - // parentSpanID here is the ORIGINAL parent (before we overrode it for child spans) - parentSpanID := tracing.ParentSpanIDFromContext(ctx) - now := time.Now().UTC() + updates := map[string]any{ + "end_time": now, + "duration_ms": int(now.Sub(start).Milliseconds()), + "status": store.SpanStatusCompleted, + "output_preview": truncate(output, 500), + } + if isError { + updates["status"] = store.SpanStatusError + updates["error"] = truncate(output, 200) + } + collector.EmitSpanUpdate(spanID, traceID, updates) +} + +// --------------------------------------------------------------------------- +// Two-phase subagent root span: start (running) + end (completed/error) +// --------------------------------------------------------------------------- + +// emitSubagentSpanStart emits a "running" root subagent span at task start. +func (sm *SubagentManager) emitSubagentSpanStart(ctx context.Context, spanID uuid.UUID, start time.Time, task *SubagentTask, model string) { + collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) + if collector == nil || traceID == uuid.Nil { + return + } + span := store.SpanData{ - ID: spanID, - TraceID: traceID, - SpanType: store.SpanTypeAgent, - Name: fmt.Sprintf("subagent:%s", task.Label), - StartTime: start, - EndTime: &now, - DurationMS: int(now.Sub(start).Milliseconds()), - Model: model, - Provider: sm.provider.Name(), - OutputPreview: truncate(output, 500), - Status: store.SpanStatusCompleted, - Level: store.SpanLevelDefault, - CreatedAt: now, + ID: spanID, + TraceID: traceID, + SpanType: store.SpanTypeAgent, + Name: fmt.Sprintf("subagent:%s", task.Label), + StartTime: start, + Status: store.SpanStatusRunning, + Level: store.SpanLevelDefault, + Model: model, + Provider: sm.provider.Name(), + InputPreview: truncate(task.Task, 500), + CreatedAt: start, } - if parentSpanID != uuid.Nil { + if parentSpanID := tracing.ParentSpanIDFromContext(ctx); parentSpanID != uuid.Nil { span.ParentSpanID = &parentSpanID } - if task.Status == TaskStatusFailed || task.Status == TaskStatusCancelled { - span.Status = store.SpanStatusError - span.Error = truncate(task.Result, 200) - } collector.EmitSpan(span) } + +// emitSubagentSpanEnd finalizes the running subagent root span. +func (sm *SubagentManager) emitSubagentSpanEnd(ctx context.Context, spanID uuid.UUID, start time.Time, task *SubagentTask, output string) { + if spanID == uuid.Nil { + return + } + collector := tracing.CollectorFromContext(ctx) + traceID := tracing.TraceIDFromContext(ctx) + if collector == nil || traceID == uuid.Nil { + return + } + + now := time.Now().UTC() + updates := map[string]any{ + "end_time": now, + "duration_ms": int(now.Sub(start).Milliseconds()), + "status": store.SpanStatusCompleted, + "output_preview": truncate(output, 500), + } + if task.Status == TaskStatusFailed || task.Status == TaskStatusCancelled { + updates["status"] = store.SpanStatusError + updates["error"] = truncate(task.Result, 200) + } + collector.EmitSpanUpdate(spanID, traceID, updates) +} diff --git a/internal/tracing/collector.go b/internal/tracing/collector.go index 4b014f58..ffe561f8 100644 --- a/internal/tracing/collector.go +++ b/internal/tracing/collector.go @@ -29,6 +29,14 @@ type SpanExporter interface { Shutdown(ctx context.Context) error } +// spanUpdate represents a deferred span field update, buffered alongside new +// spans and applied during the same flush cycle (after batch INSERT). +type spanUpdate struct { + SpanID uuid.UUID + TraceID uuid.UUID + Updates map[string]any +} + // Collector buffers spans in memory and periodically flushes them to the // TracingStore in batches. Traces are created synchronously (one per run), // while spans are buffered for async batch insert. @@ -38,9 +46,10 @@ type SpanExporter interface { type Collector struct { store store.TracingStore - spanCh chan store.SpanData - stopCh chan struct{} - wg sync.WaitGroup + spanCh chan store.SpanData + spanUpdateCh chan spanUpdate // deferred span updates (two-phase tracing) + stopCh chan struct{} + wg sync.WaitGroup // traces that need aggregate updates on flush dirtyTraces map[uuid.UUID]struct{} @@ -62,11 +71,12 @@ func NewCollector(ts store.TracingStore) *Collector { slog.Info("tracing: verbose mode enabled (GOCLAW_TRACE_VERBOSE)") } return &Collector{ - store: ts, - spanCh: make(chan store.SpanData, defaultBufferSize), - stopCh: make(chan struct{}), - dirtyTraces: make(map[uuid.UUID]struct{}), - verbose: verbose, + store: ts, + spanCh: make(chan store.SpanData, defaultBufferSize), + spanUpdateCh: make(chan spanUpdate, defaultBufferSize), + stopCh: make(chan struct{}), + dirtyTraces: make(map[uuid.UUID]struct{}), + verbose: verbose, } } @@ -132,6 +142,20 @@ func (c *Collector) EmitSpan(span store.SpanData) { } } +// EmitSpanUpdate enqueues a deferred update for an existing span. +// Used by two-phase tracing: a "running" span is emitted via EmitSpan before +// execution starts, then updated via EmitSpanUpdate when execution completes. +// Non-blocking channel send — safe to call even after ctx cancellation. +func (c *Collector) EmitSpanUpdate(spanID, traceID uuid.UUID, updates map[string]any) { + select { + case c.spanUpdateCh <- spanUpdate{SpanID: spanID, TraceID: traceID, Updates: updates}: + c.markDirty(traceID) + default: + slog.Warn("tracing: span update buffer full, dropping update", + "span_id", spanID) + } +} + // FinishTrace marks a trace as completed and schedules aggregate update. func (c *Collector) FinishTrace(ctx context.Context, traceID uuid.UUID, status string, errMsg string, outputPreview string) { now := time.Now().UTC() @@ -204,6 +228,29 @@ done: } } + // Drain and apply deferred span updates (two-phase tracing). + // Must run AFTER batch insert so that "running" spans exist before we UPDATE them. + var updates []spanUpdate + for { + select { + case u := <-c.spanUpdateCh: + updates = append(updates, u) + default: + goto doneUpdates + } + } +doneUpdates: + if len(updates) > 0 { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + for _, u := range updates { + if err := c.store.UpdateSpan(ctx, u.SpanID, u.Updates); err != nil { + slog.Warn("tracing: span update failed", "span_id", u.SpanID, "error", err) + } + } + slog.Debug("tracing: applied span updates", "count", len(updates)) + } + // Update aggregates for dirty traces c.dirtyTracesMu.Lock() dirty := c.dirtyTraces