mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 04:16:06 +00:00
tracing: nest tool calls under LLM spans, emit hook spans, fix usage_events FK (#1415)
- Reparent tool_call spans under their producing llm_call span via RunState.CurrentLLMSpanID so a model turn's trace visibly contains the tool calls it triggered. - Wire internal/hooks EmitHookSpan into the dispatcher writeExec so every hook execution (pre/post tool use, etc.) produces a trace span with input, console output, decision, error, and duration for agent troubleshooting. - Fix usage_events_span_id_fkey violations: usage events were inserted synchronously referencing a span flushed ~5s later, silently dropping token/cost data. Route them through the collector so they flush after spans in the same cycle. No schema migration; FK preserved. Co-authored-by: Bruno Clermont <bruno.clermont@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
11 files changed
+208
-46
No files matched your search
@@ -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 {
|
||||
|
||||
@@ -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()))
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 ───────────────────────────────────────────────────────────
|
||||
|
||||
+59
-28
@@ -14,16 +14,18 @@ import (
|
||||
// EmitHookSpan records a tracing span for a hook execution.
|
||||
//
|
||||
// The span name follows the plan's convention: "hook.<handlerType>.<event>"
|
||||
// (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)
|
||||
}
|
||||
@@ -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.<handlerType>.<event>".
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user