diff --git a/cmd/gateway_setup.go b/cmd/gateway_setup.go index 1facfa49..bce4ccaf 100644 --- a/cmd/gateway_setup.go +++ b/cmd/gateway_setup.go @@ -303,7 +303,7 @@ func wireTracingAndCron( ) (*tracing.Collector, *tracing.SnapshotWorker) { var traceCollector *tracing.Collector if stores.Tracing != nil { - traceCollector = tracing.NewCollector(stores.Tracing) + traceCollector = tracing.NewCollector(stores.Tracing, stores.UsageEvents) traceCollector.OnFlush = func(traceIDs []uuid.UUID) { ids := make([]string, len(traceIDs)) for i, id := range traceIDs { diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index 1c6a5828..e780fafd 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -402,6 +402,9 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c opts = append(opts, withProvider(provider.Name())) } spanID := l.emitLLMSpanStart(ctx, start, state.Iteration+1, chatReq.Messages, opts...) + if spanID != uuid.Nil { + state.CurrentLLMSpanID = &spanID + } recordUsageCapAttempt := func(reservation *usagecaps.Reservation) { if reservation != nil { opts = append(opts, withUsageCapMetadata(reservation.TraceMetadata())) diff --git a/internal/agent/loop_pipeline_tool_callbacks.go b/internal/agent/loop_pipeline_tool_callbacks.go index ee1e5368..f550e853 100644 --- a/internal/agent/loop_pipeline_tool_callbacks.go +++ b/internal/agent/loop_pipeline_tool_callbacks.go @@ -13,6 +13,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + "github.com/nextlevelbuilder/goclaw/internal/tracing" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -33,9 +34,17 @@ func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx Payload: map[string]any{"name": tc.Name, "id": tc.ID, "arguments": tc.Arguments}, }) - // Emit tool span start for tracing. + // Emit tool span start for tracing. Parent the tool span to the current + // LLM-call span so the trace tree nests tool calls under the model turn. toolStart := time.Now().UTC() - toolSpanID := l.emitToolSpanStart(ctx, toolStart, registryName, tc.ID, string(argsJSON)) + toolCtx := ctx + if state.CurrentLLMSpanID != nil { + toolCtx = tracing.WithParentSpanID(ctx, *state.CurrentLLMSpanID) + } + toolSpanID := l.emitToolSpanStart(toolCtx, toolStart, registryName, tc.ID, string(argsJSON)) + if toolSpanID != uuid.Nil { + state.CurrentToolSpanID = &toolSpanID + } // Inject agent audio snapshot so TTS tool (and any future audio consumers) // can read agent-level voice/model config without an extra DB lookup. diff --git a/internal/agent/usage_events.go b/internal/agent/usage_events.go index fa543265..5e297017 100644 --- a/internal/agent/usage_events.go +++ b/internal/agent/usage_events.go @@ -178,6 +178,20 @@ func (l *Loop) isRuntimeTool(toolName string) bool { } func (l *Loop) insertUsageEventBestEffort(ctx context.Context, event store.UsageEvent) { + // When a tracing collector is configured, buffer the event so it is flushed + // AFTER the referenced span row is written in the same flush cycle. This + // preserves the usage_events_span_id_fkey and stops rows from being dropped. + if l.traceCollector != nil { + // Ensure tenant is carried on the event: the collector flush inserts via + // InsertEvents, which scopes each row by the event's own TenantID field. + tenantID := store.TenantIDFromContext(ctx) + if tenantID != uuid.Nil { + event.TenantID = tenantID + } + l.traceCollector.EmitUsageEvent(event) + return + } + // Fallback for non-tracing contexts (e.g. some tests): direct best-effort insert. tenantID := event.TenantID go func() { bgCtx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 5*time.Second) diff --git a/internal/hooks/dispatcher.go b/internal/hooks/dispatcher.go index f4edb1b5..5a2f0ba2 100644 --- a/internal/hooks/dispatcher.go +++ b/internal/hooks/dispatcher.go @@ -2,6 +2,7 @@ package hooks import ( "context" + "encoding/json" "errors" "fmt" "log/slog" @@ -531,6 +532,24 @@ func (d *stdDispatcher) writeExec(ctx context.Context, cfg HookConfig, ev Event, if err := d.audit.Log(ctx, exec); err != nil { slog.Warn("security.hook.audit_write_failed", "err", err, "hook_id", cfg.ID) } + + // Emit a trace span for this hook execution so hook input/output/decision are + // visible alongside model and tool spans for agent troubleshooting. ctx here + // already carries the parent span ID set by the tool_stage wrappers + // (pre_tool_use → llm span; post_tool_use → tool span). For non-tool hooks + // (session start/stop), parent falls back to the agent span (ctx default). + durationMS := int(duration / time.Millisecond) + startedAt := time.Now().UTC().Add(-duration) + hookInput, _ := json.Marshal(map[string]any{ + "event": ev.HookEvent, + "tool_name": ev.ToolName, + "tool_input": ev.ToolInput, + "raw_input": ev.RawInput, + "session_id": ev.SessionID, + "agent_id": ev.AgentID, + "depth": ev.Depth, + }) + EmitHookSpan(ctx, ev.HookEvent, cfg.HandlerType, cfg, startedAt, dec, durationMS, errMsg, string(hookInput), consoleOutput) } // ── circuitBreaker ─────────────────────────────────────────────────────────── diff --git a/internal/hooks/tracing.go b/internal/hooks/tracing.go index b76fd49f..5bd18e0d 100644 --- a/internal/hooks/tracing.go +++ b/internal/hooks/tracing.go @@ -14,16 +14,18 @@ import ( // EmitHookSpan records a tracing span for a hook execution. // // The span name follows the plan's convention: "hook.." -// (e.g., "hook.command.pre_tool_use"). Duration is computed from startedAt -// to now; status is "completed" on success, "error" when errMsg is non-empty. -// The decision is persisted into Metadata as `{"decision":"allow|block|..."}` -// so dashboards can aggregate allow/block ratios per event. +// (e.g., "hook.command.pre_tool_use"). Duration is taken from durationMS +// (the caller already measured it); status is "completed" on success, +// "error" when decision == DecisionError or errMsg is non-empty. The decision +// and hook identity are persisted into Metadata so dashboards can aggregate +// allow/block/error ratios per event and hook. // // Fields lifted from ctx: // - trace id (tracing.TraceIDFromContext) // - parent span id (tracing.ParentSpanIDFromContext, omitted when nil) +// - agent id (store.AgentIDFromContext, omitted when zero) // - team id (tracing.TraceTeamIDPtrFromContext) -// - tenant id (store.TenantIDFromContext) +// - tenant id (store.TenantIDFromContext, falls back to MasterTenantID) // // No-op when ctx has no collector attached — safe in tests and for tenants // without tracing enabled. @@ -31,43 +33,68 @@ func EmitHookSpan( ctx context.Context, event HookEvent, ht HandlerType, + cfg HookConfig, startedAt time.Time, decision Decision, + durationMS int, errMsg string, + input string, + output string, ) { - collector := tracing.CollectorFromContext(ctx) - if collector == nil { - return - } + const eventPreviewLimit = 40_000 end := time.Now().UTC() - durationMS := max(int(end.Sub(startedAt)/time.Millisecond), 0) status := store.SpanStatusCompleted - if errMsg != "" { + if decision == DecisionError || errMsg != "" { status = store.SpanStatusError } - var metadata json.RawMessage - if decision != "" { - if b, err := json.Marshal(map[string]string{"decision": string(decision)}); err == nil { - metadata = b - } + // Mirror emitLLMSpanStart's ctx-extraction pattern for agent/team/tenant. + var agentID *uuid.UUID + if a := store.AgentIDFromContext(ctx); a != uuid.Nil { + agentID = &a + } + tenantID := store.TenantIDFromContext(ctx) + if tenantID == uuid.Nil { + tenantID = store.MasterTenantID + } + + errorMsg := errMsg + if len(errorMsg) > 200 { + errorMsg = errorMsg[:200] + } + + metadata := map[string]any{ + "decision": string(decision), + "hook_id": cfg.ID.String(), + "hook_name": cfg.Name, + "hook_event": string(event), + "handler_type": string(ht), + "duration_ms": durationMS, + } + var metadataJSON json.RawMessage + if b, err := json.Marshal(metadata); err == nil { + metadataJSON = b } span := store.SpanData{ - TraceID: tracing.TraceIDFromContext(ctx), - SpanType: store.SpanTypeEvent, - Name: "hook." + string(ht) + "." + string(event), - StartTime: startedAt, - EndTime: &end, - DurationMS: durationMS, - Status: status, - Error: errMsg, - Metadata: metadata, - TeamID: tracing.TraceTeamIDPtrFromContext(ctx), - TenantID: store.TenantIDFromContext(ctx), - CreatedAt: end, + TraceID: tracing.TraceIDFromContext(ctx), + SpanType: store.SpanTypeEvent, + Name: "hook." + string(ht) + "." + string(event), + StartTime: startedAt, + EndTime: &end, + DurationMS: durationMS, + Status: status, + Error: errorMsg, + Level: store.SpanLevelDefault, + InputPreview: tracing.TruncateJSON(input, eventPreviewLimit), + OutputPreview: tracing.TruncateMid(output, eventPreviewLimit), + AgentID: agentID, + TeamID: tracing.TraceTeamIDPtrFromContext(ctx), + TenantID: tenantID, + Metadata: metadataJSON, + CreatedAt: end, } // Attach parent only when present — leaving it nil avoids bogus FK edges. @@ -76,5 +103,9 @@ func EmitHookSpan( span.ParentSpanID = &p } + collector := tracing.CollectorFromContext(ctx) + if collector == nil { + return // tracing disabled — no collector attached + } collector.EmitSpan(span) } diff --git a/internal/hooks/tracing_test.go b/internal/hooks/tracing_test.go index 466f60d3..d34f0410 100644 --- a/internal/hooks/tracing_test.go +++ b/internal/hooks/tracing_test.go @@ -61,7 +61,7 @@ func newRunningCollector(t *testing.T) (*tracing.Collector, *capturingTracingSto // TestEmitHookSpan_NoCollector_NoPanic: safe no-op when ctx has no collector. func TestEmitHookSpan_NoCollector_NoPanic(t *testing.T) { hooks.EmitHookSpan(context.Background(), hooks.EventPreToolUse, hooks.HandlerCommand, - time.Now().Add(-100*time.Millisecond), hooks.DecisionAllow, "") + hooks.HookConfig{}, time.Now().Add(-100*time.Millisecond), hooks.DecisionAllow, 0, "", "", "") } // TestEmitHookSpan_NameFormat: span name = "hook..". @@ -71,7 +71,7 @@ func TestEmitHookSpan_NameFormat(t *testing.T) { ctx = tracing.WithTraceID(ctx, uuid.New()) hooks.EmitHookSpan(ctx, hooks.EventPreToolUse, hooks.HandlerCommand, - time.Now().Add(-50*time.Millisecond), hooks.DecisionAllow, "") + hooks.HookConfig{}, time.Now().Add(-50*time.Millisecond), hooks.DecisionAllow, 0, "", "", "") c.Stop() // flush synchronously @@ -89,7 +89,7 @@ func TestEmitHookSpan_NameFormat(t *testing.T) { } } -// TestEmitHookSpan_DurationAndStatus: duration is derived from startedAt; status +// TestEmitHookSpan_DurationAndStatus: duration is passed explicitly; status // reflects error presence (completed vs error). func TestEmitHookSpan_DurationAndStatus(t *testing.T) { c, cs := newRunningCollector(t) @@ -97,10 +97,10 @@ func TestEmitHookSpan_DurationAndStatus(t *testing.T) { ctx = tracing.WithTraceID(ctx, uuid.New()) startedOK := time.Now().Add(-42 * time.Millisecond) - hooks.EmitHookSpan(ctx, hooks.EventPostToolUse, hooks.HandlerHTTP, startedOK, hooks.DecisionAllow, "") + hooks.EmitHookSpan(ctx, hooks.EventPostToolUse, hooks.HandlerHTTP, hooks.HookConfig{}, startedOK, hooks.DecisionAllow, 42, "", "", "") startedErr := time.Now().Add(-17 * time.Millisecond) - hooks.EmitHookSpan(ctx, hooks.EventPostToolUse, hooks.HandlerHTTP, startedErr, hooks.DecisionBlock, "timeout") + hooks.EmitHookSpan(ctx, hooks.EventPostToolUse, hooks.HandlerHTTP, hooks.HookConfig{}, startedErr, hooks.DecisionBlock, 17, "timeout", "", "") c.Stop() @@ -117,8 +117,8 @@ func TestEmitHookSpan_DurationAndStatus(t *testing.T) { if okSpan.Error != "" { t.Errorf("ok span should have empty error; got %q", okSpan.Error) } - if okSpan.DurationMS < 40 { - t.Errorf("ok span duration_ms = %d, want >= 40", okSpan.DurationMS) + if okSpan.DurationMS != 42 { + t.Errorf("ok span duration_ms = %d, want 42", okSpan.DurationMS) } if errSpan.Status != store.SpanStatusError { t.Errorf("err span status = %q, want %q", errSpan.Status, store.SpanStatusError) @@ -135,7 +135,7 @@ func TestEmitHookSpan_DecisionInMetadata(t *testing.T) { ctx = tracing.WithTraceID(ctx, uuid.New()) hooks.EmitHookSpan(ctx, hooks.EventUserPromptSubmit, hooks.HandlerPrompt, - time.Now().Add(-10*time.Millisecond), hooks.DecisionBlock, "") + hooks.HookConfig{}, time.Now().Add(-10*time.Millisecond), hooks.DecisionBlock, 0, "", "", "") c.Stop() @@ -169,7 +169,7 @@ func TestEmitHookSpan_PropagatesTenantAndTrace(t *testing.T) { ctx = store.WithTenantID(ctx, tenantID) hooks.EmitHookSpan(ctx, hooks.EventPreToolUse, hooks.HandlerCommand, - time.Now().Add(-5*time.Millisecond), hooks.DecisionAllow, "") + hooks.HookConfig{}, time.Now().Add(-5*time.Millisecond), hooks.DecisionAllow, 0, "", "", "") c.Stop() diff --git a/internal/pipeline/run_state.go b/internal/pipeline/run_state.go index 1055d94d..98241c9f 100644 --- a/internal/pipeline/run_state.go +++ b/internal/pipeline/run_state.go @@ -3,6 +3,8 @@ package pipeline import ( "context" + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/workspace" @@ -37,6 +39,11 @@ type RunState struct { Iteration int RunID string ExitCode StageResult + + // CurrentLLMSpanID is the most recent LLM-call span in this run; tool spans parent to it. + CurrentLLMSpanID *uuid.UUID + // CurrentToolSpanID is the most recent tool-call span; post-tool-use hook spans parent to it. + CurrentToolSpanID *uuid.UUID } // NewRunState creates a RunState with identity fields set. diff --git a/internal/pipeline/tool_stage.go b/internal/pipeline/tool_stage.go index 115e89d0..f6f1f0b1 100644 --- a/internal/pipeline/tool_stage.go +++ b/internal/pipeline/tool_stage.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/hooks" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/tracing" ) const ( @@ -92,8 +93,12 @@ func (s *ToolStage) Execute(ctx context.Context, state *RunState) error { state.Tool.TotalToolCalls++ // Hook: async PostToolUse — fire and forget with detached context. + // Parent the hook span to the tool span so it nests under the tool call. if s.deps.Hooks != nil { detached := context.WithoutCancel(ctx) + if state.CurrentToolSpanID != nil { + detached = tracing.WithParentSpanID(detached, *state.CurrentToolSpanID) + } go s.deps.FireHook(detached, hooks.Event{ //nolint:errcheck EventID: uuid.NewString(), SessionID: state.Input.SessionKey, @@ -166,7 +171,13 @@ func (s *ToolStage) preflightToolCall(ctx context.Context, state *RunState, tc p } } } - r, _ := s.deps.FireHook(ctx, hooks.Event{ + // Parent the pre_tool_use hook span to the current LLM-call span so it nests + // alongside the tool call it gates. + hookCtx := ctx + if state.CurrentLLMSpanID != nil { + hookCtx = tracing.WithParentSpanID(ctx, *state.CurrentLLMSpanID) + } + r, _ := s.deps.FireHook(hookCtx, hooks.Event{ EventID: uuid.NewString(), SessionID: state.Input.SessionKey, TenantID: store.TenantIDFromContext(ctx), @@ -272,7 +283,11 @@ func (s *ToolStage) executeParallel(ctx context.Context, state *RunState, prefli results[idx] = rawResult{index: item.index, tc: item.tc, err: ctx.Err()} return } - msg, rawData, err := s.deps.ExecuteToolRaw(ctx, item.tc) + itemCtx := ctx + if state.CurrentLLMSpanID != nil { + itemCtx = tracing.WithParentSpanID(ctx, *state.CurrentLLMSpanID) + } + msg, rawData, err := s.deps.ExecuteToolRaw(itemCtx, item.tc) results[idx] = rawResult{index: item.index, tc: item.tc, msg: msg, rawData: rawData, err: err} }(i, item) } @@ -311,8 +326,15 @@ func (s *ToolStage) executeParallel(ctx context.Context, state *RunState, prefli // Hook: async PostToolUse for parallel path — fire and forget. // PreToolUse already ran in preflight before any raw I/O was scheduled. + // The parallel tool I/O is parented to the LLM span (see ExecuteToolRaw + // call site above), so parent the parallel hook span there too — the + // per-tool span ID is not plumbed back to the pipeline from the opaque + // rawData payload, so the LLM span is the correct shared ancestor. if s.deps.Hooks != nil { detached := context.WithoutCancel(ctx) + if state.CurrentLLMSpanID != nil { + detached = tracing.WithParentSpanID(detached, *state.CurrentLLMSpanID) + } go s.deps.FireHook(detached, hooks.Event{ //nolint:errcheck EventID: uuid.NewString(), SessionID: state.Input.SessionKey, diff --git a/internal/tracing/collector.go b/internal/tracing/collector.go index cb0f6a5b..9b7e2fdb 100644 --- a/internal/tracing/collector.go +++ b/internal/tracing/collector.go @@ -78,8 +78,13 @@ type pendingUpdate struct { type Collector struct { store store.TracingStore + // usageStore flushes usage events AFTER spans in the same flush cycle, + // so a usage event's referenced span row exists before the FK is checked. + usageStore store.UsageEventStore + spanCh chan store.SpanData spanUpdateCh chan spanUpdate // deferred span updates (two-phase tracing) + usageEventCh chan store.UsageEvent stopCh chan struct{} wg sync.WaitGroup @@ -106,15 +111,27 @@ type Collector struct { // NewCollector creates a new tracing collector backed by the given store. // Set GOCLAW_TRACE_VERBOSE=1 to include full LLM input in spans. -func NewCollector(ts store.TracingStore) *Collector { +// +// The optional usageStore flushes usage events AFTER spans within the same +// flush cycle, preserving the usage_events→spans FK without a schema change. +// When omitted (nil), EmitUsageEvent is a no-op and callers fall back to a +// direct best-effort insert. This keeps callers that only need tracing +// (e.g. unit/integration tests) compiling unchanged. +func NewCollector(ts store.TracingStore, usageStore ...store.UsageEventStore) *Collector { verbose := os.Getenv("GOCLAW_TRACE_VERBOSE") != "" if verbose { slog.Info("tracing: verbose mode enabled (GOCLAW_TRACE_VERBOSE)") } + var us store.UsageEventStore + if len(usageStore) > 0 { + us = usageStore[0] + } return &Collector{ store: ts, + usageStore: us, spanCh: make(chan store.SpanData, defaultBufferSize), spanUpdateCh: make(chan spanUpdate, defaultBufferSize), + usageEventCh: make(chan store.UsageEvent, defaultBufferSize), stopCh: make(chan struct{}), retryCh: make(chan pendingUpdate, retryQueueCap), dirtyTraces: make(map[uuid.UUID]struct{}), @@ -228,6 +245,21 @@ func (c *Collector) EmitSpanUpdate(spanID, traceID uuid.UUID, updates map[string } } +// EmitUsageEvent buffers a usage event to be flushed AFTER spans in the same +// flush cycle, so its referenced span row exists before the FK is checked. +// Non-blocking: drops the event if the buffer is full. No-op when no usage +// store is configured (e.g. tracing-only contexts). +func (c *Collector) EmitUsageEvent(event store.UsageEvent) { + if c == nil || c.usageStore == nil { + return + } + select { + case c.usageEventCh <- event: + default: + slog.Debug("tracing.usage_event_dropped", "resource", event.ResourceName) + } +} + // SetTraceStatus updates only the trace status and marks it dirty for re-aggregation. // Uses a detached context with retries to survive caller context cancellation. func (c *Collector) SetTraceStatus(ctx context.Context, traceID uuid.UUID, status string) { @@ -498,6 +530,32 @@ doneUpdates: slog.Debug("tracing: applied span updates", "count", len(updates)) } + // Drain and insert buffered usage events AFTER spans have been written in + // this same flush cycle. Each event carries its own TenantID (the store + // reads tenant from the row, not ctx), so the batch insert stays + // tenant-correct. Ordering here is the root-cause fix: a usage event's + // referenced span row now exists before the usage_events_span_id_fkey is + // checked, so the INSERT no longer fails and silently drops the row. + var usageEvents []store.UsageEvent + for { + select { + case ev := <-c.usageEventCh: + usageEvents = append(usageEvents, ev) + default: + goto doneUsage + } + } +doneUsage: + if len(usageEvents) > 0 { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := c.usageStore.InsertEvents(ctx, usageEvents); err != nil { + slog.Warn("tracing: usage event insert failed", "count", len(usageEvents), "error", err) + } else { + slog.Debug("tracing: flushed usage events", "count", len(usageEvents)) + } + } + // Update aggregates for dirty traces c.dirtyTracesMu.Lock() dirty := c.dirtyTraces diff --git a/tests/integration/hooks_tracing_test.go b/tests/integration/hooks_tracing_test.go index 8afde48d..71798613 100644 --- a/tests/integration/hooks_tracing_test.go +++ b/tests/integration/hooks_tracing_test.go @@ -47,7 +47,7 @@ func TestHooksTracing_EmitHookSpan(t *testing.T) { ctx := tracing.WithTraceID(tracing.WithCollector(tenantCtx(tenantID), c), traceID) started := time.Now().Add(-200 * time.Millisecond) - hooks.EmitHookSpan(ctx, hooks.EventPreToolUse, hooks.HandlerHTTP, started, hooks.DecisionAllow, "") + hooks.EmitHookSpan(ctx, hooks.EventPreToolUse, hooks.HandlerHTTP, hooks.HookConfig{}, started, hooks.DecisionAllow, 200, "", "", "") // Stop collector to flush the buffered span; Stop() drains the chan. c.Stop() @@ -161,7 +161,6 @@ func (w tracedHandler) Execute(ctx_ctx context.Context, cfg hooks.HookConfig, ev if err != nil { errMsg = err.Error() } - hooks.EmitHookSpan(ctx_ctx, ev.HookEvent, cfg.HandlerType, start, dec, errMsg) + hooks.EmitHookSpan(ctx_ctx, ev.HookEvent, cfg.HandlerType, cfg, start, dec, int(time.Since(start).Milliseconds()), errMsg, "", "") return dec, err } -