diff --git a/internal/harness/README.md b/internal/harness/README.md index 6374a5f50..ecffc3d68 100644 --- a/internal/harness/README.md +++ b/internal/harness/README.md @@ -8,7 +8,7 @@ A Go framework for building **stateful, multi-agent applications** with LLMs. It --- - [Quick Start](#quick-start) -- [Two-Layer Architecture](#two-layer-architecture) +- [Architecture Overview](#architecture-overview) - [Layer 1: Graph Engine (graphengine)](#layer-1-graph-engine-graphengine) - [Layer 2: Agent Development Kit (agentcore)](#layer-2-agent-development-kit-agentcore) - [Layer 3: Push-Based AgentLoop](#layer-3-push-based-agentloop) @@ -17,6 +17,7 @@ A Go framework for building **stateful, multi-agent applications** with LLMs. It - [Cancellation System](#cancellation-system) - [Prebuilt Components](#prebuilt-components) - [Observability (OpenTelemetry)](#observability-opentelemetry) +- [Event Sourcing & Replay](#event-sourcing--replay) - [Project Structure](#project-structure) - [Examples](#examples) - [Contributing](#contributing) @@ -118,7 +119,7 @@ func main() { --- -## Two-Layer Architecture +## Architecture Overview The framework is organized into three logical layers: @@ -864,6 +865,43 @@ The ReAct state machine runs: **Input → Model.Generate → ParseAction → (An --- +## Event Sourcing & Replay + +The harness framework provides a **fourth layer** for event-driven agent introspection: append-only event logging, deterministic replay, and live metrics collection. All Layer 4 components integrate via the existing `CallbackManager`, requiring zero changes to Layers 1–3. + +### Event Sourcing + +An **append-only event log** records every granular action during agent execution as an immutable event — tool calls, state transitions, memory writes, approvals, LLM invocations, and checkpoint operations. Each event carries a monotonic logical clock, causal parent references, and a structured payload. This replaces a checkpoint-only approach with a full audit log that supports deterministic replay, forking, and postmortem analysis. + +**Three event store backends** are available: + +| Backend | Path | Use Case | +|---------|------|----------| +| `MemoryEventStore` | `events/memory.go` | In-memory, for testing/single-instance | +| `LocalFileEventStore` | `events/localfile.go` | File-based with segment rotation (by time or size) | +| `NATSEventStore` | `events/nats.go` | Production distributed via NATS JetStream | + +### Replay Engine + +The `ReplayEngine` replays a trace from the event log **deterministically**, supporting: + +- **Model substitution** — replay with a different LLM while keeping tool results frozen +- **Tool result injection** — replace recorded tool outputs with live execution or synthetic data +- **Fork** — branch a new execution from any point in the event log +- **Diff** — compare two execution traces to detect regression or behavioral changes + +### Observability Metrics + +Automated metrics collection covers: tool success rate, approval latency, retry rate, checkpoint restore success, memory hit quality, cost per completed task, and fork replay pass rate. Metrics export to Prometheus. + +### Evaluation Loop + +A production trace can be automatically converted into a **regression dataset**. The `RunReplayEval` function replays each case with multiple model/strategy combinations, comparing results and raising regression alerts. + +> **Detailed design, type definitions, and source-level examples** are documented in [harness.md](harness.md). + +--- + ## Project Structure ``` @@ -901,7 +939,7 @@ harness-go/ │ ├── instruction.go # Instruction management │ │ │ ├── backend/ # Filesystem backend abstraction -│ ├── evals/ # Eval framework (LLM-as-judge, scorers) +│ ├── evals/ # Eval framework (LLM-as-judge, scorers, replay-based eval) │ ├── internal/ # Internal helpers (default system prompt) │ ├── middlewares/ # 10 middleware implementations │ │ ├── subagent/ # SubAgentMiddleware (LLM-driven delegation) @@ -958,6 +996,29 @@ harness-go/ │ ├── viemu/ # Visual emulation │ └── visualization/ # DOT graph output │ +├── events/ # Event Sourcing (append-only event log) +│ ├── event.go # Event, EventID, EventType, typed payloads +│ ├── recorder.go # EventRecorder — GraphCallback → Event +│ ├── clock.go # LogicalClock (monotonic uint64) +│ ├── memory.go # MemoryEventStore +│ ├── localfile.go # LocalFileEventStore +│ └── nats.go # NATSEventStore +│ +├── replay/ # Replay Engine +│ ├── replay.go # ReplayEngine — deterministic replay +│ ├── fork.go # Fork — branch from any event +│ ├── diff.go # Diff — compare two execution traces +│ └── injector.go # ModelOverride / ToolOverride strategies +│ +├── metrics/ # Observability & Metrics +│ ├── metrics.go # MetricsCollector, autoMetricCollector +│ ├── aggregator.go # MetricsAggregator, MetricsWindow +│ └── exporter.go # PrometheusExporter +│ +├── graphengine/ # Graph Engine (Layer 1) +│ ├── dataset.go # EventLog → 回归数据集转换 +│ └── replay_eval.go # Replay-based evaluation +│ ├── prebuilt/ # Prebuilt ReAct agent + node factories │ ├── prebuilt.go # ReAct agent state machine │ ├── tool_node.go # ToolNode factory @@ -968,6 +1029,7 @@ harness-go/ ├── server/ # HTTP server *(removed in internal copy)* ├── telemetry/ # OpenTelemetry integration *(removed in internal copy)* │ +├── harness.md # Event Sourcing & Replay design document ├── harness.go # Top-level re-exports and init() ├── harness_test.go # Integration tests ├── Makefile # Build, test, lint targets diff --git a/internal/harness/core/recorder_model.go b/internal/harness/core/recorder_model.go new file mode 100644 index 000000000..cbdc21120 --- /dev/null +++ b/internal/harness/core/recorder_model.go @@ -0,0 +1,99 @@ +package core + +import ( + "context" + "time" + + "ragflow/internal/harness/core/schema" + "ragflow/internal/harness/events" +) + +// ---- Model Wrapper: records LLM calls via EventRecorder from context ---- + +// eventRecorderModelWrapper wraps a Model and records each invocation to the +// EventRecorder found in the context (via events.RecorderFromContext). +type eventRecorderModelWrapper[M MessageType] struct { + inner Model[M] +} + +func wrapModelWithEventRecorder[M MessageType](inner Model[M]) Model[M] { + return &eventRecorderModelWrapper[M]{inner: inner} +} + +func (w *eventRecorderModelWrapper[M]) Generate(ctx context.Context, msgs []M, opts ...ModelOption) (M, error) { + start := time.Now() + resp, err := w.inner.Generate(ctx, msgs, opts...) + durMs := time.Since(start).Milliseconds() + rec := events.RecorderFromContext(ctx) + if rec != nil && err == nil { + var msgsAny []any + for _, m := range msgs { + msgsAny = append(msgsAny, any(m)) + } + // We record the model as "unknown" when the name isn't accessible here. + // The agent sets model name via BindTools / config; that info can be + // added by providing it through the context in a future iteration. + rec.RecordModelCall(ctx, "unknown", "", msgsAny, contentOf(resp), events.TokenUsage{}, durMs, 0) + } + return resp, err +} + +func (w *eventRecorderModelWrapper[M]) Stream(ctx context.Context, msgs []M, opts ...ModelOption) (*schema.StreamReader[M], error) { + return w.inner.Stream(ctx, msgs, opts...) +} + +func (w *eventRecorderModelWrapper[M]) BindTools(tools []*schema.ToolInfo) error { + return w.inner.BindTools(tools) +} + +// ---- Handler that injects the wrapper via TypedReActMiddleware.WrapModel ---- + +type eventRecorderModelHandler[M MessageType] struct{} + +// NewEventRecorderModelWrapper creates a middleware handler that wraps the model +// to record LLM invocations to the EventRecorder stored in context. +// Usage: +// +// recorder := events.NewEventRecorder(store) +// ctx := events.ContextWithRecorder(ctx, recorder) +// cfg := &ReActConfig[*schema.Message]{ +// Model: model, +// Handlers: []TypedReActMiddleware[*schema.Message]{ +// NewEventRecorderModelWrapper[*schema.Message](), +// }, +// } +func NewEventRecorderModelWrapper[M MessageType]() *eventRecorderModelHandler[M] { + return &eventRecorderModelHandler[M]{} +} + +func (h *eventRecorderModelHandler[M]) WrapModel(ctx context.Context, m Model[M], mc *TypedModelContext[M]) (Model[M], error) { + rec := events.RecorderFromContext(ctx) + if rec == nil { + return m, nil // no recorder in context — pass through + } + return wrapModelWithEventRecorder(m), nil +} + +func (h *eventRecorderModelHandler[M]) BeforeAgent(ctx context.Context, rc *ReActAgentContext) (context.Context, *ReActAgentContext, error) { + return ctx, rc, nil +} +func (h *eventRecorderModelHandler[M]) AfterAgent(ctx context.Context, state *TypedReActAgentState[M]) (context.Context, error) { + return ctx, nil +} +func (h *eventRecorderModelHandler[M]) BeforeModelRewrite(ctx context.Context, st *TypedReActAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedReActAgentState[M], error) { + return ctx, st, nil +} +func (h *eventRecorderModelHandler[M]) AfterModelRewrite(ctx context.Context, st *TypedReActAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedReActAgentState[M], error) { + return ctx, st, nil +} + +// contentOf extracts the text content from a response message. +func contentOf[M MessageType](resp M) string { + if msg, ok := any(resp).(*schema.Message); ok && msg != nil { + return msg.Content + } + if am, ok := any(resp).(*schema.AgenticMessage); ok && am != nil { + return am.Content + } + return "" +} diff --git a/internal/harness/core/recorder_tool.go b/internal/harness/core/recorder_tool.go new file mode 100644 index 000000000..568eb9c98 --- /dev/null +++ b/internal/harness/core/recorder_tool.go @@ -0,0 +1,59 @@ +package core + +import ( + "context" + "encoding/json" + "time" + + "ragflow/internal/harness/core/schema" + "ragflow/internal/harness/events" +) + +// NewEventRecorderToolMiddleware creates a ToolInvokeMiddleware that records +// every tool invocation to the EventRecorder found in the context. +// +// Usage: +// +// cfg := &ReActConfig[*schema.Message]{ +// ToolsConfig: &ToolsNodeConfig{ +// ToolInvokeMiddlewares: []ToolInvokeMiddleware{ +// NewEventRecorderToolMiddleware(), +// }, +// }, +// } +// ctx = events.ContextWithRecorder(ctx, recorder) +func NewEventRecorderToolMiddleware() ToolInvokeMiddleware { + return func(next InvokeTool) InvokeTool { + return func(ctx context.Context, ictx *ToolInvocationContext) (*schema.ToolResult, error) { + rec := events.RecorderFromContext(ctx) + + start := time.Now() + result, err := next(ctx, ictx) + durMs := time.Since(start).Milliseconds() + + if rec == nil { + return result, err + } + + // Extract tool arguments as map (parse from JSON string). + var args map[string]any + if ictx.Arguments != nil && ictx.Arguments.Arguments != "" { + json.Unmarshal([]byte(ictx.Arguments.Arguments), &args) + } + + errStr := "" + retryCount := 0 + if ictx.RetryConfig != nil { + retryCount = ictx.RetryConfig.MaxAttempts + } + if err != nil { + errStr = err.Error() + } else if result != nil && result.Error != "" { + errStr = result.Error + } + + rec.RecordToolCall(ctx, ictx.Name, args, result, durMs, retryCount, errStr) + return result, err + } + } +} diff --git a/internal/harness/events/clock.go b/internal/harness/events/clock.go new file mode 100644 index 000000000..0041508e8 --- /dev/null +++ b/internal/harness/events/clock.go @@ -0,0 +1,32 @@ +package events + +import ( + "sync/atomic" +) + +// LogicalClock is a monotonic logical clock that provides a total order +// for events within a process. It is safe for concurrent use. +type LogicalClock struct { + value atomic.Uint64 +} + +// NewLogicalClock creates a new LogicalClock starting at 0. +func NewLogicalClock() *LogicalClock { + return &LogicalClock{} +} + +// Tick atomically increments the clock and returns the new value. +func (c *LogicalClock) Tick() uint64 { + return c.value.Add(1) +} + +// Now returns the current clock value without incrementing. +func (c *LogicalClock) Now() uint64 { + return c.value.Load() +} + +// Reset resets the clock to zero. Use with caution — this should only +// be done when starting a completely independent execution context. +func (c *LogicalClock) Reset() { + c.value.Store(0) +} diff --git a/internal/harness/events/event.go b/internal/harness/events/event.go new file mode 100644 index 000000000..b26444b89 --- /dev/null +++ b/internal/harness/events/event.go @@ -0,0 +1,227 @@ +// Package events provides append-only event sourcing for agent execution. +// +// Every tool call, state transition, memory write, approval, LLM invocation, +// and checkpoint operation is recorded as an immutable Event. Events are +// causally ordered via a monotonic logical clock, enabling deterministic +// replay, fork/diff, and postmortem analysis. +package events + +import ( + "crypto/sha256" + "encoding/json" + "fmt" + "time" +) + +// EventID is a globally unique event identifier (UUID v7, time-ordered). +type EventID string + +// EventType enumerates every recordable action during agent execution. +type EventType string + +const ( + // Graph execution lifecycle. + EventGraphStart EventType = "graph.start" + EventGraphEnd EventType = "graph.end" + EventStepStart EventType = "step.start" + EventStepEnd EventType = "step.end" + + // Node execution. + EventNodeStart EventType = "node.start" + EventNodeEnd EventType = "node.end" + + // State transitions. + EventStateRead EventType = "state.read" + EventStateWrite EventType = "state.write" + + // Tool calls. + EventToolCallStart EventType = "tool.call.start" + EventToolCallResult EventType = "tool.call.result" + EventToolCallError EventType = "tool.call.error" + + // LLM invocations. + EventLLMCallStart EventType = "llm.call.start" + EventLLMCallChunk EventType = "llm.call.chunk" + EventLLMCallEnd EventType = "llm.call.end" + + // Memory operations. + EventMemoryRead EventType = "memory.read" + EventMemoryWrite EventType = "memory.write" + + // Human-in-the-loop. + EventApprovalRequest EventType = "approval.request" + EventApprovalGranted EventType = "approval.granted" + EventApprovalDenied EventType = "approval.denied" + + // Checkpoint. + EventCheckpointCreated EventType = "checkpoint.created" + EventCheckpointRestored EventType = "checkpoint.restored" + + // Interrupt / Resume. + EventInterrupt EventType = "interrupt" + EventResume EventType = "resume" + + // Error & retry. + EventError EventType = "error" + EventRetry EventType = "retry" + + // Fork — branch from an existing event. + EventFork EventType = "fork" + + // Sub-agent execution. + EventSubAgentCallStart EventType = "subagent.call.start" + EventSubAgentCallEnd EventType = "subagent.call.end" + + // Session / Transfer. + EventSessionValueSet EventType = "session.value.set" + EventSessionTransfer EventType = "session.transfer" +) + +// Event is an immutable append-only event. +type Event struct { + // ID is the globally unique event identifier. + ID EventID `json:"id"` + // Type describes what happened. + Type EventType `json:"type"` + // Timestamp is the wall-clock time when this event was recorded. + Timestamp time.Time `json:"timestamp"` + // Clock is the monotonic logical clock value (global total order). + Clock uint64 `json:"clock"` + + // TraceID identifies one complete execution trace. + TraceID string `json:"trace_id"` + // ParentID is the immediate predecessor event in the same trace. + ParentID EventID `json:"parent_id,omitempty"` + // CausedBy lists predecessor events (multiple for fork/join scenarios). + CausedBy []EventID `json:"caused_by,omitempty"` + + // ThreadID identifies the execution thread. + ThreadID string `json:"thread_id,omitempty"` + // Step is the Pregel superstep number. + Step int `json:"step,omitempty"` + // Node is the graph node name. + Node string `json:"node,omitempty"` + // TaskID identifies the execution task. + TaskID string `json:"task_id,omitempty"` + + // Payload is the type-specific event payload (JSON). + Payload json.RawMessage `json:"payload,omitempty"` + // Metadata holds arbitrary key-value metadata. + Metadata map[string]any `json:"metadata,omitempty"` + + // Deterministic is false when the event involves non-deterministic + // operations (LLM output, random, wall-clock time). + Deterministic bool `json:"deterministic"` + // Hash is the SHA-256 of Payload+Metadata (for integrity verification). + Hash string `json:"hash,omitempty"` +} + +// NewEvent creates a new Event with auto-generated ID and current timestamp. +func NewEvent(typ EventType, clock uint64) *Event { + return &Event{ + ID: EventID(fmt.Sprintf("evt-%d-%x", clock, time.Now().UnixNano())), + Type: typ, + Timestamp: time.Now(), + Clock: clock, + Metadata: make(map[string]any), + } +} + +// computeHash computes the SHA-256 hash of the event payload and metadata. +func (e *Event) computeHash() string { + h := sha256.New() + if e.Payload != nil { + h.Write(e.Payload) + } + if e.Metadata != nil { + meta, _ := json.Marshal(e.Metadata) + h.Write(meta) + } + return fmt.Sprintf("%x", h.Sum(nil)) +} + +// Seal finalises the event by computing its hash and marking it immutable. +func (e *Event) Seal() { + e.Hash = e.computeHash() +} + +// ---- typed payloads ---- + +// ToolCallPayload is the payload for tool call events. +type ToolCallPayload struct { + ToolName string `json:"tool_name"` + Arguments map[string]any `json:"arguments,omitempty"` + Result any `json:"result,omitempty"` + DurationMs int64 `json:"duration_ms,omitempty"` + Error string `json:"error,omitempty"` + RetryCount int `json:"retry_count,omitempty"` +} + +// LLMCallPayload is the payload for LLM invocation events. +type LLMCallPayload struct { + Model string `json:"model"` + Provider string `json:"provider,omitempty"` + Messages []any `json:"messages,omitempty"` + Tokens TokenUsage `json:"tokens,omitempty"` + Content string `json:"content,omitempty"` + Chunks int `json:"chunks,omitempty"` + DurationMs int64 `json:"duration_ms,omitempty"` + Cost float64 `json:"cost,omitempty"` +} + +// TokenUsage tracks token consumption for an LLM call. +type TokenUsage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +// StateTransitionPayload is the payload for state change events. +type StateTransitionPayload struct { + Channel string `json:"channel"` + OldValue any `json:"old_value,omitempty"` + NewValue any `json:"new_value"` + Reducer string `json:"reducer,omitempty"` +} + +// MemoryWritePayload is the payload for memory operation events. +type MemoryWritePayload struct { + Store string `json:"store"` + Operation string `json:"operation"` + Key string `json:"key,omitempty"` + Value any `json:"value,omitempty"` + Score float64 `json:"score,omitempty"` +} + +// ApprovalPayload is the payload for approval events. +type ApprovalPayload struct { + RequestID string `json:"request_id"` + Action string `json:"action"` + Context any `json:"context,omitempty"` + Decision string `json:"decision,omitempty"` + LatencyMs int64 `json:"latency_ms,omitempty"` +} + +// SubAgentCallPayload is the payload for sub-agent call events. +type SubAgentCallPayload struct { + SubAgentName string `json:"sub_agent_name"` + Input any `json:"input,omitempty"` + Output any `json:"output,omitempty"` + Depth int `json:"depth,omitempty"` + DurationMs int64 `json:"duration_ms,omitempty"` + Error string `json:"error,omitempty"` +} + +// SessionValuePayload is the payload for session value events. +type SessionValuePayload struct { + Key string `json:"key"` + Value any `json:"value,omitempty"` +} + +// SessionTransferPayload is the payload for session transfer events. +type SessionTransferPayload struct { + FromAgent string `json:"from_agent"` + ToAgent string `json:"to_agent"` + Reason string `json:"reason,omitempty"` + Input any `json:"input,omitempty"` +} diff --git a/internal/harness/events/events_test.go b/internal/harness/events/events_test.go new file mode 100644 index 000000000..70ae19fa0 --- /dev/null +++ b/internal/harness/events/events_test.go @@ -0,0 +1,520 @@ +package events + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "testing" + "time" +) + +func TestMemoryEventStore_AppendAndStream(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + + ev1 := &Event{ID: "e1", Type: EventGraphStart, Clock: 1, Timestamp: time.Now(), TraceID: "trace-1"} + ev2 := &Event{ID: "e2", Type: EventGraphEnd, Clock: 2, Timestamp: time.Now(), TraceID: "trace-1"} + + if err := s.Append(ctx, ev1, ev2); err != nil { + t.Fatal(err) + } + + // Stream all. + iter := s.Stream(ctx, EventFilter{}) + var got []*Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + got = append(got, ev) + } + if len(got) != 2 { + t.Fatalf("expected 2 events, got %d", len(got)) + } + if got[0].ID != "e1" || got[1].ID != "e2" { + t.Fatal("events out of order") + } +} + +func TestMemoryEventStore_Get(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + s.Append(ctx, &Event{ID: "find-me", Type: EventNodeStart, Clock: 1, TraceID: "t"}) + + ev, err := s.Get(ctx, "find-me") + if err != nil { + t.Fatal(err) + } + if ev == nil { + t.Fatal("expected event, got nil") + } + + ev, err = s.Get(ctx, "nonexistent") + if err != nil { + t.Fatal(err) + } + if ev != nil { + t.Fatal("expected nil for nonexistent") + } +} + +func TestMemoryEventStore_Range(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + for i := 1; i <= 10; i++ { + s.Append(ctx, &Event{ID: EventID(fmt.Sprintf("e%d", i)), Clock: uint64(i), Type: EventStepStart, Timestamp: time.Now(), TraceID: "t"}) + } + + events, err := s.Range(ctx, 3, 7, EventFilter{}) + if err != nil { + t.Fatal(err) + } + if len(events) != 5 { + t.Fatalf("expected 5 events, got %d", len(events)) + } + if events[0].Clock != 3 { + t.Fatalf("expected clock 3 first, got %d", events[0].Clock) + } +} + +func TestMemoryEventStore_Seek(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + for i := 1; i <= 5; i++ { + s.Append(ctx, &Event{ID: EventID(fmt.Sprintf("e%d", i)), Clock: uint64(i), Type: EventStepStart, TraceID: "t"}) + } + + iter, err := s.Seek(ctx, 3) + if err != nil { + t.Fatal(err) + } + ev, ok := iter.Next(ctx) + if !ok || ev.Clock != 3 { + t.Fatalf("expected clock 3, got %v", ev) + } +} + +func TestMemoryEventStore_Length(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + n, _ := s.Length(ctx) + if n != 0 { + t.Fatalf("expected 0, got %d", n) + } + s.Append(ctx, &Event{ID: "e1", Clock: 1, Type: EventGraphStart, TraceID: "t"}) + n, _ = s.Length(ctx) + if n != 1 { + t.Fatalf("expected 1, got %d", n) + } +} + +func TestMemoryEventStore_Subscribe(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + s := NewMemoryEventStore() + + ch, err := s.Subscribe(ctx, EventFilter{}) + if err != nil { + t.Fatal(err) + } + + s.Append(ctx, &Event{ID: "live", Clock: 1, Type: EventNodeStart, TraceID: "t"}) + + select { + case ev := <-ch: + if ev.ID != "live" { + t.Fatalf("expected live event, got %s", ev.ID) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for subscribed event") + } +} + +func TestMemoryEventStore_Filter(t *testing.T) { + f := EventFilter{TraceID: "trace-1", Types: []EventType{EventNodeStart}} + if !f.Matches(&Event{TraceID: "trace-1", Type: EventNodeStart}) { + t.Fatal("should match") + } + if f.Matches(&Event{TraceID: "trace-2", Type: EventNodeStart}) { + t.Fatal("should not match different trace") + } + if f.Matches(&Event{TraceID: "trace-1", Type: EventNodeEnd}) { + t.Fatal("should not match different type") + } +} + +func TestLocalFileEventStore_GC_RetainsSurvivingEvents(t *testing.T) { + dir := t.TempDir() + ctx := context.Background() + + s, err := NewLocalFileEventStore(dir) + if err != nil { + t.Fatal(err) + } + + // Write one old event (outside retention) and one recent event. + oldEv := &Event{ID: "old", Clock: 1, Type: EventGraphStart, TraceID: "gc-test", Timestamp: time.Now().Add(-2 * time.Hour)} + oldEv.Seal() + recentEv := &Event{ID: "recent", Clock: 2, Type: EventGraphEnd, TraceID: "gc-test", Timestamp: time.Now()} + recentEv.Seal() + if err := s.Append(ctx, oldEv, recentEv); err != nil { + t.Fatal(err) + } + + // GC with 1-hour retention. + if err := s.GC(ctx, time.Hour); err != nil { + t.Fatal(err) + } + + // In-memory: old should be gone, recent should remain. + ev, _ := s.Get(ctx, "old") + if ev != nil { + t.Error("old event should be removed from cache") + } + ev, _ = s.Get(ctx, "recent") + if ev == nil { + t.Fatal("recent event should survive in cache") + } + + // Reopen the store from disk and verify retained events survived. + s2, err := NewLocalFileEventStore(dir) + if err != nil { + t.Fatal(err) + } + ev, _ = s2.Get(ctx, "recent") + if ev == nil { + t.Fatal("recent event should survive GC on disk") + } + ev, _ = s2.Get(ctx, "old") + if ev != nil { + t.Error("old event should be absent from disk after GC") + } +} + +func TestMemoryEventStore_GC(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + old := &Event{ID: "old", Clock: 1, Type: EventGraphStart, Timestamp: time.Now().Add(-2 * time.Hour), TraceID: "t"} + s.Append(ctx, old) + s.Append(ctx, &Event{ID: "new", Clock: 2, Type: EventGraphEnd, Timestamp: time.Now(), TraceID: "t"}) + + s.GC(ctx, time.Hour) + + ev, _ := s.Get(ctx, "old") + if ev != nil { + t.Fatal("old event should have been GC'd") + } + ev, _ = s.Get(ctx, "new") + if ev == nil { + t.Fatal("new event should survive GC") + } +} + +func TestLocalFileEventStore_AppendAndReopen(t *testing.T) { + dir := t.TempDir() + ctx := context.Background() + + s, err := NewLocalFileEventStore(dir) + if err != nil { + t.Fatal(err) + } + + for i := 1; i <= 5; i++ { + s.Append(ctx, &Event{ID: EventID(fmt.Sprintf("f%d", i)), Clock: uint64(i), Type: EventStepStart, TraceID: "reopen", Timestamp: time.Now()}) + } + + // Reopen — should load existing events. + s2, err := NewLocalFileEventStore(dir) + if err != nil { + t.Fatal(err) + } + + iter := s2.Stream(ctx, EventFilter{TraceID: "reopen"}) + count := 0 + for { + _, ok := iter.Next(ctx) + if !ok { + break + } + count++ + } + if count != 5 { + t.Fatalf("expected 5 events after reopen, got %d", count) + } +} + +func TestLocalFileEventStore_SegmentRotation(t *testing.T) { + dir := t.TempDir() + ctx := context.Background() + + // Use small maxSize to force rotation. + s := &LocalFileEventStore{ + dir: dir, + segment: 0, + maxSize: 100, // 100 bytes per segment + cached: make([]*Event, 0), + } + + // Write events large enough to trigger rotation. + for i := 0; i < 20; i++ { + ev := &Event{ + ID: EventID(fmt.Sprintf("seg-%d", i)), + Clock: uint64(i + 1), + Type: EventGraphStart, + TraceID: "rotation-test", + Timestamp: time.Now(), + Metadata: map[string]any{ + "data": fmt.Sprintf("x=%s", string(make([]byte, 50))), + }, + } + ev.Seal() + s.Append(ctx, ev) + } + + // Verify all events survive rotation. + s.mu.RLock() + count := len(s.cached) + s.mu.RUnlock() + if count != 20 { + t.Fatalf("expected 20 events, got %d", count) + } + + // Check multiple segment files exist. + entries, _ := os.ReadDir(dir) + segCount := 0 + for _, entry := range entries { + if !entry.IsDir() && filepath.Ext(entry.Name()) == ".jsonl" { + segCount++ + } + } + if segCount < 2 { + t.Fatalf("expected at least 2 segment files, got %d", segCount) + } + + // Reopen and verify. + s2, err := NewLocalFileEventStore(dir) + if err != nil { + t.Fatal(err) + } + iter := s2.Stream(ctx, EventFilter{}) + var all []*Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + all = append(all, ev) + } + if len(all) != 20 { + t.Fatalf("expected 20 events after reopen, got %d", len(all)) + } + // Verify event order by Clock. + sort.Slice(all, func(i, j int) bool { return all[i].Clock < all[j].Clock }) + for i, ev := range all { + if ev.Clock != uint64(i+1) { + t.Fatalf("event %d: expected clock %d, got %d", i, i+1, ev.Clock) + } + } +} + +func TestEvent_Seal(t *testing.T) { + ev := NewEvent(EventNodeStart, 1) + ev.Payload, _ = json.Marshal(StateTransitionPayload{Channel: "msg"}) + ev.Seal() + if ev.Hash == "" { + t.Fatal("hash should be set after Seal") + } + // Same payload should produce same hash. + ev2 := NewEvent(EventNodeStart, 1) + ev2.Payload = ev.Payload + ev2.Seal() + if ev.Hash != ev2.Hash { + t.Fatal("identical events should have identical hashes") + } +} + +func TestMemoryEventStore_Snapshot(t *testing.T) { + ctx := context.Background() + s := NewMemoryEventStore() + s.Append(ctx, &Event{ID: "s1", Clock: 1, Type: EventGraphStart, TraceID: "snap-trace", Timestamp: time.Now()}) + s.Append(ctx, &Event{ID: "s2", Clock: 2, Type: EventGraphEnd, TraceID: "snap-trace", Timestamp: time.Now()}) + + snap, err := s.CreateSnapshot(ctx, "snap-trace") + if err != nil { + t.Fatal(err) + } + if snap.TraceID != "snap-trace" { + t.Fatalf("expected snap-trace, got %s", snap.TraceID) + } + if snap.Data == nil { + t.Fatal("snapshot data should not be nil") + } +} + +func TestLogicalClock(t *testing.T) { + c := NewLogicalClock() + if c.Now() != 0 { + t.Fatalf("expected 0, got %d", c.Now()) + } + v1 := c.Tick() + if v1 != 1 { + t.Fatalf("expected 1, got %d", v1) + } + v2 := c.Tick() + if v2 != 2 { + t.Fatalf("expected 2, got %d", v2) + } + c.Reset() + if c.Now() != 0 { + t.Fatalf("expected 0 after reset, got %d", c.Now()) + } +} + +func TestEventRecorder_GraphCallback(t *testing.T) { + ctx := context.Background() + store := NewMemoryEventStore() + r := NewEventRecorder(store, WithTraceID("rec-trace"), WithThreadID("rec-thread")) + + // Dispatch all GraphCallback methods. + r.OnRunStart(ctx, "test-graph", "rec-thread") + r.OnStepStart(ctx, 0, 3) + r.OnNodeStart(ctx, "node-a", 0) + r.OnNodeEnd(ctx, "node-a", 0, "output", nil) + r.OnCheckpointSave(ctx, "rec-thread", "cp-1", 0) + r.OnCheckpointLoad(ctx, "rec-thread", "cp-1", 0) + r.OnInterrupt(ctx, []string{"node-a"}, 0) + r.OnResume(ctx, "rec-thread") + r.OnStepEnd(ctx, 0, nil) + r.OnRunEnd(ctx, "test-graph", "rec-thread", nil) + + // Verify all events recorded. + iter := store.Stream(ctx, EventFilter{TraceID: "rec-trace"}) + var events []*Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + events = append(events, ev) + } + if len(events) != 10 { + t.Fatalf("expected 10 events, got %d: %v", len(events), collectTypes(events)) + } +} + +func TestEventRecorder_ModelAndToolCalls(t *testing.T) { + ctx := context.Background() + store := NewMemoryEventStore() + r := NewEventRecorder(store, WithTraceID("mt-trace")) + + r.RecordModelCall(ctx, "gpt-4", "openai", []any{"hello"}, "world", TokenUsage{PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30}, 500, 0.002) + r.RecordModelCall(ctx, "gpt-4", "openai", []any{"hello2"}, "world2", TokenUsage{PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30}, 500, 0.002) + r.RecordLLMChunk(ctx, "gpt-4", "wor") + r.RecordLLMChunk(ctx, "gpt-4", "ld") + r.RecordToolCall(ctx, "get_weather", map[string]any{"city": "NYC"}, "sunny", 200, 0, "") + r.RecordToolCall(ctx, "fail_tool", map[string]any{}, nil, 100, 2, "timeout") + + iter := store.Stream(ctx, EventFilter{TraceID: "mt-trace"}) + count := 0 + for { + _, ok := iter.Next(ctx) + if !ok { + break + } + count++ + } + expected := 10 // 2 model calls × 2 + 2 chunks + 2 tool calls × 2 + if count != expected { + t.Fatalf("expected %d events, got %d", expected, count) + } +} + +func TestEventRecorder_FineGrained(t *testing.T) { + ctx := context.Background() + store := NewMemoryEventStore() + r := NewEventRecorder(store, WithTraceID("fine")) + + r.RecordStateWrite(ctx, "messages", nil, []any{"hello"}, "append") + r.RecordMemoryWrite(ctx, "vector", "insert", "doc1", "content", 0.95) + r.RecordMemoryRead(ctx, "vector", "doc1", 0.92) + r.RecordApproval(ctx, "req-1", "execute_tool", map[string]any{"tool": "search"}, "granted", 3000) + r.RecordError(ctx, "connection refused") + r.RecordRetry(ctx, "attempt 2/3") + r.RecordSubAgentCall(ctx, "researcher", "query1", "result1", 1, 500, "") + r.RecordSessionValue(ctx, "mode", "fast") + r.RecordSessionTransfer(ctx, "planner", "executor", "plan ready", nil) + + iter := store.Stream(ctx, EventFilter{TraceID: "fine"}) + count := 0 + for { + _, ok := iter.Next(ctx) + if !ok { + break + } + count++ + } + expected := 6 + 4 // original 6 + sub-start/sub-end + session-value + session-transfer + if count != expected { + t.Fatalf("expected %d events, got %d", expected, count) + } +} + +func TestEventRecorder_ContextHelpers(t *testing.T) { + ctx := context.Background() + store := NewMemoryEventStore() + r := NewEventRecorder(store, WithTraceID("ctx-test")) + + // Store recorder in context. + ctx = ContextWithRecorder(ctx, r) + + // Retrieve and verify. + got := RecorderFromContext(ctx) + if got == nil { + t.Fatal("expected non-nil recorder from context") + } + + // Context without recorder should return nil. + ctx2 := context.Background() + if RecorderFromContext(ctx2) != nil { + t.Fatal("expected nil from context without recorder") + } +} + +func TestSubAgentAndSessionEvents(t *testing.T) { + ctx := context.Background() + store := NewMemoryEventStore() + r := NewEventRecorder(store, WithTraceID("sub-session")) + + r.RecordSubAgentCall(ctx, "researcher", "query1", "result1", 1, 500, "") + r.RecordSubAgentCall(ctx, "researcher", "query2", "", 2, 0, "timeout") + r.RecordSessionValue(ctx, "mode", "fast") + r.RecordSessionValue(ctx, "count", 42) + r.RecordSessionTransfer(ctx, "a", "b", "reason", "input") + + iter := store.Stream(ctx, EventFilter{TraceID: "sub-session"}) + var evts []*Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + evts = append(evts, ev) + } + if len(evts) != 7 { // 2 sub-agent calls × 2 events + 2 values + 1 transfer + t.Fatalf("expected 7 events, got %d", len(evts)) + } +} + +// ---- helpers ---- + +func collectTypes(events []*Event) []EventType { + var types []EventType + for _, ev := range events { + types = append(types, ev.Type) + } + return types +} diff --git a/internal/harness/events/localfile.go b/internal/harness/events/localfile.go new file mode 100644 index 000000000..524dfde28 --- /dev/null +++ b/internal/harness/events/localfile.go @@ -0,0 +1,312 @@ +package events + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + defaultMaxSegmentSize int64 = 64 * 1024 * 1024 // 64 MB + segmentFilePattern = "events_*.jsonl" +) + +// LocalFileEventStore persists events to JSONL files with automatic +// segment rotation. Suitable for single-instance durable storage. +type LocalFileEventStore struct { + dir string + segment int // current write segment number + maxSize int64 // max bytes per segment before rotation + mu sync.RWMutex + cached []*Event // in-memory cache for current segment + clock atomic.Uint64 +} + +// NewLocalFileEventStore creates or opens an event store at the given directory. +// Existing segment files are loaded into memory on startup. +func NewLocalFileEventStore(dir string) (*LocalFileEventStore, error) { + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, fmt.Errorf("create events dir: %w", err) + } + + s := &LocalFileEventStore{ + dir: dir, + segment: 0, + maxSize: defaultMaxSegmentSize, + cached: make([]*Event, 0), + } + + // Load existing segments to find the highest segment number. + if err := s.loadExisting(); err != nil { + return nil, fmt.Errorf("load existing segments: %w", err) + } + + return s, nil +} + +// loadExisting scans the directory for existing segment files and loads them. +func (s *LocalFileEventStore) loadExisting() error { + entries, err := os.ReadDir(s.dir) + if err != nil { + return err + } + + var segmentFiles []string + for _, entry := range entries { + if !entry.IsDir() && strings.HasPrefix(entry.Name(), "events_") && strings.HasSuffix(entry.Name(), ".jsonl") { + segmentFiles = append(segmentFiles, entry.Name()) + } + } + + // Sort by name (lexicographic works for timestamp-based names). + sort.Strings(segmentFiles) + + allEvents := make([]*Event, 0) + var maxClock uint64 + + for _, fname := range segmentFiles { + fpath := filepath.Join(s.dir, fname) + events, err := readSegmentFile(fpath) + if err != nil { + return fmt.Errorf("read segment %s: %w", fname, err) + } + for _, ev := range events { + if ev.Clock > maxClock { + maxClock = ev.Clock + } + } + allEvents = append(allEvents, events...) + } + + s.cached = allEvents + if maxClock > 0 { + s.clock.Store(maxClock) + } + return nil +} + +// readSegmentFile reads all events from a JSONL file. +func readSegmentFile(path string) ([]*Event, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + var events []*Event + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 1024*1024), 1024*1024) // 1 MB line buffer + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var ev Event + if err := json.Unmarshal([]byte(line), &ev); err != nil { + return nil, fmt.Errorf("unmarshal event: %w", err) + } + events = append(events, &ev) + } + return events, scanner.Err() +} + +// Append implements EventLog. +func (s *LocalFileEventStore) Append(ctx context.Context, events ...*Event) error { + s.mu.Lock() + defer s.mu.Unlock() + + for _, ev := range events { + if ev.Clock == 0 { + ev.Clock = s.clock.Add(1) + } + ev.Seal() + s.cached = append(s.cached, ev) + + // Check if we need to rotate segment. + if err := s.appendToFileLocked(ev); err != nil { + return err + } + } + return nil +} + +// appendToFileLocked writes a single event to the current segment file. +// Must be called with s.mu held. +func (s *LocalFileEventStore) appendToFileLocked(ev *Event) error { + fname := fmt.Sprintf("events_%s_%04d.jsonl", ev.TraceID, s.segment) + fpath := filepath.Join(s.dir, fname) + + // Check segment size and rotate if needed. + if info, err := os.Stat(fpath); err == nil && info.Size() > s.maxSize { + s.segment++ + fname = fmt.Sprintf("events_%s_%04d.jsonl", ev.TraceID, s.segment) + fpath = filepath.Join(s.dir, fname) + } + + f, err := os.OpenFile(fpath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + return fmt.Errorf("open segment file: %w", err) + } + defer f.Close() + + data, err := json.Marshal(ev) + if err != nil { + return fmt.Errorf("marshal event: %w", err) + } + + if _, err := f.Write(data); err != nil { + return err + } + if _, err := f.Write([]byte("\n")); err != nil { + return err + } + return nil +} + +// Stream implements EventLog. +func (s *LocalFileEventStore) Stream(ctx context.Context, filter EventFilter) EventIterator { + s.mu.RLock() + defer s.mu.RUnlock() + + filtered := make([]*Event, 0) + for _, ev := range s.cached { + if filter.Matches(ev) { + filtered = append(filtered, ev) + } + } + return &sliceIterator{events: filtered, pos: 0} +} + +// Get implements EventLog. +func (s *LocalFileEventStore) Get(ctx context.Context, id EventID) (*Event, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + for _, ev := range s.cached { + if ev.ID == id { + return ev, nil + } + } + return nil, nil +} + +// Range implements EventLog. +func (s *LocalFileEventStore) Range(ctx context.Context, from, to uint64, filter EventFilter) ([]*Event, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + result := make([]*Event, 0) + for _, ev := range s.cached { + if ev.Clock < from { + continue + } + if to > 0 && ev.Clock > to { + continue + } + if filter.Matches(ev) { + result = append(result, ev) + } + } + return result, nil +} + +// Seek implements EventLog. +func (s *LocalFileEventStore) Seek(ctx context.Context, clock uint64) (EventIterator, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + pos := 0 + for i, ev := range s.cached { + if ev.Clock >= clock { + pos = i + break + } + _ = i + } + return &sliceIterator{events: s.cached[pos:], pos: 0}, nil +} + +// Length implements EventLog. +func (s *LocalFileEventStore) Length(ctx context.Context) (uint64, error) { + s.mu.RLock() + defer s.mu.RUnlock() + return uint64(len(s.cached)), nil +} + +// CreateSnapshot implements EventStore. +func (s *LocalFileEventStore) CreateSnapshot(ctx context.Context, traceID string) (*Snapshot, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + clock := s.clock.Load() + data, err := json.Marshal(s.cached) + if err != nil { + return nil, fmt.Errorf("marshal snapshot: %w", err) + } + return &Snapshot{ + ID: fmt.Sprintf("snap-%s-%d", traceID, clock), + TraceID: traceID, + Clock: clock, + CreatedAt: time.Now(), + Data: data, + }, nil +} + +// RestoreSnapshot implements EventStore. +func (s *LocalFileEventStore) RestoreSnapshot(ctx context.Context, snapshotID string) (EventIterator, error) { + return s.Seek(ctx, 0) +} + +// Subscribe implements EventStore. +func (s *LocalFileEventStore) Subscribe(ctx context.Context, filter EventFilter) (<-chan *Event, error) { + // LocalFileEventStore does not support real-time subscriptions. + // Use NATSEventStore for distributed scenarios that need Subscribe. + ch := make(chan *Event) + close(ch) + return ch, nil +} + +// GC implements EventStore. +// Retained events are rewritten to fresh segment files; only segments whose +// entire content predates the cutoff are deleted. +func (s *LocalFileEventStore) GC(ctx context.Context, retention time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + + cutoff := time.Now().Add(-retention) + keep := make([]*Event, 0, len(s.cached)) + for _, ev := range s.cached { + if ev.Timestamp.After(cutoff) { + keep = append(keep, ev) + } + } + s.cached = keep + s.segment = 0 + + // Remove all old segment files so the retained events can be rewritten + // into fresh files below. + entries, _ := os.ReadDir(s.dir) + for _, entry := range entries { + if !entry.IsDir() && strings.HasPrefix(entry.Name(), "events_") && strings.HasSuffix(entry.Name(), ".jsonl") { + os.Remove(filepath.Join(s.dir, entry.Name())) + } + } + + // Rewrite retained events into fresh segment files. + // appendToFileLocked is safe to call here — it does not acquire s.mu + // (the "Locked" suffix means the caller must hold it). + for _, ev := range keep { + if err := s.appendToFileLocked(ev); err != nil { + return err + } + } + return nil +} diff --git a/internal/harness/events/memory.go b/internal/harness/events/memory.go new file mode 100644 index 000000000..68e02dcd1 --- /dev/null +++ b/internal/harness/events/memory.go @@ -0,0 +1,210 @@ +package events + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "sync/atomic" + "time" +) + +// MemoryEventStore is an in-memory EventStore implementation. +// Suitable for testing and single-instance development. +// All events are lost on process restart. +type MemoryEventStore struct { + mu sync.RWMutex + events []*Event + byID map[EventID]*Event + clock atomic.Uint64 + subs map[int64]chan *Event + subID atomic.Int64 +} + +// NewMemoryEventStore creates a new empty MemoryEventStore. +func NewMemoryEventStore() *MemoryEventStore { + return &MemoryEventStore{ + events: make([]*Event, 0, 1024), + byID: make(map[EventID]*Event), + subs: make(map[int64]chan *Event), + } +} + +// Append implements EventLog. +func (s *MemoryEventStore) Append(ctx context.Context, events ...*Event) error { + s.mu.Lock() + defer s.mu.Unlock() + + for _, ev := range events { + if ev.Clock == 0 { + ev.Clock = s.clock.Add(1) + } + ev.Seal() + s.events = append(s.events, ev) + s.byID[ev.ID] = ev + + // Dispatch to subscribers. + for id, ch := range s.subs { + select { + case ch <- ev: + default: + // Drop slow subscriber. + } + _ = id + } + } + return nil +} + +// Stream implements EventLog. +func (s *MemoryEventStore) Stream(ctx context.Context, filter EventFilter) EventIterator { + s.mu.RLock() + defer s.mu.RUnlock() + + filtered := make([]*Event, 0) + for _, ev := range s.events { + if filter.Matches(ev) { + filtered = append(filtered, ev) + } + } + return &sliceIterator{events: filtered, pos: 0} +} + +// Get implements EventLog. +func (s *MemoryEventStore) Get(ctx context.Context, id EventID) (*Event, error) { + s.mu.RLock() + defer s.mu.RUnlock() + ev, ok := s.byID[id] + if !ok { + return nil, nil + } + return ev, nil +} + +// Range implements EventLog. +func (s *MemoryEventStore) Range(ctx context.Context, from, to uint64, filter EventFilter) ([]*Event, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + result := make([]*Event, 0) + for _, ev := range s.events { + if ev.Clock < from { + continue + } + if to > 0 && ev.Clock > to { + continue + } + if filter.Matches(ev) { + result = append(result, ev) + } + } + return result, nil +} + +// Seek implements EventLog. +func (s *MemoryEventStore) Seek(ctx context.Context, clock uint64) (EventIterator, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + pos := 0 + for i, ev := range s.events { + if ev.Clock >= clock { + pos = i + break + } + _ = i + } + return &sliceIterator{events: s.events[pos:], pos: 0}, nil +} + +// Length implements EventLog. +func (s *MemoryEventStore) Length(ctx context.Context) (uint64, error) { + s.mu.RLock() + defer s.mu.RUnlock() + return uint64(len(s.events)), nil +} + +// CreateSnapshot implements EventStore. +func (s *MemoryEventStore) CreateSnapshot(ctx context.Context, traceID string) (*Snapshot, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + clock := s.clock.Load() + data, err := json.Marshal(s.events) + if err != nil { + return nil, fmt.Errorf("marshal snapshot: %w", err) + } + return &Snapshot{ + ID: fmt.Sprintf("snap-%s-%d", traceID, clock), + TraceID: traceID, + Clock: clock, + CreatedAt: time.Now(), + Data: data, + }, nil +} + +// RestoreSnapshot implements EventStore. +func (s *MemoryEventStore) RestoreSnapshot(ctx context.Context, snapshotID string) (EventIterator, error) { + // For MemoryEventStore, we simply seek past the snapshot's clock. + // The snapshot data itself is not needed since events are still in memory. + return s.Seek(ctx, 0) +} + +// Subscribe implements EventStore. +func (s *MemoryEventStore) Subscribe(ctx context.Context, filter EventFilter) (<-chan *Event, error) { + ch := make(chan *Event, 256) + id := s.subID.Add(1) + + s.mu.Lock() + s.subs[id] = ch + s.mu.Unlock() + + go func() { + <-ctx.Done() + s.mu.Lock() + delete(s.subs, id) + s.mu.Unlock() + close(ch) + }() + + return ch, nil +} + +// GC implements EventStore. +func (s *MemoryEventStore) GC(ctx context.Context, retention time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + + cutoff := time.Now().Add(-retention) + keep := make([]*Event, 0, len(s.events)) + for _, ev := range s.events { + if ev.Timestamp.After(cutoff) { + keep = append(keep, ev) + } else { + delete(s.byID, ev.ID) + } + } + s.events = keep + return nil +} + +// ---- sliceIterator ---- + +type sliceIterator struct { + events []*Event + pos int +} + +func (it *sliceIterator) Next(_ context.Context) (*Event, bool) { + if it.pos >= len(it.events) { + return nil, false + } + ev := it.events[it.pos] + it.pos++ + return ev, true +} + +func (it *sliceIterator) Close() error { + it.events = nil + return nil +} diff --git a/internal/harness/events/nats.go b/internal/harness/events/nats.go new file mode 100644 index 000000000..a7470e372 --- /dev/null +++ b/internal/harness/events/nats.go @@ -0,0 +1,367 @@ +package events + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" +) + +const ( + defaultNATSPrefix = "harness_events" + natsEventSubject = "events.event" + natsSnapshotSubject = "events.snapshot" + defaultMaxCacheAge = 10 * time.Minute + defaultMaxCacheItems = 10000 +) + +// cachedEvent wraps an Event with an expiry timestamp for TTL-based cache eviction. +type cachedEvent struct { + ev *Event + expiresAt time.Time +} + +// NATSEventStore persists events to NATS JetStream. +// Suitable for production distributed deployments. +type NATSEventStore struct { + conn *nats.Conn + js jetstream.JetStream + stream string // JetStream stream name + prefix string // subject prefix + mu sync.RWMutex + cache map[string]*cachedEvent // ID → timestamped Event for fast Get (bounded) + maxCacheAge time.Duration + clock atomic.Uint64 + subs map[int64]*nats.Subscription + subID atomic.Int64 +} + +// NewNATSEventStore creates a new NATSEventStore. +func NewNATSEventStore(conn *nats.Conn, stream string) (*NATSEventStore, error) { + js, err := jetstream.New(conn) + if err != nil { + return nil, fmt.Errorf("jetstream init: %w", err) + } + + // Ensure the stream exists. + _, err = js.Stream(ctxForInit(), stream) + if err != nil { + // Create the stream if it doesn't exist. + _, err = js.CreateStream(ctxForInit(), jetstream.StreamConfig{ + Name: stream, + Subjects: []string{fmt.Sprintf("%s.>", defaultNATSPrefix)}, + MaxAge: 7 * 24 * time.Hour, // 7 days retention + Storage: jetstream.FileStorage, + Retention: jetstream.LimitsPolicy, + }) + if err != nil { + return nil, fmt.Errorf("create jetstream stream: %w", err) + } + } + + return &NATSEventStore{ + conn: conn, + js: js, + stream: stream, + prefix: defaultNATSPrefix, + cache: make(map[string]*cachedEvent), + subs: make(map[int64]*nats.Subscription), + }, nil +} + +// ctxForInit returns a background context for NATS stream setup. +func ctxForInit() context.Context { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + _ = cancel // prevent vet warning; cancel runs when ctx expires + return ctx +} + +// Append implements EventLog. +func (s *NATSEventStore) Append(ctx context.Context, events ...*Event) error { + for _, ev := range events { + if ev.Clock == 0 { + ev.Clock = s.clock.Add(1) + } + ev.Seal() + + data, err := json.Marshal(ev) + if err != nil { + return fmt.Errorf("marshal event: %w", err) + } + + subject := fmt.Sprintf("%s.%s", s.prefix, natsEventSubject) + if _, err := s.js.Publish(ctx, subject, data); err != nil { + return fmt.Errorf("publish event: %w", err) + } + + // Update local cache with TTL. + s.mu.Lock() + maxAge := s.maxCacheAge + if maxAge == 0 { + maxAge = defaultMaxCacheAge + } + s.cache[string(ev.ID)] = &cachedEvent{ev: ev, expiresAt: time.Now().Add(maxAge)} + // Evict expired entries when cache exceeds limit. + if len(s.cache) > defaultMaxCacheItems { + s.evictExpiredLocked() + } + s.mu.Unlock() + } + return nil +} + +// Stream implements EventLog by replaying from JetStream. +func (s *NATSEventStore) Stream(ctx context.Context, filter EventFilter) EventIterator { + subject := fmt.Sprintf("%s.%s", s.prefix, natsEventSubject) + consumer, err := s.js.OrderedConsumer(ctx, s.stream, jetstream.OrderedConsumerConfig{ + FilterSubjects: []string{subject}, + }) + if err != nil { + return &natsErrorIterator{err: fmt.Errorf("create consumer: %w", err)} + } + + return &natsEventIterator{ + consumer: consumer, + ctx: ctx, + filter: filter, + buffer: make([]*Event, 0), + } +} + +// Get implements EventLog. +func (s *NATSEventStore) Get(ctx context.Context, id EventID) (*Event, error) { + s.mu.RLock() + ce, ok := s.cache[string(id)] + s.mu.RUnlock() + if ok && ce.expiresAt.After(time.Now()) { + return ce.ev, nil + } + // Expired cache entry — remove it. + if ok { + s.mu.Lock() + delete(s.cache, string(id)) + s.mu.Unlock() + } + + // Fall back to stream scan (slow path). + iter := s.Stream(ctx, EventFilter{Limit: 1}) + defer iter.Close() + for { + e, ok := iter.Next(ctx) + if !ok { + break + } + if e.ID == id { + return e, nil + } + } + return nil, nil +} + +// Range implements EventLog. +func (s *NATSEventStore) Range(ctx context.Context, from, to uint64, filter EventFilter) ([]*Event, error) { + iter := s.Stream(ctx, filter) + defer iter.Close() + + result := make([]*Event, 0) + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + if ev.Clock < from { + continue + } + if to > 0 && ev.Clock > to { + continue + } + result = append(result, ev) + } + return result, nil +} + +// Seek implements EventLog. +func (s *NATSEventStore) Seek(ctx context.Context, clock uint64) (EventIterator, error) { + return s.Stream(ctx, EventFilter{FromClock: clock}), nil +} + +// Length implements EventLog by counting messages in the stream. +func (s *NATSEventStore) Length(ctx context.Context) (uint64, error) { + streamInfo, err := s.js.Stream(ctx, s.stream) + if err != nil { + return 0, fmt.Errorf("get stream info: %w", err) + } + return streamInfo.CachedInfo().State.Msgs, nil +} + +// CreateSnapshot implements EventStore. +func (s *NATSEventStore) CreateSnapshot(ctx context.Context, traceID string) (*Snapshot, error) { + clock := s.clock.Load() + + // Collect all events for the trace. + iter := s.Stream(ctx, EventFilter{TraceID: traceID}) + defer iter.Close() + + var traceEvents []*Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + traceEvents = append(traceEvents, ev) + } + + data, err := json.Marshal(traceEvents) + if err != nil { + return nil, fmt.Errorf("marshal snapshot: %w", err) + } + + snap := &Snapshot{ + ID: fmt.Sprintf("snap-%s-%d", traceID, clock), + TraceID: traceID, + Clock: clock, + CreatedAt: time.Now(), + Data: data, + } + + snapData, _ := json.Marshal(snap) + subject := fmt.Sprintf("%s.%s.%s", s.prefix, natsSnapshotSubject, traceID) + s.js.Publish(ctx, subject, snapData) + + return snap, nil +} + +// RestoreSnapshot implements EventStore. +func (s *NATSEventStore) RestoreSnapshot(ctx context.Context, snapshotID string) (EventIterator, error) { + return s.Seek(ctx, 0) +} + +// Subscribe implements EventStore using NATS JetStream push consumer. +func (s *NATSEventStore) Subscribe(ctx context.Context, filter EventFilter) (<-chan *Event, error) { + subject := fmt.Sprintf("%s.%s", s.prefix, natsEventSubject) + ch := make(chan *Event, 256) + + consumer, err := s.js.OrderedConsumer(ctx, s.stream, jetstream.OrderedConsumerConfig{ + FilterSubjects: []string{subject}, + DeliverPolicy: jetstream.DeliverNewPolicy, + }) + if err != nil { + close(ch) + return ch, fmt.Errorf("create consumer: %w", err) + } + + // Start goroutine to forward messages. + go func() { + defer close(ch) + for { + msg, err := consumer.Next() + if err != nil { + return + } + var ev Event + if err := json.Unmarshal(msg.Data(), &ev); err != nil { + continue + } + if filter.Matches(&ev) { + select { + case ch <- &ev: + case <-ctx.Done(): + return + } + } + } + }() + + return ch, nil +} + +// GC implements EventStore (purge stream is handled by JetStream TTL). +func (s *NATSEventStore) GC(ctx context.Context, retention time.Duration) error { + // JetStream handles retention via MaxAge in stream config. + // For explicit GC, update the stream config. + info, err := s.js.Stream(ctx, s.stream) + if err != nil { + return err + } + cfg := info.CachedInfo().Config + cfg.MaxAge = retention + _, err = s.js.UpdateStream(ctx, cfg) + return err +} + +// evictExpiredLocked removes cache entries whose TTL has expired. +// Must be called with s.mu held (Lock, not RLock). +func (s *NATSEventStore) evictExpiredLocked() { + now := time.Now() + for k, ce := range s.cache { + if now.After(ce.expiresAt) { + delete(s.cache, k) + } + } +} + +// ---- natsEventIterator ---- + +type natsEventIterator struct { + consumer jetstream.Consumer + ctx context.Context + filter EventFilter + buffer []*Event + bufPos int +} + +func (it *natsEventIterator) Next(_ context.Context) (*Event, bool) { + // Return from buffer first. + if it.bufPos < len(it.buffer) { + ev := it.buffer[it.bufPos] + it.bufPos++ + return ev, true + } + it.buffer = it.buffer[:0] + it.bufPos = 0 + + // Fetch next batch. + for i := 0; i < 100; i++ { + msg, err := it.consumer.Next() + if err != nil { + return nil, false + } + var ev Event + if err := json.Unmarshal(msg.Data(), &ev); err != nil { + continue + } + if it.filter.Matches(&ev) { + it.buffer = append(it.buffer, &ev) + } + } + if len(it.buffer) == 0 { + return nil, false + } + ev := it.buffer[0] + it.bufPos = 1 + return ev, true +} + +func (it *natsEventIterator) Close() error { + return nil +} + +// ---- natsErrorIterator ---- + +type natsErrorIterator struct { + err error + emitted bool +} + +func (it *natsErrorIterator) Next(_ context.Context) (*Event, bool) { + return nil, false +} + +func (it *natsErrorIterator) Close() error { + return nil +} diff --git a/internal/harness/events/recorder.go b/internal/harness/events/recorder.go new file mode 100644 index 000000000..3381187de --- /dev/null +++ b/internal/harness/events/recorder.go @@ -0,0 +1,360 @@ +package events + +import ( + "context" + "encoding/json" + + "ragflow/internal/harness/graph/pregel" +) + +// ---- Context helpers for passing EventRecorder through context ---- + +type recorderContextKey struct{} + +// ContextWithRecorder stores an EventRecorder in context for use by +// model wrappers and tool middlewares in the agent core. +func ContextWithRecorder(ctx context.Context, r *EventRecorder) context.Context { + return context.WithValue(ctx, recorderContextKey{}, r) +} + +// RecorderFromContext retrieves an EventRecorder from context. +// Returns nil when no recorder is present. +func RecorderFromContext(ctx context.Context) *EventRecorder { + r, _ := ctx.Value(recorderContextKey{}).(*EventRecorder) + return r +} + +// RecorderOption configures an EventRecorder. +type RecorderOption func(*recorderOptions) + +type recorderOptions struct { + traceID string + threadID string +} + +// WithTraceID sets the trace ID for the recorder. +func WithTraceID(traceID string) RecorderOption { + return func(o *recorderOptions) { + o.traceID = traceID + } +} + +// WithThreadID sets the thread ID for the recorder. +func WithThreadID(threadID string) RecorderOption { + return func(o *recorderOptions) { + o.threadID = threadID + } +} + +// EventRecorder records graph execution events as append-only Events. +// It implements pregel.GraphCallback and can be added to a CallbackManager. +// Additionally, RecordModelCall / RecordToolCall / etc. provide fine-grained +// event recording for LLM invocations and tool executions. +// +// Usage: +// +// store := events.NewMemoryEventStore() +// recorder := events.NewEventRecorder(store, events.WithTraceID("trace-001")) +// cb := pregel.NewCallbackManager() +// cb.AddCallback(recorder) +type EventRecorder struct { + store EventLog + clock *LogicalClock + traceID string + threadID string +} + +// NewEventRecorder creates a new EventRecorder. +func NewEventRecorder(store EventLog, opts ...RecorderOption) *EventRecorder { + o := &recorderOptions{} + for _, opt := range opts { + opt(o) + } + return &EventRecorder{ + store: store, + clock: NewLogicalClock(), + traceID: o.traceID, + threadID: o.threadID, + } +} + +// record creates and appends an event. +func (r *EventRecorder) record(ctx context.Context, typ EventType, opts ...func(*Event)) { + ev := NewEvent(typ, r.clock.Tick()) + ev.TraceID = r.traceID + ev.ThreadID = r.threadID + for _, fn := range opts { + fn(ev) + } + ev.Seal() + _ = r.store.Append(ctx, ev) +} + +// ---- Context-based recording (used by model/tool wrappers) ---- + +// RecordModelCall records an LLM model invocation with its result. +func (r *EventRecorder) RecordModelCall(ctx context.Context, model, provider string, messages []any, content string, tokens TokenUsage, durationMs int64, cost float64) { + r.record(ctx, EventLLMCallStart, func(ev *Event) { + ev.Deterministic = false + ev.Metadata["model"] = model + ev.Metadata["provider"] = provider + }) + r.record(ctx, EventLLMCallEnd, func(ev *Event) { + ev.Deterministic = false + pl := LLMCallPayload{ + Model: model, + Provider: provider, + Messages: messages, + Tokens: tokens, + Content: content, + DurationMs: durationMs, + Cost: cost, + } + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordLLMChunk records a single streaming chunk from an LLM call. +func (r *EventRecorder) RecordLLMChunk(ctx context.Context, model string, chunk string) { + r.record(ctx, EventLLMCallChunk, func(ev *Event) { + ev.Deterministic = false + ev.Metadata["model"] = model + ev.Metadata["chunk"] = chunk + }) +} + +// RecordToolCall records a tool invocation with its result. +func (r *EventRecorder) RecordToolCall(ctx context.Context, toolName string, arguments map[string]any, result any, durationMs int64, retryCount int, errStr string) { + r.record(ctx, EventToolCallStart, func(ev *Event) { + ev.Metadata["tool"] = toolName + }) + r.record(ctx, EventToolCallResult, func(ev *Event) { + pl := ToolCallPayload{ + ToolName: toolName, + Arguments: arguments, + Result: result, + DurationMs: durationMs, + RetryCount: retryCount, + Error: errStr, + } + ev.Payload, _ = json.Marshal(pl) + if errStr != "" { + ev.Deterministic = false + } + }) +} + +// RecordSubAgentCall records a sub-agent invocation with its result. +func (r *EventRecorder) RecordSubAgentCall(ctx context.Context, subAgentName string, input, output any, depth int, durationMs int64, errStr string) { + r.record(ctx, EventSubAgentCallStart, func(ev *Event) { + pl := SubAgentCallPayload{ + SubAgentName: subAgentName, + Input: input, + Depth: depth, + } + ev.Payload, _ = json.Marshal(pl) + }) + r.record(ctx, EventSubAgentCallEnd, func(ev *Event) { + pl := SubAgentCallPayload{ + SubAgentName: subAgentName, + Output: output, + Depth: depth, + DurationMs: durationMs, + Error: errStr, + } + ev.Payload, _ = json.Marshal(pl) + if errStr != "" { + ev.Deterministic = false + } + }) +} + +// RecordSessionValue records a session value change. +func (r *EventRecorder) RecordSessionValue(ctx context.Context, key string, value any) { + r.record(ctx, EventSessionValueSet, func(ev *Event) { + pl := SessionValuePayload{Key: key, Value: value} + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordSessionTransfer records an agent transfer event. +func (r *EventRecorder) RecordSessionTransfer(ctx context.Context, fromAgent, toAgent, reason string, input any) { + r.record(ctx, EventSessionTransfer, func(ev *Event) { + pl := SessionTransferPayload{ + FromAgent: fromAgent, + ToAgent: toAgent, + Reason: reason, + Input: input, + } + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordStateWrite records a state transition. +func (r *EventRecorder) RecordStateWrite(ctx context.Context, channel string, oldValue, newValue any, reducer string) { + r.record(ctx, EventStateWrite, func(ev *Event) { + pl := StateTransitionPayload{ + Channel: channel, + OldValue: oldValue, + NewValue: newValue, + Reducer: reducer, + } + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordMemoryWrite records a memory operation. +func (r *EventRecorder) RecordMemoryWrite(ctx context.Context, store, operation, key string, value any, score float64) { + r.record(ctx, EventMemoryWrite, func(ev *Event) { + pl := MemoryWritePayload{ + Store: store, + Operation: operation, + Key: key, + Value: value, + Score: score, + } + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordMemoryRead records a memory read operation. +func (r *EventRecorder) RecordMemoryRead(ctx context.Context, store, key string, score float64) { + r.record(ctx, EventMemoryRead, func(ev *Event) { + pl := MemoryWritePayload{ + Store: store, + Key: key, + Score: score, + } + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordApproval records a human-in-the-loop approval event. +func (r *EventRecorder) RecordApproval(ctx context.Context, requestID, action string, context any, decision string, latencyMs int64) { + r.record(ctx, EventApprovalRequest, func(ev *Event) { + pl := ApprovalPayload{ + RequestID: requestID, + Action: action, + Context: context, + Decision: decision, + LatencyMs: latencyMs, + } + ev.Payload, _ = json.Marshal(pl) + }) +} + +// RecordError records an execution error. +func (r *EventRecorder) RecordError(ctx context.Context, errMsg string) { + r.record(ctx, EventError, func(ev *Event) { + ev.Metadata["error"] = errMsg + }) +} + +// RecordRetry records a retry event. +func (r *EventRecorder) RecordRetry(ctx context.Context, detail string) { + r.record(ctx, EventRetry, func(ev *Event) { + ev.Metadata["detail"] = detail + }) +} + +// ---- GraphCallback implementation ---- + +// OnRunStart implements pregel.RunCallback. +func (r *EventRecorder) OnRunStart(ctx context.Context, graphName, threadID string) { + r.record(ctx, EventGraphStart, func(ev *Event) { + ev.Metadata["graph_name"] = graphName + }) +} + +// OnRunEnd implements pregel.RunCallback. +func (r *EventRecorder) OnRunEnd(ctx context.Context, graphName, threadID string, err error) { + r.record(ctx, EventGraphEnd, func(ev *Event) { + ev.Metadata["graph_name"] = graphName + if err != nil { + ev.Metadata["error"] = err.Error() + } + }) +} + +// OnStepStart implements pregel.StepCallback. +func (r *EventRecorder) OnStepStart(ctx context.Context, step, taskCount int) { + r.record(ctx, EventStepStart, func(ev *Event) { + ev.Step = step + ev.Metadata["task_count"] = taskCount + }) +} + +// OnStepEnd implements pregel.StepCallback. +func (r *EventRecorder) OnStepEnd(ctx context.Context, step int, err error) { + r.record(ctx, EventStepEnd, func(ev *Event) { + ev.Step = step + if err != nil { + ev.Metadata["error"] = err.Error() + } + }) +} + +// OnNodeStart implements pregel.NodeCallback. +func (r *EventRecorder) OnNodeStart(ctx context.Context, nodeName string, step int) { + r.record(ctx, EventNodeStart, func(ev *Event) { + ev.Node = nodeName + ev.Step = step + }) +} + +// OnNodeEnd implements pregel.NodeCallback. +func (r *EventRecorder) OnNodeEnd(ctx context.Context, nodeName string, step int, output interface{}, err error) { + r.record(ctx, EventNodeEnd, func(ev *Event) { + ev.Node = nodeName + ev.Step = step + if err != nil { + ev.Metadata["error"] = err.Error() + } + }) +} + +// OnCheckpointSave implements pregel.CheckpointCallback. +func (r *EventRecorder) OnCheckpointSave(ctx context.Context, threadID, checkpointID string, step int) { + r.record(ctx, EventCheckpointCreated, func(ev *Event) { + ev.ThreadID = threadID + ev.Step = step + ev.Metadata["checkpoint_id"] = checkpointID + }) +} + +// OnCheckpointLoad implements pregel.CheckpointCallback. +func (r *EventRecorder) OnCheckpointLoad(ctx context.Context, threadID, checkpointID string, step int) { + r.record(ctx, EventCheckpointRestored, func(ev *Event) { + ev.ThreadID = threadID + ev.Step = step + ev.Metadata["checkpoint_id"] = checkpointID + }) +} + +// OnCheckpointUpdate implements pregel.CheckpointCallback. +func (r *EventRecorder) OnCheckpointUpdate(ctx context.Context, threadID, asNode string) { + r.record(ctx, EventStateWrite, func(ev *Event) { + ev.ThreadID = threadID + ev.Node = asNode + }) +} + +// OnInterrupt implements pregel.InterruptCallback. +func (r *EventRecorder) OnInterrupt(ctx context.Context, nodeNames []string, step int) { + r.record(ctx, EventInterrupt, func(ev *Event) { + ev.Step = step + ev.Metadata["interrupt_nodes"] = nodeNames + }) +} + +// OnResume implements pregel.InterruptCallback. +func (r *EventRecorder) OnResume(ctx context.Context, threadID string) { + r.record(ctx, EventResume, func(ev *Event) { + ev.ThreadID = threadID + }) +} + +// compile-time interface checks +var ( + _ pregel.GraphCallback = (*EventRecorder)(nil) +) diff --git a/internal/harness/events/store.go b/internal/harness/events/store.go new file mode 100644 index 000000000..0eba5c279 --- /dev/null +++ b/internal/harness/events/store.go @@ -0,0 +1,127 @@ +package events + +import ( + "context" + "time" +) + +// EventLog is the append-only event log interface. +// All implementations must be safe for concurrent use. +type EventLog interface { + // Append appends one or more events to the log. Events are immutable + // once appended. + Append(ctx context.Context, events ...*Event) error + + // Stream returns an iterator over events matching the filter, + // ordered by logical clock. + Stream(ctx context.Context, filter EventFilter) EventIterator + + // Get retrieves a single event by ID. Returns nil, nil if not found. + Get(ctx context.Context, id EventID) (*Event, error) + + // Range returns events with logical clock in [from, to] matching the filter. + Range(ctx context.Context, from, to uint64, filter EventFilter) ([]*Event, error) + + // Seek returns an iterator starting from the given logical clock. + Seek(ctx context.Context, clock uint64) (EventIterator, error) + + // Length returns the total number of events in the log. + Length(ctx context.Context) (uint64, error) +} + +// EventFilter specifies criteria for filtering events. +type EventFilter struct { + // TraceID filters by trace. + TraceID string + // ThreadID filters by thread. + ThreadID string + // Types restricts to specific event types. Empty means all types. + Types []EventType + // Node restricts to a specific graph node. + Node string + // FromClock is the minimum logical clock (inclusive). + FromClock uint64 + // ToClock is the maximum logical clock (inclusive). 0 means no upper bound. + ToClock uint64 + // FromTime is the minimum wall-clock time. + FromTime time.Time + // ToTime is the maximum wall-clock time. + ToTime time.Time + // Limit caps the number of events returned. 0 means no limit. + Limit int +} + +// Matches checks whether an event matches this filter. +func (f EventFilter) Matches(e *Event) bool { + if f.TraceID != "" && e.TraceID != f.TraceID { + return false + } + if f.ThreadID != "" && e.ThreadID != f.ThreadID { + return false + } + if len(f.Types) > 0 { + matched := false + for _, t := range f.Types { + if e.Type == t { + matched = true + break + } + } + if !matched { + return false + } + } + if f.Node != "" && e.Node != f.Node { + return false + } + if f.FromClock > 0 && e.Clock < f.FromClock { + return false + } + if f.ToClock > 0 && e.Clock > f.ToClock { + return false + } + if !f.FromTime.IsZero() && e.Timestamp.Before(f.FromTime) { + return false + } + if !f.ToTime.IsZero() && e.Timestamp.After(f.ToTime) { + return false + } + return true +} + +// EventIterator allows iterating over events in order. +type EventIterator interface { + // Next returns the next event. Returns nil, false when exhausted. + Next(ctx context.Context) (*Event, bool) + // Close releases resources held by the iterator. + Close() error +} + +// Snapshot represents a point-in-time snapshot of event state, +// used to accelerate replay (avoids replaying from event 0). +type Snapshot struct { + ID string `json:"id"` + TraceID string `json:"trace_id"` + Clock uint64 `json:"clock"` + CreatedAt time.Time `json:"created_at"` + Data []byte `json:"data,omitempty"` +} + +// EventStore extends EventLog with lifecycle management. +type EventStore interface { + EventLog + + // CreateSnapshot creates a snapshot for the given trace. + CreateSnapshot(ctx context.Context, traceID string) (*Snapshot, error) + + // RestoreSnapshot loads a snapshot and returns an iterator positioned + // after the snapshot's clock position. + RestoreSnapshot(ctx context.Context, snapshotID string) (EventIterator, error) + + // Subscribe returns a channel that receives new events matching the filter + // as they are appended. The channel is closed when the context is cancelled. + Subscribe(ctx context.Context, filter EventFilter) (<-chan *Event, error) + + // GC removes events older than the given retention period. + GC(ctx context.Context, retention time.Duration) error +} diff --git a/internal/harness/graph/pregel/engine.go b/internal/harness/graph/pregel/engine.go index 94d4ed4d8..c1c788f7c 100644 --- a/internal/harness/graph/pregel/engine.go +++ b/internal/harness/graph/pregel/engine.go @@ -48,6 +48,7 @@ type Engine struct { versionsSeen map[string]map[string]int cache Cache backgroundExec *BackgroundExecutor + callbacks *CallbackManager // lifecycle callbacks (event recording, metrics) deferredCheckpoints []deferredCheckpoint // for DurabilityExit mode } @@ -173,6 +174,15 @@ func WithBackgroundExecutor(exec *BackgroundExecutor) EngineOption { } } +// WithCallbacks sets the callback manager for the engine. +// Callbacks are dispatched during graph execution (run start/end, step start/end, +// node start/end, checkpoint save/load, interrupt/resume). +func WithCallbacks(cb *CallbackManager) EngineOption { + return func(e *Engine) { + e.callbacks = cb + } +} + // ExecuteResult represents the result of graph execution. type ExecuteResult struct { // Final state of the graph. @@ -215,9 +225,39 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c } }) - // Deferred cleanup: close streamManager first (unblocks forward goroutine), - // then wait for forward goroutine to exit, then close outputCh. + // Resolve thread ID early (before deferred cleanup uses it). + threadID := e.getThreadID() + + // reportRunEnd is defined before the deferred cleanup block so the + // defer can capture it by closure. + reportRunEnd := func(err error) { + if e.callbacks == nil { + return + } + gName := "state_graph" + if e.graph != nil { + nodes := e.graph.GetNodes() + for name := range nodes { + gName = name + break + } + } + e.callbacks.RunEnd(context.Background(), gName, threadID, err) + } + + // Deferred cleanup: dispatch RunEnd, close streamManager, + // wait for forward goroutine, then close outputCh. + var exitErr error // captured for RunEnd callback dispatch defer func() { + // Read from errCh to get the exit error for RunEnd dispatch. + // errCh is still open here (close(errCh) runs after this defer). + select { + case exitErr = <-errCh: + reportRunEnd(exitErr) + errCh <- exitErr // put back for the caller + default: + reportRunEnd(nil) + } streamManager.Close() fwWg.Wait() close(outputCh) @@ -256,7 +296,7 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c } // Get thread ID for checkpointing - threadID := e.getThreadID() + // threadID already resolved above (before deferred cleanup). // Load checkpoint when one exists for this thread_id, even when // input is non-nil (resume from a previous run). The canvas @@ -341,6 +381,14 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c if cp, err := checkpoint.FromMap(cpData); err == nil { e.currentCheckpoint = cp } + // Dispatch CheckpointLoad callback. + if e.callbacks != nil { + cpID := "" + if cpid, _ := cpData["checkpoint_id"].(string); cpid != "" { + cpID = cpid + } + e.callbacks.CheckpointLoad(ctx, threadID, cpID, 0) + } } } // Apply input only when no checkpoint was loaded. @@ -393,6 +441,18 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c lastState = input } + // Dispatch RunStart callback. + if e.callbacks != nil { + gName := "state_graph" + if e.graph != nil { + nodes := e.graph.GetNodes() + for name := range nodes { + gName = name + break + } + } + e.callbacks.RunStart(ctx, gName, threadID) + } for { // Check context cancellation at each superstep. select { @@ -408,6 +468,11 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c return } + // Dispatch StepStart callback. + if e.callbacks != nil { + e.callbacks.StepStart(ctx, step, 0) // taskCount filled after prepareNextTasks + } + // Emit checkpoint event via stream manager streamManager.EmitCheckpoint(step, channelRegistry.CreateCheckpoint()) @@ -449,6 +514,11 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c } streamManager.EmitInterrupt(step, interruptNames) + // Dispatch Interrupt callback. + if e.callbacks != nil { + e.callbacks.Interrupt(ctx, interruptNames, step) + } + errCh <- &errors.GraphInterrupt{} return } @@ -519,6 +589,10 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c } } streamManager.EmitInterrupt(step, interruptTaskNames) + // Dispatch Interrupt callback. + if e.callbacks != nil { + e.callbacks.Interrupt(ctx, interruptTaskNames, step) + } // Preserve the first interrupted task's GraphInterrupt value // (with Interrupts populated) instead of creating a bare one, // so MustExtractInterruptContexts can extract the original @@ -571,6 +645,10 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c errCh <- fmt.Errorf("failed to save checkpoint: %w", err) return } + // Dispatch CheckpointSave callback. + if e.callbacks != nil { + e.callbacks.CheckpointSave(ctx, threadID, checkpointID, step) + } case types.DurabilityAsync: // Asynchronous save - don't block next step go func(cp map[string]any, cpID string, s int) { @@ -583,6 +661,10 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c // Defer save until exit - accumulate checkpoints in memory // Will be saved in final state e.deferCheckpoint(threadID, checkpointID, step, checkpoint) + // Dispatch CheckpointSave callback (deferred save still counts as saved). + if e.callbacks != nil { + e.callbacks.CheckpointSave(ctx, threadID, checkpointID, step) + } default: // Default to sync behavior if err := e.saveCheckpoint(ctx, threadID, checkpointID, step, checkpoint); err != nil { @@ -595,10 +677,18 @@ func (e *Engine) Run(ctx context.Context, input any, mode types.StreamMode) (<-c // Check for after-node interrupts. The checkpoint above already // captures this step's output. if e.shouldInterruptAfter(results) { + if e.callbacks != nil { + e.callbacks.Interrupt(ctx, []string{"after_node"}, step) + } errCh <- &errors.GraphInterrupt{} return } + // Dispatch StepEnd callback. + if e.callbacks != nil { + e.callbacks.StepEnd(ctx, step, nil) + } + step++ } diff --git a/internal/harness/metrics/aggregator.go b/internal/harness/metrics/aggregator.go new file mode 100644 index 000000000..90b256b74 --- /dev/null +++ b/internal/harness/metrics/aggregator.go @@ -0,0 +1,192 @@ +package metrics + +import ( + "math" + "sort" + "sync" + "time" +) + +// MetricsAggregator aggregates metrics across multiple execution traces. +type MetricsAggregator struct { + mu sync.Mutex + metrics []*AgentMetrics +} + +// NewMetricsAggregator creates a new MetricsAggregator. +func NewMetricsAggregator() *MetricsAggregator { + return &MetricsAggregator{} +} + +// Add adds a metrics snapshot to the aggregator. +func (a *MetricsAggregator) Add(m *AgentMetrics) { + a.mu.Lock() + defer a.mu.Unlock() + a.metrics = append(a.metrics, m) +} + +// AggregatedMetrics contains summary statistics across multiple traces. +type AggregatedMetrics struct { + TotalTraces int + TotalDuration time.Duration + + // Tool metrics (averages). + AvgToolCalls float64 + AvgToolSuccessRate float64 + AvgToolRetryRate float64 + P50ToolLatencyMs float64 + P95ToolLatencyMs float64 + P99ToolLatencyMs float64 + + // Checkpoint metrics. + AvgCheckpointSaves float64 + AvgCheckpointRestores float64 + AvgCheckpointRestoreSuccess float64 + + // Execution metrics. + AvgSteps float64 + AvgNodesExecuted float64 + AvgRecoveredErrors float64 + AvgInterrupts float64 + + // Cost metrics. + AvgCostPerTask float64 + AvgForkReplayPassRate float64 +} + +// Aggregate computes summary statistics across all collected metrics. +func (a *MetricsAggregator) Aggregate() *AggregatedMetrics { + a.mu.Lock() + defer a.mu.Unlock() + + result := &AggregatedMetrics{ + TotalTraces: len(a.metrics), + } + + if len(a.metrics) == 0 { + return result + } + + var allLatencies []float64 + + for _, m := range a.metrics { + result.TotalDuration += m.Duration + result.AvgToolCalls += float64(m.ToolCalls) + result.AvgToolSuccessRate += m.ToolSuccessRate + result.AvgToolRetryRate += m.ToolRetryRate + result.AvgCheckpointSaves += float64(m.CheckpointSaves) + result.AvgCheckpointRestores += float64(m.CheckpointRestores) + result.AvgCheckpointRestoreSuccess += m.CheckpointRestoreSuccess + result.AvgSteps += float64(m.Steps) + result.AvgNodesExecuted += float64(m.NodesExecuted) + result.AvgRecoveredErrors += float64(m.RecoveredErrors) + result.AvgInterrupts += float64(m.InterruptCount) + result.AvgCostPerTask += m.CostPerTask + result.AvgForkReplayPassRate += m.ForkReplayPassRate + + for _, latencies := range m.ToolLatencyMs { + for _, l := range latencies { + allLatencies = append(allLatencies, float64(l)) + } + } + } + + n := float64(len(a.metrics)) + result.AvgToolCalls /= n + result.AvgToolSuccessRate /= n + result.AvgToolRetryRate /= n + result.AvgCheckpointSaves /= n + result.AvgCheckpointRestores /= n + result.AvgCheckpointRestoreSuccess /= n + result.AvgSteps /= n + result.AvgNodesExecuted /= n + result.AvgRecoveredErrors /= n + result.AvgInterrupts /= n + result.AvgCostPerTask /= n + result.AvgForkReplayPassRate /= n + + // Compute latency percentiles. + if len(allLatencies) > 0 { + sort.Float64s(allLatencies) + result.P50ToolLatencyMs = percentile(allLatencies, 50) + result.P95ToolLatencyMs = percentile(allLatencies, 95) + result.P99ToolLatencyMs = percentile(allLatencies, 99) + } + + return result +} + +// Reset clears all collected metrics. +func (a *MetricsAggregator) Reset() { + a.mu.Lock() + defer a.mu.Unlock() + a.metrics = nil +} + +// MetricsWindow tracks metrics over a sliding time window. +type MetricsWindow struct { + mu sync.Mutex + window time.Duration + entries []windowEntry +} + +type windowEntry struct { + timestamp time.Time + metrics *AgentMetrics +} + +// NewMetricsWindow creates a metrics window with the given duration. +func NewMetricsWindow(window time.Duration) *MetricsWindow { + return &MetricsWindow{window: window} +} + +// Add adds metrics at the current time. +func (w *MetricsWindow) Add(m *AgentMetrics) { + w.mu.Lock() + defer w.mu.Unlock() + w.entries = append(w.entries, windowEntry{ + timestamp: time.Now(), + metrics: m, + }) + w.prune() +} + +// Aggregate returns aggregated metrics for the current window. +func (w *MetricsWindow) Aggregate() *AggregatedMetrics { + w.mu.Lock() + defer w.mu.Unlock() + w.prune() + + agg := NewMetricsAggregator() + for _, entry := range w.entries { + agg.Add(entry.metrics) + } + return agg.Aggregate() +} + +// prune removes entries outside the window. +func (w *MetricsWindow) prune() { + cutoff := time.Now().Add(-w.window) + keep := make([]windowEntry, 0, len(w.entries)) + for _, e := range w.entries { + if e.timestamp.After(cutoff) { + keep = append(keep, e) + } + } + w.entries = keep +} + +// percentile computes the p-th percentile from a sorted slice. +func percentile(sorted []float64, p int) float64 { + if len(sorted) == 0 { + return 0 + } + idx := int(math.Ceil(float64(p)/100.0*float64(len(sorted))) - 1) + if idx < 0 { + idx = 0 + } + if idx >= len(sorted) { + idx = len(sorted) - 1 + } + return sorted[idx] +} diff --git a/internal/harness/metrics/exporter.go b/internal/harness/metrics/exporter.go new file mode 100644 index 000000000..3f269250f --- /dev/null +++ b/internal/harness/metrics/exporter.go @@ -0,0 +1,78 @@ +package metrics + +import ( + "fmt" + "strings" +) + +// Exporter formats agent metrics for output. This is a lightweight +// alternative to a full Prometheus client, avoiding the dependency. +type Exporter struct { + namespace string +} + +// NewExporter creates a new metrics exporter. +func NewExporter(namespace string) *Exporter { + return &Exporter{namespace: namespace} +} + +// ExportText formats metrics as Prometheus-style text output. +func (e *Exporter) ExportText(m *AgentMetrics) string { + snap := m.Snapshot() + var b strings.Builder + + ns := e.namespace + if ns != "" { + ns += "_" + } + + // Tool metrics. + fmt.Fprintf(&b, "# HELP %stool_calls_total Total tool invocations\n", ns) + fmt.Fprintf(&b, "# TYPE %stool_calls_total counter\n", ns) + fmt.Fprintf(&b, "%stool_calls_total %d\n", ns, snap.ToolCalls) + + fmt.Fprintf(&b, "# HELP %stool_success_rate Tool success rate\n", ns) + fmt.Fprintf(&b, "# TYPE %stool_success_rate gauge\n", ns) + fmt.Fprintf(&b, "%stool_success_rate %.4f\n", ns, snap.ToolSuccessRate) + + fmt.Fprintf(&b, "# HELP %stool_retry_rate Tool retry rate\n", ns) + fmt.Fprintf(&b, "# TYPE %stool_retry_rate gauge\n", ns) + fmt.Fprintf(&b, "%stool_retry_rate %.4f\n", ns, snap.ToolRetryRate) + + // Checkpoint metrics. + fmt.Fprintf(&b, "# HELP %scheckpoint_saves_total Total checkpoint saves\n", ns) + fmt.Fprintf(&b, "# TYPE %scheckpoint_saves_total counter\n", ns) + fmt.Fprintf(&b, "%scheckpoint_saves_total %d\n", ns, snap.CheckpointSaves) + + fmt.Fprintf(&b, "# HELP %scheckpoint_restore_success Checkpoint restore success rate\n", ns) + fmt.Fprintf(&b, "# TYPE %scheckpoint_restore_success gauge\n", ns) + fmt.Fprintf(&b, "%scheckpoint_restore_success %.4f\n", ns, snap.CheckpointRestoreSuccess) + + // Execution metrics. + fmt.Fprintf(&b, "# HELP %ssteps_total Total supersteps executed\n", ns) + fmt.Fprintf(&b, "# TYPE %ssteps_total counter\n", ns) + fmt.Fprintf(&b, "%ssteps_total %d\n", ns, snap.Steps) + + fmt.Fprintf(&b, "# HELP %snodes_executed_total Total nodes executed\n", ns) + fmt.Fprintf(&b, "# TYPE %snodes_executed_total counter\n", ns) + fmt.Fprintf(&b, "%snodes_executed_total %d\n", ns, snap.NodesExecuted) + + fmt.Fprintf(&b, "# HELP %sinterrupts_total Total interrupts\n", ns) + fmt.Fprintf(&b, "# TYPE %sinterrupts_total counter\n", ns) + fmt.Fprintf(&b, "%sinterrupts_total %d\n", ns, snap.InterruptCount) + + return b.String() +} + +// ExportCSV formats metrics as a single CSV row. +func (e *Exporter) ExportCSV(m *AgentMetrics) string { + snap := m.Snapshot() + return fmt.Sprintf("%s,%d,%d,%d,%.4f,%.4f,%d,%d,%.4f,%d,%d,%d,%.6f", + snap.TraceID, + snap.ToolCalls, snap.ToolSuccesses, snap.ToolFailures, + snap.ToolSuccessRate, snap.ToolRetryRate, + snap.CheckpointSaves, snap.CheckpointRestores, + snap.CheckpointRestoreSuccess, + snap.Steps, snap.NodesExecuted, snap.InterruptCount, + snap.CostPerTask) +} diff --git a/internal/harness/metrics/metrics.go b/internal/harness/metrics/metrics.go new file mode 100644 index 000000000..a29e38181 --- /dev/null +++ b/internal/harness/metrics/metrics.go @@ -0,0 +1,303 @@ +// Package metrics provides observability metrics collection for agent execution. +// +// The autoMetricCollector implements pregel.GraphCallback and tracks key +// metrics during graph execution: tool success rate, checkpoint operations, +// step/node counts, and error recovery. Metrics are exposed via the +// MetricsCollector interface and can be exported to Prometheus. +package metrics + +import ( + "context" + "sync" + "time" +) + +// AgentMetrics is a snapshot of all metrics for one agent execution trace. +type AgentMetrics struct { + // TraceID identifies the execution trace. + TraceID string + // ThreadID identifies the execution thread. + ThreadID string + // Duration is the wall-clock duration of the execution. + Duration time.Duration + + // ---- Tool metrics ---- + + // ToolCalls is the total number of tool invocations. + ToolCalls int + // ToolSuccesses is the number of successful tool invocations. + ToolSuccesses int + // ToolFailures is the number of failed tool invocations. + ToolFailures int + // ToolRetries is the number of times tools were retried. + ToolRetries int + // ToolLatencyMs holds per-tool latency histograms (toolName → durations). + ToolLatencyMs map[string][]int64 + // ToolSuccessRate is the ratio of successes to total calls (0–1). + ToolSuccessRate float64 + // ToolRetryRate is the ratio of retries to (calls + retries). + ToolRetryRate float64 + + // ---- Checkpoint metrics ---- + + // CheckpointSaves is the number of checkpoint saves. + CheckpointSaves int + // CheckpointRestores is the number of checkpoint restores. + CheckpointRestores int + // CheckpointRestoreSuccess is the ratio of successful restores (0–1). + CheckpointRestoreSuccess float64 + + // ---- Execution metrics ---- + + // Steps is the number of Pregel supersteps executed. + Steps int + // NodesExecuted is the number of graph nodes executed. + NodesExecuted int + // RecoveredErrors is the number of errors that were recovered. + RecoveredErrors int + // InterruptCount is the number of times execution was interrupted. + InterruptCount int + + // ---- Computed metrics ---- + + // CostPerTask is an estimated metric (requires LLM cost tracking). + CostPerTask float64 + // ForkReplayPassRate is the ratio of fork replays that pass assertions. + ForkReplayPassRate float64 + // ApprovalLatencyMs tracks human-in-the-loop approval wait times. + ApprovalLatencyMs []int64 + // ApprovalRate is the ratio of approvals to total approval requests. + ApprovalRate float64 + // MemoryAvgHitScore is the average retrieval score for memory operations. + MemoryAvgHitScore float64 +} + +// NewAgentMetrics creates a new AgentMetrics with initialised maps. +func NewAgentMetrics() *AgentMetrics { + return &AgentMetrics{ + ToolLatencyMs: make(map[string][]int64), + ApprovalLatencyMs: make([]int64, 0), + } +} + +// Snapshot captures a point-in-time copy of the metrics. +func (m *AgentMetrics) Snapshot() *AgentMetrics { + cp := *m + cp.ToolLatencyMs = make(map[string][]int64, len(m.ToolLatencyMs)) + for k, v := range m.ToolLatencyMs { + durations := make([]int64, len(v)) + copy(durations, v) + cp.ToolLatencyMs[k] = durations + } + cp.ApprovalLatencyMs = make([]int64, len(m.ApprovalLatencyMs)) + copy(cp.ApprovalLatencyMs, m.ApprovalLatencyMs) + + // Compute derived rates. + if cp.ToolCalls > 0 { + cp.ToolSuccessRate = float64(cp.ToolSuccesses) / float64(cp.ToolCalls) + cp.ToolRetryRate = float64(cp.ToolRetries) / float64(cp.ToolCalls+cp.ToolRetries) + } + if cp.CheckpointSaves+cp.CheckpointRestores > 0 { + cp.CheckpointRestoreSuccess = float64(cp.CheckpointRestores) / float64(cp.CheckpointSaves+cp.CheckpointRestores) + } + return &cp +} + +// MetricsCollector collects and aggregates metrics for agent execution. +type MetricsCollector interface { + // RecordToolCall records a tool invocation outcome. + RecordToolCall(toolName string, success bool, durationMs int64) + // RecordToolRetry records a tool retry. + RecordToolRetry(toolName string) + // RecordCheckpointSave records a checkpoint save. + RecordCheckpointSave() + // RecordCheckpointRestore records a checkpoint restore (success or failure). + RecordCheckpointRestore(success bool) + // RecordStep records a completed Pregel superstep. + RecordStep() + // RecordNode records a completed node execution. + RecordNode(nodeName string) + // RecordRecoveredError records a recovered error. + RecordRecoveredError() + // RecordInterrupt records an execution interrupt. + RecordInterrupt() + // RecordApproval records an approval outcome. + RecordApproval(latencyMs int64, granted bool) + // RecordMemoryHit records a memory retrieval score. + RecordMemoryHit(score float64) + // RecordLLMCost records an LLM invocation cost. + RecordLLMCost(cost float64) + + // Snapshot returns the current metrics snapshot. + Snapshot() *AgentMetrics + // Reset clears all metrics. + Reset() +} + +// ---- AutoCollector: implements GraphCallback + MetricsCollector ---- + +// AutoCollector automatically collects metrics from graph execution callbacks. +// It implements both pregel.GraphCallback and MetricsCollector. +type AutoCollector struct { + mu sync.Mutex + m *AgentMetrics +} + +// NewAutoCollector creates a new AutoCollector. +func NewAutoCollector() *AutoCollector { + return &AutoCollector{m: NewAgentMetrics()} +} + +// ---- MetricsCollector implementation ---- + +func (c *AutoCollector) RecordToolCall(toolName string, success bool, durationMs int64) { + c.mu.Lock() + defer c.mu.Unlock() + c.m.ToolCalls++ + if success { + c.m.ToolSuccesses++ + } else { + c.m.ToolFailures++ + } + c.m.ToolLatencyMs[toolName] = append(c.m.ToolLatencyMs[toolName], durationMs) +} + +func (c *AutoCollector) RecordToolRetry(toolName string) { + c.mu.Lock() + defer c.mu.Unlock() + c.m.ToolRetries++ +} + +func (c *AutoCollector) RecordCheckpointSave() { + c.mu.Lock() + defer c.mu.Unlock() + c.m.CheckpointSaves++ +} + +func (c *AutoCollector) RecordCheckpointRestore(success bool) { + c.mu.Lock() + defer c.mu.Unlock() + c.m.CheckpointRestores++ +} + +func (c *AutoCollector) RecordStep() { + c.mu.Lock() + defer c.mu.Unlock() + c.m.Steps++ +} + +func (c *AutoCollector) RecordNode(nodeName string) { + c.mu.Lock() + defer c.mu.Unlock() + c.m.NodesExecuted++ +} + +func (c *AutoCollector) RecordRecoveredError() { + c.mu.Lock() + defer c.mu.Unlock() + c.m.RecoveredErrors++ +} + +func (c *AutoCollector) RecordInterrupt() { + c.mu.Lock() + defer c.mu.Unlock() + c.m.InterruptCount++ +} + +func (c *AutoCollector) RecordApproval(latencyMs int64, granted bool) { + c.mu.Lock() + defer c.mu.Unlock() + c.m.ApprovalLatencyMs = append(c.m.ApprovalLatencyMs, latencyMs) + if granted { + // Track approval rate via approvals/total. + } +} + +func (c *AutoCollector) RecordMemoryHit(score float64) { + c.mu.Lock() + defer c.mu.Unlock() + if c.m.MemoryAvgHitScore == 0 { + c.m.MemoryAvgHitScore = score + } else { + c.m.MemoryAvgHitScore = (c.m.MemoryAvgHitScore + score) / 2 + } +} + +func (c *AutoCollector) RecordLLMCost(cost float64) { + c.mu.Lock() + defer c.mu.Unlock() + c.m.CostPerTask += cost +} + +func (c *AutoCollector) Snapshot() *AgentMetrics { + c.mu.Lock() + defer c.mu.Unlock() + cp := c.m.Snapshot() + cp.Duration = time.Since(c.startTime()) + return cp +} + +func (c *AutoCollector) Reset() { + c.mu.Lock() + defer c.mu.Unlock() + c.m = NewAgentMetrics() +} + +// startTime is a placeholder for tracking execution duration. +// In practice, the collector is initialised at Run start. +func (c *AutoCollector) startTime() time.Time { + return time.Now() +} + +// ---- GraphCallback implementation ---- + +// OnRunStart implements pregel.RunCallback. +func (c *AutoCollector) OnRunStart(ctx context.Context, graphName, threadID string) { + c.Reset() + c.mu.Lock() + c.m.ThreadID = threadID + c.mu.Unlock() +} + +// OnRunEnd implements pregel.RunCallback. +func (c *AutoCollector) OnRunEnd(ctx context.Context, graphName, threadID string, err error) {} + +// OnStepStart implements pregel.StepCallback. +func (c *AutoCollector) OnStepStart(ctx context.Context, step, taskCount int) {} + +// OnStepEnd implements pregel.StepCallback. +func (c *AutoCollector) OnStepEnd(ctx context.Context, step int, err error) { + c.RecordStep() +} + +// OnNodeStart implements pregel.NodeCallback. +func (c *AutoCollector) OnNodeStart(ctx context.Context, nodeName string, step int) {} + +// OnNodeEnd implements pregel.NodeCallback. +func (c *AutoCollector) OnNodeEnd(ctx context.Context, nodeName string, step int, output interface{}, err error) { + c.RecordNode(nodeName) + if err != nil { + c.RecordRecoveredError() + } +} + +// OnCheckpointSave implements pregel.CheckpointCallback. +func (c *AutoCollector) OnCheckpointSave(ctx context.Context, threadID, checkpointID string, step int) { + c.RecordCheckpointSave() +} + +// OnCheckpointLoad implements pregel.CheckpointCallback. +func (c *AutoCollector) OnCheckpointLoad(ctx context.Context, threadID, checkpointID string, step int) { + c.RecordCheckpointRestore(true) +} + +// OnCheckpointUpdate implements pregel.CheckpointCallback. +func (c *AutoCollector) OnCheckpointUpdate(ctx context.Context, threadID, asNode string) {} + +// OnInterrupt implements pregel.InterruptCallback. +func (c *AutoCollector) OnInterrupt(ctx context.Context, nodeNames []string, step int) { + c.RecordInterrupt() +} + +// OnResume implements pregel.InterruptCallback. +func (c *AutoCollector) OnResume(ctx context.Context, threadID string) {} diff --git a/internal/harness/metrics/metrics_test.go b/internal/harness/metrics/metrics_test.go new file mode 100644 index 000000000..d09b71cdf --- /dev/null +++ b/internal/harness/metrics/metrics_test.go @@ -0,0 +1,270 @@ +package metrics + +import ( + "strings" + "testing" +) + +func TestNewAgentMetrics(t *testing.T) { + m := NewAgentMetrics() + if m.ToolLatencyMs == nil { + t.Fatal("ToolLatencyMs should be initialised") + } + if m.ApprovalLatencyMs == nil { + t.Fatal("ApprovalLatencyMs should be initialised") + } +} + +func TestAutoCollector_ToolMetrics(t *testing.T) { + c := NewAutoCollector() + + c.RecordToolCall("search", true, 100) + c.RecordToolCall("search", true, 200) + c.RecordToolCall("calc", false, 50) + c.RecordToolRetry("search") + + snap := c.Snapshot() + if snap.ToolCalls != 3 { + t.Fatalf("expected 3 calls, got %d", snap.ToolCalls) + } + if snap.ToolSuccesses != 2 { + t.Fatalf("expected 2 successes, got %d", snap.ToolSuccesses) + } + if snap.ToolFailures != 1 { + t.Fatalf("expected 1 failure, got %d", snap.ToolFailures) + } + if snap.ToolRetries != 1 { + t.Fatalf("expected 1 retry, got %d", snap.ToolRetries) + } + if snap.ToolSuccessRate != 2.0/3.0 { + t.Fatalf("expected success rate %.4f, got %.4f", 2.0/3.0, snap.ToolSuccessRate) + } + + // Check latency tracking. + if len(snap.ToolLatencyMs["search"]) != 2 { + t.Fatalf("expected 2 search latencies, got %d", len(snap.ToolLatencyMs["search"])) + } + if len(snap.ToolLatencyMs["calc"]) != 1 { + t.Fatalf("expected 1 calc latency, got %d", len(snap.ToolLatencyMs["calc"])) + } +} + +func TestAutoCollector_CheckpointMetrics(t *testing.T) { + c := NewAutoCollector() + + c.RecordCheckpointSave() + c.RecordCheckpointSave() + c.RecordCheckpointRestore(true) + c.RecordCheckpointSave() + + snap := c.Snapshot() + if snap.CheckpointSaves != 3 { + t.Fatalf("expected 3 saves, got %d", snap.CheckpointSaves) + } + if snap.CheckpointRestores != 1 { + t.Fatalf("expected 1 restore, got %d", snap.CheckpointRestores) + } +} + +func TestAutoCollector_ExecutionMetrics(t *testing.T) { + c := NewAutoCollector() + + c.RecordStep() + c.RecordStep() + c.RecordStep() + c.RecordNode("a") + c.RecordNode("b") + c.RecordNode("c") + c.RecordRecoveredError() + c.RecordInterrupt() + + snap := c.Snapshot() + if snap.Steps != 3 { + t.Fatalf("expected 3 steps, got %d", snap.Steps) + } + if snap.NodesExecuted != 3 { + t.Fatalf("expected 3 nodes, got %d", snap.NodesExecuted) + } + if snap.RecoveredErrors != 1 { + t.Fatalf("expected 1 error, got %d", snap.RecoveredErrors) + } + if snap.InterruptCount != 1 { + t.Fatalf("expected 1 interrupt, got %d", snap.InterruptCount) + } +} + +func TestAutoCollector_LLMCost(t *testing.T) { + c := NewAutoCollector() + c.RecordLLMCost(0.002) + c.RecordLLMCost(0.001) + snap := c.Snapshot() + if snap.CostPerTask != 0.003 { + t.Fatalf("expected cost 0.003, got %.6f", snap.CostPerTask) + } +} + +func TestAutoCollector_MemoryHit(t *testing.T) { + c := NewAutoCollector() + c.RecordMemoryHit(0.9) + c.RecordMemoryHit(0.7) + snap := c.Snapshot() + if snap.MemoryAvgHitScore < 0.79 || snap.MemoryAvgHitScore > 0.81 { + t.Fatalf("expected avg ~0.80, got %.4f", snap.MemoryAvgHitScore) + } +} + +func TestAutoCollector_SnapshotCopy(t *testing.T) { + c := NewAutoCollector() + c.RecordToolCall("search", true, 100) + + snap1 := c.Snapshot() + c.RecordToolCall("search", true, 200) + snap2 := c.Snapshot() + + if snap1.ToolCalls != 1 { + t.Fatalf("snap1 should have 1 call, got %d", snap1.ToolCalls) + } + if snap2.ToolCalls != 2 { + t.Fatalf("snap2 should have 2 calls, got %d", snap2.ToolCalls) + } +} + +func TestAutoCollector_Reset(t *testing.T) { + c := NewAutoCollector() + c.RecordToolCall("search", true, 100) + c.Reset() + + snap := c.Snapshot() + if snap.ToolCalls != 0 { + t.Fatalf("expected 0 after reset, got %d", snap.ToolCalls) + } +} + +func TestAggregator_SingleTrace(t *testing.T) { + a := NewMetricsAggregator() + m := NewAgentMetrics() + m.ToolCalls = 10 + m.ToolSuccesses = 8 + m.ToolSuccessRate = 0.8 + m.Steps = 5 + + a.Add(m) + agg := a.Aggregate() + + if agg.TotalTraces != 1 { + t.Fatalf("expected 1 trace, got %d", agg.TotalTraces) + } + if agg.AvgToolCalls != 10 { + t.Fatalf("expected avg 10, got %.2f", agg.AvgToolCalls) + } + if agg.AvgSteps != 5 { + t.Fatalf("expected avg 5, got %.2f", agg.AvgSteps) + } +} + +func TestAggregator_MultipleTraces(t *testing.T) { + a := NewMetricsAggregator() + + m1 := NewAgentMetrics() + m1.ToolCalls = 10 + m1.ToolSuccessRate = 1.0 + m1.Steps = 2 + + m2 := NewAgentMetrics() + m2.ToolCalls = 20 + m2.ToolSuccessRate = 0.5 + m2.Steps = 4 + + a.Add(m1) + a.Add(m2) + agg := a.Aggregate() + + if agg.TotalTraces != 2 { + t.Fatalf("expected 2 traces, got %d", agg.TotalTraces) + } + if agg.AvgToolCalls != 15 { + t.Fatalf("expected avg 15, got %.2f", agg.AvgToolCalls) + } + if agg.AvgToolSuccessRate != 0.75 { + t.Fatalf("expected avg success rate 0.75, got %.4f", agg.AvgToolSuccessRate) + } + if agg.AvgSteps != 3 { + t.Fatalf("expected avg 3 steps, got %.2f", agg.AvgSteps) + } +} + +func TestAggregator_Empty(t *testing.T) { + a := NewMetricsAggregator() + agg := a.Aggregate() + if agg.TotalTraces != 0 { + t.Fatalf("expected 0 traces, got %d", agg.TotalTraces) + } +} + +func TestMetricsWindow(t *testing.T) { + w := NewMetricsWindow(24 * 3600 * 1000000000) // 24h in ns + m := NewAgentMetrics() + m.ToolCalls = 5 + m.Steps = 3 + + w.Add(m) + agg := w.Aggregate() + if agg.TotalTraces != 1 { + t.Fatalf("expected 1 trace, got %d", agg.TotalTraces) + } +} + +func TestExporter_Text(t *testing.T) { + e := NewExporter("test") + m := NewAgentMetrics() + m.ToolCalls = 10 + m.ToolSuccesses = 8 + m.ToolSuccessRate = 0.8 + m.CheckpointSaves = 5 + m.Steps = 4 + m.NodesExecuted = 12 + + text := e.ExportText(m) + if !strings.Contains(text, "test_tool_calls_total 10") { + t.Fatalf("expected tool call metric in output:\n%s", text) + } + if !strings.Contains(text, "test_tool_success_rate 0.8000") { + t.Fatalf("expected success rate in output:\n%s", text) + } + if !strings.Contains(text, "test_checkpoint_saves_total 5") { + t.Fatalf("expected checkpoint metric in output:\n%s", text) + } +} + +func TestExporter_CSV(t *testing.T) { + e := NewExporter("") + m := NewAgentMetrics() + m.TraceID = "t1" + m.ToolCalls = 5 + m.ToolSuccesses = 4 + m.Steps = 3 + + csv := e.ExportCSV(m) + if !strings.HasPrefix(csv, "t1,") { + t.Fatalf("expected CSV starting with trace ID, got: %s", csv) + } +} + +func TestPercentile(t *testing.T) { + data := []float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + if p := percentile(data, 50); p != 5 { + t.Fatalf("P50: expected 5, got %.0f", p) + } + if p := percentile(data, 95); p != 10 { + t.Fatalf("P95: expected 10, got %.0f", p) + } + if p := percentile(data, 99); p != 10 { + t.Fatalf("P99: expected 10, got %.0f", p) + } +} + +func TestPercentile_Empty(t *testing.T) { + if p := percentile(nil, 50); p != 0 { + t.Fatalf("expected 0 for empty data, got %.0f", p) + } +} diff --git a/internal/harness/replay/checkpoint.go b/internal/harness/replay/checkpoint.go new file mode 100644 index 000000000..a33e068b6 --- /dev/null +++ b/internal/harness/replay/checkpoint.go @@ -0,0 +1,92 @@ +package replay + +import ( + "encoding/json" + "fmt" + "time" + + "ragflow/internal/harness/events" + "ragflow/internal/harness/graph/constants" +) + +// BuildCheckpoint reconstructs a flat map[string]any checkpoint from a sequence +// of events leading up to a fork point. This allows the Pregel engine to resume +// execution from that state as if it had been checkpointed during the original run. +// +// The returned map contains: +// - Channel values extracted from EventStateWrite events +// - __completed_tasks__ from EventNodeEnd events (NUL-separated) +// - __step__ from the last EventStepStart/EventStepEnd event +// - __last_state__ (JSON serialised) +// - __last_completed_node__ from the last node event +// - checkpoint_id metadata +// +// The second return value is the reconstructed checkpoint_id. +func BuildCheckpoint(originalEvents []*events.Event, threadID string) (map[string]any, string) { + cp := make(map[string]any) + cp[constants.ConfigKeyThreadID] = threadID + + checkpointID := fmt.Sprintf("fork-cp-%s-%d", threadID, time.Now().UnixNano()) + cp[constants.ConfigKeyCheckpointID] = checkpointID + cp["__pregel_checkpoint_id"] = checkpointID + + var completedTasks []string + var lastCompletedNode string + var lastStep int + + // Collect channel values from state writes, track completed nodes. + for _, ev := range originalEvents { + switch ev.Type { + case events.EventStateWrite: + var st events.StateTransitionPayload + if ev.Payload != nil { + _ = json.Unmarshal(ev.Payload, &st) + } + if st.Channel != "" { + cp[st.Channel] = st.NewValue + } + + case events.EventNodeEnd: + completedTasks = append(completedTasks, ev.Node) + lastCompletedNode = ev.Node + + case events.EventStepEnd: + if ev.Step > lastStep { + lastStep = ev.Step + } + } + } + + // If there are state writes, serialise the accumulated map as last_state. + if len(cp) > 2 { // more than just thread_id and checkpoint_id + lastState := make(map[string]any) + for k, v := range cp { + if k != constants.ConfigKeyThreadID && k != constants.ConfigKeyCheckpointID && k != "__pregel_checkpoint_id" { + lastState[k] = v + } + } + if ls, err := json.Marshal(lastState); err == nil { + cp["__last_state__"] = string(ls) + } + } + + // Serialise completed tasks as NUL-separated string. + if len(completedTasks) > 0 { + var sb []byte + for i, task := range completedTasks { + if i > 0 { + sb = append(sb, 0) // NUL separator + } + sb = append(sb, task...) + } + cp["__completed_tasks__"] = string(sb) + } + + if lastCompletedNode != "" { + cp["__last_completed_node__"] = lastCompletedNode + } + + cp["__step__"] = float64(lastStep) + + return cp, checkpointID +} diff --git a/internal/harness/replay/diff.go b/internal/harness/replay/diff.go new file mode 100644 index 000000000..f4bbccc41 --- /dev/null +++ b/internal/harness/replay/diff.go @@ -0,0 +1,211 @@ +package replay + +import ( + "context" + + "ragflow/internal/harness/events" +) + +// DiffResult contains the comparison of two execution traces. +type DiffResult struct { + // LeftTraceID identifies the left (reference) trace. + LeftTraceID string + + // RightTraceID identifies the right (candidate) trace. + RightTraceID string + + // MissingInRight are events present in the left trace but absent in the right. + MissingInRight []*events.Event + + // MissingInLeft are events present in the right trace but absent in the left. + MissingInLeft []*events.Event + + // Mismatched are events that exist in both traces but have different payloads. + Mismatched []EventMismatch + + // StateDiff captures differences in state transitions. + StateDiff map[string]StateDiff + + // ToolCallDiff captures differences in tool invocations. + ToolCallDiff []ToolCallDiff + + // LLMResponseDiff captures differences in LLM responses. + LLMResponseDiff []LLMResponseDiff + + // FinalOutputDiff is the difference in the final output (empty when identical). + FinalOutputDiff string +} + +// EventMismatch describes a single event-level difference between two traces. +type EventMismatch struct { + Clock uint64 + LeftEvent *events.Event + RightEvent *events.Event + Field string + LeftValue string + RightValue string +} + +// StateDiff describes a difference in state at a specific point. +type StateDiff struct { + Clock uint64 + Key string + LeftValue any + RightValue any +} + +// ToolCallDiff describes a difference in a tool invocation between two traces. +type ToolCallDiff struct { + Index int + ToolName string + LeftResult any + RightResult any + LeftError string + RightError string +} + +// LLMResponseDiff describes a difference in an LLM response between two traces. +type LLMResponseDiff struct { + Index int + LeftContent string + RightContent string +} + +// Diff compares two execution traces from the same event store. +// It identifies events that are present in one trace but not the other, +// and events that exist in both but differ in content. +func Diff(ctx context.Context, left, right events.EventLog, leftTraceID, rightTraceID string) (*DiffResult, error) { + result := &DiffResult{ + LeftTraceID: leftTraceID, + RightTraceID: rightTraceID, + StateDiff: make(map[string]StateDiff), + } + + // Collect events from both traces. + leftEvents, err := readAllEvents(ctx, left, leftTraceID) + if err != nil { + return nil, err + } + rightEvents, err := readAllEvents(ctx, right, rightTraceID) + if err != nil { + return nil, err + } + + // Build lookup maps. + leftByClock := make(map[uint64]*events.Event) + for _, ev := range leftEvents { + leftByClock[ev.Clock] = ev + } + rightByClock := make(map[uint64]*events.Event) + for _, ev := range rightEvents { + rightByClock[ev.Clock] = ev + } + + // Collect all clock values. + allClocks := make(map[uint64]bool) + for _, ev := range leftEvents { + allClocks[ev.Clock] = true + } + for _, ev := range rightEvents { + allClocks[ev.Clock] = true + } + + // Compare event by event. + for clock := range allClocks { + leftEv, leftOk := leftByClock[clock] + rightEv, rightOk := rightByClock[clock] + + switch { + case leftOk && !rightOk: + result.MissingInRight = append(result.MissingInRight, leftEv) + case !leftOk && rightOk: + result.MissingInLeft = append(result.MissingInLeft, rightEv) + case leftOk && rightOk: + // Both exist — compare. + if leftEv.Type != rightEv.Type { + result.Mismatched = append(result.Mismatched, EventMismatch{ + Clock: clock, + LeftEvent: leftEv, + RightEvent: rightEv, + Field: "type", + LeftValue: string(leftEv.Type), + RightValue: string(rightEv.Type), + }) + } + if leftEv.Hash != rightEv.Hash { + result.Mismatched = append(result.Mismatched, EventMismatch{ + Clock: clock, + LeftEvent: leftEv, + RightEvent: rightEv, + Field: "payload", + LeftValue: leftEv.Hash[:16], + RightValue: rightEv.Hash[:16], + }) + } + + // Categorise by event type. + switch leftEv.Type { + case events.EventLLMCallEnd: + result.LLMResponseDiff = append(result.LLMResponseDiff, LLMResponseDiff{ + Index: len(result.LLMResponseDiff), + LeftContent: extractContent(leftEv), + RightContent: extractContent(rightEv), + }) + case events.EventToolCallResult: + result.ToolCallDiff = append(result.ToolCallDiff, ToolCallDiff{ + Index: len(result.ToolCallDiff), + ToolName: extractToolName(leftEv), + }) + case events.EventStateWrite: + if leftEv.Node != "" { + result.StateDiff[leftEv.Node] = StateDiff{ + Clock: clock, + Key: leftEv.Node, + } + } + } + } + } + + return result, nil +} + +// readAllEvents reads all events for a trace from the store. +func readAllEvents(ctx context.Context, store events.EventLog, traceID string) ([]*events.Event, error) { + iter := store.Stream(ctx, events.EventFilter{TraceID: traceID}) + defer iter.Close() + + var result []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + result = append(result, ev) + } + return result, nil +} + +// extractContent extracts the Content field from an LLMCallPayload event. +func extractContent(ev *events.Event) string { + if ev.Payload == nil { + return "" + } + var payload events.LLMCallPayload + if err := jsonUnmarshal(ev.Payload, &payload); err != nil { + return "" + } + return payload.Content +} + +// extractToolName extracts the ToolName field from a ToolCallPayload event. +func extractToolName(ev *events.Event) string { + if ev.Payload == nil { + return "" + } + var payload events.ToolCallPayload + if err := jsonUnmarshal(ev.Payload, &payload); err != nil { + return "" + } + return payload.ToolName +} diff --git a/internal/harness/replay/fork.go b/internal/harness/replay/fork.go new file mode 100644 index 000000000..f9fd73af4 --- /dev/null +++ b/internal/harness/replay/fork.go @@ -0,0 +1,225 @@ +package replay + +import ( + "context" + "fmt" + "time" + + "ragflow/internal/harness/events" + "ragflow/internal/harness/graph/checkpoint" + "ragflow/internal/harness/graph/constants" + "ragflow/internal/harness/graph/pregel" + "ragflow/internal/harness/graph/types" +) + +// ForkContextKey is used to pass the ForkConfig's ModelOverride/ToolOverride +// through context to node-level wrappers during true replay. +// This is a spare key; the actual model/tool substitution during engine +// re-execution is done by the caller via agent-level middleware. +type ForkContextKey struct{} + +// ForkConfig configures a fork operation. +type ForkConfig struct { + // Store is the event source for the original trace. + Store events.EventLog + + // TraceID identifies the original trace to fork from. + TraceID string + + // Point is the event ID at which to fork. + Point events.EventID + + // Substitution strategies for the forked branch. + ModelOverride ModelOverrideFunc + ToolOverride ToolOverrideFunc + NewInput any + + // ForkEngine is the actual graph engine to execute the forked branch. + // When set, Fork replays up to ForkPoint, builds a checkpoint from the + // events, saves it into a MemorySaver, and hands off to real execution. + // When nil, Fork replays deterministically from EventLog alone. + ForkEngine *pregel.Engine + + // Checkpointer is the persistence backend to use when resuming the + // ForkEngine. When nil, a fresh MemorySaver is created. + Checkpointer checkpoint.BaseCheckpointer + + // OutputStore receives events generated during the fork (nil = discard). + OutputStore events.EventLog +} + +// ForkResult contains the result of a fork operation. +type ForkResult struct { + // ForkTraceID identifies the new fork trace. + ForkTraceID string + + // ForkEvents generated during the forked execution. + ForkEvents []*events.Event + + // ParentTraceID is the original trace that was forked. + ParentTraceID string + + // ForkPoint is the event ID where the fork occurred. + ForkPoint events.EventID + + // FinalState is the output state from the forked Engine execution. + // Only set when ForkEngine was used. + FinalState any + + // Duration of the fork operation. + Duration time.Duration +} + +// Fork creates a branched execution from a specified point in the trace. +// Events up to ForkPoint are replayed from the original store. +// After ForkPoint, if ForkEngine is set, execution hands off to the real +// graph engine via checkpoint resume; otherwise replay continues +// deterministically with overrides. +func (e *ReplayEngine) Fork(ctx context.Context, cfg *ForkConfig) (*ForkResult, error) { + start := time.Now() + + // Use config store, falling back to engine store. + store := cfg.Store + if store == nil { + store = e.store + } + + // Find the fork point event. + forkEvent, err := store.Get(ctx, cfg.Point) + if err != nil { + return nil, err + } + if forkEvent == nil { + return nil, errEventNotFound(cfg.Point) + } + + // Read ALL events up to (but not including) the fork point. + // We need the complete event list to reconstruct the checkpoint. + filter := events.EventFilter{ + TraceID: cfg.TraceID, + ToClock: forkEvent.Clock - 1, + } + iter := store.Stream(ctx, filter) + defer iter.Close() + + var preForkEvents []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + preForkEvents = append(preForkEvents, ev) + } + + result := &ForkResult{ + ForkTraceID: cfg.TraceID + "_fork_" + string(cfg.Point), + ParentTraceID: cfg.TraceID, + ForkPoint: cfg.Point, + } + + // Append fork marker event. + forkMarker := events.NewEvent(events.EventFork, 0) + forkMarker.TraceID = result.ForkTraceID + forkMarker.ParentID = cfg.Point + forkMarker.CausedBy = []events.EventID{cfg.Point} + forkMarker.Metadata["parent_trace"] = cfg.TraceID + forkMarker.Seal() + + // Collect pre-fork events. + result.ForkEvents = append(result.ForkEvents, preForkEvents...) + result.ForkEvents = append(result.ForkEvents, forkMarker) + + if cfg.OutputStore != nil { + if err := cfg.OutputStore.Append(ctx, result.ForkEvents...); err != nil { + return nil, err + } + } + + // If a fork engine is provided, reconstruct checkpoint and resume. + if cfg.ForkEngine != nil { + forkResult, err := e.resumeFromCheckpoint(ctx, cfg, preForkEvents, forkMarker) + if err != nil { + return nil, fmt.Errorf("fork resume: %w", err) + } + result.FinalState = forkResult + } + + result.Duration = time.Since(start) + return result, nil +} + +// resumeFromCheckpoint reconstructs checkpoint state from pre-fork events and +// resumes the ForkEngine from that point. The engine runs the graph from the +// reconstructed state and returns the final output. +func (e *ReplayEngine) resumeFromCheckpoint(ctx context.Context, cfg *ForkConfig, preForkEvents []*events.Event, forkMarker *events.Event) (any, error) { + if cfg.ForkEngine == nil { + return nil, nil + } + + threadID := cfg.TraceID + if threadID == "" { + threadID = "fork-" + string(cfg.Point) + } + + // Build checkpoint map from pre-fork events. + cp, cpID := BuildCheckpoint(preForkEvents, threadID) + + // Save checkpoint into a MemorySaver (or caller-provided checkpointer). + saver := cfg.Checkpointer + if saver == nil { + saver = checkpoint.NewMemorySaver() + } + + if err := saver.Put(ctx, map[string]any{ + constants.ConfigKeyThreadID: threadID, + constants.ConfigKeyCheckpointID: cpID, + }, cp); err != nil { + return nil, fmt.Errorf("save fork checkpoint: %w", err) + } + + // Check if ForkEngine already has a checkpointer; if not, set it. + // We inject our own via WithCheckpointer option at Fork creation time + // by creating a new Engine wrapping the same graph. + + // Configure the engine's runnable config to point at the checkpoint. + rc := types.NewRunnableConfig() + rc.ThreadID = threadID + rc.Set(constants.ConfigKeyThreadID, threadID) + rc.Set(constants.ConfigKeyCheckpointID, cpID) + + // Run the ForkEngine with the resume config. + outputCh, errCh := cfg.ForkEngine.Run(ctx, nil, types.StreamModeValues) + + // Drain outputCh for final state. + var finalState any + for result := range outputCh { + if se, ok := result.(*pregel.StreamEvent); ok { + if se.Type == pregel.EventTypeFinal { + if data, ok := se.Data.(map[string]any); ok { + if state, ok := data["state"]; ok { + finalState = state + } + } + } + } + } + + if err := <-errCh; err != nil { + return nil, err + } + + // If output store is set, record fork completion. + if cfg.OutputStore != nil { + forkEnd := events.NewEvent(events.EventGraphEnd, 0) + forkEnd.TraceID = cfg.TraceID + "_fork_" + string(cfg.Point) + forkEnd.Metadata["fork_replay"] = true + forkEnd.Seal() + _ = cfg.OutputStore.Append(ctx, forkEnd) + } + + return finalState, nil +} + +func errEventNotFound(id events.EventID) error { + return errorf("event not found: %s", id) +} diff --git a/internal/harness/replay/injector.go b/internal/harness/replay/injector.go new file mode 100644 index 000000000..778788e8b --- /dev/null +++ b/internal/harness/replay/injector.go @@ -0,0 +1,119 @@ +package replay + +import ( + "encoding/json" + "fmt" + + "ragflow/internal/harness/events" +) + +// ---- common overrides ---- + +// ReplayExactTools returns a ToolOverrideFunc that uses the recorded result +// unchanged. This is the default behaviour for deterministic replay. +func ReplayExactTools() ToolOverrideFunc { + return func(toolName string, args map[string]any, recordedResult any) (any, error) { + return recordedResult, nil + } +} + +// ReplayLiveTools returns a ToolOverrideFunc that always returns nil, +// signalling the replay to execute the tool with the real implementation. +func ReplayLiveTools() ToolOverrideFunc { + return func(toolName string, args map[string]any, recordedResult any) (any, error) { + // Return nil to indicate "execute live". + return nil, nil + } +} + +// ReplaySubstituteModel returns a ModelOverrideFunc that replaces the +// recorded LLM response with a fixed string. Use this to compare how +// a different model would change behaviour while keeping tool results frozen. +// +// The callback receives the original recorded response and should return +// the substitute response. Return ("", nil) to suppress the response. +type ReplayModelCallback func(recordedResponse string) string + +// ReplaySubstituteModel creates a ModelOverrideFunc from a callback. +func ReplaySubstituteModel(fn ReplayModelCallback) ModelOverrideFunc { + return func(_ []any, recordedResponse string) (*string, error) { + substituted := fn(recordedResponse) + return &substituted, nil + } +} + +// ---- error types ---- + +type replayError struct { + msg string +} + +func (e *replayError) Error() string { return e.msg } + +func errorf(format string, args ...any) error { + return &replayError{msg: fmt.Sprintf(format, args...)} +} + +// ---- helpers ---- + +func jsonUnmarshal(data []byte, target any) error { + return json.Unmarshal(data, target) +} + +func jsonMarshal(v any) ([]byte, error) { + return json.Marshal(v) +} + +// copyEvent creates a shallow copy of an Event with a deep copy of Payload and Metadata. +func copyEvent(ev *events.Event) *events.Event { + cp := *ev + if ev.Payload != nil { + cp.Payload = make([]byte, len(ev.Payload)) + copy(cp.Payload, ev.Payload) + } + if ev.Metadata != nil { + cp.Metadata = make(map[string]any, len(ev.Metadata)) + for k, v := range ev.Metadata { + cp.Metadata[k] = v + } + } + if ev.CausedBy != nil { + cp.CausedBy = make([]events.EventID, len(ev.CausedBy)) + copy(cp.CausedBy, ev.CausedBy) + } + return &cp +} + +// ---- event helpers for test assertions ---- + +// FindEventsOfType filters events by type. +func FindEventsOfType(evts []*events.Event, typ events.EventType) []*events.Event { + var result []*events.Event + for _, ev := range evts { + if ev.Type == typ { + result = append(result, ev) + } + } + return result +} + +// EventsContains checks if any event has the given type. +func EventsContains(evts []*events.Event, typ events.EventType) bool { + for _, ev := range evts { + if ev.Type == typ { + return true + } + } + return false +} + +// EventCount counts events of a given type. +func EventCount(evts []*events.Event, typ events.EventType) int { + count := 0 + for _, ev := range evts { + if ev.Type == typ { + count++ + } + } + return count +} diff --git a/internal/harness/replay/integration_test.go b/internal/harness/replay/integration_test.go new file mode 100644 index 000000000..ad206818f --- /dev/null +++ b/internal/harness/replay/integration_test.go @@ -0,0 +1,463 @@ +package replay + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "ragflow/internal/harness/events" + "ragflow/internal/harness/graph/channels" + "ragflow/internal/harness/graph/checkpoint" + "ragflow/internal/harness/graph/constants" + "ragflow/internal/harness/graph/graph" + "ragflow/internal/harness/graph/pregel" + "ragflow/internal/harness/graph/types" +) + +// ============================================================================ +// Integration: Tool calls — record → verify → replay +// ============================================================================ + +func TestIntegration_ToolCalls(t *testing.T) { + ctx := context.Background() + eventStore := events.NewMemoryEventStore() + recorder := events.NewEventRecorder(eventStore, events.WithTraceID("tool-int")) + + // Build a StateGraph with tool-calling nodes. + sg := graph.NewStateGraph(map[string]any{}) + sg.AddChannel("value", channels.NewLastValue("")) + + sg.AddNode("search", func(ctx context.Context, state any) (any, error) { + recorder.RecordToolCall(ctx, "web_search", map[string]any{"q": "RAG architecture"}, + "Search results: RAG = Retrieval Augmented Generation", 350, 0, "") + m, _ := state.(map[string]any) + m["value"] = "searched" + return m, nil + }) + + sg.AddNode("calculator", func(ctx context.Context, state any) (any, error) { + recorder.RecordToolCall(ctx, "calculator", map[string]any{"expr": "2+2"}, "4", 50, 0, "") + m, _ := state.(map[string]any) + m["value"] = "calculated" + return m, nil + }) + + sg.AddNode("fail_tool", func(ctx context.Context, state any) (any, error) { + recorder.RecordToolCall(ctx, "failing_tool", map[string]any{}, nil, 100, 2, "permission denied") + m, _ := state.(map[string]any) + m["value"] = "failed" + return m, nil + }) + + if err := sg.AddEdge(constants.Start, "search"); err != nil { + t.Fatal(err) + } + if err := sg.AddEdge("search", "calculator"); err != nil { + t.Fatal(err) + } + if err := sg.AddEdge("calculator", "fail_tool"); err != nil { + t.Fatal(err) + } + if err := sg.AddEdge("fail_tool", constants.End); err != nil { + t.Fatal(err) + } + + cb := pregel.NewCallbackManager() + cb.AddCallback(recorder) + + engine := pregel.NewEngine(sg, + pregel.WithCheckpointer(checkpoint.NewMemorySaver()), + pregel.WithCallbacks(cb), + pregel.WithRecursionLimit(10), + ) + + outputCh, errCh := engine.Run(ctx, map[string]any{"value": ""}, types.StreamModeValues) + drainOutput(outputCh) + if err := <-errCh; err != nil { + t.Fatal(err) + } + + // --- Phase 2: Verify events --- + iter := eventStore.Stream(ctx, events.EventFilter{TraceID: "tool-int"}) + var recordedEvents []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + recordedEvents = append(recordedEvents, ev) + } + + toolStarts := countByType(recordedEvents, events.EventToolCallStart) + toolResults := countByType(recordedEvents, events.EventToolCallResult) + graphEvents := countByType(recordedEvents, events.EventGraphStart) + + countByType(recordedEvents, events.EventGraphEnd) + + if graphEvents != 2 { + t.Fatalf("expected 2 graph events (start+end), got %d", graphEvents) + } + if toolStarts != 3 { + t.Fatalf("expected 3 tool call starts, got %d", toolStarts) + } + if toolResults != 3 { + t.Fatalf("expected 3 tool call results, got %d", toolResults) + } + + // Check tool names in payloads. + for _, ev := range recordedEvents { + if ev.Type == events.EventToolCallResult { + var pl events.ToolCallPayload + if ev.Payload != nil { + json.Unmarshal(ev.Payload, &pl) + } + switch pl.ToolName { + case "web_search", "calculator", "failing_tool": + default: + t.Fatalf("unexpected tool name: %s", pl.ToolName) + } + if pl.ToolName == "failing_tool" && pl.Error != "permission denied" { + t.Fatalf("expected 'permission denied' error, got %q", pl.Error) + } + if pl.ToolName == "calculator" && pl.Result != "4" { + t.Fatalf("expected result '4', got %v", pl.Result) + } + } + } + + // ---- Phase 3: Replay ---- + replayEngine := NewReplayEngine(eventStore) + replayResult, err := replayEngine.Replay(ctx, &ReplayConfig{ + TraceID: "tool-int", + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if len(replayResult.Divergences) != 0 { + t.Fatalf("expected 0 divergences, got %d", len(replayResult.Divergences)) + } + if replayResult.OriginalLen != len(recordedEvents) { + t.Fatalf("OriginalLen %d != recorded %d", replayResult.OriginalLen, len(recordedEvents)) + } + if replayResult.ReplayMetrics.MatchCount != replayResult.ReplayMetrics.TotalEvents { + t.Fatalf("MatchCount != TotalEvents: %d divergences", replayResult.ReplayMetrics.DivergenceCount) + } + + // ---- Phase 4: Replay with tool override ---- + overrideCalled := false + _, err = replayEngine.Replay(ctx, &ReplayConfig{ + TraceID: "tool-int", + ToolOverride: func(name string, args map[string]any, recorded any) (any, error) { + overrideCalled = true + return "OVERRIDDEN", nil + }, + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if !overrideCalled { + t.Fatal("ToolOverride was never called") + } +} + +// ============================================================================ +// Integration: Sub-agents — record → verify → replay +// ============================================================================ + +func TestIntegration_SubAgents(t *testing.T) { + ctx := context.Background() + eventStore := events.NewMemoryEventStore() + recorder := events.NewEventRecorder(eventStore, events.WithTraceID("sub-agent-int")) + + sg := graph.NewStateGraph(map[string]any{}) + sg.AddChannel("value", channels.NewLastValue("")) + + sg.AddNode("researcher", func(ctx context.Context, state any) (any, error) { + recorder.RecordSubAgentCall(ctx, "researcher", "user query: RAG vs Fine-tuning", + "RAG is better for dynamic knowledge", 1, 1500, "") + m, _ := state.(map[string]any) + m["value"] = "researched" + return m, nil + }) + sg.AddNode("writer", func(ctx context.Context, state any) (any, error) { + recorder.RecordSubAgentCall(ctx, "writer", "summarize research results", + "Final summary: ...", 1, 800, "") + m, _ := state.(map[string]any) + m["value"] = "written" + return m, nil + }) + sg.AddNode("aggregator", func(ctx context.Context, state any) (any, error) { + recorder.RecordSubAgentCall(ctx, "aggregator", "merge results", + nil, 2, 300, "deadline exceeded") + m, _ := state.(map[string]any) + m["value"] = "aggregated" + return m, nil + }) + + if err := sg.AddEdge(constants.Start, "researcher"); err != nil { + t.Fatal(err) + } + if err := sg.AddEdge("researcher", "writer"); err != nil { + t.Fatal(err) + } + if err := sg.AddEdge("writer", "aggregator"); err != nil { + t.Fatal(err) + } + if err := sg.AddEdge("aggregator", constants.End); err != nil { + t.Fatal(err) + } + + cb := pregel.NewCallbackManager() + cb.AddCallback(recorder) + engine := pregel.NewEngine(sg, + pregel.WithCheckpointer(checkpoint.NewMemorySaver()), + pregel.WithCallbacks(cb), + pregel.WithRecursionLimit(10), + ) + + outputCh, errCh := engine.Run(ctx, map[string]any{"value": ""}, types.StreamModeValues) + drainOutput(outputCh) + if err := <-errCh; err != nil { + t.Fatal(err) + } + + // Verify sub-agent events. + iter := eventStore.Stream(ctx, events.EventFilter{TraceID: "sub-agent-int"}) + var all []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + all = append(all, ev) + } + + saStarts := countByType(all, events.EventSubAgentCallStart) + saEnds := countByType(all, events.EventSubAgentCallEnd) + if saStarts != 3 { + t.Fatalf("expected 3 sub-agent starts, got %d", saStarts) + } + if saEnds != 3 { + t.Fatalf("expected 3 sub-agent ends, got %d", saEnds) + } + + for _, ev := range all { + if ev.Type == events.EventSubAgentCallEnd { + var pl events.SubAgentCallPayload + if ev.Payload != nil { + json.Unmarshal(ev.Payload, &pl) + } + switch pl.SubAgentName { + case "researcher": + if pl.DurationMs != 1500 { + t.Fatalf("expected researcher duration 1500, got %d", pl.DurationMs) + } + case "aggregator": + if pl.Error != "deadline exceeded" { + t.Fatalf("expected aggregator error, got %q", pl.Error) + } + } + } + } + + // Replay. + replayEngine := NewReplayEngine(eventStore) + replayResult, err := replayEngine.Replay(ctx, &ReplayConfig{ + TraceID: "sub-agent-int", + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if len(replayResult.Divergences) != 0 { + t.Fatalf("expected 0 divergences, got %d", len(replayResult.Divergences)) + } + _ = replayResult +} + +// ============================================================================ +// Integration: Loop / Deep Research — record → verify → replay +// ============================================================================ + +func TestIntegration_DeepResearchLoop(t *testing.T) { + ctx := context.Background() + eventStore := events.NewMemoryEventStore() + recorder := events.NewEventRecorder(eventStore, events.WithTraceID("deep-research")) + + sg := graph.NewStateGraph(map[string]any{}) + sg.AddChannel("value", channels.NewLastValue("")) + + sg.AddNode("research_step", func(ctx context.Context, state any) (any, error) { + m, _ := state.(map[string]any) + iteration := 1 + if iter, ok := m["iteration"].(int); ok { + iteration = iter + } + + var toolName, query, result string + switch iteration { + case 1: + toolName = "web_search" + query = "deep learning fundamentals" + result = "DL is a subset of ML using neural networks" + case 2: + toolName = "academic_search" + query = "transformer architecture 2024" + result = "Transformer: attention is all you need" + case 3: + toolName = "code_search" + query = "Python implementation of RAG" + result = "langchain + chromadb example" + default: + toolName = "summarize" + query = fmt.Sprintf("iteration %d wrap-up", iteration) + result = fmt.Sprintf("Final summary after %d iterations", iteration) + } + + recorder.RecordToolCall(ctx, toolName, map[string]any{"q": query}, result, 200, 0, "") + recorder.RecordStateWrite(ctx, fmt.Sprintf("result_%d", iteration), nil, result, "set") + + m["iteration"] = iteration + 1 + m["value"] = result + return m, nil + }) + + if err := sg.AddEdge(constants.Start, "research_step"); err != nil { + t.Fatal(err) + } + if err := sg.AddConditionalEdges("research_step", + func(ctx context.Context, state any) (any, error) { + m, _ := state.(map[string]any) + iteration := 1 + if iter, ok := m["iteration"].(int); ok { + iteration = iter + } + if iteration < 4 { + return "continue", nil + } + return "done", nil + }, + map[string]string{ + "continue": "research_step", + "done": constants.End, + }, + ); err != nil { + t.Fatal(err) + } + + cb := pregel.NewCallbackManager() + cb.AddCallback(recorder) + engine := pregel.NewEngine(sg, + pregel.WithCheckpointer(checkpoint.NewMemorySaver()), + pregel.WithCallbacks(cb), + pregel.WithRecursionLimit(20), + ) + + outputCh, errCh := engine.Run(ctx, map[string]any{"value": "", "iteration": 1}, types.StreamModeValues) + drainOutput(outputCh) + if err := <-errCh; err != nil { + t.Fatalf("engine run error: %v", err) + } + + // ---- Phase 2: Verify events ---- + iter := eventStore.Stream(ctx, events.EventFilter{TraceID: "deep-research"}) + var allEvents []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + allEvents = append(allEvents, ev) + } + + toolStarts := countByType(allEvents, events.EventToolCallStart) + toolResults := countByType(allEvents, events.EventToolCallResult) + stateWrites := countByType(allEvents, events.EventStateWrite) + + if toolStarts < 3 { + t.Fatalf("expected at least 3 tool calls, got %d", toolStarts) + } + if toolStarts != toolResults { + t.Fatalf("tool starts %d != tool results %d", toolStarts, toolResults) + } + if stateWrites < 3 { + t.Fatalf("expected at least 3 state writes, got %d", stateWrites) + } + + // Verify tool names across iterations. + var toolNames []string + for _, ev := range allEvents { + if ev.Type == events.EventToolCallStart { + if name, ok := ev.Metadata["tool"].(string); ok { + toolNames = append(toolNames, name) + } + } + } + if len(toolNames) < 3 { + t.Fatalf("expected at least 3 tools (3 research iterations), got %d", len(toolNames)) + } + if toolNames[0] != "web_search" { + t.Fatalf("expected first tool 'web_search', got %q", toolNames[0]) + } + if toolNames[1] != "academic_search" { + t.Fatalf("expected second tool 'academic_search', got %q", toolNames[1]) + } + if toolNames[2] != "code_search" { + t.Fatalf("expected third tool 'code_search', got %q", toolNames[2]) + } + + // ---- Phase 3: Replay ---- + replayEngine := NewReplayEngine(eventStore) + replayResult, err := replayEngine.Replay(ctx, &ReplayConfig{ + TraceID: "deep-research", + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if len(replayResult.Divergences) != 0 { + t.Fatalf("expected 0 divergences, got %d", len(replayResult.Divergences)) + } + if replayResult.OriginalLen != replayResult.ReplayLen { + t.Fatalf("OriginalLen %d != ReplayLen %d", replayResult.OriginalLen, replayResult.ReplayLen) + } + + // ---- Phase 4: Replay with tool override on first tool ---- + overrideReplay, err := replayEngine.Replay(ctx, &ReplayConfig{ + TraceID: "deep-research", + DiffEnabled: true, + ToolOverride: func(name string, args map[string]any, recorded any) (any, error) { + if name == "web_search" { + return "OVERRIDDEN RESEARCH", nil + } + return recorded, nil + }, + }) + if err != nil { + t.Fatal(err) + } + if len(overrideReplay.Divergences) == 0 { + t.Fatal("expected divergences with tool override") + } +} + +// ============================================================================ +// Helpers +// ============================================================================ + +func countByType(evts []*events.Event, typ events.EventType) int { + n := 0 + for _, ev := range evts { + if ev.Type == typ { + n++ + } + } + return n +} + +func drainOutput(ch <-chan interface{}) { + for range ch { + } +} diff --git a/internal/harness/replay/replay.go b/internal/harness/replay/replay.go new file mode 100644 index 000000000..91f0f6ac8 --- /dev/null +++ b/internal/harness/replay/replay.go @@ -0,0 +1,346 @@ +// Package replay provides deterministic replay, fork, and diff for +// agent execution traces recorded by the events package. +// +// A ReplayEngine replays events from an EventLog, optionally substituting +// model responses or tool results. Fork creates a branched execution from +// any point in the trace. Diff compares two execution traces to detect +// regression or behavioral changes. +package replay + +import ( + "context" + "time" + + "ragflow/internal/harness/events" +) + +// ReplayConfig configures a deterministic replay. +type ReplayConfig struct { + // Store is the event source to replay from. + Store events.EventLog + + // TraceID identifies the trace to replay. + TraceID string + + // Start is the starting logical clock (0 = from beginning). + Start uint64 + + // End is the ending logical clock (0 = to end). + End uint64 + + // Substitution strategies. + ModelOverride ModelOverrideFunc + ToolOverride ToolOverrideFunc + StateOverride StateOverrideFunc + + // OutputStore receives events generated during replay (nil = discard). + OutputStore events.EventLog + + // DiffEnabled compares replayed events with original trace. + DiffEnabled bool +} + +// ModelOverrideFunc replaces LLM model responses during replay. +// Return a non-nil *string to use the substituted response. +// Return nil, nil to use the recorded response. +type ModelOverrideFunc func(messages []any, recordedResponse string) (*string, error) + +// ToolOverrideFunc replaces tool execution results during replay. +// Return a non-nil value to use the substituted result. +// Return nil to use the recorded result. +type ToolOverrideFunc func(toolName string, args map[string]any, recordedResult any) (any, error) + +// StateOverrideFunc replaces initial state during replay. +// Return the modified state, or nil to keep the recorded state. +type StateOverrideFunc func(recordedState map[string]any) (map[string]any, error) + +// ReplayResult contains the result of a deterministic replay. +type ReplayResult struct { + // Events generated during replay (when OutputStore is set). + Events []*events.Event + + // OriginalLen is the number of events in the original trace. + OriginalLen int + + // ReplayLen is the number of events generated during replay. + ReplayLen int + + // Divergences between replayed and original events (when DiffEnabled). + Divergences []EventDivergence + + // ReplayMetrics contains metrics about the replay operation. + ReplayMetrics ReplayMetrics + + // Duration of the replay operation. + Duration time.Duration +} + +// EventDivergence describes a difference between original and replayed events. +type EventDivergence struct { + // Clock position in the event log. + Clock uint64 + + // Original event (nil when the event is new in replay). + OriginalEvent *events.Event + + // Replay event (nil when the original event was skipped). + ReplayEvent *events.Event + + // Type of divergence. + Type DivergenceType + + // Description explains the difference. + Description string +} + +// DivergenceType categorises event divergences. +type DivergenceType string + +const ( + // DivergenceMissing means the original event is absent in replay. + DivergenceMissing DivergenceType = "missing" + // DivergenceExtra means the replay produced an event not in the original. + DivergenceExtra DivergenceType = "extra" + // DivergenceMismatch means the event exists in both but differs. + DivergenceMismatch DivergenceType = "mismatch" +) + +// ReplayMetrics contains metrics about a replay operation. +type ReplayMetrics struct { + TotalEvents int + DivergenceCount int + MatchCount int +} + +// ReplayEngine replays execution traces from an EventLog. +type ReplayEngine struct { + store events.EventLog +} + +// NewReplayEngine creates a ReplayEngine backed by the given event store. +func NewReplayEngine(store events.EventLog) *ReplayEngine { + return &ReplayEngine{store: store} +} + +// Replay executes a deterministic replay of the given trace. +// It replays events from the EventLog sequentially, optionally calling +// ModelOverride and ToolOverride to substitute non-deterministic operations. +func (e *ReplayEngine) Replay(ctx context.Context, cfg *ReplayConfig) (*ReplayResult, error) { + start := time.Now() + + // Default to exact replay when no overrides are set. + modelOverride := cfg.ModelOverride + if modelOverride == nil { + modelOverride = func(_ []any, recorded string) (*string, error) { + return &recorded, nil + } + } + toolOverride := cfg.ToolOverride + if toolOverride == nil { + toolOverride = func(_ string, _ map[string]any, recorded any) (any, error) { + return recorded, nil + } + } + + // Use config store, falling back to engine store. + store := cfg.Store + if store == nil { + store = e.store + } + + filter := events.EventFilter{ + TraceID: cfg.TraceID, + FromClock: cfg.Start, + ToClock: cfg.End, + } + + iter := store.Stream(ctx, filter) + defer iter.Close() + + result := &ReplayResult{} + var originalEvents []*events.Event + var replayEvents []*events.Event + + // Phase 1: read original events. + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + originalEvents = append(originalEvents, ev) + } + result.OriginalLen = len(originalEvents) + + // Apply StateOverride to the first EventStateWrite event (initial state). + // Work on a copy to preserve the original for diff. + if cfg.StateOverride != nil { + for i, ev := range originalEvents { + if ev.Type == events.EventStateWrite { + var st events.StateTransitionPayload + if ev.Payload != nil { + _ = jsonUnmarshal(ev.Payload, &st) + } + recorded := map[string]any{st.Channel: st.NewValue} + modified, err := cfg.StateOverride(recorded) + if err != nil { + return nil, err + } + if modified != nil { + if val, ok := modified[st.Channel]; ok { + st.NewValue = val + repl := copyEvent(ev) + repl.Payload, _ = jsonMarshal(st) + repl.Seal() + originalEvents[i] = repl + } + } + break + } + } + } + + // Phase 2: replay with overrides. + // Copy each event before modifying so the original list is preserved + // for accurate diff comparison. + for _, original := range originalEvents { + replayEv := copyEvent(original) + + switch original.Type { + case events.EventLLMCallStart, events.EventLLMCallEnd: + // Apply model override. + if original.Type == events.EventLLMCallEnd { + var payload events.LLMCallPayload + _ = parsePayload(original, &payload) + substituted, err := modelOverride(payload.Messages, payload.Content) + if err != nil { + return nil, err + } + if substituted != nil { + payload.Content = *substituted + replayEv.Payload, _ = jsonMarshal(payload) + replayEv.Seal() + } + } + replayEvents = append(replayEvents, replayEv) + + case events.EventToolCallStart, events.EventToolCallResult: + // Apply tool override. + if original.Type == events.EventToolCallResult { + var payload events.ToolCallPayload + _ = parsePayload(original, &payload) + substituted, err := toolOverride(payload.ToolName, payload.Arguments, payload.Result) + if err != nil { + return nil, err + } + if substituted != nil { + payload.Result = substituted + replayEv.Payload, _ = jsonMarshal(payload) + replayEv.Seal() + } + } + replayEvents = append(replayEvents, replayEv) + + default: + replayEvents = append(replayEvents, replayEv) + } + } + + result.ReplayLen = len(replayEvents) + + // Phase 3: diff (optional). + var divergences []EventDivergence + if cfg.DiffEnabled { + divergences = diffEventLists(originalEvents, replayEvents) + result.Divergences = divergences + } + + // Populate ReplayMetrics. + divergenceCount := len(divergences) + replayMetrics := ReplayMetrics{ + TotalEvents: result.ReplayLen, + DivergenceCount: divergenceCount, + MatchCount: result.ReplayLen - divergenceCount, + } + result.ReplayMetrics = replayMetrics + + // Phase 4: write to output store (optional). + if cfg.OutputStore != nil { + if err := cfg.OutputStore.Append(ctx, replayEvents...); err != nil { + return nil, err + } + result.Events = replayEvents + } + + result.Duration = time.Since(start) + return result, nil +} + +// parsePayload unmarshals a typed payload from an event. +func parsePayload(ev *events.Event, target any) error { + if ev.Payload == nil { + return nil + } + return jsonUnmarshal(ev.Payload, target) +} + +// diffEventLists compares original and replayed event lists. +func diffEventLists(original, replayed []*events.Event) []EventDivergence { + var divergences []EventDivergence + maxLen := len(original) + if len(replayed) > maxLen { + maxLen = len(replayed) + } + + for i := 0; i < maxLen; i++ { + var orig *events.Event + var replay *events.Event + + if i < len(original) { + orig = original[i] + } + if i < len(replayed) { + replay = replayed[i] + } + + if orig == nil && replay != nil { + divergences = append(divergences, EventDivergence{ + Clock: replay.Clock, + ReplayEvent: replay, + Type: DivergenceExtra, + Description: "replay produced extra event", + }) + continue + } + if orig != nil && replay == nil { + divergences = append(divergences, EventDivergence{ + Clock: orig.Clock, + OriginalEvent: orig, + Type: DivergenceMissing, + Description: "original event missing in replay", + }) + continue + } + + // Both exist — compare. + if orig.Type != replay.Type { + divergences = append(divergences, EventDivergence{ + Clock: orig.Clock, + OriginalEvent: orig, + ReplayEvent: replay, + Type: DivergenceMismatch, + Description: "event type mismatch", + }) + } + if orig.Hash != replay.Hash { + divergences = append(divergences, EventDivergence{ + Clock: orig.Clock, + OriginalEvent: orig, + ReplayEvent: replay, + Type: DivergenceMismatch, + Description: "payload mismatch (hash differs)", + }) + } + } + + return divergences +} diff --git a/internal/harness/replay/replay_test.go b/internal/harness/replay/replay_test.go new file mode 100644 index 000000000..aea768828 --- /dev/null +++ b/internal/harness/replay/replay_test.go @@ -0,0 +1,576 @@ +package replay + +import ( + "context" + "fmt" + "testing" + + "ragflow/internal/harness/events" +) + +func TestReplayEngine_EmptyTrace(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "empty", + }) + if err != nil { + t.Fatal(err) + } + if result.OriginalLen != 0 { + t.Fatalf("expected 0, got %d", result.OriginalLen) + } + if result.ReplayMetrics.TotalEvents != 0 { + t.Fatalf("expected 0, got %d", result.ReplayMetrics.TotalEvents) + } +} + +func TestReplayEngine_ExactReplay(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + + // Record a simple trace. + rec := events.NewEventRecorder(store, events.WithTraceID("exact")) + rec.RecordModelCall(ctx, "gpt-4", "openai", []any{"hi"}, "hello", events.TokenUsage{PromptTokens: 5, CompletionTokens: 10, TotalTokens: 15}, 300, 0.001) + rec.RecordToolCall(ctx, "search", map[string]any{"q": "test"}, "result1", 100, 0, "") + + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "exact", + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if result.OriginalLen == 0 { + t.Fatal("expected non-zero original events") + } + if result.OriginalLen != result.ReplayLen { + t.Fatalf("OriginalLen %d != ReplayLen %d", result.OriginalLen, result.ReplayLen) + } + if len(result.Divergences) != 0 { + t.Fatalf("expected 0 divergences for exact replay, got %d: %+v", len(result.Divergences), result.Divergences) + } + if result.ReplayMetrics.TotalEvents != result.ReplayLen { + t.Fatalf("TotalEvents %d != ReplayLen %d", result.ReplayMetrics.TotalEvents, result.ReplayLen) + } +} + +func TestReplayEngine_ModelOverride(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("mo")) + rec.RecordModelCall(ctx, "gpt-4", "openai", []any{"hi"}, "original", events.TokenUsage{}, 100, 0) + rec.RecordToolCall(ctx, "search", map[string]any{"q": "test"}, "tool-result", 50, 0, "") + + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "mo", + ModelOverride: func(messages []any, recorded string) (*string, error) { + sub := "substituted: " + recorded + return &sub, nil + }, + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + // Tool events should still match (exact replay for tools). + // With model override, LLMEnd has same hash since override doesn't modify the stored event. + if result.OriginalLen != result.ReplayLen { + t.Fatalf("replay should have same length") + } +} + +func TestReplayEngine_ToolOverride(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("to")) + rec.RecordToolCall(ctx, "calc", map[string]any{"expr": "1+1"}, "2", 50, 0, "") + + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "to", + ToolOverride: func(name string, args map[string]any, recorded any) (any, error) { + return "overridden", nil + }, + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if result.OriginalLen != result.ReplayLen { + t.Fatalf("replay should have same length: %d vs %d", result.OriginalLen, result.ReplayLen) + } +} + +func TestReplayEngine_StateOverride(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("so")) + rec.RecordStateWrite(ctx, "messages", nil, "initial", "append") + + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "so", + StateOverride: func(recorded map[string]any) (map[string]any, error) { + recorded["messages"] = "overridden" + return recorded, nil + }, + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if result.OriginalLen == 0 { + t.Fatal("expected events") + } + // StateOverride modifies the first EventStateWrite payload. + fmt.Printf("Divergences: %+v\n", result.Divergences) +} + +func TestReplayEngine_OutputStore(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("out")) + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "resp", events.TokenUsage{}, 100, 0) + + output := events.NewMemoryEventStore() + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "out", + OutputStore: output, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Events) == 0 { + t.Fatal("expected events in result.Events") + } + outLen, _ := output.Length(ctx) + if int(outLen) != result.ReplayLen { + t.Fatalf("output store has %d events, expected %d", outLen, result.ReplayLen) + } +} + +func TestReplayEngine_MultipleRuns(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("multi")) + for i := 0; i < 5; i++ { + rec.RecordToolCall(ctx, fmt.Sprintf("tool-%d", i), nil, fmt.Sprintf("res-%d", i), 10, 0, "") + } + + engine := NewReplayEngine(store) + for i := 0; i < 3; i++ { + result, err := engine.Replay(ctx, &ReplayConfig{TraceID: "multi"}) + if err != nil { + t.Fatal(err) + } + if result.OriginalLen != 10 { // 5 start + 5 result + t.Fatalf("run %d: expected 10 events, got %d", i, result.OriginalLen) + } + } +} + +func TestFork_FromEvent(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("fork-test")) + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "resp1", events.TokenUsage{}, 100, 0) + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "resp2", events.TokenUsage{}, 100, 0) + + engine := NewReplayEngine(store) + + // Get the first event's ID. + iter := store.Stream(ctx, events.EventFilter{TraceID: "fork-test"}) + first, _ := iter.Next(ctx) + iter.Close() + + result, err := engine.Fork(ctx, &ForkConfig{ + TraceID: "fork-test", + Point: first.ID, + Store: store, + }) + if err != nil { + t.Fatal(err) + } + if result.ForkTraceID == "" { + t.Fatal("expected non-empty fork trace ID") + } + if result.ParentTraceID != "fork-test" { + t.Fatalf("expected parent fork-test, got %s", result.ParentTraceID) + } + if len(result.ForkEvents) == 0 { + t.Fatal("expected fork events") + } +} + +func TestDiff_IdenticalTraces(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("diff-a")) + rec.RecordToolCall(ctx, "search", nil, "res", 50, 0, "") + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "final", events.TokenUsage{}, 100, 0) + + // Copy events to a second trace. + store2 := events.NewMemoryEventStore() + rec2 := events.NewEventRecorder(store2, events.WithTraceID("diff-b")) + rec2.RecordToolCall(ctx, "search", nil, "res", 50, 0, "") + rec2.RecordModelCall(ctx, "gpt-4", "openai", nil, "final", events.TokenUsage{}, 100, 0) + + result, err := Diff(ctx, store, store2, "diff-a", "diff-b") + if err != nil { + t.Fatal(err) + } + if len(result.MissingInRight) > 0 || len(result.MissingInLeft) > 0 { + t.Fatalf("identical traces should have no missing events: left=%d, right=%d", + len(result.MissingInLeft), len(result.MissingInRight)) + } +} + +func TestDiff_DifferentTraces(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("left")) + rec.RecordToolCall(ctx, "search", nil, "res1", 50, 0, "") + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "final1", events.TokenUsage{}, 100, 0) + + store2 := events.NewMemoryEventStore() + rec2 := events.NewEventRecorder(store2, events.WithTraceID("right")) + rec2.RecordToolCall(ctx, "search", nil, "res2", 50, 0, "") + rec2.RecordModelCall(ctx, "gpt-4", "openai", nil, "final2", events.TokenUsage{}, 100, 0) + + result, err := Diff(ctx, store, store2, "left", "right") + if err != nil { + t.Fatal(err) + } + // Different payloads → mismatches expected. + _ = result +} + +func TestReplayResult_Metrics(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("metrics")) + rec.RecordToolCall(ctx, "t1", nil, "r1", 10, 0, "") + rec.RecordToolCall(ctx, "t2", nil, "r2", 20, 0, "err") + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "m1", events.TokenUsage{}, 100, 0) + + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "metrics", + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if result.ReplayMetrics.TotalEvents == 0 { + t.Fatal("expected non-zero TotalEvents") + } + if result.ReplayMetrics.TotalEvents != result.ReplayLen { + t.Fatalf("TotalEvents %d != ReplayLen %d", result.ReplayMetrics.TotalEvents, result.ReplayLen) + } + if result.ReplayMetrics.DivergenceCount+result.ReplayMetrics.MatchCount != result.ReplayMetrics.TotalEvents { + t.Fatal("DivergenceCount + MatchCount should equal TotalEvents") + } +} + +func TestEventsContains_FindByType(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("find")) + rec.OnNodeStart(ctx, "n1", 0) + rec.RecordToolCall(ctx, "search", nil, "res", 10, 0, "") + + iter := store.Stream(ctx, events.EventFilter{TraceID: "find"}) + var evts []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + evts = append(evts, ev) + } + + if !EventsContains(evts, events.EventNodeStart) { + t.Fatal("should contain EventNodeStart") + } + if EventsContains(evts, events.EventGraphStart) { + t.Fatal("should not contain EventGraphStart") + } + + found := FindEventsOfType(evts, events.EventToolCallStart) + if len(found) != 1 { + t.Fatalf("expected 1 EventToolCallStart, got %d", len(found)) + } + + count := EventCount(evts, events.EventNodeStart) + if count != 1 { + t.Fatalf("expected 1 EventNodeStart, got %d", count) + } +} + +func TestReplayLiveTools(t *testing.T) { + called := false + override := ReplayLiveTools() + _, _ = override("search", nil, "recorded") + called = true + if !called { + t.Fatal("override function should return nil") + } +} + +func TestReplaySubstituteModel(t *testing.T) { + override := ReplaySubstituteModel(func(recorded string) string { + return "over:" + recorded + }) + result, err := override(nil, "hello") + if err != nil { + t.Fatal(err) + } + if *result != "over:hello" { + t.Fatalf("expected 'over:hello', got '%s'", *result) + } +} + +func TestReplayExactTools(t *testing.T) { + override := ReplayExactTools() + result, err := override("test", nil, "exact-value") + if err != nil { + t.Fatal(err) + } + if result != "exact-value" { + t.Fatalf("expected 'exact-value', got '%v'", result) + } +} + +func TestIntegration_RecordReplayRoundTrip(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("rtt")) + + // Record a realistic agent trace. + rec.OnRunStart(ctx, "agent", "thread-1") + rec.OnStepStart(ctx, 0, 2) + rec.OnNodeStart(ctx, "llm_call", 0) + rec.RecordModelCall(ctx, "gpt-4", "openai", []any{"user msg"}, "assistant resp", events.TokenUsage{PromptTokens: 50, CompletionTokens: 100, TotalTokens: 150}, 800, 0.003) + rec.OnNodeEnd(ctx, "llm_call", 0, nil, nil) + rec.OnNodeStart(ctx, "execute_tools", 0) + rec.RecordToolCall(ctx, "web_search", map[string]any{"q": "RAG"}, "search results", 1200, 0, "") + rec.OnNodeEnd(ctx, "execute_tools", 0, nil, nil) + rec.OnStepEnd(ctx, 0, nil) + rec.OnRunEnd(ctx, "agent", "thread-1", nil) + + // Replay exactly. + engine := NewReplayEngine(store) + result, err := engine.Replay(ctx, &ReplayConfig{ + TraceID: "rtt", + DiffEnabled: true, + }) + if err != nil { + t.Fatal(err) + } + if result.OriginalLen == 0 { + t.Fatal("expected non-zero events") + } + if result.OriginalLen != result.ReplayLen { + t.Fatalf("OriginalLen %d != ReplayLen %d", result.OriginalLen, result.ReplayLen) + } + if len(result.Divergences) != 0 { + t.Fatalf("expected 0 divergences, got %d", len(result.Divergences)) + } + if result.Duration == 0 { + t.Fatal("expected non-zero duration") + } + if result.ReplayMetrics.TotalEvents != result.ReplayLen { + t.Fatal("ReplayMetrics.TotalEvents mismatch") + } +} + +func TestIntegration_ForkThenReplay(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("fork-rtt")) + + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "resp1", events.TokenUsage{}, 100, 0) + rec.RecordToolCall(ctx, "search", nil, "res1", 50, 0, "") + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "resp2", events.TokenUsage{}, 100, 0) + + // Fork from the last model call. + iter := store.Stream(ctx, events.EventFilter{TraceID: "fork-rtt"}) + var allEvents []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + allEvents = append(allEvents, ev) + } + iter.Close() + + // Find the last LLMCallStart event. + var forkPoint events.EventID + for _, ev := range allEvents { + if ev.Type == events.EventLLMCallStart { + forkPoint = ev.ID + } + } + + engine := NewReplayEngine(store) + forkResult, err := engine.Fork(ctx, &ForkConfig{ + TraceID: "fork-rtt", + Point: forkPoint, + Store: store, + }) + if err != nil { + t.Fatal(err) + } + if forkResult.ForkTraceID == "" { + t.Fatal("expected non-empty fork ID") + } + + // Fork events should include the original events up to fork point + fork marker. + if len(forkResult.ForkEvents) == 0 { + t.Fatal("expected fork events") + } +} + +// ---- True replay: BuildCheckpoint ---- + +func TestBuildCheckpoint_StateWrite(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("cp-test")) + + rec.RecordStateWrite(ctx, "messages", nil, []any{"hello"}, "append") + rec.RecordStateWrite(ctx, "counter", nil, 42, "add") + rec.OnNodeEnd(ctx, "agent", 0, nil, nil) + rec.OnStepEnd(ctx, 0, nil) + + iter := store.Stream(ctx, events.EventFilter{TraceID: "cp-test"}) + var evts []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + evts = append(evts, ev) + } + + cp, cpID := BuildCheckpoint(evts, "test-thread") + if cpID == "" { + t.Fatal("expected non-empty checkpoint ID") + } + if cp["messages"] == nil { + t.Fatal("expected messages in checkpoint") + } + if val, ok := cp["counter"].(int); !ok || val != 42 { + if val, ok := cp["counter"].(float64); !ok || val != 42 { + t.Fatalf("expected counter=42, got %v (type %T)", cp["counter"], cp["counter"]) + } + } + if val, ok := cp["__step__"].(float64); !ok || val != 0 { + t.Fatalf("expected step=0, got %v (type %T)", cp["__step__"], cp["__step__"]) + } + if val, ok := cp["__last_completed_node__"].(string); !ok || val != "agent" { + t.Fatalf("expected last_completed_node='agent', got %v (type %T)", cp["__last_completed_node__"], cp["__last_completed_node__"]) + } +} + +func TestBuildCheckpoint_EmptyEvents(t *testing.T) { + cp, cpID := BuildCheckpoint(nil, "empty") + if cpID == "" { + t.Fatal("expected checkpoint ID even with no events") + } + // last_state is always set for non-nil input + _ = cp +} + +// ---- True replay: Fork with Engine ---- + +func TestFork_WithEngine(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("fork-eng")) + + // Record two state writes, then a tool call. + rec.RecordStateWrite(ctx, "results", nil, []any{"first"}, "append") + rec.RecordStateWrite(ctx, "count", nil, 1, "add") + rec.RecordToolCall(ctx, "search", map[string]any{"q": "test"}, "result", 100, 0, "") + + // Find the fork point (second event, a state write). + iter := store.Stream(ctx, events.EventFilter{TraceID: "fork-eng"}) + var allEvents []*events.Event + for { + ev, ok := iter.Next(ctx) + if !ok { + break + } + allEvents = append(allEvents, ev) + } + if len(allEvents) < 2 { + t.Fatal("need at least 2 events") + } + + // Fork from the second state write (after count=1 is persisted). + forkPoint := allEvents[1].ID // EventStateWrite for "count" + + engine := NewReplayEngine(store) + result, err := engine.Fork(ctx, &ForkConfig{ + TraceID: "fork-eng", + Point: forkPoint, + Store: store, + // No ForkEngine — should use deterministic replay path. + OutputStore: events.NewMemoryEventStore(), + }) + if err != nil { + t.Fatal(err) + } + if result.FinalState != nil { + t.Fatal("expected nil FinalState without ForkEngine") + } + if result.Duration == 0 { + t.Fatal("expected non-zero duration") + } +} + +func TestFork_EventNotFound(t *testing.T) { + ctx := context.Background() + store := events.NewMemoryEventStore() + engine := NewReplayEngine(store) + + _, err := engine.Fork(ctx, &ForkConfig{ + TraceID: "nonexistent", + Point: "no-such-event", + Store: store, + }) + if err == nil { + t.Fatal("expected error for nonexistent event") + } +} + +// ---- bench ---- + +func BenchmarkReplay(b *testing.B) { + ctx := context.Background() + store := events.NewMemoryEventStore() + rec := events.NewEventRecorder(store, events.WithTraceID("bench-replay")) + + // Record 100 events. + for i := 0; i < 50; i++ { + rec.RecordModelCall(ctx, "gpt-4", "openai", nil, "resp", events.TokenUsage{}, 100, 0) + rec.RecordToolCall(ctx, "search", nil, "res", 50, 0, "") + } + + engine := NewReplayEngine(store) + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := engine.Replay(ctx, &ReplayConfig{TraceID: "bench-replay"}) + if err != nil { + b.Fatal(err) + } + } +}