From fd24e3191267a47924dc128f7e8a5ca57fcf5445 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 23 Feb 2026 12:18:13 +0800 Subject: [PATCH 1/3] Update Makefile for enhanced testing and add new event documentation - Modify unit test commands in the Makefile to include additional skip patterns for memory leak tests, improving test coverage and accuracy. - Expand benchmark and memory leak detection to include the event module, ensuring comprehensive testing across all components. - Add new design and TODO documentation files for the event module to facilitate future development. --- .gitignore | 2 + Makefile | 16 +- event/README.md | 184 ++++++++++++++++++ event/bench_test.go | 382 +++++++++++++++++++++++++++++++++++++ event/bus.go | 159 ++++++++++++++++ event/bus_test.go | 320 +++++++++++++++++++++++++++++++ event/leak_test.go | 348 ++++++++++++++++++++++++++++++++++ event/listener.go | 141 ++++++++++++++ event/listener_test.go | 222 ++++++++++++++++++++++ event/option.go | 51 +++++ event/queue.go | 208 ++++++++++++++++++++ event/queue_test.go | 387 ++++++++++++++++++++++++++++++++++++++ event/service.go | 221 ++++++++++++++++++++++ event/service_test.go | 229 ++++++++++++++++++++++ event/sub.go | 101 ++++++++++ event/sub_test.go | 181 ++++++++++++++++++ event/types/interfaces.go | 23 +++ event/types/types.go | 102 ++++++++++ event/types/types_test.go | 267 ++++++++++++++++++++++++++ event/worker.go | 94 +++++++++ event/worker_test.go | 213 +++++++++++++++++++++ 21 files changed, 3843 insertions(+), 8 deletions(-) create mode 100644 event/README.md create mode 100644 event/bench_test.go create mode 100644 event/bus.go create mode 100644 event/bus_test.go create mode 100644 event/leak_test.go create mode 100644 event/listener.go create mode 100644 event/listener_test.go create mode 100644 event/option.go create mode 100644 event/queue.go create mode 100644 event/queue_test.go create mode 100644 event/service.go create mode 100644 event/service_test.go create mode 100644 event/sub.go create mode 100644 event/sub_test.go create mode 100644 event/types/interfaces.go create mode 100644 event/types/types.go create mode 100644 event/types/types_test.go create mode 100644 event/worker.go create mode 100644 event/worker_test.go diff --git a/.gitignore b/.gitignore index d67186a4..e899992a 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,5 @@ release/* sandbox/TODO-VNC.md sandbox/docker/chrome/PLAN.md sandbox/DESIGN-REMOTE.md +event/DESIGN.md +event/TODO.md diff --git a/Makefile b/Makefile index b852bd93..cae35a18 100644 --- a/Makefile +++ b/Makefile @@ -30,7 +30,7 @@ TESTTAGS ?= "" unit-test: echo "mode: count" > coverage.out for d in $(TESTFOLDER); do \ - $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal' $$d > tmp.out; \ + $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' $$d > tmp.out; \ cat tmp.out; \ if grep -q "^--- FAIL" tmp.out; then \ rm tmp.out; \ @@ -56,7 +56,7 @@ unit-test: unit-test-core: echo "mode: count" > coverage.out for d in $(TESTFOLDER_CORE); do \ - $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal' $$d > tmp.out; \ + $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' $$d > tmp.out; \ cat tmp.out; \ if grep -q "^--- FAIL" tmp.out; then \ rm tmp.out; \ @@ -224,9 +224,9 @@ unit-test-sandbox: benchmark: @echo "" @echo "=============================================" - @echo "Running Benchmark Tests (agent & trace)..." + @echo "Running Benchmark Tests (agent, trace, event)..." @echo "=============================================" - @for d in $$($(GO) list ./agent/... ./trace/...); do \ + @for d in $$($(GO) list ./agent/... ./trace/... ./event/...); do \ if $(GO) test -list=Benchmark $$d 2>/dev/null | grep -q "^Benchmark"; then \ echo ""; \ echo "πŸ“Š Benchmarking: $$d"; \ @@ -244,14 +244,14 @@ benchmark: memory-leak: @echo "" @echo "=============================================" - @echo "Running Memory Leak Detection (agent & trace)..." + @echo "Running Memory Leak Detection (agent, trace, event)..." @echo "=============================================" - @for d in $$($(GO) list ./agent/... ./trace/...); do \ - if $(GO) test -list='TestMemoryLeak|TestIsolateDisposal|TestGoroutineLeak' $$d 2>/dev/null | grep -qE "^Test(MemoryLeak|IsolateDisposal|GoroutineLeak)"; then \ + @for d in $$($(GO) list ./agent/... ./trace/... ./event/...); do \ + if $(GO) test -list='TestMemoryLeak|TestIsolateDisposal|TestGoroutineLeak|TestLeak_|TestScenario_' $$d 2>/dev/null | grep -qE "^Test(MemoryLeak|IsolateDisposal|GoroutineLeak|Leak_|Scenario_)"; then \ echo ""; \ echo "πŸ” Memory Leak Detection: $$d"; \ echo "---------------------------------------------"; \ - $(GO) test -run='TestMemoryLeak|TestIsolateDisposal|TestGoroutineLeak' -v -timeout=5m $$d || exit 1; \ + $(GO) test -run='TestMemoryLeak|TestIsolateDisposal|TestGoroutineLeak|TestLeak_|TestScenario_' -v -timeout=5m $$d || exit 1; \ fi; \ done @echo "" diff --git a/event/README.md b/event/README.md new file mode 100644 index 00000000..36a4f660 --- /dev/null +++ b/event/README.md @@ -0,0 +1,184 @@ +# event β€” Yao In-Process Event Bus + +Global event service for async/sync event routing, serial queue processing, and real-time subscriptions. All operations are goroutine-safe. + +## Import + +```go +import ( + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) +``` + +## Core Concepts + +| Concept | Description | +|---|---| +| **Push** | Async fire-and-forget delivery. Returns event ID immediately. | +| **Call** | Sync request-response. Blocks until handler writes to `resp`. | +| **Handler** | One per prefix (e.g. `"trace"`). Processes `Push` and `Call` events. | +| **Queue** | FIFO serial processing per entity (e.g. per traceID). Events in same queue never run concurrently. | +| **Listener** | Persistent background consumer (registered at startup). Gets a copy of every matching event. | +| **Subscriber** | Dynamic subscription (e.g. SSE/WebSocket). Non-blocking; skips if channel full. | + +## Lifecycle + +```go +// 1. Register handlers and listeners (before Start, typically in init()) +event.Register("trace", traceHandler, event.MaxWorkers(512), event.ReservedWorkers(20)) +event.Register("job", jobHandler) +event.Listen("trace.*", traceListener) + +// 2. Start +event.Start() + +// 3. Use (from any goroutine) +event.Push(ctx, "trace.add", payload, event.Queue(traceQueueID)) +id, data, err := event.Call(ctx, "trace.get", req, event.Queue(traceQueueID)) + +// 4. Stop (during shutdown) +event.Stop(ctx) +``` + +## Handler + +Implement `types.Handler`: + +```go +type TraceHandler struct{} + +func (h *TraceHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + var p TracePayload + if err := ev.Should(&p); err != nil { + if ev.IsCall { resp <- types.Result{Err: err} } + return + } + // ... business logic ... + if ev.IsCall { + resp <- types.Result{Data: result} + } +} + +func (h *TraceHandler) Shutdown(ctx context.Context) error { return nil } +``` + +- `ctx`: non-cancellable for Push; caller's context for Call. +- `resp`: always non-nil. Write exactly once for Call; ignore for Push. +- `ev.Should(&target)`: type-safe payload extraction. +- Panics are recovered automatically; `ErrHandlerPanic` is returned to Call. + +## Queue + +```go +queueID, err := event.QueueCreate("trace") // auto-generated ID +queueID, err := event.QueueCreate("trace", "my-id") // custom ID + +event.Push(ctx, "trace.add", data, event.Queue(queueID)) // serial +event.Call(ctx, "trace.get", req, event.Queue(queueID)) // serial, same queue + +event.QueueRelease(queueID) // graceful: drain pending, reject new +event.QueueAbort(queueID) // forceful: discard pending, reject new +``` + +## Listener + +Implement `types.Listener`: + +```go +type MailListener struct{} +func (l *MailListener) OnEvent(ev *types.Event) { /* ... */ } +func (l *MailListener) Shutdown(ctx context.Context) error { return nil } + +// Register before Start +event.Listen("mail.*", &MailListener{}, event.Filter(fn), event.BufferSize(4096)) +``` + +- Each listener runs in its own goroutine. +- Non-blocking: if buffer full, event is skipped (logged as warning). + +## Subscriber + +```go +ch := make(chan *types.Event, 256) +subID := event.Subscribe("trace.*", ch, event.Filter(fn)) +defer event.Unsubscribe(subID) + +for ev := range ch { + // push to SSE / WebSocket +} +``` + +- Non-blocking: if `ch` full, event is skipped silently. +- Call `Unsubscribe` when client disconnects. + +## Context Propagation + +```go +ctx = event.WithSID(ctx, sessionID) +ctx = event.WithAuth(ctx, &types.AuthorizedInfo{UserID: "u-1"}) + +// Inside handler: +sid := ev.SID +auth := ev.Auth // may be nil +``` + +SID and Auth are extracted from `ctx` automatically when calling `Push`/`Call`. + +## Pattern Matching + +Used by `Listen` and `Subscribe`: + +| Pattern | Matches | +|---|---| +| `"*"` | Everything | +| `"trace.*"` | `"trace.add"`, `"trace.get"`, etc. | +| `"trace.add"` | Exact match only | + +## Handler Options + +| Option | Default | Description | +|---|---|---| +| `MaxWorkers(n)` | 512 | Max concurrent goroutines for this handler | +| `ReservedWorkers(n)` | 10 | Slots reserved for Call (Push can use Maxβˆ’Reserved) | +| `QueueSize(n)` | 8192 | Per-queue buffered channel capacity | + +## Errors + +| Error | When | +|---|---| +| `ErrNotStarted` | Push/Call before Start or after Stop | +| `ErrNoHandler` | No handler registered for event prefix | +| `ErrQueueFull` | Queue buffer at capacity | +| `ErrQueueNotFound` | Queue ID never created | +| `ErrQueueReleased` | Queue already released/aborted | +| `ErrQueueExists` | QueueCreate with duplicate ID | +| `ErrHandlerPanic` | Handler panicked (recovered) | + +## Performance (M2 Max, 12 cores) + +| Metric | Value | +|---|---| +| Push (no queue) | ~860K ops/sec, 456 B/op | +| Call (no queue) | ~1.2M ops/sec, 440 B/op | +| Push (with queue) | ~2.9M ops/sec, 341 B/op | +| 1000-user scenario (2000 queues, 27K events) | ~100K events/sec, 280ms total | +| Steady-state memory (1000 users) | ~27 MB | +| Goroutine leaks | Zero | + +## File Structure + +``` +event/ +β”œβ”€β”€ types/ +β”‚ β”œβ”€β”€ types.go # Event, Result, HandlerEntry, FilterEntry, options +β”‚ └── interfaces.go # Handler, Listener interfaces +β”œβ”€β”€ service.go # Register, Start, Stop, Reload, global state +β”œβ”€β”€ bus.go # Push, Call, QueueCreate/Release/Abort +β”œβ”€β”€ queue.go # FIFO queue + queue manager +β”œβ”€β”€ worker.go # Worker pool (two-tier semaphore) +β”œβ”€β”€ listener.go # Listener manager + pattern matching +β”œβ”€β”€ sub.go # Subscriber manager +β”œβ”€β”€ option.go # Option functions +└── README.md +``` diff --git a/event/bench_test.go b/event/bench_test.go new file mode 100644 index 00000000..4538af32 --- /dev/null +++ b/event/bench_test.go @@ -0,0 +1,382 @@ +package event_test + +import ( + "context" + "fmt" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --------------------------------------------------------------------------- +// Shared bench handler: lightweight, simulates minimal real work. +// --------------------------------------------------------------------------- + +type benchHandler struct { + processed atomic.Int64 +} + +func (h *benchHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + h.processed.Add(1) + if ev.IsCall { + resp <- types.Result{Data: "ok"} + } +} + +func (h *benchHandler) Shutdown(ctx context.Context) error { return nil } + +// benchListener counts received events. +type benchListener struct { + received atomic.Int64 +} + +func (l *benchListener) OnEvent(ev *types.Event) { l.received.Add(1) } +func (l *benchListener) Shutdown(ctx context.Context) error { return nil } + +// --------------------------------------------------------------------------- +// Benchmark: Push throughput (no queue, pure worker dispatch) +// --------------------------------------------------------------------------- + +func BenchmarkPush_NoQueue(b *testing.B) { + event.Reset() + h := &benchHandler{} + event.Register("bench", h, event.MaxWorkers(512)) + _ = event.Start() + b.Cleanup(func() { _ = event.Stop(context.Background()); event.Reset() }) + + ctx := context.Background() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + _, _ = event.Push(ctx, "bench.work", nil) + } + }) + b.StopTimer() + + // Drain workers + time.Sleep(100 * time.Millisecond) + b.ReportMetric(float64(h.processed.Load()), "events_handled") +} + +// --------------------------------------------------------------------------- +// Benchmark: Call throughput (no queue, synchronous round-trip) +// --------------------------------------------------------------------------- + +func BenchmarkCall_NoQueue(b *testing.B) { + event.Reset() + h := &benchHandler{} + event.Register("bench", h, event.MaxWorkers(512)) + _ = event.Start() + b.Cleanup(func() { _ = event.Stop(context.Background()); event.Reset() }) + + ctx := context.Background() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + _, _, _ = event.Call(ctx, "bench.get", nil) + } + }) +} + +// --------------------------------------------------------------------------- +// Benchmark: Push throughput with Queue (serial per queue) +// --------------------------------------------------------------------------- + +func BenchmarkPush_WithQueue(b *testing.B) { + event.Reset() + h := &benchHandler{} + event.Register("bench", h, event.MaxWorkers(512), event.QueueSize(8192)) + _ = event.Start() + b.Cleanup(func() { _ = event.Stop(context.Background()); event.Reset() }) + + qID, _ := event.QueueCreate("bench") + b.Cleanup(func() { event.QueueRelease(qID) }) + + ctx := context.Background() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = event.Push(ctx, "bench.work", nil, event.Queue(qID)) + } + b.StopTimer() + + event.QueueRelease(qID) + time.Sleep(200 * time.Millisecond) +} + +// --------------------------------------------------------------------------- +// Scenario: 1000 concurrent users, each with trace + job queues. +// +// Simulates: +// - 1000 users Γ— 2 queues (trace + job) = 2000 queues +// - Each user pushes 20 trace events + 5 job events + 1 Call per queue +// - 200 SSE subscribers watching "trace.*" and "job.*" +// - 2 Listeners (trace.* + job.*) +// +// Reports: total duration, events/sec, memory delta. +// --------------------------------------------------------------------------- + +func TestScenario_1000Users(t *testing.T) { + event.Reset() + defer event.Reset() + + traceH := &benchHandler{} + jobH := &benchHandler{} + event.Register("trace", traceH, event.MaxWorkers(512), event.ReservedWorkers(20), event.QueueSize(8192)) + event.Register("job", jobH, event.MaxWorkers(256), event.ReservedWorkers(10), event.QueueSize(4096)) + + traceL := &benchListener{} + jobL := &benchListener{} + event.Listen("trace.*", traceL) + event.Listen("job.*", jobL) + + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + const ( + numUsers = 1000 + tracePushPerUser = 20 + jobPushPerUser = 5 + callsPerQueue = 1 + numSubscribers = 200 + subscriberBufSize = 256 + ) + + // --- Subscribers --- + subChans := make([]chan *types.Event, numSubscribers) + subIDs := make([]string, numSubscribers) + for i := 0; i < numSubscribers; i++ { + ch := make(chan *types.Event, subscriberBufSize) + subChans[i] = ch + pattern := "trace.*" + if i%2 == 1 { + pattern = "job.*" + } + subIDs[i] = event.Subscribe(pattern, ch) + } + defer func() { + for _, id := range subIDs { + event.Unsubscribe(id) + } + }() + + // Drain subscribers in background + var subReceived atomic.Int64 + subDone := make(chan struct{}) + go func() { + defer close(subDone) + for _, ch := range subChans { + go func(c chan *types.Event) { + for range c { + subReceived.Add(1) + } + }(ch) + } + }() + + // --- Memory before --- + runtime.GC() + var memBefore runtime.MemStats + runtime.ReadMemStats(&memBefore) + + // --- Run --- + start := time.Now() + var wg sync.WaitGroup + + for u := 0; u < numUsers; u++ { + wg.Add(1) + go func(userID int) { + defer wg.Done() + ctx := event.WithSID(context.Background(), fmt.Sprintf("sess-%d", userID)) + ctx = event.WithAuth(ctx, &types.AuthorizedInfo{UserID: fmt.Sprintf("u-%d", userID)}) + + // Create trace queue + traceQID, err := event.QueueCreate("trace") + if err != nil { + t.Errorf("user %d: trace QueueCreate: %v", userID, err) + return + } + + // Create job queue + jobQID, err := event.QueueCreate("job") + if err != nil { + t.Errorf("user %d: job QueueCreate: %v", userID, err) + event.QueueRelease(traceQID) + return + } + + // Push trace events + for i := 0; i < tracePushPerUser; i++ { + _, _ = event.Push(ctx, "trace.add", i, event.Queue(traceQID)) + } + + // Push job events + for i := 0; i < jobPushPerUser; i++ { + _, _ = event.Push(ctx, "job.progress", i, event.Queue(jobQID)) + } + + // Call on each queue + for i := 0; i < callsPerQueue; i++ { + callCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + _, _, _ = event.Call(callCtx, "trace.get", nil, event.Queue(traceQID)) + cancel() + + callCtx2, cancel2 := context.WithTimeout(ctx, 5*time.Second) + _, _, _ = event.Call(callCtx2, "job.status", nil, event.Queue(jobQID)) + cancel2() + } + + // Release queues + event.QueueRelease(traceQID) + event.QueueRelease(jobQID) + }(u) + } + + wg.Wait() + elapsed := time.Since(start) + + // Wait for queues to drain + time.Sleep(500 * time.Millisecond) + + // --- Memory after --- + runtime.GC() + var memAfter runtime.MemStats + runtime.ReadMemStats(&memAfter) + + // --- Results --- + totalPush := int64(numUsers) * int64(tracePushPerUser+jobPushPerUser) + totalCall := int64(numUsers) * int64(callsPerQueue) * 2 + totalEvents := totalPush + totalCall + traceProcessed := traceH.processed.Load() + jobProcessed := jobH.processed.Load() + listenerTrace := traceL.received.Load() + listenerJob := jobL.received.Load() + memDeltaMB := float64(memAfter.TotalAlloc-memBefore.TotalAlloc) / 1024 / 1024 + + t.Logf("=== 1000-User Scenario Results ===") + t.Logf("Users: %d", numUsers) + t.Logf("Queues created: %d (trace: %d, job: %d)", numUsers*2, numUsers, numUsers) + t.Logf("Subscribers: %d", numSubscribers) + t.Logf("Total events: %d (push: %d, call: %d)", totalEvents, totalPush, totalCall) + t.Logf("Trace processed: %d", traceProcessed) + t.Logf("Job processed: %d", jobProcessed) + t.Logf("Listener trace: %d", listenerTrace) + t.Logf("Listener job: %d", listenerJob) + t.Logf("Sub received: %d", subReceived.Load()) + t.Logf("Elapsed: %v", elapsed) + t.Logf("Throughput: %.0f events/sec", float64(totalEvents)/elapsed.Seconds()) + t.Logf("Memory delta: %.2f MB (TotalAlloc)", memDeltaMB) + + // --- Assertions --- + expectedProcessed := totalPush + totalCall + actualProcessed := traceProcessed + jobProcessed + if actualProcessed < expectedProcessed { + t.Errorf("processed %d < expected %d (some events lost)", actualProcessed, expectedProcessed) + } + + if elapsed > 30*time.Second { + t.Errorf("scenario took %v, expected < 30s", elapsed) + } +} + +// --------------------------------------------------------------------------- +// Benchmark: Queue create/release churn (lifecycle overhead) +// --------------------------------------------------------------------------- + +func BenchmarkQueueCreateRelease(b *testing.B) { + event.Reset() + h := &benchHandler{} + event.Register("bench", h, event.QueueSize(64)) + _ = event.Start() + b.Cleanup(func() { _ = event.Stop(context.Background()); event.Reset() }) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + qID, err := event.QueueCreate("bench") + if err != nil { + b.Fatalf("QueueCreate: %v", err) + } + event.QueueRelease(qID) + } + }) +} + +// --------------------------------------------------------------------------- +// Benchmark: Subscriber notify throughput (fanout to 200 subscribers) +// --------------------------------------------------------------------------- + +func BenchmarkSubscriberFanout(b *testing.B) { + event.Reset() + h := &benchHandler{} + event.Register("bench", h, event.MaxWorkers(512)) + _ = event.Start() + b.Cleanup(func() { _ = event.Stop(context.Background()); event.Reset() }) + + const numSubs = 200 + for i := 0; i < numSubs; i++ { + ch := make(chan *types.Event, 1024) + event.Subscribe("bench.*", ch) + go func(c chan *types.Event) { + for range c { + } + }(ch) + } + + ctx := context.Background() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + _, _ = event.Push(ctx, "bench.work", nil) + } + }) +} + +// --------------------------------------------------------------------------- +// Benchmark: Mixed Push/Call with 2000 queues (1000 users Γ— 2) +// --------------------------------------------------------------------------- + +func BenchmarkMixed_2000Queues(b *testing.B) { + event.Reset() + h := &benchHandler{} + event.Register("mix", h, event.MaxWorkers(512), event.ReservedWorkers(20), event.QueueSize(4096)) + _ = event.Start() + b.Cleanup(func() { _ = event.Stop(context.Background()); event.Reset() }) + + const numQueues = 2000 + queueIDs := make([]string, numQueues) + for i := 0; i < numQueues; i++ { + qID, err := event.QueueCreate("mix") + if err != nil { + b.Fatalf("QueueCreate %d: %v", i, err) + } + queueIDs[i] = qID + } + b.Cleanup(func() { + for _, qID := range queueIDs { + event.QueueRelease(qID) + } + time.Sleep(200 * time.Millisecond) + }) + + ctx := context.Background() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + i := 0 + for pb.Next() { + qID := queueIDs[i%numQueues] + if i%10 == 0 { + callCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + _, _, _ = event.Call(callCtx, "mix.get", nil, event.Queue(qID)) + cancel() + } else { + _, _ = event.Push(ctx, "mix.work", nil, event.Queue(qID)) + } + i++ + } + }) +} diff --git a/event/bus.go b/event/bus.go new file mode 100644 index 00000000..34517c3b --- /dev/null +++ b/event/bus.go @@ -0,0 +1,159 @@ +package event + +import ( + "context" + "fmt" + "strings" + "sync/atomic" + + "github.com/yaoapp/yao/event/types" +) + +var eventIDCounter atomic.Uint64 + +func nextEventID() string { + id := eventIDCounter.Add(1) + return fmt.Sprintf("ev-%d", id) +} + +// prefixOf extracts the handler prefix from an event type. +// "trace.add" -> "trace", "job.progress" -> "job" +func prefixOf(typ string) string { + if i := strings.IndexByte(typ, '.'); i >= 0 { + return typ[:i] + } + return typ +} + +// Push delivers an event asynchronously (fire-and-forget). +// SID and Auth are extracted from ctx automatically. +// Returns the auto-generated event ID. +func Push(ctx context.Context, typ string, payload any, opts ...types.PushOption) (string, error) { + prefix := prefixOf(typ) + entry, pool, err := getHandler(prefix) + if err != nil { + return "", err + } + _ = entry // used for queue config lookup + + ev := &types.Event{ + Type: typ, + ID: nextEventID(), + IsCall: false, + Payload: payload, + SID: SIDFrom(ctx), + Auth: AuthFrom(ctx), + } + for _, opt := range opts { + opt(ev) + } + + // Notify listeners and subscribers (non-blocking, before handler) + svc.lmgr.notify(ev) + svc.smgr.notify(ev) + + // Route to queue or direct dispatch + if ev.Queue != "" { + q, err := svc.queues.get(ev.Queue) + if err != nil { + return ev.ID, err + } + discard := make(chan types.Result, 1) + if err := q.enqueue(ctx, ev, discard); err != nil { + return ev.ID, err + } + return ev.ID, nil + } + + // No queue: direct dispatch with discard channel + discard := make(chan types.Result, 1) + pushCtx := context.WithoutCancel(ctx) + if _, err := pool.dispatch(pushCtx, ev, discard); err != nil { + return ev.ID, fmt.Errorf("event push: worker unavailable: %w", err) + } + return ev.ID, nil +} + +// Call delivers an event synchronously and blocks until the handler responds. +// SID and Auth are extracted from ctx automatically. +// Returns the auto-generated event ID and the handler's result. +func Call(ctx context.Context, typ string, payload any, opts ...types.PushOption) (string, any, error) { + prefix := prefixOf(typ) + _, pool, err := getHandler(prefix) + if err != nil { + return "", nil, err + } + + ev := &types.Event{ + Type: typ, + ID: nextEventID(), + IsCall: true, + Payload: payload, + SID: SIDFrom(ctx), + Auth: AuthFrom(ctx), + } + for _, opt := range opts { + opt(ev) + } + + // Notify listeners and subscribers + svc.lmgr.notify(ev) + svc.smgr.notify(ev) + + resp := make(chan types.Result, 1) + + if ev.Queue != "" { + q, err := svc.queues.get(ev.Queue) + if err != nil { + return ev.ID, nil, err + } + if err := q.enqueue(ctx, ev, resp); err != nil { + return ev.ID, nil, err + } + } else { + if _, err := pool.dispatch(ctx, ev, resp); err != nil { + return ev.ID, nil, fmt.Errorf("event call: worker unavailable: %w", err) + } + } + + // Wait for handler result or context cancellation + select { + case result := <-resp: + return ev.ID, result.Data, result.Err + case <-ctx.Done(): + return ev.ID, nil, ctx.Err() + } +} + +// QueueCreate creates a new event queue bound to a handler prefix. +// Returns the queue ID. If no id is provided, one is auto-generated. +func QueueCreate(prefix string, id ...string) (string, error) { + entry, pool, err := getHandler(prefix) + if err != nil { + return "", err + } + + queueID := "" + if len(id) > 0 && id[0] != "" { + queueID = id[0] + } else { + queueID = fmt.Sprintf("q-%s-%d", prefix, eventIDCounter.Add(1)) + } + + if err := svc.queues.create(prefix, queueID, entry.QueueSize, pool); err != nil { + return "", err + } + return queueID, nil +} + +// QueueRelease gracefully releases a queue (async). +// Rejects new events immediately; existing events are drained internally. +func QueueRelease(queueID string) { + svc.queues.release(queueID) +} + +// QueueAbort forcefully releases a queue (async). +// Rejects new events, discards pending events, waits for in-flight to finish. +func QueueAbort(queueID string) { + svc.queues.abortOne(queueID) +} diff --git a/event/bus_test.go b/event/bus_test.go new file mode 100644 index 00000000..97de5f05 --- /dev/null +++ b/event/bus_test.go @@ -0,0 +1,320 @@ +package event_test + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --- Test handler --- + +type recordHandler struct { + mu sync.Mutex + calls []string // records ev.Type for each Handle call + shutdown bool +} + +func (h *recordHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + h.mu.Lock() + h.calls = append(h.calls, ev.Type) + h.mu.Unlock() + + if ev.IsCall { + var p string + if err := ev.Should(&p); err == nil { + resp <- types.Result{Data: "echo:" + p} + } else { + resp <- types.Result{Data: "echo:" + ev.Type} + } + } +} + +func (h *recordHandler) Shutdown(ctx context.Context) error { + h.mu.Lock() + defer h.mu.Unlock() + h.shutdown = true + return nil +} + +func (h *recordHandler) getCalls() []string { + h.mu.Lock() + defer h.mu.Unlock() + cp := make([]string, len(h.calls)) + copy(cp, h.calls) + return cp +} + +// --- Phase 3: Push / Call basic routing (no queue) --- + +func TestPush_NoQueue(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &recordHandler{} + event.Register("foo", h) + if err := event.Start(); err != nil { + t.Fatal(err) + } + defer func() { _ = event.Stop(context.Background()) }() + + id, err := event.Push(context.Background(), "foo.bar", "payload1") + if err != nil { + t.Fatalf("Push failed: %v", err) + } + if id == "" { + t.Fatal("expected non-empty event ID") + } + + // Wait for async handler + time.Sleep(50 * time.Millisecond) + calls := h.getCalls() + if len(calls) != 1 || calls[0] != "foo.bar" { + t.Fatalf("expected [foo.bar], got %v", calls) + } +} + +func TestCall_NoQueue(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &recordHandler{} + event.Register("foo", h) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + id, data, err := event.Call(context.Background(), "foo.get", "hello") + if err != nil { + t.Fatalf("Call failed: %v", err) + } + if id == "" { + t.Fatal("expected non-empty event ID") + } + if data != "echo:hello" { + t.Fatalf("expected echo:hello, got %v", data) + } +} + +func TestPush_UnregisteredPrefix(t *testing.T) { + event.Reset() + defer event.Reset() + + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, err := event.Push(context.Background(), "unknown.thing", nil) + if err != event.ErrNoHandler { + t.Fatalf("expected ErrNoHandler, got %v", err) + } +} + +func TestPush_NotStarted(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + _, err := event.Push(context.Background(), "foo.bar", nil) + if err != event.ErrNotStarted { + t.Fatalf("expected ErrNotStarted, got %v", err) + } +} + +func TestPush_SIDAndAuth(t *testing.T) { + event.Reset() + defer event.Reset() + + var captured *types.Event + var mu sync.Mutex + + h := &captureHandler{onHandle: func(ev *types.Event) { + mu.Lock() + captured = ev + mu.Unlock() + }} + event.Register("foo", h) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ctx := event.WithSID(context.Background(), "sess-abc") + ctx = event.WithAuth(ctx, &types.AuthorizedInfo{UserID: "u-1"}) + + _, err := event.Push(ctx, "foo.bar", "data") + if err != nil { + t.Fatalf("Push failed: %v", err) + } + + time.Sleep(50 * time.Millisecond) + mu.Lock() + defer mu.Unlock() + if captured == nil { + t.Fatal("handler was not called") + } + if captured.SID != "sess-abc" { + t.Fatalf("expected SID sess-abc, got %s", captured.SID) + } + if captured.Auth == nil || captured.Auth.UserID != "u-1" { + t.Fatalf("expected Auth.UserID u-1, got %+v", captured.Auth) + } +} + +// captureHandler captures the event for inspection. +type captureHandler struct { + onHandle func(*types.Event) +} + +func (h *captureHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + if h.onHandle != nil { + h.onHandle(ev) + } + if ev.IsCall { + resp <- types.Result{Data: "ok"} + } +} + +func (h *captureHandler) Shutdown(ctx context.Context) error { return nil } + +// --- Coverage: prefixOf without dot --- + +func TestPush_TypeWithoutDot(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &recordHandler{} + event.Register("nodot", h) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + id, err := event.Push(context.Background(), "nodot", "payload") + if err != nil { + t.Fatalf("Push failed: %v", err) + } + if id == "" { + t.Fatal("expected non-empty event ID") + } + time.Sleep(50 * time.Millisecond) + calls := h.getCalls() + if len(calls) != 1 || calls[0] != "nodot" { + t.Fatalf("expected [nodot], got %v", calls) + } +} + +// --- Coverage: Call unregistered prefix --- + +func TestCall_UnregisteredPrefix(t *testing.T) { + event.Reset() + defer event.Reset() + + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, _, err := event.Call(context.Background(), "unknown.thing", nil) + if err != event.ErrNoHandler { + t.Fatalf("expected ErrNoHandler, got %v", err) + } +} + +// --- Coverage: Call with queue (happy path) --- + +func TestCall_WithQueue(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &recordHandler{} + event.Register("foo", h) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, err := event.QueueCreate("foo") + if err != nil { + t.Fatalf("QueueCreate failed: %v", err) + } + defer event.QueueRelease(qID) + + id, data, err := event.Call(context.Background(), "foo.get", "hello", event.Queue(qID)) + if err != nil { + t.Fatalf("Call with queue failed: %v", err) + } + if id == "" { + t.Fatal("expected non-empty event ID") + } + if data != "echo:hello" { + t.Fatalf("expected echo:hello, got %v", data) + } +} + +// --- Coverage: Call with non-existent queue --- + +func TestCall_QueueNotFound(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, _, err := event.Call(context.Background(), "foo.get", nil, event.Queue("no-such-queue")) + if err != event.ErrQueueNotFound { + t.Fatalf("expected ErrQueueNotFound, got %v", err) + } +} + +// --- Coverage: Call ctx timeout --- + +func TestCall_CtxTimeout(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &captureHandler{onHandle: func(ev *types.Event) { + time.Sleep(500 * time.Millisecond) + }} + event.Register("slow", h) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + _, _, err := event.Call(ctx, "slow.op", nil) + if err == nil { + t.Fatal("expected timeout error") + } +} + +// --- Coverage: Call no-queue dispatch failure (ctx cancelled) --- + +func TestCall_NoQueue_DispatchFail(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &concurrencyHandler{ + peak: &atomic.Int32{}, + current: &atomic.Int32{}, + delay: 200 * time.Millisecond, + } + event.Register("tiny", h, event.MaxWorkers(1), event.ReservedWorkers(0)) + _ = event.Start() + + // Saturate the single total slot with a Call in background + bgDone := make(chan struct{}) + go func() { + defer close(bgDone) + _, _, _ = event.Call(context.Background(), "tiny.work", nil) + }() + time.Sleep(10 * time.Millisecond) + + // Another Call with already-cancelled context should fail at dispatch + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, _, err := event.Call(ctx, "tiny.op", nil) + if err == nil { + t.Fatal("expected error for cancelled ctx call") + } + + // Wait for background goroutine to finish before Stop + <-bgDone + _ = event.Stop(context.Background()) +} diff --git a/event/leak_test.go b/event/leak_test.go new file mode 100644 index 00000000..244419d3 --- /dev/null +++ b/event/leak_test.go @@ -0,0 +1,348 @@ +package event_test + +import ( + "context" + "fmt" + "runtime" + "sync" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --------------------------------------------------------------------------- +// Helper: snapshot goroutine count after GC stabilization. +// --------------------------------------------------------------------------- + +func stableGoroutineCount() int { + // Let runtime settle: GC + finalizers + scheduler + for i := 0; i < 5; i++ { + runtime.GC() + runtime.Gosched() + time.Sleep(10 * time.Millisecond) + } + return runtime.NumGoroutine() +} + +// leakHandler is a no-op handler for leak tests. +type leakHandler struct{} + +func (h *leakHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + if ev.IsCall { + resp <- types.Result{Data: "ok"} + } +} + +func (h *leakHandler) Shutdown(ctx context.Context) error { return nil } + +// leakListener is a no-op listener for leak tests. +type leakListener struct{} + +func (l *leakListener) OnEvent(ev *types.Event) {} +func (l *leakListener) Shutdown(ctx context.Context) error { return nil } + +// --------------------------------------------------------------------------- +// Test: 1000 Queue create/release cycles leak no goroutines. +// --------------------------------------------------------------------------- + +func TestLeak_QueueCreateRelease(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("leak", &leakHandler{}, event.QueueSize(64)) + _ = event.Start() + + before := stableGoroutineCount() + + const cycles = 1000 + for i := 0; i < cycles; i++ { + qID, err := event.QueueCreate("leak") + if err != nil { + t.Fatalf("cycle %d: QueueCreate: %v", i, err) + } + // Push a few events to exercise consumer goroutine + for j := 0; j < 3; j++ { + _, _ = event.Push(context.Background(), "leak.work", j, event.Queue(qID)) + } + event.QueueRelease(qID) + } + + // Let all consumer goroutines drain and exit + time.Sleep(500 * time.Millisecond) + after := stableGoroutineCount() + + _ = event.Stop(context.Background()) + + leaked := after - before + t.Logf("goroutines: before=%d after=%d delta=%d (over %d cycles)", before, after, leaked, cycles) + + // Allow a small margin for runtime jitter (GC, timers, etc.) + if leaked > 5 { + t.Errorf("goroutine leak: %d goroutines accumulated over %d queue cycles", leaked, cycles) + } +} + +// --------------------------------------------------------------------------- +// Test: 1000 Queue create/abort cycles leak no goroutines. +// --------------------------------------------------------------------------- + +func TestLeak_QueueCreateAbort(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("leak", &leakHandler{}, event.QueueSize(64)) + _ = event.Start() + + before := stableGoroutineCount() + + const cycles = 1000 + for i := 0; i < cycles; i++ { + qID, err := event.QueueCreate("leak") + if err != nil { + t.Fatalf("cycle %d: QueueCreate: %v", i, err) + } + for j := 0; j < 3; j++ { + _, _ = event.Push(context.Background(), "leak.work", j, event.Queue(qID)) + } + event.QueueAbort(qID) + } + + time.Sleep(500 * time.Millisecond) + after := stableGoroutineCount() + + _ = event.Stop(context.Background()) + + leaked := after - before + t.Logf("goroutines: before=%d after=%d delta=%d (over %d cycles)", before, after, leaked, cycles) + + if leaked > 5 { + t.Errorf("goroutine leak: %d goroutines accumulated over %d abort cycles", leaked, cycles) + } +} + +// --------------------------------------------------------------------------- +// Test: Subscriber create/unsubscribe cycles leak no goroutines or memory. +// --------------------------------------------------------------------------- + +func TestLeak_SubscriberLifecycle(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("leak", &leakHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + before := stableGoroutineCount() + runtime.GC() + var memBefore runtime.MemStats + runtime.ReadMemStats(&memBefore) + + const cycles = 1000 + for i := 0; i < cycles; i++ { + ch := make(chan *types.Event, 16) + subID := event.Subscribe("leak.*", ch) + + _, _ = event.Push(context.Background(), "leak.work", nil) + time.Sleep(time.Microsecond) // let notify propagate + + event.Unsubscribe(subID) + } + + time.Sleep(200 * time.Millisecond) + after := stableGoroutineCount() + + runtime.GC() + var memAfter runtime.MemStats + runtime.ReadMemStats(&memAfter) + + leaked := after - before + memDeltaMB := float64(int64(memAfter.HeapInuse)-int64(memBefore.HeapInuse)) / 1024 / 1024 + + t.Logf("goroutines: before=%d after=%d delta=%d", before, after, leaked) + t.Logf("heap in-use delta: %.2f MB", memDeltaMB) + + if leaked > 3 { + t.Errorf("goroutine leak: %d goroutines after %d sub/unsub cycles", leaked, cycles) + } +} + +// --------------------------------------------------------------------------- +// Test: Start/Stop cycles leak no goroutines. +// --------------------------------------------------------------------------- + +func TestLeak_StartStopCycles(t *testing.T) { + before := stableGoroutineCount() + + const cycles = 20 + for i := 0; i < cycles; i++ { + event.Reset() + event.Register("leak", &leakHandler{}) + event.Listen("leak.*", &leakListener{}) + _ = event.Start() + + ctx := context.Background() + for j := 0; j < 10; j++ { + _, _ = event.Push(ctx, "leak.work", j) + } + time.Sleep(10 * time.Millisecond) + + _ = event.Stop(ctx) + } + event.Reset() + + time.Sleep(300 * time.Millisecond) + after := stableGoroutineCount() + + leaked := after - before + t.Logf("goroutines: before=%d after=%d delta=%d (over %d start/stop cycles)", before, after, leaked, cycles) + + if leaked > 3 { + t.Errorf("goroutine leak: %d goroutines after %d start/stop cycles", leaked, cycles) + } +} + +// --------------------------------------------------------------------------- +// Test: 1000 concurrent users creating/using/releasing queues, verify +// no goroutine leak when everything settles. +// --------------------------------------------------------------------------- + +func TestLeak_1000Users_FullCycle(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("trace", &leakHandler{}, event.MaxWorkers(512), event.QueueSize(8192)) + event.Register("job", &leakHandler{}, event.MaxWorkers(256), event.QueueSize(4096)) + event.Listen("trace.*", &leakListener{}) + _ = event.Start() + + before := stableGoroutineCount() + + const numUsers = 1000 + var wg sync.WaitGroup + for u := 0; u < numUsers; u++ { + wg.Add(1) + go func(uid int) { + defer wg.Done() + ctx := event.WithSID(context.Background(), fmt.Sprintf("s-%d", uid)) + + tqID, err := event.QueueCreate("trace") + if err != nil { + return + } + jqID, err := event.QueueCreate("job") + if err != nil { + event.QueueRelease(tqID) + return + } + + for i := 0; i < 10; i++ { + _, _ = event.Push(ctx, "trace.add", i, event.Queue(tqID)) + } + for i := 0; i < 3; i++ { + _, _ = event.Push(ctx, "job.progress", i, event.Queue(jqID)) + } + + callCtx, cancel := context.WithTimeout(ctx, 3*time.Second) + _, _, _ = event.Call(callCtx, "trace.get", nil, event.Queue(tqID)) + cancel() + + event.QueueRelease(tqID) + event.QueueRelease(jqID) + }(u) + } + + wg.Wait() + time.Sleep(1 * time.Second) // let all consumers drain + + after := stableGoroutineCount() + + _ = event.Stop(context.Background()) + + // Final check after full stop + afterStop := stableGoroutineCount() + + leaked := after - before + leakedAfterStop := afterStop - before + + t.Logf("goroutines: before=%d after_drain=%d after_stop=%d", before, after, afterStop) + t.Logf("delta after drain: %d, delta after stop: %d", leaked, leakedAfterStop) + + if leaked > 10 { + t.Errorf("goroutine leak after drain: %d (1000 users Γ— 2 queues)", leaked) + } + if leakedAfterStop > 3 { + t.Errorf("goroutine leak after stop: %d", leakedAfterStop) + } +} + +// --------------------------------------------------------------------------- +// Test: Memory stability under sustained load. +// Push 100k events through 100 queues, measure heap growth. +// --------------------------------------------------------------------------- + +func TestLeak_MemoryStability(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("mem", &leakHandler{}, event.MaxWorkers(256), event.QueueSize(8192)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + const ( + numQueues = 100 + eventsPerQueue = 1000 + totalEvents = numQueues * eventsPerQueue + ) + + queueIDs := make([]string, numQueues) + for i := 0; i < numQueues; i++ { + qID, _ := event.QueueCreate("mem") + queueIDs[i] = qID + } + + runtime.GC() + var memBefore runtime.MemStats + runtime.ReadMemStats(&memBefore) + + ctx := context.Background() + var wg sync.WaitGroup + for q := 0; q < numQueues; q++ { + wg.Add(1) + go func(qIdx int) { + defer wg.Done() + qID := queueIDs[qIdx] + for i := 0; i < eventsPerQueue; i++ { + _, _ = event.Push(ctx, "mem.work", i, event.Queue(qID)) + } + }(q) + } + wg.Wait() + + // Release all and wait + for _, qID := range queueIDs { + event.QueueRelease(qID) + } + time.Sleep(1 * time.Second) + + runtime.GC() + var memAfter runtime.MemStats + runtime.ReadMemStats(&memAfter) + + // Use signed arithmetic to handle GC reclaiming memory between snapshots. + heapDeltaMB := float64(int64(memAfter.HeapInuse)-int64(memBefore.HeapInuse)) / 1024 / 1024 + allocDeltaMB := float64(memAfter.TotalAlloc-memBefore.TotalAlloc) / 1024 / 1024 + + t.Logf("=== Memory Stability ===") + t.Logf("Events: %d (%d queues Γ— %d events)", totalEvents, numQueues, eventsPerQueue) + t.Logf("HeapInuse delta: %.2f MB", heapDeltaMB) + t.Logf("TotalAlloc: %.2f MB", allocDeltaMB) + t.Logf("Alloc/event: %.0f bytes", allocDeltaMB*1024*1024/float64(totalEvents)) + + // After drain, heap should not retain significant memory. + // Allow generous 50 MB for 100k events (runtime overhead, GC timing). + if heapDeltaMB > 50 { + t.Errorf("heap grew %.2f MB after %d events, possible leak", heapDeltaMB, totalEvents) + } +} diff --git a/event/listener.go b/event/listener.go new file mode 100644 index 00000000..371088c7 --- /dev/null +++ b/event/listener.go @@ -0,0 +1,141 @@ +package event + +import ( + "context" + "strings" + "sync" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/event/types" +) + +// listenerEntry holds a registered listener with its filter configuration. +type listenerEntry struct { + pattern string + listener types.Listener + filter func(*types.Event) bool + bufferSize int + ch chan *types.Event + done chan struct{} +} + +// listenerManager manages all registered listeners. +type listenerManager struct { + mu sync.RWMutex + entries []*listenerEntry + started bool +} + +func newListenerManager() *listenerManager { + return &listenerManager{} +} + +// register adds a listener. Must be called before start(). +func (lm *listenerManager) register(pattern string, listener types.Listener, opts ...types.FilterOption) { + fe := &types.FilterEntry{ + Pattern: pattern, + BufferSize: types.DefaultBufferSize, + } + for _, opt := range opts { + opt(fe) + } + + lm.mu.Lock() + defer lm.mu.Unlock() + lm.entries = append(lm.entries, &listenerEntry{ + pattern: pattern, + listener: listener, + filter: fe.Filter, + bufferSize: fe.BufferSize, + }) +} + +// start creates channels and goroutines for each listener. +func (lm *listenerManager) start() { + lm.mu.Lock() + defer lm.mu.Unlock() + + for _, entry := range lm.entries { + entry.ch = make(chan *types.Event, entry.bufferSize) + entry.done = make(chan struct{}) + go lm.consume(entry) + } + lm.started = true +} + +// consume is the goroutine that reads from a listener's channel. +func (lm *listenerManager) consume(entry *listenerEntry) { + defer close(entry.done) + for ev := range entry.ch { + func() { + defer func() { + if r := recover(); r != nil { + log.Error("event listener panic: pattern=%s type=%s err=%v", entry.pattern, ev.Type, r) + } + }() + entry.listener.OnEvent(ev) + }() + } +} + +// notify sends an event to all matching listeners (non-blocking). +func (lm *listenerManager) notify(ev *types.Event) { + lm.mu.RLock() + defer lm.mu.RUnlock() + + if !lm.started { + return + } + + for _, entry := range lm.entries { + if !matchPattern(entry.pattern, ev.Type) { + continue + } + if entry.filter != nil && !entry.filter(ev) { + continue + } + select { + case entry.ch <- ev: + default: + log.Warn("event listener buffer full: pattern=%s type=%s id=%s (skipped)", entry.pattern, ev.Type, ev.ID) + } + } +} + +// stop shuts down all listeners. +func (lm *listenerManager) stop(ctx context.Context) { + lm.mu.Lock() + lm.started = false + entries := lm.entries + lm.mu.Unlock() + + for _, entry := range entries { + close(entry.ch) + } + for _, entry := range entries { + <-entry.done + _ = entry.listener.Shutdown(ctx) + } +} + +// matchPattern matches an event type against a listener/subscriber pattern. +// - "*" matches everything +// - "foo.*" matches any type starting with "foo." +// - "foo.bar" matches exactly "foo.bar" +func matchPattern(pattern, eventType string) bool { + if pattern == "*" { + return true + } + if strings.HasSuffix(pattern, ".*") { + prefix := strings.TrimSuffix(pattern, "*") + return strings.HasPrefix(eventType, prefix) + } + return pattern == eventType +} + +// Listen registers a persistent listener. Must be called before Start. +func Listen(pattern string, listener types.Listener, opts ...types.FilterOption) { + svc.mu.Lock() + defer svc.mu.Unlock() + svc.lmgr.register(pattern, listener, opts...) +} diff --git a/event/listener_test.go b/event/listener_test.go new file mode 100644 index 00000000..37bbee0e --- /dev/null +++ b/event/listener_test.go @@ -0,0 +1,222 @@ +package event_test + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --- Phase 6: Listener tests --- + +// collectListener collects received events. +type collectListener struct { + mu sync.Mutex + events []*types.Event + shut bool +} + +func (l *collectListener) OnEvent(ev *types.Event) { + l.mu.Lock() + defer l.mu.Unlock() + l.events = append(l.events, ev) +} + +func (l *collectListener) Shutdown(ctx context.Context) error { + l.mu.Lock() + defer l.mu.Unlock() + l.shut = true + return nil +} + +func (l *collectListener) getEvents() []*types.Event { + l.mu.Lock() + defer l.mu.Unlock() + cp := make([]*types.Event, len(l.events)) + copy(cp, l.events) + return cp +} + +func TestListener_PatternMatch(t *testing.T) { + event.Reset() + defer event.Reset() + + allL := &collectListener{} + fooL := &collectListener{} + exactL := &collectListener{} + + event.Listen("*", allL) + event.Listen("foo.*", fooL) + event.Listen("foo.exact", exactL) + + h := &recordHandler{} + event.Register("foo", h) + event.Register("bar", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, _ = event.Push(context.Background(), "foo.exact", nil) + _, _ = event.Push(context.Background(), "foo.other", nil) + _, _ = event.Push(context.Background(), "bar.thing", nil) + + time.Sleep(100 * time.Millisecond) + + allEvents := allL.getEvents() + fooEvents := fooL.getEvents() + exactEvents := exactL.getEvents() + + if len(allEvents) != 3 { + t.Fatalf("all listener expected 3, got %d", len(allEvents)) + } + if len(fooEvents) != 2 { + t.Fatalf("foo.* listener expected 2, got %d", len(fooEvents)) + } + if len(exactEvents) != 1 { + t.Fatalf("foo.exact listener expected 1, got %d", len(exactEvents)) + } + if exactEvents[0].Type != "foo.exact" { + t.Fatalf("expected foo.exact, got %s", exactEvents[0].Type) + } +} + +func TestListener_Filter(t *testing.T) { + event.Reset() + defer event.Reset() + + filtered := &collectListener{} + event.Listen("foo.*", filtered, event.Filter(func(ev *types.Event) bool { + return ev.Type == "foo.keep" + })) + + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, _ = event.Push(context.Background(), "foo.keep", nil) + _, _ = event.Push(context.Background(), "foo.drop", nil) + + time.Sleep(100 * time.Millisecond) + events := filtered.getEvents() + if len(events) != 1 || events[0].Type != "foo.keep" { + t.Fatalf("filter should only pass foo.keep, got %v", events) + } +} + +func TestListener_BufferFull_Skip(t *testing.T) { + event.Reset() + defer event.Reset() + + // Use buffer size 2, listener that blocks + blocking := &blockingListener{unblock: make(chan struct{})} + event.Listen("foo.*", blocking, event.BufferSize(2)) + + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + // Push 5 events; 1 being processed + 2 buffered = 3, rest skipped + for i := 0; i < 5; i++ { + _, _ = event.Push(context.Background(), "foo.item", i) + } + + time.Sleep(50 * time.Millisecond) + close(blocking.unblock) // unblock listener + time.Sleep(100 * time.Millisecond) + + count := blocking.count.Load() + if count > 3 { + t.Fatalf("expected at most 3 events with buffer=2, got %d", count) + } + if count < 1 { + t.Fatal("expected at least 1 event") + } +} + +type blockingListener struct { + unblock chan struct{} + count atomic.Int32 +} + +func (l *blockingListener) OnEvent(ev *types.Event) { + <-l.unblock + l.count.Add(1) +} + +func (l *blockingListener) Shutdown(ctx context.Context) error { return nil } + +func TestListener_Shutdown(t *testing.T) { + event.Reset() + defer event.Reset() + + listener := &collectListener{} + event.Listen("foo.*", listener) + event.Register("foo", &recordHandler{}) + _ = event.Start() + _ = event.Stop(context.Background()) + + if !listener.shut { + t.Fatal("listener Shutdown should have been called") + } +} + +func TestListener_PanicRecovery(t *testing.T) { + event.Reset() + defer event.Reset() + + var afterPanic atomic.Int32 + pl := &panicListener{afterPanic: &afterPanic} + event.Listen("foo.*", pl) + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, _ = event.Push(context.Background(), "foo.panic", nil) + _, _ = event.Push(context.Background(), "foo.ok", nil) + time.Sleep(100 * time.Millisecond) + + if afterPanic.Load() < 1 { + t.Fatal("listener should recover from panic and process next event") + } +} + +type panicListener struct { + afterPanic *atomic.Int32 + first atomic.Bool +} + +func (l *panicListener) OnEvent(ev *types.Event) { + if !l.first.Load() { + l.first.Store(true) + panic("listener panic") + } + l.afterPanic.Add(1) +} + +func (l *panicListener) Shutdown(ctx context.Context) error { return nil } + +// --- Coverage: notify when listener manager not started --- + +func TestListener_NotifyBeforeStart(t *testing.T) { + event.Reset() + defer event.Reset() + + listener := &collectListener{} + event.Listen("foo.*", listener) + + // Register handler but do NOT start service; Push will fail with ErrNotStarted. + // Instead, we test that listener.notify returns silently before start. + event.Register("foo", &recordHandler{}) + + // Manually start and immediately stop to verify no events leaked + _ = event.Start() + _ = event.Stop(context.Background()) + + events := listener.getEvents() + if len(events) != 0 { + t.Fatalf("expected 0 events before any push, got %d", len(events)) + } +} diff --git a/event/option.go b/event/option.go new file mode 100644 index 00000000..ff578318 --- /dev/null +++ b/event/option.go @@ -0,0 +1,51 @@ +package event + +import "github.com/yaoapp/yao/event/types" + +// MaxWorkers sets the max concurrent worker goroutines for a Handler. +// Default is 512. Workers are fire-and-forget (goroutine ends after task). +func MaxWorkers(n int) types.HandlerOption { + return func(e *types.HandlerEntry) { + e.MaxWorkers = n + } +} + +// ReservedWorkers sets the number of workers reserved for Call events. +// Default is 10. Push can use MaxWorkers - Reserved; Call can use MaxWorkers. +func ReservedWorkers(n int) types.HandlerOption { + return func(e *types.HandlerEntry) { + e.ReservedWorkers = n + } +} + +// QueueSize sets the per-queue capacity. Default is 8192. +// When a queue is full, Push/Call returns ErrQueueFull immediately. +func QueueSize(n int) types.HandlerOption { + return func(e *types.HandlerEntry) { + e.QueueSize = n + } +} + +// Queue sets the queue key for a Push/Call invocation. +// Events with the same queue key are processed serially (FIFO). +func Queue(key string) types.PushOption { + return func(ev *types.Event) { + ev.Queue = key + } +} + +// Filter sets a custom filter function for Listen or Subscribe. +// Events that do not pass the filter are skipped. +func Filter(fn func(*types.Event) bool) types.FilterOption { + return func(e *types.FilterEntry) { + e.Filter = fn + } +} + +// BufferSize sets the Listener channel buffer size. Default is 8192. +// Only effective for Listen; ignored by Subscribe. +func BufferSize(n int) types.FilterOption { + return func(e *types.FilterEntry) { + e.BufferSize = n + } +} diff --git a/event/queue.go b/event/queue.go new file mode 100644 index 00000000..676d98dc --- /dev/null +++ b/event/queue.go @@ -0,0 +1,208 @@ +package event + +import ( + "context" + "sync" + + "github.com/yaoapp/yao/event/types" +) + +// queueItem wraps an event with its execution context and response channel. +type queueItem struct { + ctx context.Context + ev *types.Event + resp chan<- types.Result +} + +// eventQueue is a single FIFO queue bound to a specific handler prefix. +// Events are enqueued and consumed serially by a dedicated goroutine. +type eventQueue struct { + id string + prefix string + ch chan queueItem + released bool + aborted bool + mu sync.Mutex + done chan struct{} // closed when consumer goroutine exits +} + +// enqueue adds an event to the queue. Returns error if full, released, or aborted. +// The send to q.ch is performed while holding q.mu to prevent a race with +// release()/abort() closing the channel between the flag check and the send. +func (q *eventQueue) enqueue(ctx context.Context, ev *types.Event, resp chan<- types.Result) error { + q.mu.Lock() + defer q.mu.Unlock() + + if q.released || q.aborted { + return ErrQueueReleased + } + + select { + case q.ch <- queueItem{ctx: ctx, ev: ev, resp: resp}: + return nil + default: + return ErrQueueFull + } +} + +// release gracefully stops the queue: rejects new events, drains existing ones. +func (q *eventQueue) release() { + q.mu.Lock() + if q.released || q.aborted { + q.mu.Unlock() + return + } + q.released = true + close(q.ch) + q.mu.Unlock() +} + +// abort forcefully stops the queue: rejects new events, discards pending. +// The consumer goroutine detects the aborted flag and skips remaining items. +func (q *eventQueue) abort() { + q.mu.Lock() + if q.aborted { + q.mu.Unlock() + return + } + wasReleased := q.released + q.aborted = true + q.released = true + if !wasReleased { + close(q.ch) + } + q.mu.Unlock() +} + +// consumer is the goroutine that processes queued events serially. +func (q *eventQueue) consumer(pool *workerPool) { + defer close(q.done) + for item := range q.ch { + q.mu.Lock() + aborted := q.aborted + q.mu.Unlock() + if aborted { + continue + } + + // For Push events, use a non-cancellable context so that queued + // fire-and-forget events are not dropped when the caller's ctx expires. + // For Call events, preserve the caller's ctx for deadline/cancellation. + dispatchCtx := item.ctx + if !item.ev.IsCall { + dispatchCtx = context.WithoutCancel(item.ctx) + } + + done, err := pool.dispatch(dispatchCtx, item.ev, item.resp) + if err != nil { + select { + case item.resp <- types.Result{Err: err}: + default: + } + continue + } + <-done + } +} + +// queueManager manages all active queues. +type queueManager struct { + mu sync.RWMutex + queues map[string]*eventQueue + released map[string]struct{} // tracks IDs that have been released/aborted +} + +func newQueueManager() *queueManager { + return &queueManager{ + queues: make(map[string]*eventQueue), + released: make(map[string]struct{}), + } +} + +// create creates a new queue bound to a handler prefix. +func (qm *queueManager) create(prefix string, queueID string, queueSize int, pool *workerPool) error { + qm.mu.Lock() + defer qm.mu.Unlock() + + if _, exists := qm.queues[queueID]; exists { + return ErrQueueExists + } + + q := &eventQueue{ + id: queueID, + prefix: prefix, + ch: make(chan queueItem, queueSize), + done: make(chan struct{}), + } + qm.queues[queueID] = q + go q.consumer(pool) + return nil +} + +// get returns a queue by ID. +// Returns ErrQueueNotFound if the queue was never created, +// or ErrQueueReleased if it has been released/aborted. +func (qm *queueManager) get(queueID string) (*eventQueue, error) { + qm.mu.RLock() + defer qm.mu.RUnlock() + + q, ok := qm.queues[queueID] + if !ok { + if _, wasReleased := qm.released[queueID]; wasReleased { + return nil, ErrQueueReleased + } + return nil, ErrQueueNotFound + } + return q, nil +} + +// release gracefully releases a queue. +func (qm *queueManager) release(queueID string) { + qm.mu.Lock() + q, ok := qm.queues[queueID] + if !ok { + qm.mu.Unlock() + return + } + delete(qm.queues, queueID) + qm.released[queueID] = struct{}{} + qm.mu.Unlock() + + q.release() + go func() { <-q.done }() +} + +// abortOne forcefully releases a single queue. +func (qm *queueManager) abortOne(queueID string) { + qm.mu.Lock() + q, ok := qm.queues[queueID] + if !ok { + qm.mu.Unlock() + return + } + delete(qm.queues, queueID) + qm.released[queueID] = struct{}{} + qm.mu.Unlock() + + q.abort() + go func() { <-q.done }() +} + +// abortAll forcefully releases all queues. Used during Stop. +func (qm *queueManager) abortAll() { + qm.mu.Lock() + queues := make([]*eventQueue, 0, len(qm.queues)) + for id, q := range qm.queues { + queues = append(queues, q) + qm.released[id] = struct{}{} + } + qm.queues = make(map[string]*eventQueue) + qm.mu.Unlock() + + for _, q := range queues { + q.abort() + } + for _, q := range queues { + <-q.done + } +} diff --git a/event/queue_test.go b/event/queue_test.go new file mode 100644 index 00000000..e2028eed --- /dev/null +++ b/event/queue_test.go @@ -0,0 +1,387 @@ +package event_test + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --- Phase 4: Queue tests --- + +// orderHandler records the order of payload values to verify FIFO. +type orderHandler struct { + mu sync.Mutex + order []int +} + +func (h *orderHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + var v int + if err := ev.Should(&v); err == nil { + h.mu.Lock() + h.order = append(h.order, v) + h.mu.Unlock() + } + if ev.IsCall { + resp <- types.Result{Data: v} + } +} + +func (h *orderHandler) Shutdown(ctx context.Context) error { return nil } + +func (h *orderHandler) getOrder() []int { + h.mu.Lock() + defer h.mu.Unlock() + cp := make([]int, len(h.order)) + copy(cp, h.order) + return cp +} + +func TestQueueCreate_Release_FIFO(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &orderHandler{} + event.Register("seq", h, event.QueueSize(100)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, err := event.QueueCreate("seq") + if err != nil { + t.Fatalf("QueueCreate failed: %v", err) + } + if qID == "" { + t.Fatal("expected non-empty queue ID") + } + + n := 20 + for i := 0; i < n; i++ { + _, err := event.Push(context.Background(), "seq.append", i, event.Queue(qID)) + if err != nil { + t.Fatalf("Push %d failed: %v", i, err) + } + } + + // Release and wait for drain + event.QueueRelease(qID) + time.Sleep(200 * time.Millisecond) + + order := h.getOrder() + if len(order) != n { + t.Fatalf("expected %d events, got %d", n, len(order)) + } + for i, v := range order { + if v != i { + t.Fatalf("FIFO violation at index %d: expected %d, got %d", i, i, v) + } + } +} + +func TestQueueCreate_CustomID(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, err := event.QueueCreate("seq", "my-custom-id") + if err != nil { + t.Fatalf("QueueCreate failed: %v", err) + } + if qID != "my-custom-id" { + t.Fatalf("expected my-custom-id, got %s", qID) + } + event.QueueRelease(qID) +} + +func TestQueueCreate_Duplicate(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, _ = event.QueueCreate("seq", "dup-id") + _, err := event.QueueCreate("seq", "dup-id") + if err != event.ErrQueueExists { + t.Fatalf("expected ErrQueueExists, got %v", err) + } + event.QueueRelease("dup-id") +} + +func TestQueueCreate_UnregisteredPrefix(t *testing.T) { + event.Reset() + defer event.Reset() + + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, err := event.QueueCreate("nonexist") + if err != event.ErrNoHandler { + t.Fatalf("expected ErrNoHandler, got %v", err) + } +} + +func TestPush_QueueNotFound(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, err := event.Push(context.Background(), "seq.append", 1, event.Queue("no-such-queue")) + if err != event.ErrQueueNotFound { + t.Fatalf("expected ErrQueueNotFound, got %v", err) + } +} + +func TestPush_QueueReleased(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, _ := event.QueueCreate("seq") + event.QueueRelease(qID) + time.Sleep(50 * time.Millisecond) + + _, err := event.Push(context.Background(), "seq.append", 1, event.Queue(qID)) + if err != event.ErrQueueReleased { + t.Fatalf("expected ErrQueueReleased after release, got %v", err) + } +} + +func TestQueueAbort_DiscardsPending(t *testing.T) { + event.Reset() + defer event.Reset() + + // slowHandler delays processing to let events pile up + var processed atomic.Int32 + slow := &slowHandler{delay: 50 * time.Millisecond, counter: &processed} + event.Register("slow", slow, event.QueueSize(100)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, _ := event.QueueCreate("slow") + + // Push 10 events; first will start processing, rest queue up + for i := 0; i < 10; i++ { + _, _ = event.Push(context.Background(), "slow.work", i, event.Queue(qID)) + } + + time.Sleep(30 * time.Millisecond) // let first event start + event.QueueAbort(qID) + time.Sleep(200 * time.Millisecond) + + count := processed.Load() + if count >= 10 { + t.Fatalf("abort should discard pending events, but %d were processed", count) + } +} + +// slowHandler processes events with a delay. +type slowHandler struct { + delay time.Duration + counter *atomic.Int32 +} + +func (h *slowHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + time.Sleep(h.delay) + h.counter.Add(1) + if ev.IsCall { + resp <- types.Result{Data: "done"} + } +} + +func (h *slowHandler) Shutdown(ctx context.Context) error { return nil } + +func TestQueue_CallInsideQueue_Serial(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &orderHandler{} + event.Register("seq", h, event.QueueSize(100)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, _ := event.QueueCreate("seq") + defer event.QueueRelease(qID) + + // Push 5, then Call, then Push 5 more + for i := 0; i < 5; i++ { + _, _ = event.Push(context.Background(), "seq.append", i, event.Queue(qID)) + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _, data, err := event.Call(ctx, "seq.append", 99, event.Queue(qID)) + if err != nil { + t.Fatalf("Call failed: %v", err) + } + if data != 99 { + t.Fatalf("expected 99, got %v", data) + } + + for i := 5; i < 10; i++ { + _, _ = event.Push(context.Background(), "seq.append", i, event.Queue(qID)) + } + + time.Sleep(200 * time.Millisecond) + order := h.getOrder() + + // The Call (99) should appear after the first 5 and before the last 5 + found := false + for i, v := range order { + if v == 99 { + if i < 5 { + t.Fatalf("Call should be after first 5 pushes, found at index %d", i) + } + found = true + break + } + } + if !found { + t.Fatalf("Call result (99) not found in order: %v", order) + } +} + +func TestQueueFull(t *testing.T) { + event.Reset() + defer event.Reset() + + var processed atomic.Int32 + slow := &slowHandler{delay: 100 * time.Millisecond, counter: &processed} + event.Register("tiny", slow, event.QueueSize(2)) + _ = event.Start() + + qID, _ := event.QueueCreate("tiny") + + // Fill the queue (size=2) + _, err1 := event.Push(context.Background(), "tiny.work", 1, event.Queue(qID)) + _, err2 := event.Push(context.Background(), "tiny.work", 2, event.Queue(qID)) + + // These may or may not succeed depending on timing, but eventually one should fail + var fullErr error + for i := 0; i < 10; i++ { + _, err := event.Push(context.Background(), "tiny.work", i+3, event.Queue(qID)) + if err == event.ErrQueueFull { + fullErr = err + break + } + } + + if err1 != nil { + t.Fatalf("first push should succeed: %v", err1) + } + if err2 != nil { + t.Fatalf("second push should succeed: %v", err2) + } + if fullErr == nil { + t.Log("warning: queue never reported full (handler may be too fast)") + } + + // Wait for queued events to finish before Stop to avoid race between + // consumer goroutine (dispatch/wg.Add) and Stop (pool.wait/wg.Wait). + event.QueueRelease(qID) + time.Sleep(500 * time.Millisecond) + _ = event.Stop(context.Background()) +} + +// --- Coverage: QueueRelease idempotent (release non-existent queue) --- + +func TestQueueRelease_NonExistent(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + // Should not panic + event.QueueRelease("never-created") +} + +// --- Coverage: QueueAbort idempotent (abort non-existent queue) --- + +func TestQueueAbort_NonExistent(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + // Should not panic + event.QueueAbort("never-created") +} + +// --- Coverage: QueueAbort after already released --- + +func TestQueueAbort_AfterRelease(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, _ := event.QueueCreate("seq") + event.QueueRelease(qID) + time.Sleep(50 * time.Millisecond) + + // Abort after release should not panic (already removed from map) + event.QueueAbort(qID) +} + +// --- Coverage: Stop with active queues (abortAll path) --- + +func TestStop_WithActiveQueues(t *testing.T) { + event.Reset() + defer event.Reset() + + var processed atomic.Int32 + slow := &slowHandler{delay: 30 * time.Millisecond, counter: &processed} + event.Register("bg", slow, event.QueueSize(100)) + _ = event.Start() + + qID, _ := event.QueueCreate("bg") + for i := 0; i < 5; i++ { + _, _ = event.Push(context.Background(), "bg.work", i, event.Queue(qID)) + } + time.Sleep(10 * time.Millisecond) + + // Stop should abort all queues and wait + err := event.Stop(context.Background()) + if err != nil { + t.Fatalf("Stop failed: %v", err) + } +} + +// --- Coverage: Call with queue enqueue failure (queue released) --- + +func TestCall_QueueReleased(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("seq", &orderHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + qID, _ := event.QueueCreate("seq") + event.QueueRelease(qID) + time.Sleep(50 * time.Millisecond) + + _, _, err := event.Call(context.Background(), "seq.get", nil, event.Queue(qID)) + if err != event.ErrQueueReleased { + t.Fatalf("expected ErrQueueReleased, got %v", err) + } +} diff --git a/event/service.go b/event/service.go new file mode 100644 index 00000000..d9dfc7a2 --- /dev/null +++ b/event/service.go @@ -0,0 +1,221 @@ +package event + +import ( + "context" + "errors" + "sync" + + "github.com/yaoapp/yao/event/types" +) + +// Sentinel errors. +var ( + ErrNotStarted = errors.New("event: service not started") + ErrAlreadyStart = errors.New("event: service already started") + ErrQueueFull = errors.New("event: queue is full") + ErrQueueNotFound = errors.New("event: queue not found") + ErrQueueExists = errors.New("event: queue already exists") + ErrQueueReleased = errors.New("event: queue already released") + ErrNoHandler = errors.New("event: no handler registered for prefix") + ErrHandlerPanic = errors.New("event: handler panicked") +) + +// Context keys for SID and Auth propagation. +type ctxKey int + +const ( + ctxKeySID ctxKey = iota + ctxKeyAuth +) + +// WithSID returns a context carrying the given session ID. +func WithSID(ctx context.Context, sid string) context.Context { + return context.WithValue(ctx, ctxKeySID, sid) +} + +// SIDFrom extracts the session ID from ctx. Returns empty string if not set. +func SIDFrom(ctx context.Context) string { + if v, ok := ctx.Value(ctxKeySID).(string); ok { + return v + } + return "" +} + +// WithAuth returns a context carrying the given authorized info. +func WithAuth(ctx context.Context, auth *types.AuthorizedInfo) context.Context { + return context.WithValue(ctx, ctxKeyAuth, auth) +} + +// AuthFrom extracts the authorized info from ctx. Returns nil if not set. +func AuthFrom(ctx context.Context) *types.AuthorizedInfo { + if v, ok := ctx.Value(ctxKeyAuth).(*types.AuthorizedInfo); ok { + return v + } + return nil +} + +// service holds all global state for the event bus. +type service struct { + mu sync.RWMutex + started bool + handlers map[string]*types.HandlerEntry // prefix -> registration + pools map[string]*workerPool // prefix -> worker pool + queues *queueManager // queue lifecycle + lmgr *listenerManager // listener manager + smgr *subManager // subscriber manager +} + +var svc = &service{} + +func init() { + svc.reset() +} + +// Register registers a handler for the given prefix. +// Must be called before Start (typically in init()). +func Register(prefix string, handler types.Handler, opts ...types.HandlerOption) { + entry := &types.HandlerEntry{ + Prefix: prefix, + Handler: handler, + MaxWorkers: types.DefaultMaxWorkers, + ReservedWorkers: types.DefaultReservedWorkers, + QueueSize: types.DefaultQueueSize, + } + for _, opt := range opts { + opt(entry) + } + + svc.mu.Lock() + defer svc.mu.Unlock() + svc.handlers[prefix] = entry +} + +// Start initializes and starts the event service. +// Called during engine startup, after runtime is ready. +func Start() error { + svc.mu.Lock() + defer svc.mu.Unlock() + + if svc.started { + return ErrAlreadyStart + } + + // Create worker pools for each registered handler + for prefix, entry := range svc.handlers { + svc.pools[prefix] = newWorkerPool(entry) + } + + // Start listener manager + svc.lmgr.start() + + svc.started = true + return nil +} + +// Stop gracefully shuts down the event service. +// Waits for in-flight events to finish, discards pending queue items, +// and calls Shutdown on all handlers and listeners. +// +// The lock is released before waiting for workers so that in-flight handlers +// calling Push/Call (which acquire RLock via getHandler) do not deadlock. +// Once started=false, getHandler returns ErrNotStarted for any new calls. +func Stop(ctx context.Context) error { + svc.mu.Lock() + if !svc.started { + svc.mu.Unlock() + return nil + } + svc.started = false + + // Snapshot references under lock, then release. + queues := svc.queues + pools := make([]*workerPool, 0, len(svc.pools)) + for _, p := range svc.pools { + pools = append(pools, p) + } + handlers := make([]*types.HandlerEntry, 0, len(svc.handlers)) + for _, e := range svc.handlers { + handlers = append(handlers, e) + } + lmgr := svc.lmgr + smgr := svc.smgr + svc.mu.Unlock() + + // From here on, started=false prevents any new Push/Call/QueueCreate. + // Existing in-flight workers may still call getHandler and get ErrNotStarted, + // which is the correct behavior during shutdown. + + // Abort all queues (discard pending, wait for in-flight) + queues.abortAll() + + // Wait for all worker pools to drain + for _, pool := range pools { + pool.wait() + } + + // Shutdown all handlers + for _, entry := range handlers { + if entry.Handler != nil { + _ = entry.Handler.Shutdown(ctx) + } + } + + // Stop listener manager + lmgr.stop(ctx) + + // Clear subscribers + smgr.clear() + + return nil +} + +// Reload performs a hot-reload. Preserves queues and in-flight events, +// reloads dynamic configuration only. +func Reload() error { + svc.mu.RLock() + defer svc.mu.RUnlock() + + if !svc.started { + return ErrNotStarted + } + return nil +} + +// IsStarted reports whether the service is currently running. +func IsStarted() bool { + svc.mu.RLock() + defer svc.mu.RUnlock() + return svc.started +} + +// getHandler returns the handler entry and its worker pool for the given prefix. +func getHandler(prefix string) (*types.HandlerEntry, *workerPool, error) { + svc.mu.RLock() + defer svc.mu.RUnlock() + + if !svc.started { + return nil, nil, ErrNotStarted + } + entry, ok := svc.handlers[prefix] + if !ok { + return nil, nil, ErrNoHandler + } + pool := svc.pools[prefix] + return entry, pool, nil +} + +// Reset clears all state. For testing only. +func Reset() { + svc.mu.Lock() + defer svc.mu.Unlock() + svc.reset() +} + +func (s *service) reset() { + s.started = false + s.handlers = make(map[string]*types.HandlerEntry) + s.pools = make(map[string]*workerPool) + s.queues = newQueueManager() + s.lmgr = newListenerManager() + s.smgr = newSubManager() +} diff --git a/event/service_test.go b/event/service_test.go new file mode 100644 index 00000000..9a8b7196 --- /dev/null +++ b/event/service_test.go @@ -0,0 +1,229 @@ +package event_test + +import ( + "context" + "testing" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// stubHandler is a minimal Handler for testing registration and lifecycle. +type stubHandler struct { + shutdownCalled bool +} + +func (h *stubHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) {} + +func (h *stubHandler) Shutdown(ctx context.Context) error { + h.shutdownCalled = true + return nil +} + +// --- Register + Start/Stop lifecycle --- + +func TestStartStop_Basic(t *testing.T) { + event.Reset() + defer event.Reset() + + if event.IsStarted() { + t.Fatal("service should not be started initially") + } + + if err := event.Start(); err != nil { + t.Fatalf("Start failed: %v", err) + } + if !event.IsStarted() { + t.Fatal("service should be started after Start") + } + + if err := event.Stop(context.Background()); err != nil { + t.Fatalf("Stop failed: %v", err) + } + if event.IsStarted() { + t.Fatal("service should not be started after Stop") + } +} + +func TestStart_Double(t *testing.T) { + event.Reset() + defer event.Reset() + + if err := event.Start(); err != nil { + t.Fatalf("Start failed: %v", err) + } + + err := event.Start() + if err != event.ErrAlreadyStart { + t.Fatalf("expected ErrAlreadyStart, got: %v", err) + } + + _ = event.Stop(context.Background()) +} + +func TestStop_WhenNotStarted(t *testing.T) { + event.Reset() + defer event.Reset() + + if err := event.Stop(context.Background()); err != nil { + t.Fatalf("Stop on non-started service should succeed, got: %v", err) + } +} + +func TestReload_WhenNotStarted(t *testing.T) { + event.Reset() + defer event.Reset() + + err := event.Reload() + if err != event.ErrNotStarted { + t.Fatalf("expected ErrNotStarted, got: %v", err) + } +} + +func TestReload_WhenStarted(t *testing.T) { + event.Reset() + defer event.Reset() + + _ = event.Start() + if err := event.Reload(); err != nil { + t.Fatalf("Reload failed: %v", err) + } + _ = event.Stop(context.Background()) +} + +// --- Register + options --- + +func TestRegister_DefaultOptions(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &stubHandler{} + event.Register("test", h) + + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + if !event.IsStarted() { + t.Fatal("service should be started") + } +} + +func TestRegister_CustomOptions(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &stubHandler{} + event.Register("test", h, + event.MaxWorkers(128), + event.ReservedWorkers(5), + event.QueueSize(2048), + ) + + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + if !event.IsStarted() { + t.Fatal("service should be started") + } +} + +// --- Stop calls Shutdown on handlers --- + +func TestStop_CallsHandlerShutdown(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &stubHandler{} + event.Register("test", h) + _ = event.Start() + + if err := event.Stop(context.Background()); err != nil { + t.Fatalf("Stop failed: %v", err) + } + if !h.shutdownCalled { + t.Fatal("Handler.Shutdown should have been called on Stop") + } +} + +func TestStop_MultipleHandlersShutdown(t *testing.T) { + event.Reset() + defer event.Reset() + + h1 := &stubHandler{} + h2 := &stubHandler{} + event.Register("alpha", h1) + event.Register("bravo", h2) + _ = event.Start() + + if err := event.Stop(context.Background()); err != nil { + t.Fatalf("Stop failed: %v", err) + } + if !h1.shutdownCalled || !h2.shutdownCalled { + t.Fatal("all handlers should have been shut down") + } +} + +// --- Context SID/Auth propagation --- + +func TestWithSID_SIDFrom(t *testing.T) { + ctx := event.WithSID(context.Background(), "sess-123") + got := event.SIDFrom(ctx) + if got != "sess-123" { + t.Fatalf("expected sess-123, got %s", got) + } +} + +func TestSIDFrom_Empty(t *testing.T) { + got := event.SIDFrom(context.Background()) + if got != "" { + t.Fatalf("expected empty, got %s", got) + } +} + +func TestWithAuth_AuthFrom(t *testing.T) { + auth := &types.AuthorizedInfo{UserID: "u-1", TeamID: "t-1"} + ctx := event.WithAuth(context.Background(), auth) + got := event.AuthFrom(ctx) + if got == nil { + t.Fatal("expected non-nil auth") + } + if got.UserID != "u-1" || got.TeamID != "t-1" { + t.Fatalf("unexpected auth: %+v", got) + } +} + +func TestAuthFrom_Nil(t *testing.T) { + got := event.AuthFrom(context.Background()) + if got != nil { + t.Fatal("expected nil auth from bare context") + } +} + +func TestWithSIDAndAuth_Combined(t *testing.T) { + auth := &types.AuthorizedInfo{UserID: "u-2"} + ctx := event.WithSID(context.Background(), "sess-456") + ctx = event.WithAuth(ctx, auth) + + if event.SIDFrom(ctx) != "sess-456" { + t.Fatal("SID mismatch") + } + if event.AuthFrom(ctx).UserID != "u-2" { + t.Fatal("Auth mismatch") + } +} + +// --- Reset --- + +func TestReset_ClearsState(t *testing.T) { + event.Reset() + + h := &stubHandler{} + event.Register("test", h) + _ = event.Start() + + event.Reset() + + if event.IsStarted() { + t.Fatal("service should not be started after Reset") + } +} diff --git a/event/sub.go b/event/sub.go new file mode 100644 index 00000000..089ad94d --- /dev/null +++ b/event/sub.go @@ -0,0 +1,101 @@ +package event + +import ( + "fmt" + "sync" + "sync/atomic" + + "github.com/yaoapp/yao/event/types" +) + +var subIDCounter atomic.Uint64 + +func nextSubID() string { + id := subIDCounter.Add(1) + return fmt.Sprintf("sub-%d", id) +} + +// subEntry holds a dynamic subscriber registration. +type subEntry struct { + id string + pattern string + filter func(*types.Event) bool + ch chan<- *types.Event +} + +// subManager manages dynamic subscribers. +type subManager struct { + mu sync.RWMutex + entries map[string]*subEntry // id -> entry +} + +func newSubManager() *subManager { + return &subManager{ + entries: make(map[string]*subEntry), + } +} + +// subscribe adds a dynamic subscriber. Returns the subscription ID. +func (sm *subManager) subscribe(pattern string, ch chan<- *types.Event, opts ...types.FilterOption) string { + fe := &types.FilterEntry{Pattern: pattern} + for _, opt := range opts { + opt(fe) + } + + id := nextSubID() + sm.mu.Lock() + defer sm.mu.Unlock() + sm.entries[id] = &subEntry{ + id: id, + pattern: pattern, + filter: fe.Filter, + ch: ch, + } + return id +} + +// unsubscribe removes a subscriber by ID. +func (sm *subManager) unsubscribe(id string) { + sm.mu.Lock() + defer sm.mu.Unlock() + delete(sm.entries, id) +} + +// notify sends an event to all matching subscribers (non-blocking). +func (sm *subManager) notify(ev *types.Event) { + sm.mu.RLock() + defer sm.mu.RUnlock() + + for _, entry := range sm.entries { + if !matchPattern(entry.pattern, ev.Type) { + continue + } + if entry.filter != nil && !entry.filter(ev) { + continue + } + select { + case entry.ch <- ev: + default: + // Subscriber chan full, skip (non-blocking) + } + } +} + +// clear removes all subscribers. Used during Stop. +func (sm *subManager) clear() { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.entries = make(map[string]*subEntry) +} + +// Subscribe dynamically subscribes to events matching the given pattern. +// Returns the subscription ID for later unsubscription. +// Event delivery is non-blocking: if ch is full, the event is skipped. +func Subscribe(pattern string, ch chan<- *types.Event, opts ...types.FilterOption) string { + return svc.smgr.subscribe(pattern, ch, opts...) +} + +// Unsubscribe removes a dynamic subscription by ID. +func Unsubscribe(id string) { + svc.smgr.unsubscribe(id) +} diff --git a/event/sub_test.go b/event/sub_test.go new file mode 100644 index 00000000..abdde4f9 --- /dev/null +++ b/event/sub_test.go @@ -0,0 +1,181 @@ +package event_test + +import ( + "context" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --- Phase 7: Subscriber tests --- + +func TestSubscribe_Basic(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ch := make(chan *types.Event, 10) + subID := event.Subscribe("foo.*", ch) + if subID == "" { + t.Fatal("expected non-empty subscription ID") + } + defer event.Unsubscribe(subID) + + _, _ = event.Push(context.Background(), "foo.bar", "payload") + _, _ = event.Push(context.Background(), "foo.baz", "payload2") + + received := drainChan(ch, 2, 200*time.Millisecond) + if len(received) != 2 { + t.Fatalf("expected 2 events, got %d", len(received)) + } + if received[0].Type != "foo.bar" { + t.Fatalf("expected foo.bar, got %s", received[0].Type) + } + if received[1].Type != "foo.baz" { + t.Fatalf("expected foo.baz, got %s", received[1].Type) + } +} + +func TestSubscribe_PatternFilter(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + event.Register("bar", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ch := make(chan *types.Event, 10) + subID := event.Subscribe("foo.*", ch, event.Filter(func(ev *types.Event) bool { + return ev.Type == "foo.keep" + })) + defer event.Unsubscribe(subID) + + _, _ = event.Push(context.Background(), "foo.keep", nil) + _, _ = event.Push(context.Background(), "foo.drop", nil) + _, _ = event.Push(context.Background(), "bar.thing", nil) + + received := drainChan(ch, 1, 200*time.Millisecond) + if len(received) != 1 { + t.Fatalf("expected 1 filtered event, got %d", len(received)) + } + if received[0].Type != "foo.keep" { + t.Fatalf("expected foo.keep, got %s", received[0].Type) + } +} + +func TestSubscribe_Unsubscribe(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ch := make(chan *types.Event, 10) + subID := event.Subscribe("foo.*", ch) + + _, _ = event.Push(context.Background(), "foo.first", nil) + time.Sleep(50 * time.Millisecond) + + event.Unsubscribe(subID) + + _, _ = event.Push(context.Background(), "foo.second", nil) + time.Sleep(50 * time.Millisecond) + + received := drainChan(ch, 10, 100*time.Millisecond) + for _, ev := range received { + if ev.Type == "foo.second" { + t.Fatal("should not receive events after Unsubscribe") + } + } +} + +func TestSubscribe_ChanFull_Skip(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ch := make(chan *types.Event, 1) // tiny buffer + subID := event.Subscribe("foo.*", ch) + defer event.Unsubscribe(subID) + + // Push multiple events quickly; only 1 should fit in buffer + for i := 0; i < 5; i++ { + _, _ = event.Push(context.Background(), "foo.item", i) + } + + time.Sleep(100 * time.Millisecond) + + // Should have at most 1 in channel (rest skipped) + count := len(ch) + if count > 1 { + t.Fatalf("expected at most 1 buffered event, got %d", count) + } +} + +func TestSubscribe_WildcardAll(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + event.Register("bar", &recordHandler{}) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + ch := make(chan *types.Event, 10) + subID := event.Subscribe("*", ch) + defer event.Unsubscribe(subID) + + _, _ = event.Push(context.Background(), "foo.one", nil) + _, _ = event.Push(context.Background(), "bar.two", nil) + + received := drainChan(ch, 2, 200*time.Millisecond) + if len(received) != 2 { + t.Fatalf("wildcard * should receive all events, got %d", len(received)) + } +} + +func TestSubscribe_StopClearsSubscribers(t *testing.T) { + event.Reset() + defer event.Reset() + + event.Register("foo", &recordHandler{}) + _ = event.Start() + + ch := make(chan *types.Event, 10) + _ = event.Subscribe("foo.*", ch) + + _ = event.Stop(context.Background()) + + // After Stop, Push should fail + _, err := event.Push(context.Background(), "foo.bar", nil) + if err != event.ErrNotStarted { + t.Fatalf("expected ErrNotStarted after Stop, got %v", err) + } +} + +// drainChan reads up to n events from ch within timeout. +func drainChan(ch chan *types.Event, n int, timeout time.Duration) []*types.Event { + var result []*types.Event + timer := time.NewTimer(timeout) + defer timer.Stop() + + for range n { + select { + case ev := <-ch: + result = append(result, ev) + case <-timer.C: + return result + } + } + return result +} diff --git a/event/types/interfaces.go b/event/types/interfaces.go new file mode 100644 index 00000000..9cf17294 --- /dev/null +++ b/event/types/interfaces.go @@ -0,0 +1,23 @@ +package types + +import "context" + +// Handler processes events for a given prefix (registered at startup, one per prefix). +// +// Handle is invoked by the WorkerPool. +// - ctx: for Call, this carries the caller's deadline/cancellation; for Push, a non-cancellable context. +// - resp is always non-nil. For Push the framework passes a discard channel; for Call it waits for a read. +// Use ev.IsCall to decide whether to write a meaningful result. +type Handler interface { + Handle(ctx context.Context, ev *Event, resp chan<- Result) + Shutdown(ctx context.Context) error +} + +// Listener receives matched events in a dedicated goroutine (registered at startup). +// +// OnEvent is called in the Listener's own goroutine; it does not block other +// Listeners or Subscribers. +type Listener interface { + OnEvent(ev *Event) + Shutdown(ctx context.Context) error +} diff --git a/event/types/types.go b/event/types/types.go new file mode 100644 index 00000000..aaed979a --- /dev/null +++ b/event/types/types.go @@ -0,0 +1,102 @@ +package types + +import ( + "fmt" + "reflect" + + "github.com/yaoapp/gou/process" +) + +// AuthorizedInfo is an alias for gou/process.AuthorizedInfo. +type AuthorizedInfo = process.AuthorizedInfo + +// Event represents a single event in the event bus. +type Event struct { + Type string // Event type, e.g. "trace.add", "job.progress" + ID string // Auto-generated event ID + Queue string // Queue key for serial processing; empty means no queue + IsCall bool // true = synchronous Call, false = asynchronous Push + Payload any // Business data; concrete type is determined by event type + SID string // Session ID, extracted from caller context + Auth *AuthorizedInfo // Authorized info, extracted from caller context; may be nil +} + +// Should asserts the Payload to the target pointer type. +// target must be a non-nil pointer. Returns an error if the type does not match. +// +// Usage: +// +// var p MyPayload +// if err := ev.Should(&p); err != nil { ... } +func (ev *Event) Should(target any) error { + if target == nil { + return fmt.Errorf("event.Should: target must be a non-nil pointer") + } + + rv := reflect.ValueOf(target) + if rv.Kind() != reflect.Ptr || rv.IsNil() { + return fmt.Errorf("event.Should: target must be a non-nil pointer, got %T", target) + } + + if ev.Payload == nil { + return fmt.Errorf("event.Should: payload is nil") + } + + // Direct assignment: payload is already the expected pointer type + payloadVal := reflect.ValueOf(ev.Payload) + targetElem := rv.Elem() + + // If payload is a pointer, dereference it + if payloadVal.Kind() == reflect.Ptr { + if payloadVal.IsNil() { + return fmt.Errorf("event.Should: payload is nil pointer") + } + payloadVal = payloadVal.Elem() + } + + if !payloadVal.Type().AssignableTo(targetElem.Type()) { + return fmt.Errorf("event.Should: payload type %T is not assignable to %s", ev.Payload, targetElem.Type()) + } + + targetElem.Set(payloadVal) + return nil +} + +// Result holds the response from a synchronous Call. +type Result struct { + Data any + Err error +} + +// HandlerOption configures a Handler registration. +type HandlerOption func(*HandlerEntry) + +// HandlerEntry is the internal registration record for a Handler. +type HandlerEntry struct { + Prefix string + Handler Handler + MaxWorkers int // Max concurrent workers, default 512 + ReservedWorkers int // Workers reserved for Call, default 10 + QueueSize int // Per-queue capacity, default 8192 +} + +// FilterOption configures a Listener or Subscriber registration. +type FilterOption func(*FilterEntry) + +// FilterEntry is the internal registration record for a Listener/Subscriber. +type FilterEntry struct { + Pattern string + Filter func(*Event) bool // Custom filter function + BufferSize int // Listener chan buffer size, default 8192; only for Listen +} + +// PushOption configures a Push or Call invocation. +type PushOption func(*Event) + +// Default configuration values. +const ( + DefaultMaxWorkers = 512 + DefaultReservedWorkers = 10 + DefaultQueueSize = 8192 + DefaultBufferSize = 8192 +) diff --git a/event/types/types_test.go b/event/types/types_test.go new file mode 100644 index 00000000..8fd99b5a --- /dev/null +++ b/event/types/types_test.go @@ -0,0 +1,267 @@ +package types_test + +import ( + "fmt" + "testing" + + "github.com/yaoapp/yao/event/types" +) + +// samplePayload is a test-only struct with no business semantics. +type samplePayload struct { + Name string + Value int + Tags []string +} + +// --- Should: basic struct assignment --- + +func TestShould_StructValue(t *testing.T) { + ev := &types.Event{ + Payload: samplePayload{Name: "alpha", Value: 1, Tags: []string{"a", "b"}}, + } + + var got samplePayload + if err := ev.Should(&got); err != nil { + t.Fatalf("Should returned error: %v", err) + } + if got.Name != "alpha" || got.Value != 1 || len(got.Tags) != 2 { + t.Fatalf("unexpected payload: %+v", got) + } +} + +// --- Should: pointer payload --- + +func TestShould_PointerPayload(t *testing.T) { + ev := &types.Event{ + Payload: &samplePayload{Name: "beta", Value: 2}, + } + + var got samplePayload + if err := ev.Should(&got); err != nil { + t.Fatalf("Should returned error: %v", err) + } + if got.Name != "beta" || got.Value != 2 { + t.Fatalf("unexpected payload: %+v", got) + } +} + +// --- Should: primitive payloads --- + +func TestShould_StringPayload(t *testing.T) { + ev := &types.Event{ + Payload: "hello world", + } + + var got string + if err := ev.Should(&got); err != nil { + t.Fatalf("Should returned error: %v", err) + } + if got != "hello world" { + t.Fatalf("unexpected string: %s", got) + } +} + +func TestShould_IntPayload(t *testing.T) { + ev := &types.Event{ + Payload: 42, + } + + var got int + if err := ev.Should(&got); err != nil { + t.Fatalf("Should returned error: %v", err) + } + if got != 42 { + t.Fatalf("unexpected int: %d", got) + } +} + +// --- Should: error cases --- + +func TestShould_NilTarget(t *testing.T) { + ev := &types.Event{Payload: "data"} + if err := ev.Should(nil); err == nil { + t.Fatal("expected error for nil target") + } +} + +func TestShould_NonPointerTarget(t *testing.T) { + ev := &types.Event{Payload: "data"} + var s string + if err := ev.Should(s); err == nil { + t.Fatal("expected error for non-pointer target") + } +} + +func TestShould_NilPayload(t *testing.T) { + ev := &types.Event{Payload: nil} + var got string + if err := ev.Should(&got); err == nil { + t.Fatal("expected error for nil payload") + } +} + +func TestShould_NilPointerPayload(t *testing.T) { + ev := &types.Event{Payload: (*samplePayload)(nil)} + var got samplePayload + if err := ev.Should(&got); err == nil { + t.Fatal("expected error for nil pointer payload") + } +} + +func TestShould_TypeMismatch(t *testing.T) { + ev := &types.Event{ + Payload: "wrong type", + } + var got samplePayload + if err := ev.Should(&got); err == nil { + t.Fatal("expected error for type mismatch") + } +} + +// --- Event fields --- + +func TestEvent_NilAuth(t *testing.T) { + ev := &types.Event{ + Type: "x.y", + ID: "ev-100", + Auth: nil, + } + if ev.Auth != nil { + t.Fatal("Auth should be nil") + } + if ev.Type != "x.y" || ev.ID != "ev-100" { + t.Fatalf("unexpected Type/ID: %s/%s", ev.Type, ev.ID) + } +} + +func TestEvent_WithAuth(t *testing.T) { + ev := &types.Event{ + Type: "x.y", + ID: "ev-101", + SID: "sess-abc", + Auth: &types.AuthorizedInfo{ + UserID: "u-1", + TeamID: "t-1", + }, + } + if ev.Type != "x.y" || ev.ID != "ev-101" { + t.Fatalf("unexpected Type/ID: %s/%s", ev.Type, ev.ID) + } + if ev.SID != "sess-abc" { + t.Fatalf("unexpected SID: %s", ev.SID) + } + if ev.Auth.UserID != "u-1" || ev.Auth.TeamID != "t-1" { + t.Fatalf("unexpected Auth: %+v", ev.Auth) + } +} + +func TestEvent_QueueAndIsCall(t *testing.T) { + push := &types.Event{Queue: "q-1", IsCall: false} + call := &types.Event{Queue: "q-1", IsCall: true} + + if push.IsCall { + t.Fatal("Push event should not be IsCall") + } + if !call.IsCall { + t.Fatal("Call event should be IsCall") + } + if push.Queue != "q-1" || call.Queue != "q-1" { + t.Fatal("Queue key mismatch") + } +} + +// --- Result --- + +func TestResult_Success(t *testing.T) { + r := types.Result{Data: map[string]string{"k": "v"}, Err: nil} + if r.Err != nil { + t.Fatal("expected nil error") + } + m, ok := r.Data.(map[string]string) + if !ok || m["k"] != "v" { + t.Fatal("unexpected result data") + } +} + +func TestResult_Error(t *testing.T) { + r := types.Result{Data: nil, Err: fmt.Errorf("something failed")} + if r.Err == nil { + t.Fatal("expected error") + } + if r.Err.Error() != "something failed" { + t.Fatalf("unexpected error message: %s", r.Err.Error()) + } + if r.Data != nil { + t.Fatal("expected nil data") + } +} + +// --- HandlerEntry defaults --- + +func TestHandlerEntry_Defaults(t *testing.T) { + entry := types.HandlerEntry{} + if entry.MaxWorkers != 0 { + t.Fatal("zero value should be 0 before applying options") + } + + if entry.MaxWorkers == 0 { + entry.MaxWorkers = types.DefaultMaxWorkers + } + if entry.ReservedWorkers == 0 { + entry.ReservedWorkers = types.DefaultReservedWorkers + } + if entry.QueueSize == 0 { + entry.QueueSize = types.DefaultQueueSize + } + + if entry.MaxWorkers != 512 { + t.Fatalf("expected MaxWorkers 512, got %d", entry.MaxWorkers) + } + if entry.ReservedWorkers != 10 { + t.Fatalf("expected ReservedWorkers 10, got %d", entry.ReservedWorkers) + } + if entry.QueueSize != 8192 { + t.Fatalf("expected QueueSize 8192, got %d", entry.QueueSize) + } +} + +// --- FilterEntry --- + +func TestFilterEntry_WithFilter(t *testing.T) { + called := false + entry := types.FilterEntry{ + Pattern: "x.*", + Filter: func(ev *types.Event) bool { + called = true + return ev.Type == "x.hit" + }, + BufferSize: 4096, + } + + if entry.Pattern != "x.*" { + t.Fatalf("unexpected Pattern: %s", entry.Pattern) + } + if !entry.Filter(&types.Event{Type: "x.hit"}) { + t.Fatal("filter should match x.hit") + } + if !called { + t.Fatal("filter was not called") + } + if entry.Filter(&types.Event{Type: "x.miss"}) { + t.Fatal("filter should not match x.miss") + } + if entry.BufferSize != 4096 { + t.Fatalf("unexpected BufferSize: %d", entry.BufferSize) + } +} + +func TestFilterEntry_NilFilter(t *testing.T) { + entry := types.FilterEntry{Pattern: "y.*"} + if entry.Pattern != "y.*" { + t.Fatalf("unexpected Pattern: %s", entry.Pattern) + } + if entry.Filter != nil { + t.Fatal("Filter should be nil when not set") + } +} diff --git a/event/worker.go b/event/worker.go new file mode 100644 index 00000000..3fb752b8 --- /dev/null +++ b/event/worker.go @@ -0,0 +1,94 @@ +package event + +import ( + "context" + "sync" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/event/types" +) + +// workerPool manages goroutine-based workers for a single Handler. +// Workers are fire-and-forget: each goroutine processes one task then exits. +// MaxWorkers limits total concurrent goroutines. +// ReservedWorkers reserves slots for Call events so Push cannot starve them. +type workerPool struct { + handler types.Handler + + // semTotal is a buffered channel of size MaxWorkers. + semTotal chan struct{} + + // semPush is a buffered channel of size (MaxWorkers - ReservedWorkers). + // Push events must acquire from both semPush and semTotal. + // Call events only acquire from semTotal. + semPush chan struct{} + + wg sync.WaitGroup +} + +func newWorkerPool(entry *types.HandlerEntry) *workerPool { + pushSlots := entry.MaxWorkers - entry.ReservedWorkers + if pushSlots < 1 { + pushSlots = 1 + } + return &workerPool{ + handler: entry.Handler, + semTotal: make(chan struct{}, entry.MaxWorkers), + semPush: make(chan struct{}, pushSlots), + } +} + +// dispatch runs the handler for one event in a new goroutine. +// Returns a done channel that is closed when the handler finishes. +// Blocks until a worker slot is available or ctx is cancelled. +func (wp *workerPool) dispatch(ctx context.Context, ev *types.Event, resp chan<- types.Result) (done <-chan struct{}, err error) { + isPush := !ev.IsCall + + if isPush { + select { + case wp.semPush <- struct{}{}: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + + select { + case wp.semTotal <- struct{}{}: + case <-ctx.Done(): + if isPush { + <-wp.semPush + } + return nil, ctx.Err() + } + + ch := make(chan struct{}) + wp.wg.Add(1) + go func() { + defer close(ch) + defer wp.wg.Done() + defer func() { <-wp.semTotal }() + if isPush { + defer func() { <-wp.semPush }() + } + defer wp.recoverPanic(ev, resp) + + wp.handler.Handle(ctx, ev, resp) + }() + + return ch, nil +} + +func (wp *workerPool) recoverPanic(ev *types.Event, resp chan<- types.Result) { + if r := recover(); r != nil { + log.Error("event worker panic: type=%s id=%s err=%v", ev.Type, ev.ID, r) + select { + case resp <- types.Result{Err: ErrHandlerPanic}: + default: + } + } +} + +// wait blocks until all active workers finish. Used during Stop. +func (wp *workerPool) wait() { + wp.wg.Wait() +} diff --git a/event/worker_test.go b/event/worker_test.go new file mode 100644 index 00000000..d2106bf6 --- /dev/null +++ b/event/worker_test.go @@ -0,0 +1,213 @@ +package event_test + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/event/types" +) + +// --- Phase 5: Worker pool tests --- + +func TestWorker_MaxConcurrency(t *testing.T) { + event.Reset() + defer event.Reset() + + var peak atomic.Int32 + var current atomic.Int32 + + h := &concurrencyHandler{peak: &peak, current: ¤t, delay: 30 * time.Millisecond} + event.Register("conc", h, event.MaxWorkers(4), event.ReservedWorkers(1)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + var wg sync.WaitGroup + for i := 0; i < 20; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, _ = event.Push(context.Background(), "conc.work", i) + }(i) + } + wg.Wait() + time.Sleep(300 * time.Millisecond) + + p := peak.Load() + if p > 4 { + t.Fatalf("peak concurrency %d exceeded MaxWorkers 4", p) + } + if p < 2 { + t.Fatalf("peak concurrency %d seems too low, expected at least 2", p) + } +} + +type concurrencyHandler struct { + peak *atomic.Int32 + current *atomic.Int32 + delay time.Duration +} + +func (h *concurrencyHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + c := h.current.Add(1) + for { + old := h.peak.Load() + if c <= old || h.peak.CompareAndSwap(old, c) { + break + } + } + time.Sleep(h.delay) + h.current.Add(-1) + if ev.IsCall { + resp <- types.Result{Data: "ok"} + } +} + +func (h *concurrencyHandler) Shutdown(ctx context.Context) error { return nil } + +func TestWorker_CallReservation(t *testing.T) { + event.Reset() + defer event.Reset() + + // MaxWorkers=4, ReservedWorkers=2 => Push can use 2, Call can use 4 + var pushActive atomic.Int32 + var callDone atomic.Int32 + + h := &reservationHandler{pushActive: &pushActive, callDone: &callDone} + event.Register("res", h, event.MaxWorkers(4), event.ReservedWorkers(2)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + // Saturate push slots (only 2 available for push) + for i := 0; i < 4; i++ { + _, _ = event.Push(context.Background(), "res.work", i) + } + time.Sleep(20 * time.Millisecond) // let pushes start + + // Call should still work (reserved slots) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _, data, err := event.Call(ctx, "res.get", "ping") + if err != nil { + t.Fatalf("Call should succeed with reserved workers: %v", err) + } + if data != "pong" { + t.Fatalf("expected pong, got %v", data) + } +} + +type reservationHandler struct { + pushActive *atomic.Int32 + callDone *atomic.Int32 +} + +func (h *reservationHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + if ev.IsCall { + resp <- types.Result{Data: "pong"} + h.callDone.Add(1) + return + } + h.pushActive.Add(1) + time.Sleep(100 * time.Millisecond) + h.pushActive.Add(-1) +} + +func (h *reservationHandler) Shutdown(ctx context.Context) error { return nil } + +func TestWorker_PanicRecovery(t *testing.T) { + event.Reset() + defer event.Reset() + + var afterPanic atomic.Bool + + h := &panicHandler{afterPanic: &afterPanic} + event.Register("pan", h) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + // First push panics + _, _ = event.Push(context.Background(), "pan.crash", "boom") + time.Sleep(50 * time.Millisecond) + + // Second push should still work + _, _ = event.Push(context.Background(), "pan.ok", "fine") + time.Sleep(50 * time.Millisecond) + + if !afterPanic.Load() { + t.Fatal("handler should have processed event after panic recovery") + } +} + +type panicHandler struct { + afterPanic *atomic.Bool +} + +func (h *panicHandler) Handle(ctx context.Context, ev *types.Event, resp chan<- types.Result) { + if ev.Type == "pan.crash" { + panic("test panic") + } + h.afterPanic.Store(true) + if ev.IsCall { + resp <- types.Result{Data: "ok"} + } +} + +func (h *panicHandler) Shutdown(ctx context.Context) error { return nil } + +// --- Coverage: ReservedWorkers >= MaxWorkers (pushSlots clamped to 1) --- + +func TestWorker_ReservedExceedsMax(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &recordHandler{} + event.Register("edge", h, event.MaxWorkers(2), event.ReservedWorkers(5)) + _ = event.Start() + defer func() { _ = event.Stop(context.Background()) }() + + _, err := event.Push(context.Background(), "edge.work", "data") + if err != nil { + t.Fatalf("Push failed: %v", err) + } + time.Sleep(50 * time.Millisecond) + + calls := h.getCalls() + if len(calls) != 1 { + t.Fatalf("expected 1 call, got %d", len(calls)) + } +} + +// --- Coverage: dispatch Call ctx cancel while waiting for semTotal --- + +func TestWorker_Call_CtxCancel_SemTotal(t *testing.T) { + event.Reset() + defer event.Reset() + + h := &concurrencyHandler{ + peak: &atomic.Int32{}, + current: &atomic.Int32{}, + delay: 200 * time.Millisecond, + } + event.Register("lim", h, event.MaxWorkers(1), event.ReservedWorkers(0)) + _ = event.Start() + + bgDone := make(chan struct{}) + go func() { + defer close(bgDone) + _, _, _ = event.Call(context.Background(), "lim.work", nil) + }() + time.Sleep(10 * time.Millisecond) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + _, _, err := event.Call(ctx, "lim.op", nil) + if err == nil { + t.Fatal("expected error for call with saturated pool") + } + + <-bgDone + _ = event.Stop(context.Background()) +} From af981dc0b070456bb1783f8d6ff98c189f253bd3 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 23 Feb 2026 19:22:31 +0800 Subject: [PATCH 2/3] Update dependencies and enhance logging in the Assistant module - Add new indirect dependencies including various Charmbracelet packages for improved UI handling. - Enhance logging in the Assistant module by adding tool completion and start logging for better traceability of tool calls. - Modify context handling in the RequestLogger to support a stack-based assistant ID management, improving the logging structure for agent requests. - Implement event service integration for better trace management and debugging capabilities. --- agent/assistant/agent.go | 1 + agent/assistant/mcp.go | 27 +- agent/caller/orchestrator.go | 48 ++ agent/context/context.go | 28 +- agent/context/log.go | 393 ++++++++++++----- agent/context/tui.go | 799 ++++++++++++++++++++++++++++++++++ agent/context/tui_msg.go | 90 ++++ cmd/start.go | 42 +- engine/load.go | 15 + go.mod | 14 + go.sum | 29 ++ test/utils.go | 19 +- trace/event_listener.go | 25 ++ trace/handler.go | 28 ++ trace/manager.go | 154 +++---- trace/node.go | 23 +- trace/pubsub/pubsub.go | 142 ------ trace/pubsub/subscriber.go | 62 --- trace/space.go | 20 +- trace/state.go | 362 ++++----------- trace/subscription.go | 78 +++- trace/trace.go | 105 +---- trace/trace_basic_test.go | 13 +- trace/trace_lifecycle_test.go | 25 +- 24 files changed, 1678 insertions(+), 864 deletions(-) create mode 100644 agent/context/tui.go create mode 100644 agent/context/tui_msg.go create mode 100644 trace/event_listener.go create mode 100644 trace/handler.go delete mode 100644 trace/pubsub/pubsub.go delete mode 100644 trace/pubsub/subscriber.go diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 5d9f070d..e1fcdda0 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -103,6 +103,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Log end of request ctx.Logger.End(finalStatus == context.StepStatusCompleted, finalError) + ctx.Logger.RestoreAssistantID() }() // Determine stream handler diff --git a/agent/assistant/mcp.go b/agent/assistant/mcp.go index 1e544e00..4696f230 100644 --- a/agent/assistant/mcp.go +++ b/agent/assistant/mcp.go @@ -319,6 +319,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Content = result.Error.Error() result.IsRetryableError = true // Argument parsing error is retryable by LLM ctx.Logger.Error("Failed to parse arguments: %v", err) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -333,6 +334,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Content = result.Error.Error() result.IsRetryableError = true // Type error is retryable by LLM ctx.Logger.Error("Arguments must be an object, got %T", parsed) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -346,6 +348,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Content = result.Error.Error() result.IsRetryableError = true // Validation error is retryable by LLM ctx.Logger.Error("Argument validation failed: %v", err) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -393,7 +396,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall } result.Content = string(contentBytes) - ctx.Logger.ToolComplete(toolName, true) + ctx.Logger.ToolComplete(toolCall.Function.Name, true) if toolNode != nil { toolNode.Complete(map[string]any{ @@ -563,6 +566,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context Arguments: args, }) callMap[toolName] = tc + ctx.Logger.ToolStart(tc.Function.Name) // Add trace input for this tool parallelInputs = append(parallelInputs, types.TraceParallelInput{ @@ -598,11 +602,15 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls, ctx) if err != nil { ctx.Logger.Error("Parallel call failed: %v", err) - // Mark all trace nodes as failed - for _, node := range toolNodes { + for i, node := range toolNodes { if node != nil { node.Fail(err) } + if i < len(mcpCalls) { + if tc, ok := callMap[mcpCalls[i].Name]; ok { + ctx.Logger.ToolComplete(tc.Function.Name, false) + } + } } return nil, true } @@ -631,6 +639,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context result.Content = result.Error.Error() result.IsRetryableError = false // Serialization error is not retryable hasErrors = true + ctx.Logger.ToolComplete(originalCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -643,11 +652,12 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context result.IsRetryableError = isRetryableToolError(result.Error) hasErrors = true ctx.Logger.Error("Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError) + ctx.Logger.ToolComplete(originalCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } } else { - // Success + ctx.Logger.ToolComplete(originalCall.Function.Name, true) if toolNode != nil { toolNode.Complete(map[string]any{ "result": mcpResult.Content, @@ -670,6 +680,8 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte ctx.Logger.Debug("Calling %d tools sequentially on server '%s'", len(toolCalls), serverID) for _, tc := range toolCalls { + ctx.Logger.ToolStart(tc.Function.Name) + _, toolName, ok := ParseMCPToolName(tc.Function.Name) if !ok { results = append(results, ToolCallResult{ @@ -678,6 +690,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte Content: fmt.Sprintf("Invalid tool name format: %s", tc.Function.Name), Error: fmt.Errorf("invalid tool name format"), }) + ctx.Logger.ToolComplete(tc.Function.Name, false) hasErrors = true continue } @@ -727,6 +740,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } results = append(results, result) hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -747,6 +761,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } results = append(results, result) hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -765,6 +780,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } results = append(results, result) hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -788,6 +804,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte result.IsRetryableError = isRetryableToolError(err) hasErrors = true ctx.Logger.Error("Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError) + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -806,11 +823,13 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte result.Content = fmt.Sprintf("Failed to serialize result: %v", err) result.IsRetryableError = false // Serialization error is not retryable hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } } else { result.Content = string(contentBytes) + ctx.Logger.ToolComplete(tc.Function.Name, !mcpResult.IsError) if toolNode != nil { toolNode.Complete(map[string]any{ "result": mcpResult.Content, diff --git a/agent/caller/orchestrator.go b/agent/caller/orchestrator.go index 65589c0a..fdf78655 100644 --- a/agent/caller/orchestrator.go +++ b/agent/caller/orchestrator.go @@ -5,6 +5,7 @@ import ( "sync" agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/trace/types" ) // Orchestrator handles parallel agent calls with different concurrency patterns @@ -257,13 +258,60 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ ctxOpts.OnMessage = req.Handler } + // Add trace node for A2A call using the ORIGINAL parent context's trace + // (forked contexts have nil trace and nil Stack, so ctx.Trace() would create a new orphan trace) + parentTrace, _ := o.ctx.Trace() + var a2aNode types.Node + if parentTrace != nil { + a2aNode, _ = parentTrace.Add( + map[string]any{ + "agent_id": req.AgentID, + "referer": string(ctx.Referer), + }, + types.TraceNodeOption{ + Label: fmt.Sprintf("Agent: %s", req.AgentID), + Type: "agent_call", + Icon: "smart_toy", + Description: fmt.Sprintf("A2A call to '%s'", req.AgentID), + }, + ) + } + + // Notify TUI of A2A call start (use parent requestID so it appears in parent panel) + parentRequestID := o.ctx.RequestID() + agentContext.SendTUI(agentContext.AgentEventMsg{ + RequestID: parentRequestID, + Event: agentContext.EventA2AStart, + Data: map[string]interface{}{"target": req.AgentID}, + }) + // Execute the agent call with the provided context // The agent.Stream method will use the context's Writer for output resp, err := agent.Stream(ctx, req.Messages, ctxOpts) if err != nil { + if a2aNode != nil { + a2aNode.Fail(err) + } + agentContext.SendTUI(agentContext.AgentEventMsg{ + RequestID: parentRequestID, + Event: agentContext.EventA2ADone, + Data: map[string]interface{}{"target": req.AgentID, "error": err.Error()}, + }) return NewResult(req.AgentID, nil, fmt.Errorf("agent call failed: %w", err)) } + if a2aNode != nil { + a2aNode.Complete(map[string]any{ + "agent_id": req.AgentID, + "status": "completed", + }) + } + agentContext.SendTUI(agentContext.AgentEventMsg{ + RequestID: parentRequestID, + Event: agentContext.EventA2ADone, + Data: map[string]interface{}{"target": req.AgentID}, + }) + return NewResult(req.AgentID, resp, nil) } diff --git a/agent/context/context.go b/agent/context/context.go index b419cf49..05582029 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -71,19 +71,21 @@ func (ctx *Context) Release() { ctx.Interrupt = nil } - // Complete and release trace if exists + // Complete and release trace if exists. + // Only the root context (non-forked) owns the trace lifecycle. + // Forked contexts share the same trace manager but must not release it. if ctx.trace != nil && ctx.Stack != nil && ctx.Stack.TraceID != "" { - if ctx.Logger != nil { - ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID) - } - - // Check if context is cancelled - if so, mark as cancelled instead of complete - if ctx.Context != nil && ctx.Context.Err() != nil { - trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()) - trace.Release(ctx.Stack.TraceID) - } else { - ctx.trace.MarkComplete() - trace.Release(ctx.Stack.TraceID) + if ctx.ForkParent == nil { + if ctx.Logger != nil { + ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID) + } + if ctx.Context != nil && ctx.Context.Err() != nil { + trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()) + trace.Release(ctx.Stack.TraceID) + } else { + ctx.trace.MarkComplete() + trace.Release(ctx.Stack.TraceID) + } } ctx.trace = nil } @@ -191,7 +193,7 @@ func (ctx *Context) Fork() *Context { // Create independent resources to avoid race conditions IDGenerator: message.NewIDGenerator(), - Logger: NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID), + Logger: NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID, WithParentID(ctx.ID)), messageMetadata: newMessageMetadataStore(), // Inherit context metadata diff --git a/agent/context/log.go b/agent/context/log.go index 16658e42..503fc3fe 100644 --- a/agent/context/log.go +++ b/agent/context/log.go @@ -72,11 +72,12 @@ type LogEntry struct { // RequestLogger provides request-scoped async logging type RequestLogger struct { - assistantID string - chatID string - requestID string - shortID string // Short version of requestID for display - startTime time.Time + assistantIDStack []string // Stack-based: delegate calls push, pop on exit; top = current + chatID string + requestID string + shortID string // Short version of requestID for display + parentID string // Parent request ID for A2A tree structure + startTime time.Time ch chan LogEntry done chan struct{} @@ -86,19 +87,33 @@ type RequestLogger struct { mu sync.RWMutex } +// LoggerOption configures a RequestLogger +type LoggerOption func(*RequestLogger) + +// WithParentID sets the parent request ID for A2A tree structure +func WithParentID(parentID string) LoggerOption { + return func(l *RequestLogger) { + l.parentID = parentID + } +} + // noopLogger is a shared no-op logger instance var noopLogger = &RequestLogger{noop: true} // NewRequestLogger creates a new request-scoped logger with async processing -func NewRequestLogger(assistantID, chatID, requestID string) *RequestLogger { +func NewRequestLogger(assistantID, chatID, requestID string, opts ...LoggerOption) *RequestLogger { l := &RequestLogger{ - assistantID: assistantID, - chatID: chatID, - requestID: requestID, - shortID: shortID(requestID), - startTime: time.Now(), - ch: make(chan LogEntry, 100), // Buffered channel - done: make(chan struct{}), + assistantIDStack: []string{assistantID}, + chatID: chatID, + requestID: requestID, + shortID: shortID(requestID), + startTime: time.Now(), + ch: make(chan LogEntry, 100), // Buffered channel + done: make(chan struct{}), + } + + for _, opt := range opts { + opt(l) } // Start consumer goroutine @@ -112,12 +127,35 @@ func Noop() *RequestLogger { return noopLogger } -// SetAssistantID sets the assistant ID (called when entering Stream) +// SetAssistantID pushes a new assistant ID onto the stack (called when entering Stream). +// Each SetAssistantID must be paired with a RestoreAssistantID on exit. func (l *RequestLogger) SetAssistantID(id string) { if l.noop { return } - l.assistantID = id + l.mu.Lock() + l.assistantIDStack = append(l.assistantIDStack, id) + l.mu.Unlock() +} + +// RestoreAssistantID pops the current assistant ID, reverting to the previous one. +// Safe to call even if the stack has only one entry (the initial ID is never removed). +func (l *RequestLogger) RestoreAssistantID() { + if l.noop { + return + } + l.mu.Lock() + if len(l.assistantIDStack) > 1 { + l.assistantIDStack = l.assistantIDStack[:len(l.assistantIDStack)-1] + } + l.mu.Unlock() +} + +func (l *RequestLogger) currentAssistantID() string { + if len(l.assistantIDStack) == 0 { + return "" + } + return l.assistantIDStack[len(l.assistantIDStack)-1] } // Close closes the logger and waits for all entries to be processed @@ -148,13 +186,17 @@ func (l *RequestLogger) consume() { func (l *RequestLogger) processEntry(entry LogEntry) { if config.IsDevelopment() { l.printDev(entry) + l.writeLog(entry, true) } else { - l.printProd(entry) + l.writeLog(entry, false) } } -// printDev prints colorful output for development mode +// printDev sends to TUI if available, otherwise prints colored output to stdout func (l *RequestLogger) printDev(entry LogEntry) { + if GetTUIProgram() != nil { + return + } switch entry.Level { case LogLevelTrace: fmt.Printf("%s β†’ %s%s\n", colorGray, entry.Message, colorReset) @@ -169,10 +211,13 @@ func (l *RequestLogger) printDev(entry LogEntry) { } } -// printProd logs to kun/log for production mode -func (l *RequestLogger) printProd(entry LogEntry) { +// writeLog writes structured events to kun/log +func (l *RequestLogger) writeLog(entry LogEntry, devMode bool) { prefix := fmt.Sprintf("[AGENT] %s ", l.shortID) - + if devMode { + kunlog.Trace("%s%s", prefix, entry.Message) + return + } switch entry.Level { case LogLevelTrace: kunlog.Trace("%s%s", prefix, entry.Message) @@ -263,18 +308,29 @@ func (l *RequestLogger) Start() { return } + kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s", + l.shortID, l.currentAssistantID(), shortID(l.chatID), shortID(l.requestID)) + if !config.IsDevelopment() { - kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s", - l.shortID, l.assistantID, shortID(l.chatID), shortID(l.requestID)) return } - // Development: colorful output (direct print, not through channel for immediate display) + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + ParentID: l.parentID, + AssistantID: l.currentAssistantID(), + Event: EventRequestStart, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Println() fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset) - fmt.Printf("%s πŸš€ AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset) + fmt.Printf("%s AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset) fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("─", 60), colorReset) - fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset) + fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.currentAssistantID(), colorReset) fmt.Printf("%s Chat ID: %s%s%s\n", colorGray, colorWhite, l.chatID, colorReset) fmt.Printf("%s Request: %s%s%s\n", colorGray, colorWhite, l.requestID, colorReset) fmt.Printf("%s Time: %s%s%s\n", colorGray, colorWhite, l.startTime.Format("15:04:05.000"), colorReset) @@ -289,28 +345,43 @@ func (l *RequestLogger) End(success bool, err error) { duration := time.Since(l.startTime) + if success { + kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v", + l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond)) + } else { + kunlog.Error("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v", + l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond), err) + } + if !config.IsDevelopment() { - if success { - kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v", - l.shortID, l.assistantID, duration.Round(time.Millisecond)) - } else { - kunlog.Trace("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v", - l.shortID, l.assistantID, duration.Round(time.Millisecond), err) - } return } - // Development: colorful output (direct print for immediate display) + data := map[string]interface{}{"duration": duration.Round(time.Millisecond)} + if err != nil { + data["error"] = err.Error() + } + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + AssistantID: l.currentAssistantID(), + Event: EventRequestEnd, + Data: data, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) if success { - fmt.Printf("%s βœ… REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset) + fmt.Printf("%s REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset) } else { - fmt.Printf("%s ❌ REQUEST %s FAILED%s\n", colorBoldRed, l.shortID, colorReset) + fmt.Printf("%s REQUEST %s FAILED%s\n", colorBoldRed, l.shortID, colorReset) if err != nil { fmt.Printf("%s Error: %s%v%s\n", colorGray, colorRed, err, colorReset) } } - fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset) + fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.currentAssistantID(), colorReset) fmt.Printf("%s Duration: %s%v%s\n", colorGray, colorWhite, duration.Round(time.Millisecond), colorReset) fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) fmt.Println() @@ -323,12 +394,22 @@ func (l *RequestLogger) Phase(name string) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s β–Ά %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventPhase, + Data: map[string]interface{}{"name": name, "elapsed": elapsed}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s > %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset) } // PhaseComplete logs the completion of a phase @@ -338,12 +419,22 @@ func (l *RequestLogger) PhaseComplete(name string) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s βœ“ %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventPhaseDone, + Data: map[string]interface{}{"name": name, "elapsed": elapsed}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s + %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset) } // PhaseSkip logs a skipped phase (development only) @@ -351,9 +442,21 @@ func (l *RequestLogger) PhaseSkip(name, reason string) { if l.noop { return } - if config.IsDevelopment() { - fmt.Printf("%s ⊘ %s (%s)%s\n", colorGray, name, reason, colorReset) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventPhaseSkip, + Data: map[string]interface{}{"name": name, "reason": reason}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s - %s (%s)%s\n", colorGray, name, reason, colorReset) } // LLMStart logs the start of an LLM call @@ -363,17 +466,31 @@ func (l *RequestLogger) LLMStart(connector, model string, messageCount int) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s πŸ€– LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset) - fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset) - if model != "" { - fmt.Printf("%s Model: %s%s%s\n", colorGray, colorWhite, model, colorReset) - } - fmt.Printf("%s Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset) - } else { - kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventLLMCall, + Data: map[string]interface{}{ + "connector": connector, + "model": model, + "messages": messageCount, + }, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset) + fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset) + if model != "" { + fmt.Printf("%s Model: %s%s%s\n", colorGray, colorWhite, model, colorReset) + } + fmt.Printf("%s Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset) } // LLMComplete logs the completion of an LLM call @@ -388,15 +505,30 @@ func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) { status = "tool_calls" } - if config.IsDevelopment() { - fmt.Printf("%s βœ“ LLM Response (%s)%s", colorGreen, status, colorReset) - if tokens > 0 { - fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset) - } - fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed) + kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventLLMDone, + Data: map[string]interface{}{ + "detail": fmt.Sprintf("%s [tokens:%d, %v]", status, tokens, elapsed), + "tokens": tokens, + "status": status, + }, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s + LLM Response (%s)%s", colorGreen, status, colorReset) + if tokens > 0 { + fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset) + } + fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset) } // ToolStart logs the start of tool execution @@ -405,11 +537,22 @@ func (l *RequestLogger) ToolStart(toolName string) { return } - if config.IsDevelopment() { - fmt.Printf("%s πŸ”§ Tool: %s%s\n", colorYellow, toolName, colorReset) - } else { - kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName) + kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventToolCall, + Data: map[string]interface{}{"name": toolName}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s Tool: %s%s\n", colorYellow, toolName, colorReset) } // ToolComplete logs the completion of tool execution @@ -418,18 +561,29 @@ func (l *RequestLogger) ToolComplete(toolName string, success bool) { return } - if config.IsDevelopment() { - if success { - fmt.Printf("%s βœ“ %s completed%s\n", colorGreen, toolName, colorReset) - } else { - fmt.Printf("%s βœ— %s failed%s\n", colorRed, toolName, colorReset) - } + if success { + kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName) } else { - if success { - kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName) - } else { - kunlog.Trace("[AGENT] %s Tool failed: %s", l.shortID, toolName) - } + kunlog.Error("[AGENT] %s Tool failed: %s", l.shortID, toolName) + } + + if !config.IsDevelopment() { + return + } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventToolDone, + Data: map[string]interface{}{"name": toolName, "success": success}, + }) + + if GetTUIProgram() != nil { + return + } + if success { + fmt.Printf("%s + %s completed%s\n", colorGreen, toolName, colorReset) + } else { + fmt.Printf("%s x %s failed%s\n", colorRed, toolName, colorReset) } } @@ -440,12 +594,22 @@ func (l *RequestLogger) HookStart(hookName string) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s πŸͺ Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventHook, + Data: map[string]interface{}{"name": hookName}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset) } // HookComplete logs the completion of a hook @@ -454,11 +618,22 @@ func (l *RequestLogger) HookComplete(hookName string) { return } - if config.IsDevelopment() { - fmt.Printf("%s βœ“ %s done%s\n", colorGreen, hookName, colorReset) - } else { - kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName) + kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventHookDone, + Data: map[string]interface{}{"name": hookName}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s + %s done%s\n", colorGreen, hookName, colorReset) } // Cleanup logs resource cleanup @@ -467,11 +642,12 @@ func (l *RequestLogger) Cleanup(resource string) { return } - if config.IsDevelopment() { - fmt.Printf("%s βœ“ %s%s\n", colorGray, resource, colorReset) - } else { - kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource) + kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource) + + if !config.IsDevelopment() || GetTUIProgram() != nil { + return } + fmt.Printf("%s + %s%s\n", colorGray, resource, colorReset) } // HistoryLoad logs history loading @@ -480,11 +656,12 @@ func (l *RequestLogger) HistoryLoad(count, maxSize int) { return } - if config.IsDevelopment() { - fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset) - } else { - kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize) + kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize) + + if !config.IsDevelopment() || GetTUIProgram() != nil { + return } + fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset) } // HistoryOverlap logs overlap detection @@ -494,11 +671,12 @@ func (l *RequestLogger) HistoryOverlap(overlapCount int) { } if overlapCount > 0 { - if config.IsDevelopment() { - fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset) - } else { - kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount) + kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount) + + if !config.IsDevelopment() || GetTUIProgram() != nil { + return } + fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset) } } @@ -508,11 +686,22 @@ func (l *RequestLogger) Release() { return } - if config.IsDevelopment() { - fmt.Printf("%s 🧹 RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.assistantID, colorReset) - } else { - kunlog.Trace("[AGENT] %s Release started", l.shortID) + kunlog.Trace("[AGENT] %s Release started", l.shortID) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventContextRelease, + Data: map[string]interface{}{"assistant": l.currentAssistantID()}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.currentAssistantID(), colorReset) } // ============================================================================= diff --git a/agent/context/tui.go b/agent/context/tui.go new file mode 100644 index 00000000..5062bcc9 --- /dev/null +++ b/agent/context/tui.go @@ -0,0 +1,799 @@ +package context + +import ( + "fmt" + "strings" + "sync" + "time" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" +) + +var ( + tuiProgram *tea.Program + tuiProgramMu sync.RWMutex +) + +// SetTUIProgram sets the global TUI program (called from start.go after HTTP READY) +func SetTUIProgram(p *tea.Program) { + tuiProgramMu.Lock() + tuiProgram = p + tuiProgramMu.Unlock() +} + +// GetTUIProgram returns the global TUI program (nil if not in TUI mode) +func GetTUIProgram() *tea.Program { + tuiProgramMu.RLock() + defer tuiProgramMu.RUnlock() + return tuiProgram +} + +// SendTUI sends a message to the TUI program if available +func SendTUI(msg tea.Msg) { + if p := GetTUIProgram(); p != nil { + p.Send(msg) + } +} + +// TUILogWriter implements io.Writer to bridge gou DevWriter -> TUI AppLogMsg +type TUILogWriter struct { + Program *tea.Program +} + +func (w *TUILogWriter) Write(p []byte) (n int, err error) { + content := strings.TrimRight(string(p), "\n") + if content == "" { + return len(p), nil + } + w.Program.Send(AppLogMsg{Content: content}) + return len(p), nil +} + +// ─── Styles ─────────────────────────────────────────────────────────────────── + +var ( + boxRunning = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("33")). + PaddingLeft(1).PaddingRight(1) + + boxDone = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("240")). + PaddingLeft(1).PaddingRight(1) + + boxFailed = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("31")). + PaddingLeft(1).PaddingRight(1) + + boxAppLog = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("240")). + PaddingLeft(1).PaddingRight(1) + + sRunning = lipgloss.NewStyle().Foreground(lipgloss.Color("33")) + sDone = lipgloss.NewStyle().Foreground(lipgloss.Color("34")) + sFailed = lipgloss.NewStyle().Foreground(lipgloss.Color("31")) + sDim = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + sBold = lipgloss.NewStyle().Bold(true) + sYellow = lipgloss.NewStyle().Foreground(lipgloss.Color("33")) + sRed = lipgloss.NewStyle().Foreground(lipgloss.Color("31")) + sBlue = lipgloss.NewStyle().Foreground(lipgloss.Color("34")) + sMagenta = lipgloss.NewStyle().Foreground(lipgloss.Color("35")) + sTree = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) +) + +// ─── Data ───────────────────────────────────────────────────────────────────── + +// RequestPanel represents a single top-level agent request +type RequestPanel struct { + RequestID string + ShortID string + AssistantID string + StartTime time.Time + EndTime time.Time // set when done/failed, freezes elapsed display + Status PanelStatus + Nodes []TreeNode + ParentID string + Collapsed bool + viewRow int // Y offset of the header line (for mouse click) +} + +// TreeNode represents a step within a request panel +type TreeNode struct { + Kind NodeKind + Label string + Status NodeStatus + Detail string + Children []*TreeNode + StartTime time.Time + EndTime time.Time + Collapsed bool +} + +// AgentTUIModel is the bubbletea Model for agent request visualization +type AgentTUIModel struct { + panels []*RequestPanel + panelIndex map[string]int // requestID -> index in panels (first registration wins) + appLogs []AppLogEntry + appLogExpand bool + appLogRow int // Y offset of app log header + cursor int + width int + height int + scrollOffset int + autoFollow bool // auto-scroll to bottom when new content arrives + mouseOn bool + quitting bool +} + +// NewAgentTUIModel creates a new TUI model +func NewAgentTUIModel() AgentTUIModel { + return AgentTUIModel{ + panels: []*RequestPanel{}, + panelIndex: map[string]int{}, + appLogs: []AppLogEntry{}, + width: 80, + height: 24, + autoFollow: true, + } +} + +func (m AgentTUIModel) Init() tea.Cmd { + return tickCmd() +} + +func tickCmd() tea.Cmd { + return tea.Tick(500*time.Millisecond, func(t time.Time) tea.Msg { + return TickMsg(t) + }) +} + +// ─── Update ─────────────────────────────────────────────────────────────────── + +func (m AgentTUIModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.height = msg.Height + return m, nil + + case tea.KeyMsg: + return m.handleKey(msg) + + case tea.MouseMsg: + return m.handleMouse(msg) + + case AgentEventMsg: + return m.handleAgentEvent(msg), nil + + case AppLogMsg: + m.appLogs = append(m.appLogs, AppLogEntry{ + Content: msg.Content, + Time: time.Now(), + }) + return m, nil + + case TickMsg: + return m, tickCmd() + } + return m, nil +} + +func (m AgentTUIModel) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { + topPanels := m.topLevelPanels() + total := len(topPanels) + 1 // +1 for app log + viewH := m.viewHeight() + + switch msg.String() { + case "q", "ctrl+c": + m.quitting = true + return m, tea.Quit + + // Scrolling + case "j", "down": + m.scrollOffset++ + m.autoFollow = false + case "k", "up": + if m.scrollOffset > 0 { + m.scrollOffset-- + } + m.autoFollow = false + case "pgdown", "ctrl+d": + m.scrollOffset += viewH / 2 + m.autoFollow = false + case "pgup", "ctrl+u": + m.scrollOffset -= viewH / 2 + if m.scrollOffset < 0 { + m.scrollOffset = 0 + } + m.autoFollow = false + case "G", "end": + m.autoFollow = true + case "g", "home": + m.scrollOffset = 0 + m.autoFollow = false + + // Cursor navigation for panel selection (wraps around) + case "tab": + m.cursor = (m.cursor + 1) % total + m.scrollToCursor(topPanels) + case "shift+tab": + m.cursor = (m.cursor - 1 + total) % total + m.scrollToCursor(topPanels) + + case "enter", " ": + if m.cursor < len(topPanels) { + topPanels[m.cursor].Collapsed = !topPanels[m.cursor].Collapsed + } else { + m.appLogExpand = !m.appLogExpand + } + case "c": + m.appLogExpand = !m.appLogExpand + case "a": + for _, p := range m.panels { + p.Collapsed = false + } + m.appLogExpand = true + case "A": + for _, p := range m.panels { + p.Collapsed = true + } + m.appLogExpand = false + case "m": + m.mouseOn = !m.mouseOn + if m.mouseOn { + return m, tea.EnableMouseCellMotion + } + return m, tea.DisableMouse + } + return m, nil +} + +func (m AgentTUIModel) viewHeight() int { + h := m.height - 2 // reserve for status bar + if h < 4 { + h = 4 + } + return h +} + +func (m *AgentTUIModel) scrollToCursor(topPanels []*RequestPanel) { + targetRow := 0 + if m.cursor < len(topPanels) { + targetRow = topPanels[m.cursor].viewRow + } else { + targetRow = m.appLogRow + } + viewH := m.viewHeight() + if targetRow < m.scrollOffset { + m.scrollOffset = targetRow + } else if targetRow >= m.scrollOffset+viewH { + m.scrollOffset = targetRow - viewH + 3 + } +} + +func (m AgentTUIModel) handleMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) { + switch { + case msg.Button == tea.MouseButtonWheelUp: + m.scrollOffset -= 3 + if m.scrollOffset < 0 { + m.scrollOffset = 0 + } + m.autoFollow = false + return m, nil + case msg.Button == tea.MouseButtonWheelDown: + m.scrollOffset += 3 + m.autoFollow = false + return m, nil + } + + if msg.Button != tea.MouseButtonLeft || msg.Action != tea.MouseActionRelease { + return m, nil + } + y := msg.Y + m.scrollOffset + + // Check app log header + if y == m.appLogRow { + m.appLogExpand = !m.appLogExpand + return m, nil + } + + // Check panel headers + for _, p := range m.panels { + if p.ParentID != "" { + continue + } + if y == p.viewRow { + p.Collapsed = !p.Collapsed + return m, nil + } + } + return m, nil +} + +// ─── Agent Events ───────────────────────────────────────────────────────────── + +func (m *AgentTUIModel) handleAgentEvent(msg AgentEventMsg) tea.Model { + switch msg.Event { + case EventRequestStart: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + // Delegate sub-call: same requestID, different assistantID. + // Add as a tree node inside the existing panel instead of creating a new one. + p := m.panels[idx] + p.Nodes = append(p.Nodes, TreeNode{ + Kind: NodeA2A, + Label: msg.AssistantID, + Status: NodeRunning, + StartTime: time.Now(), + }) + return m + } + + panel := &RequestPanel{ + RequestID: msg.RequestID, + ShortID: shortID(msg.RequestID), + AssistantID: msg.AssistantID, + StartTime: time.Now(), + Status: PanelRunning, + ParentID: msg.ParentID, + } + m.panelIndex[msg.RequestID] = len(m.panels) + m.panels = append(m.panels, panel) + + case EventRequestEnd: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + p := m.panels[idx] + + // Only mark panel done if the ending assistantID matches the panel's original assistantID + // (delegate sub-calls End with a different assistantID, they update their tree node instead) + if msg.AssistantID == p.AssistantID || msg.AssistantID == "" { + if errVal, has := msg.Data["error"]; has && errVal != nil { + p.Status = PanelFailed + } else { + p.Status = PanelSuccess + } + p.EndTime = time.Now() + p.Collapsed = true + + // Finalize any still-running child nodes (e.g. hook interrupted mid-execution) + finalStatus := NodeDone + if p.Status == PanelFailed { + finalStatus = NodeFailed + } + for i := range p.Nodes { + if p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = finalStatus + p.Nodes[i].EndTime = p.EndTime + } + } + } else { + // Delegate sub-call finished: mark its tree node as done + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeA2A && p.Nodes[i].Label == msg.AssistantID && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + } + + case EventLLMCall: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeLLM, Label: "LLM", Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventLLMDone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeLLM && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + if d, has := msg.Data["detail"]; has { + p.Nodes[i].Detail = fmt.Sprintf("%v", d) + } + break + } + } + } + + case EventToolCall: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + name := dataStr(msg.Data, "name") + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeTool, Label: name, Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventToolDone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + name := dataStr(msg.Data, "name") + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeTool && p.Nodes[i].Label == name && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + + case EventHook: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + name := dataStr(msg.Data, "name") + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeHook, Label: name, Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventHookDone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeHook && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + + case EventA2AStart: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + target := dataStr(msg.Data, "target") + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeA2A, Label: target, Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventA2ADone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + target := dataStr(msg.Data, "target") + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeA2A && p.Nodes[i].Status == NodeRunning && (target == "" || p.Nodes[i].Label == target) { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + } + return m +} + +// ─── View ───────────────────────────────────────────────────────────────────── + +func (m AgentTUIModel) View() string { + if m.quitting { + return "" + } + + boxW := m.width - 2 + if boxW < 40 { + boxW = 40 + } + + // Render full content + var sb strings.Builder + row := 0 + topIdx := 0 + + for _, panel := range m.panels { + if panel.ParentID != "" { + continue + } + selected := (topIdx == m.cursor) + rendered := m.renderPanelBox(panel, boxW, selected, &row) + sb.WriteString(rendered) + sb.WriteString("\n") + row++ + topIdx++ + } + + // App Log + m.appLogRow = row + sb.WriteString(m.renderAppLogBox(boxW, topIdx == m.cursor, &row)) + + fullContent := sb.String() + lines := strings.Split(fullContent, "\n") + totalLines := len(lines) + viewH := m.viewHeight() + + // Auto-follow: snap to bottom + if m.autoFollow { + m.scrollOffset = totalLines - viewH + } + + // Clamp scroll offset + maxScroll := totalLines - viewH + if maxScroll < 0 { + maxScroll = 0 + } + if m.scrollOffset > maxScroll { + m.scrollOffset = maxScroll + } + if m.scrollOffset < 0 { + m.scrollOffset = 0 + } + + // Slice visible lines + end := m.scrollOffset + viewH + if end > totalLines { + end = totalLines + } + visible := lines[m.scrollOffset:end] + + // Build output + var out strings.Builder + out.WriteString(strings.Join(visible, "\n")) + + // Status bar with scroll indicator + mouseLabel := "off" + if m.mouseOn { + mouseLabel = "on" + } + scrollInfo := "" + if totalLines > viewH { + pct := 100 + if maxScroll > 0 { + pct = m.scrollOffset * 100 / maxScroll + } + scrollInfo = fmt.Sprintf(" [%d%%]", pct) + } + followLabel := "" + if m.autoFollow { + followLabel = " AUTO" + } + hint := sDim.Render(fmt.Sprintf(" j/k:scroll tab:select space:toggle a/A:all G:bottom g:top m:mouse(%s)%s%s q:quit", + mouseLabel, scrollInfo, followLabel)) + out.WriteString("\n" + hint) + + return out.String() +} + +func (m AgentTUIModel) renderPanelBox(panel *RequestPanel, boxW int, selected bool, row *int) string { + // Record header row for mouse + panel.viewRow = *row + + elapsed := m.panelElapsed(panel) + icon, statusText, style := panelStatusDisplay(panel.Status, elapsed) + + // Title line + collapser := "β–Ύ" + if panel.Collapsed { + collapser = "β–Έ" + } + cursor := " " + if selected { + cursor = "β€Ί" + } + title := fmt.Sprintf("%s %s %s %s %s", + sDim.Render(cursor), + sDim.Render(collapser), + sBold.Render(panel.ShortID), + panel.AssistantID, + style.Render(icon+" "+statusText), + ) + + if panel.Collapsed { + box := boxForStatus(panel.Status).Width(boxW) + result := box.Render(title) + *row += strings.Count(result, "\n") + 1 + return result + } + + // Build body + var body strings.Builder + body.WriteString(title + "\n") + + for _, node := range panel.Nodes { + body.WriteString(m.renderTreeNode(node, " ", false, panel)) + } + + // Render fork children (different requestID, parentID matches) + children := m.childPanels(panel.RequestID) + for i, child := range children { + isLast := (i == len(children)-1) + body.WriteString(m.renderChildSummary(child, " ", isLast)) + } + + box := boxForStatus(panel.Status).Width(boxW) + result := box.Render(body.String()) + *row += strings.Count(result, "\n") + 1 + return result +} + +func (m AgentTUIModel) renderTreeNode(node TreeNode, prefix string, isChild bool, panel *RequestPanel) string { + panelEnded := panel != nil && panel.Status != PanelRunning + displayNode := node + if panelEnded && displayNode.Status == NodeRunning { + displayNode.Status = NodeFailed + } + icon, statusText := nodeStatusDisplay(displayNode) + elapsed := m.nodeElapsed(node, panelEnded, panel.EndTime) + + label := "" + switch node.Kind { + case NodeHook: + label = sMagenta.Render("Hook: "+node.Label) + " " + statusText + case NodeLLM: + detail := "" + if node.Detail != "" { + detail = " " + sDim.Render("["+node.Detail+"]") + } + label = sBlue.Render("LLM") + " " + statusText + detail + case NodeTool: + label = sTree.Render("β”œ ") + sYellow.Render(node.Label) + " " + statusText + case NodeA2A: + label = sTree.Render("β€· ") + sBold.Render(node.Label) + " " + statusText + case NodePhase: + label = node.Label + " " + statusText + default: + label = node.Label + " " + statusText + } + + _ = icon + line := prefix + label + if elapsed != "" { + line += " " + sDim.Render(elapsed) + } + return line + "\n" +} + +func (m AgentTUIModel) renderChildSummary(panel *RequestPanel, prefix string, isLast bool) string { + elapsed := m.panelElapsed(panel) + icon, statusText, style := panelStatusDisplay(panel.Status, elapsed) + + branch := sTree.Render("β”œβ”€ ") + if isLast { + branch = sTree.Render("└─ ") + } + return fmt.Sprintf("%s%s%s %s %s\n", + prefix, branch, + sBold.Render(panel.ShortID+" "+panel.AssistantID), + style.Render(icon+" "+statusText), + sDim.Render(elapsed), + ) +} + +func (m AgentTUIModel) renderAppLogBox(boxW int, selected bool, row *int) string { + cursor := " " + if selected { + cursor = "β€Ί" + } + collapser := "β–Έ" + if m.appLogExpand { + collapser = "β–Ύ" + } + + count := len(m.appLogs) + title := fmt.Sprintf("%s %s %s (%d)", + sDim.Render(cursor), + sDim.Render(collapser), + sBold.Render("App Output"), + count, + ) + + if !m.appLogExpand || count == 0 { + result := boxAppLog.Width(boxW).Render(title) + *row += strings.Count(result, "\n") + 1 + return result + } + + var body strings.Builder + body.WriteString(title + "\n") + + start := 0 + if count > 50 { + start = count - 50 + } + for _, entry := range m.appLogs[start:] { + body.WriteString(" " + entry.Content + "\n") + } + + result := boxAppLog.Width(boxW).Render(body.String()) + *row += strings.Count(result, "\n") + 1 + return result +} + +// ─── Helpers ────────────────────────────────────────────────────────────────── + +func (m AgentTUIModel) topLevelPanels() []*RequestPanel { + var result []*RequestPanel + for _, p := range m.panels { + if p.ParentID == "" { + result = append(result, p) + } + } + return result +} + +func (m AgentTUIModel) childPanels(parentRequestID string) []*RequestPanel { + var result []*RequestPanel + for _, p := range m.panels { + if p.ParentID == parentRequestID { + result = append(result, p) + } + } + return result +} + +func (m AgentTUIModel) panelElapsed(p *RequestPanel) string { + if p.Status != PanelRunning && !p.EndTime.IsZero() { + return fmtDuration(p.EndTime.Sub(p.StartTime)) + } + return fmtDuration(time.Since(p.StartTime)) +} + +func (m AgentTUIModel) nodeElapsed(n TreeNode, panelEnded bool, panelEndTime time.Time) string { + if n.Status == NodeDone || n.Status == NodeFailed { + if !n.EndTime.IsZero() { + return fmtDuration(n.EndTime.Sub(n.StartTime)) + } + } + if n.Status == NodeRunning { + if panelEnded && !panelEndTime.IsZero() { + return fmtDuration(panelEndTime.Sub(n.StartTime)) + } + return fmtDuration(time.Since(n.StartTime)) + } + return "" +} + +func panelStatusDisplay(status PanelStatus, elapsed string) (icon string, text string, style lipgloss.Style) { + switch status { + case PanelRunning: + return "⟳", "running " + elapsed, sRunning + case PanelSuccess: + return "βœ“", "done " + elapsed, sDone + case PanelFailed: + return "βœ—", "failed " + elapsed, sFailed + } + return "", "", sDim +} + +func nodeStatusDisplay(n TreeNode) (icon string, text string) { + switch n.Status { + case NodePending: + return "…", sDim.Render("…") + case NodeRunning: + return "⟳", sRunning.Render("⟳") + case NodeDone: + return "βœ“", sDone.Render("βœ“") + case NodeFailed: + return "βœ—", sFailed.Render("βœ—") + } + return "", "" +} + +func boxForStatus(status PanelStatus) lipgloss.Style { + switch status { + case PanelRunning: + return boxRunning + case PanelFailed: + return boxFailed + default: + return boxDone + } +} + +func dataStr(data map[string]interface{}, key string) string { + if v, ok := data[key]; ok { + return fmt.Sprintf("%v", v) + } + return "" +} + +func fmtDuration(d time.Duration) string { + if d < time.Second { + return fmt.Sprintf("%dms", d.Milliseconds()) + } + return fmt.Sprintf("%.1fs", d.Seconds()) +} diff --git a/agent/context/tui_msg.go b/agent/context/tui_msg.go new file mode 100644 index 00000000..f910df44 --- /dev/null +++ b/agent/context/tui_msg.go @@ -0,0 +1,90 @@ +package context + +import "time" + +// EventType represents the type of agent lifecycle event +type EventType int + +const ( + EventRequestStart EventType = iota + EventPhase + EventPhaseDone + EventPhaseSkip + EventLLMCall + EventLLMDone + EventToolCall + EventToolDone + EventHook + EventHookDone + EventA2AStart + EventA2ADone + EventRequestEnd + EventContextFork + EventContextRelease +) + +// AgentEventMsg is sent from RequestLogger to the TUI Program +type AgentEventMsg struct { + RequestID string + ParentID string + AssistantID string + Event EventType + Data map[string]interface{} +} + +// AppLogLevel represents the severity of application-side output +type AppLogLevel int + +const ( + AppLogLevelLog AppLogLevel = iota + AppLogLevelInfo + AppLogLevelWarn + AppLogLevelError + AppLogLevelException +) + +// AppLogMsg is sent from the DevWriter (gou layer) to the TUI Program +type AppLogMsg struct { + Level AppLogLevel + Content string +} + +// AppLogEntry stores a single application output entry +type AppLogEntry struct { + Level AppLogLevel + Content string + Time time.Time +} + +// PanelStatus represents the lifecycle state of a request panel +type PanelStatus int + +const ( + PanelRunning PanelStatus = iota + PanelSuccess + PanelFailed +) + +// NodeKind represents the type of a tree node within a request panel +type NodeKind int + +const ( + NodePhase NodeKind = iota + NodeLLM + NodeTool + NodeHook + NodeA2A +) + +// NodeStatus represents the state of a tree node +type NodeStatus int + +const ( + NodePending NodeStatus = iota + NodeRunning + NodeDone + NodeFailed +) + +// TickMsg triggers periodic UI refresh for elapsed time display +type TickMsg time.Time diff --git a/cmd/start.go b/cmd/start.go index 9b221474..d663505b 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -8,11 +8,14 @@ import ( "strings" "syscall" + tea "github.com/charmbracelet/bubbletea" "github.com/fatih/color" + "github.com/mattn/go-isatty" "github.com/spf13/cobra" "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/fs" + "github.com/yaoapp/gou/helper" "github.com/yaoapp/gou/mcp" "github.com/yaoapp/gou/plugin" "github.com/yaoapp/gou/schedule" @@ -20,7 +23,9 @@ import ( "github.com/yaoapp/gou/store" "github.com/yaoapp/gou/task" "github.com/yaoapp/gou/websocket" + "github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/log" + agentcontext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/engine" "github.com/yaoapp/yao/openapi" @@ -33,6 +38,7 @@ import ( var startDebug = false var startDisableWatching = false +var startTUI = false var startCmd = &cobra.Command{ Use: "start", @@ -265,8 +271,9 @@ var startCmd = &cobra.Command{ switch v { case http.READY: - fmt.Println(color.GreenString(L("✨Server is up and running..."))) - fmt.Println(color.GreenString("✨Ctrl+C to stop")) + fmt.Println(color.GreenString(L("Server is up and running..."))) + fmt.Println(color.GreenString("Ctrl+C to stop")) + initAgentTUI() break case http.CLOSED: @@ -627,7 +634,38 @@ func colorMehtod(method string) string { } } +// initAgentTUI initializes the TUI for agent request visualization in dev mode. +// Must be called after HTTP READY to avoid interfering with startup messages. +func initAgentTUI() { + if !config.IsDevelopment() { + return + } + + if !startTUI && os.Getenv("YAO_TUI") != "on" { + return + } + + if !isatty.IsTerminal(os.Stdout.Fd()) { + return + } + + model := agentcontext.NewAgentTUIModel() + p := tea.NewProgram(model, tea.WithoutSignalHandler()) + + agentcontext.SetTUIProgram(p) + tuiWriter := &agentcontext.TUILogWriter{Program: p} + helper.SetDevWriter(tuiWriter) + exception.SetWriter(tuiWriter) + + go func() { + if _, err := p.Run(); err != nil { + log.Error("TUI error: %s", err.Error()) + } + }() +} + func init() { startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode")) startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching")) + startCmd.PersistentFlags().BoolVarP(&startTUI, "tui", "", false, L("Enable TUI for agent request visualization")) } diff --git a/engine/load.go b/engine/load.go index 411d19a5..f96f2090 100644 --- a/engine/load.go +++ b/engine/load.go @@ -1,6 +1,7 @@ package engine import ( + "context" "fmt" "log" "os" @@ -24,6 +25,7 @@ import ( "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/connector" "github.com/yaoapp/yao/data" + "github.com/yaoapp/yao/event" "github.com/yaoapp/yao/flow" "github.com/yaoapp/yao/fs" "github.com/yaoapp/yao/i18n" @@ -48,6 +50,8 @@ import ( "github.com/yaoapp/yao/websocket" "github.com/yaoapp/yao/widget" "github.com/yaoapp/yao/widgets" + + _ "github.com/yaoapp/yao/trace" // register trace handler/listener via init() ) // LoadHooks used to load custom widgets/processes @@ -212,6 +216,14 @@ func Load(cfg config.Config, options LoadOption, progressCallback ...func(string warnings = append(warnings, Warning{Widget: "Store", Error: err}) } + // Start Event Service (handlers registered via init(), e.g. trace) + err = loadStep("Event", func() error { + return event.Start() + }, callback) + if err != nil { + warnings = append(warnings, Warning{Widget: "Event", Error: err}) + } + // Load Uploaders err = loadStep("Uploader", func() error { return attachment.Load(cfg) @@ -425,6 +437,9 @@ func Unload() (err error) { } } + // Stop Event Service (before runtime, so in-flight handlers can still use V8) + event.Stop(context.Background()) + // Stop Runtime err = runtime.Stop() diff --git a/go.mod b/go.mod index 4b25c0cf..55089598 100644 --- a/go.mod +++ b/go.mod @@ -64,11 +64,18 @@ require ( github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.15 // indirect github.com/aws/smithy-go v1.22.3 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/blang/semver/v4 v4.0.0 // indirect github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect github.com/bytedance/sonic v1.13.2 // indirect github.com/bytedance/sonic/loader v0.2.4 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/charmbracelet/bubbletea v1.3.10 // indirect + github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect + github.com/charmbracelet/lipgloss v1.1.0 // indirect + github.com/charmbracelet/x/ansi v0.10.1 // indirect + github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect + github.com/charmbracelet/x/term v0.2.1 // indirect github.com/cloudwego/base64x v0.1.5 // indirect github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect @@ -78,6 +85,7 @@ require ( github.com/dlclark/regexp2 v1.11.5 // indirect github.com/docker/go-units v0.5.0 // indirect github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // indirect + github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/gabriel-vasile/mimetype v1.4.9 // indirect github.com/gin-contrib/sse v1.1.0 // indirect @@ -115,9 +123,11 @@ require ( github.com/klauspost/cpuid/v2 v2.2.10 // indirect github.com/leodido/go-urn v1.4.0 // indirect github.com/lib/pq v1.10.9 // indirect + github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mark3labs/mcp-go v0.32.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-localereader v0.0.1 // indirect github.com/mattn/go-runewidth v0.0.16 // indirect github.com/mattn/go-sqlite3 v1.14.28 // indirect github.com/miekg/dns v1.1.66 // indirect @@ -126,6 +136,9 @@ require ( github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/montanaflynn/stats v0.7.1 // indirect + github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/muesli/termenv v0.16.0 // indirect github.com/neo4j/neo4j-go-driver/v5 v5.28.1 // indirect github.com/oklog/run v1.1.0 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect @@ -159,6 +172,7 @@ require ( github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.1.2 // indirect github.com/xdg-go/stringprep v1.0.4 // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/xuri/efp v0.0.1 // indirect github.com/xuri/nfp v0.0.1 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/go.sum b/go.sum index 52182b71..0ff37529 100644 --- a/go.sum +++ b/go.sum @@ -39,6 +39,8 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3 h1:BRXS0U76Z8wfF+bnkilA2QwpIch6U github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3/go.mod h1:bNXKFFyaiVvWuR6O16h/I1724+aXe/tAkA9/QS01t5k= github.com/aws/smithy-go v1.22.3 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k= github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ= github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk= github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= @@ -62,6 +64,18 @@ github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqy github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw= +github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4= +github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs= +github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk= +github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY= +github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30= +github.com/charmbracelet/x/ansi v0.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ= +github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE= +github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8= +github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs= +github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= +github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= @@ -98,6 +112,8 @@ github.com/emersion/go-sasl v0.0.0-20200509203442-7bfe0ed36a21/go.mod h1:iL2twTe github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 h1:oP4q0fw+fOSWn3DfFi4EXdT+B+gTtzx8GC9xsc26Znk= github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6/go.mod h1:iL2twTeMvZnrg54ZoPDNfJaJaqy0xIQFuBdrLsmspwQ= github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594/go.mod h1:aqO8z8wPrjkscevZJFVE1wXJrLpC5LtJG7fqLOsPb2U= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/evanw/esbuild v0.25.4 h1:k1bTSim+usBG27w7BfOCorhgx3tO+6bAfMj5pR+6SKg= github.com/evanw/esbuild v0.25.4/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48= github.com/expr-lang/expr v1.17.7 h1:Q0xY/e/2aCIp8g9s/LGvMDCC5PxYlvHgDZRQ4y16JX8= @@ -231,6 +247,8 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= +github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= +github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mark3labs/mcp-go v0.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8= github.com/mark3labs/mcp-go v0.32.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4= github.com/matoous/go-nanoid/v2 v2.1.0 h1:P64+dmq21hhWdtvZfEAofnvJULaRR1Yib0+PnU669bE= @@ -243,6 +261,8 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= +github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= @@ -269,6 +289,12 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/mozillazg/go-pinyin v0.20.0 h1:BtR3DsxpApHfKReaPO1fCqF4pThRwH9uwvXzm+GnMFQ= github.com/mozillazg/go-pinyin v0.20.0/go.mod h1:iR4EnMMRXkfpFVV5FMi4FNB6wGq9NV6uDWbUuPhP4Yc= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= +github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/neo4j/neo4j-go-driver/v5 v5.28.1 h1:RKWQW7wTgYAY2fU9S+9LaJ9OwRPbRc0I17tlT7nDmAY= github.com/neo4j/neo4j-go-driver/v5 v5.28.1/go.mod h1:Vff8OwT7QpLm7L2yYr85XNWe9Rbqlbeb9asNXJTHO4k= github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= @@ -379,6 +405,8 @@ github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY= github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3kKLN4= github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8= github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/xuri/efp v0.0.1 h1:fws5Rv3myXyYni8uwj2qKjVaRP30PdjeYe2Y6FDsCL8= github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI= github.com/xuri/excelize/v2 v2.9.1 h1:VdSGk+rraGmgLHGFaGG9/9IWu1nj4ufjJ7uwMDtj8Qw= @@ -471,6 +499,7 @@ golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/test/utils.go b/test/utils.go index 11bf75f5..8b3efb80 100644 --- a/test/utils.go +++ b/test/utils.go @@ -66,12 +66,14 @@ // 7. Starts V8 JavaScript runtime // 8. Registers query engines for database operations // 9. Creates temporary data directories for test isolation +// 10. Starts the Event Service (handlers registered via init(), e.g. trace) // // WHAT test.Clean() DOES: -// 1. Stops V8 runtime and releases resources -// 2. Closes all database connections -// 3. Removes temporary test data stores -// 4. Resets global state to prevent test interference +// 1. Stops the Event Service (drains in-flight events) +// 2. Stops V8 runtime and releases resources +// 3. Closes all database connections +// 4. Removes temporary test data stores +// 5. Resets global state to prevent test interference // // WHAT test.Start() DOES: // 1. Creates Gin HTTP server with API routes @@ -159,6 +161,7 @@ package test import ( + "context" "fmt" "os" "path/filepath" @@ -184,11 +187,14 @@ import ( "github.com/yaoapp/xun/capsule" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/data" + "github.com/yaoapp/yao/event" "github.com/yaoapp/yao/fs" "github.com/yaoapp/yao/helper" "github.com/yaoapp/yao/runtime" "github.com/yaoapp/yao/share" "github.com/yaoapp/yao/utils" + + _ "github.com/yaoapp/yao/trace" // register trace event handler via init() ) var testServer *http.Server = nil @@ -488,10 +494,15 @@ func Prepare(t *testing.T, cfg config.Config, opts ...interface{}) { load(t, cfg) startRuntime(t, cfg) + // Start event service (trace handler registered via blank import above) + if err := event.Start(); err != nil { + t.Fatalf("Failed to start event service: %v", err) + } } // Clean the test environment func Clean() { + event.Stop(context.Background()) dbclose() runtime.Stop() diff --git a/trace/event_listener.go b/trace/event_listener.go new file mode 100644 index 00000000..f699595d --- /dev/null +++ b/trace/event_listener.go @@ -0,0 +1,25 @@ +package trace + +import ( + "context" + + "github.com/yaoapp/yao/event" + eventTypes "github.com/yaoapp/yao/event/types" +) + +// traceUpdateListener receives trace update events for cross-cutting concerns +// (e.g., audit logging, metrics). Trace updates are broadcast via event.Push +// and delivered to this listener and any dynamic subscribers. +type traceUpdateListener struct{} + +func (l *traceUpdateListener) OnEvent(ev *eventTypes.Event) {} + +func (l *traceUpdateListener) Shutdown(ctx context.Context) error { + return nil +} + +func init() { + event.Listen("trace.*", &traceUpdateListener{}, + event.BufferSize(4096), + ) +} diff --git a/trace/handler.go b/trace/handler.go new file mode 100644 index 00000000..e73feafd --- /dev/null +++ b/trace/handler.go @@ -0,0 +1,28 @@ +package trace + +import ( + "context" + + "github.com/yaoapp/yao/event" + eventTypes "github.com/yaoapp/yao/event/types" +) + +// traceHandler processes trace events dispatched through the event service. +// It enables event.Push routing for trace.* events (used by addUpdateAndBroadcast). +type traceHandler struct{} + +func (h *traceHandler) Handle(ctx context.Context, ev *eventTypes.Event, resp chan<- eventTypes.Result) { + resp <- eventTypes.Result{} +} + +func (h *traceHandler) Shutdown(ctx context.Context) error { + return nil +} + +func init() { + event.Register("trace", &traceHandler{}, + event.MaxWorkers(256), + event.ReservedWorkers(32), + event.QueueSize(4096), + ) +} diff --git a/trace/manager.go b/trace/manager.go index 8aaf6b25..2b4989e4 100644 --- a/trace/manager.go +++ b/trace/manager.go @@ -3,63 +3,56 @@ package trace import ( "context" "fmt" + "sync" "time" gonanoid "github.com/matoous/go-nanoid/v2" "github.com/yaoapp/kun/log" - "github.com/yaoapp/yao/trace/pubsub" + "github.com/yaoapp/yao/event" "github.com/yaoapp/yao/trace/types" ) -// manager implements the Manager interface with channel-based state management +// manager implements the Manager interface. +// State is protected by a mutex, replacing the old channel-based state worker. +// This eliminates the context-cancel bug while maintaining thread safety. type manager struct { - ctx context.Context - cancel context.CancelFunc - traceID string - driver types.Driver - stateCmdChan chan stateCommand // Single channel for all state mutations - closed int32 // Atomic flag: 1 = closed, safeSend rejects new commands - autoArchive bool // Auto-archive on complete/fail - pubsub *pubsub.PubSub // Reference to independent pubsub service (for publishing only, doesn't own it) + mu sync.Mutex + traceID string + driver types.Driver + state *managerState + autoArchive bool } -// NewManager creates a new trace manager instance -// pubsubService: reference to independent pubsub service (manager doesn't own it, just publishes to it) -func NewManager(ctx context.Context, traceID string, driver types.Driver, pubsubService *pubsub.PubSub, option *types.TraceOption) (types.Manager, error) { - // Create a cancellable context for the manager - managerCtx, cancel := context.WithCancel(ctx) - - // Determine auto-archive setting +// NewManager creates a new trace manager instance. +func NewManager(ctx context.Context, traceID string, driver types.Driver, option *types.TraceOption) (types.Manager, error) { autoArchive := false if option != nil { autoArchive = option.AutoArchive } m := &manager{ - ctx: managerCtx, - cancel: cancel, - traceID: traceID, - driver: driver, - stateCmdChan: make(chan stateCommand, 100), // Buffered channel for performance - autoArchive: autoArchive, - pubsub: pubsubService, // Reference only, doesn't manage lifecycle + traceID: traceID, + driver: driver, + autoArchive: autoArchive, + state: &managerState{ + spaces: make(map[string]*types.TraceSpace), + traceStatus: types.TraceStatusPending, + updates: make([]*types.TraceUpdate, 0, 100), + }, } - // Start state worker goroutine - go m.startStateWorker() - - // Try to load existing updates from driver (for resumed traces) + // Load existing updates from driver (for resumed traces). + // Safe to access m.state directly here β€” no Queue yet, single goroutine. if existingUpdates, err := driver.LoadUpdates(ctx, traceID, 0); err == nil && len(existingUpdates) > 0 { log.Trace("[MANAGER] NewManager: loaded %d existing updates from driver for trace %s", len(existingUpdates), traceID) - m.stateSetUpdates(existingUpdates) - // Check if trace was already completed + m.state.updates = existingUpdates for _, update := range existingUpdates { if update.Type == types.UpdateTypeComplete { log.Trace("[MANAGER] NewManager: trace %s was already completed, marking as completed", traceID) - m.stateMarkCompleted() + m.state.completed = true if data, ok := update.Data.(*types.TraceCompleteData); ok { log.Trace("[MANAGER] NewManager: setting trace status to %s", data.Status) - m.stateSetTraceStatus(data.Status) + m.state.traceStatus = data.Status } break } @@ -70,7 +63,6 @@ func NewManager(ctx context.Context, traceID string, driver types.Driver, pubsub } else { log.Trace("[MANAGER] NewManager: no existing updates found for trace %s, creating new trace", traceID) } - // New trace - create and broadcast init event now := time.Now().UnixMilli() m.addUpdateAndBroadcast(&types.TraceUpdate{ Type: types.UpdateTypeInit, @@ -89,36 +81,24 @@ func genNodeID() string { return id } -// addUpdateAndBroadcast persists, adds to history, and publishes an update +// addUpdateAndBroadcast persists, adds to history, and broadcasts via event service. func (m *manager) addUpdateAndBroadcast(update *types.TraceUpdate) { - // Persist to driver (synchronous - no race) if err := m.driver.SaveUpdate(context.Background(), m.traceID, update); err != nil { log.Trace("[MANAGER] addUpdateAndBroadcast: failed to save update type=%s for trace %s: %v", update.Type, m.traceID, err) } - // else { - // log.Trace("[MANAGER] addUpdateAndBroadcast: successfully saved update type=%s for trace %s", update.Type, m.traceID) - // } - // Add to in-memory history m.stateAddUpdate(update) - // Publish to independent PubSub service (manager just publishes, doesn't manage pubsub lifecycle) - if m.pubsub != nil { - m.pubsub.Publish(update) - } + // Broadcast to subscribers via event service (fire-and-forget, non-blocking). + // Uses Push with the update as payload so event.Subscribe filters can match by traceID. + event.Push(context.Background(), "trace.update", update) } -// checkContext checks if context is cancelled +// checkContext checks if the trace has been completed/released. +// With event-based state management, the manager no longer binds a context. +// Lifecycle is controlled by QueueCreate/QueueRelease. func (m *manager) checkContext() error { - select { - case <-m.ctx.Done(): - // Context cancelled - just return the error - // Don't call handleCancellation here to avoid deadlock - // handleCancellation should be called explicitly when needed - return m.ctx.Err() - default: - return nil - } + return nil } // Add creates next sequential node - auto-joins if currently in parallel state @@ -147,7 +127,7 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ } // Save root node - if err := m.driver.SaveNode(m.ctx, m.traceID, rootNode); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, rootNode); err != nil { return nil, fmt.Errorf("failed to save root node: %w", err) } @@ -210,14 +190,14 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ // Add to each parent's children for _, parent := range currentNodes { parent.Children = append(parent.Children, newNodeData) - if err := m.driver.SaveNode(m.ctx, m.traceID, parent); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, parent); err != nil { // Log error but continue m.Error("Failed to update parent node %s: %v", parent.ID, err) } } // Save new node - if err := m.driver.SaveNode(m.ctx, m.traceID, newNodeData); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, newNodeData); err != nil { return nil, err } @@ -310,7 +290,7 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N // Save all nodes in batch - collect errors var saveErrors []error for _, data := range nodeData { - if err := m.driver.SaveNode(m.ctx, m.traceID, data); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, data); err != nil { saveErrors = append(saveErrors, fmt.Errorf("failed to save node %s: %w", data.ID, err)) } } @@ -321,7 +301,7 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N } // Save parent node - if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, parentNode); err != nil { return nil, fmt.Errorf("failed to save parent node: %w", err) } @@ -380,7 +360,7 @@ func (m *manager) log(level string, message string, args ...any) { NodeID: node.ID, } // Save log (ignore errors for non-critical logging) - _ = m.driver.SaveLog(m.ctx, m.traceID, log) + _ = m.driver.SaveLog(context.Background(), m.traceID, log) // Broadcast log event m.addUpdateAndBroadcast(&types.TraceUpdate{ @@ -404,7 +384,7 @@ func (m *manager) SetOutput(output types.TraceOutput) error { for _, node := range nodes { node.Output = output node.UpdatedAt = now - if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, node); err != nil { return err } @@ -434,7 +414,7 @@ func (m *manager) SetMetadata(key string, value any) error { } node.Metadata[key] = value node.UpdatedAt = now - if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, node); err != nil { return err } @@ -485,7 +465,7 @@ func (m *manager) Complete(output ...types.TraceOutput) error { node.Status = types.StatusCompleted node.EndTime = now node.UpdatedAt = now - if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, node); err != nil { return err } @@ -516,7 +496,7 @@ func (m *manager) Fail(err error) error { node.Status = types.StatusFailed node.EndTime = now node.UpdatedAt = now - if saveErr := m.driver.SaveNode(m.ctx, m.traceID, node); saveErr != nil { + if saveErr := m.driver.SaveNode(context.Background(), m.traceID, node); saveErr != nil { return saveErr } @@ -545,7 +525,7 @@ func (m *manager) GetRootNode() (*types.TraceNode, error) { // GetNode returns a node by ID func (m *manager) GetNode(id string) (*types.TraceNode, error) { - return m.driver.LoadNode(m.ctx, m.traceID, id) + return m.driver.LoadNode(context.Background(), m.traceID, id) } // GetCurrentNodes returns current active nodes @@ -581,7 +561,7 @@ func (m *manager) MarkComplete() error { // Auto-archive if enabled if m.autoArchive { - if err := m.driver.Archive(m.ctx, m.traceID); err != nil { + if err := m.driver.Archive(context.Background(), m.traceID); err != nil { // Log error but don't fail the complete operation m.Debug("Failed to auto-archive trace", map[string]any{ "trace_id": m.traceID, @@ -610,7 +590,7 @@ func (m *manager) CreateSpace(option types.TraceSpaceOption) (*types.TraceSpace, } // Save to driver - if err := m.driver.SaveSpace(m.ctx, m.traceID, space); err != nil { + if err := m.driver.SaveSpace(context.Background(), m.traceID, space); err != nil { return nil, err } @@ -637,7 +617,7 @@ func (m *manager) GetSpace(id string) (*types.TraceSpace, error) { } // Load from driver - space, err := m.driver.LoadSpace(m.ctx, m.traceID, id) + space, err := m.driver.LoadSpace(context.Background(), m.traceID, id) if err != nil { return nil, err } @@ -658,7 +638,7 @@ func (m *manager) HasSpace(id string) bool { } // Check in driver - space, _ := m.driver.LoadSpace(m.ctx, m.traceID, id) + space, _ := m.driver.LoadSpace(context.Background(), m.traceID, id) return space != nil } @@ -674,7 +654,7 @@ func (m *manager) DeleteSpace(id string) error { m.stateDeleteSpace(id) // Delete from driver - if err := m.driver.DeleteSpace(m.ctx, m.traceID, id); err != nil { + if err := m.driver.DeleteSpace(context.Background(), m.traceID, id); err != nil { return err } @@ -693,7 +673,7 @@ func (m *manager) DeleteSpace(id string) error { // ListSpaces returns all spaces func (m *manager) ListSpaces() []*types.TraceSpace { // Load from driver to ensure we have all spaces - spaceIDs, err := m.driver.ListSpaces(m.ctx, m.traceID) + spaceIDs, err := m.driver.ListSpaces(context.Background(), m.traceID) if err != nil { // Fallback to cached spaces return m.stateGetAllSpaces() @@ -727,13 +707,13 @@ func (m *manager) SetSpaceValue(spaceID, key string, value any) error { // Set value in driver (through state worker for concurrent safety) err = m.stateExecuteSpaceOp(spaceID, func() error { - if err := m.driver.SetSpaceKey(m.ctx, m.traceID, spaceID, key, value); err != nil { + if err := m.driver.SetSpaceKey(context.Background(), m.traceID, spaceID, key, value); err != nil { return err } // Update space timestamp space.UpdatedAt = now - if err := m.driver.SaveSpace(m.ctx, m.traceID, space); err != nil { + if err := m.driver.SaveSpace(context.Background(), m.traceID, space); err != nil { return err } @@ -761,7 +741,7 @@ func (m *manager) GetSpaceValue(spaceID, key string) (any, error) { var result any err := m.stateExecuteSpaceOp(spaceID, func() error { var err error - result, err = m.driver.GetSpaceKey(m.ctx, m.traceID, spaceID, key) + result, err = m.driver.GetSpaceKey(context.Background(), m.traceID, spaceID, key) return err }) return result, err @@ -771,7 +751,7 @@ func (m *manager) GetSpaceValue(spaceID, key string) (any, error) { func (m *manager) HasSpaceValue(spaceID, key string) bool { var result bool _ = m.stateExecuteSpaceOp(spaceID, func() error { - result = m.driver.HasSpaceKey(m.ctx, m.traceID, spaceID, key) + result = m.driver.HasSpaceKey(context.Background(), m.traceID, spaceID, key) return nil }) return result @@ -787,7 +767,7 @@ func (m *manager) DeleteSpaceValue(spaceID, key string) error { // Delete value from driver (through state worker for concurrent safety) err := m.stateExecuteSpaceOp(spaceID, func() error { - return m.driver.DeleteSpaceKey(m.ctx, m.traceID, spaceID, key) + return m.driver.DeleteSpaceKey(context.Background(), m.traceID, spaceID, key) }) if err != nil { @@ -816,7 +796,7 @@ func (m *manager) ClearSpaceValues(spaceID string) error { // Clear values from driver (through state worker for concurrent safety) err := m.stateExecuteSpaceOp(spaceID, func() error { - return m.driver.ClearSpaceKeys(m.ctx, m.traceID, spaceID) + return m.driver.ClearSpaceKeys(context.Background(), m.traceID, spaceID) }) if err != nil { @@ -840,7 +820,7 @@ func (m *manager) ListSpaceKeys(spaceID string) []string { var keys []string _ = m.stateExecuteSpaceOp(spaceID, func() error { var err error - keys, err = m.driver.ListSpaceKeys(m.ctx, m.traceID, spaceID) + keys, err = m.driver.ListSpaceKeys(context.Background(), m.traceID, spaceID) return err }) return keys @@ -865,7 +845,7 @@ func (m *manager) GetTraceInfo() (*types.TraceInfo, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadTraceInfo(m.ctx, m.traceID) + return m.driver.LoadTraceInfo(context.Background(), m.traceID) } // GetAllNodes retrieves all nodes from storage @@ -875,7 +855,7 @@ func (m *manager) GetAllNodes() ([]*types.TraceNode, error) { } // Load the root node tree from storage - rootNode, err := m.driver.LoadTrace(m.ctx, m.traceID) + rootNode, err := m.driver.LoadTrace(context.Background(), m.traceID) if err != nil { return nil, err } @@ -906,7 +886,7 @@ func (m *manager) GetNodeByID(nodeID string) (*types.TraceNode, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadNode(m.ctx, m.traceID, nodeID) + return m.driver.LoadNode(context.Background(), m.traceID, nodeID) } // GetAllLogs retrieves all logs from storage @@ -914,7 +894,7 @@ func (m *manager) GetAllLogs() ([]*types.TraceLog, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadLogs(m.ctx, m.traceID, "") + return m.driver.LoadLogs(context.Background(), m.traceID, "") } // GetLogsByNode retrieves logs for a specific node from storage @@ -922,7 +902,7 @@ func (m *manager) GetLogsByNode(nodeID string) ([]*types.TraceLog, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadLogs(m.ctx, m.traceID, nodeID) + return m.driver.LoadLogs(context.Background(), m.traceID, nodeID) } // GetAllSpaces retrieves all spaces from storage @@ -932,7 +912,7 @@ func (m *manager) GetAllSpaces() ([]*types.TraceSpace, error) { } // Get all space IDs from driver - spaceIDs, err := m.driver.ListSpaces(m.ctx, m.traceID) + spaceIDs, err := m.driver.ListSpaces(context.Background(), m.traceID) if err != nil { return nil, err } @@ -940,7 +920,7 @@ func (m *manager) GetAllSpaces() ([]*types.TraceSpace, error) { // Load all spaces spaces := make([]*types.TraceSpace, 0, len(spaceIDs)) for _, spaceID := range spaceIDs { - space, err := m.driver.LoadSpace(m.ctx, m.traceID, spaceID) + space, err := m.driver.LoadSpace(context.Background(), m.traceID, spaceID) if err != nil { continue // Skip spaces that fail to load } @@ -959,7 +939,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) { } // Load space metadata - space, err := m.driver.LoadSpace(m.ctx, m.traceID, spaceID) + space, err := m.driver.LoadSpace(context.Background(), m.traceID, spaceID) if err != nil { return nil, err } @@ -968,7 +948,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) { } // Load all keys in the space - keys, err := m.driver.ListSpaceKeys(m.ctx, m.traceID, spaceID) + keys, err := m.driver.ListSpaceKeys(context.Background(), m.traceID, spaceID) if err != nil { return nil, err } @@ -976,7 +956,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) { // Load all key-value pairs data := make(map[string]any) for _, key := range keys { - value, err := m.driver.GetSpaceKey(m.ctx, m.traceID, spaceID, key) + value, err := m.driver.GetSpaceKey(context.Background(), m.traceID, spaceID, key) if err != nil { continue // Skip keys that fail to load } diff --git a/trace/node.go b/trace/node.go index cd09f14e..aca80e6c 100644 --- a/trace/node.go +++ b/trace/node.go @@ -1,6 +1,7 @@ package trace import ( + "context" "time" "github.com/yaoapp/yao/trace/types" @@ -60,7 +61,7 @@ func (n *node) log(level string, message string, args ...any) *types.TraceLog { NodeID: n.data.ID, } // Save log (ignore errors for non-critical logging) - _ = n.manager.driver.SaveLog(n.manager.ctx, n.manager.traceID, log) + _ = n.manager.driver.SaveLog(context.Background(), n.manager.traceID, log) return log } @@ -85,10 +86,10 @@ func (n *node) Add(input types.TraceInput, option types.TraceNodeOption) (types. n.data.Children = append(n.data.Children, childNodeData) // Save both nodes - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, childNodeData); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, childNodeData); err != nil { return nil, err } - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data); err != nil { return nil, err } @@ -120,7 +121,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node n.data.Children = append(n.data.Children, childNodeData) // Save node - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, childNodeData); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, childNodeData); err != nil { return nil, err } @@ -132,7 +133,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node } // Save parent node - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data); err != nil { return nil, err } @@ -165,7 +166,7 @@ func (n *node) Join(nodes []*types.TraceNode, input types.TraceInput, option typ } // Save join node - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, joinNodeData); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, joinNodeData); err != nil { return nil, err } @@ -185,7 +186,7 @@ func (n *node) ID() string { func (n *node) SetOutput(output types.TraceOutput) error { n.data.Output = output n.data.UpdatedAt = time.Now().UnixMilli() - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // SetMetadata sets node metadata @@ -195,14 +196,14 @@ func (n *node) SetMetadata(key string, value any) error { } n.data.Metadata[key] = value n.data.UpdatedAt = time.Now().UnixMilli() - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // SetStatus sets the node status func (n *node) SetStatus(status string) error { n.data.Status = types.NodeStatus(status) n.data.UpdatedAt = time.Now().UnixMilli() - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // Complete marks the node as completed (public method, broadcasts event) @@ -236,7 +237,7 @@ func (n *node) complete(output ...types.TraceOutput) error { n.data.Status = types.StatusCompleted n.data.EndTime = now n.data.UpdatedAt = now - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // Fail marks the node as failed (public method, broadcasts event) @@ -269,5 +270,5 @@ func (n *node) fail(err error) error { n.data.EndTime = now n.data.UpdatedAt = now - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } diff --git a/trace/pubsub/pubsub.go b/trace/pubsub/pubsub.go deleted file mode 100644 index 9f8ee309..00000000 --- a/trace/pubsub/pubsub.go +++ /dev/null @@ -1,142 +0,0 @@ -package pubsub - -import ( - "sync" - - gonanoid "github.com/matoous/go-nanoid/v2" - "github.com/yaoapp/kun/log" - "github.com/yaoapp/yao/trace/types" -) - -// PubSub is an independent publish-subscribe service for trace updates -// It acts as a message broker between trace writers and readers -type PubSub struct { - eventBus chan *types.TraceUpdate // Event bus for incoming events - subscribers map[string]chan *types.TraceUpdate // Active subscribers - mu sync.RWMutex // Protects subscribers map - stopCh chan struct{} // Signal to stop the service - stopped bool // Whether service is stopped -} - -// New creates a new PubSub service -func New() *PubSub { - ps := &PubSub{ - eventBus: make(chan *types.TraceUpdate, 1000), // Buffered event bus - subscribers: make(map[string]chan *types.TraceUpdate), - stopCh: make(chan struct{}), - stopped: false, - } - - // Start forwarding service - go ps.forward() - - return ps -} - -// forward continuously forwards events from eventBus to all subscribers -// This runs in a dedicated goroutine -func (ps *PubSub) forward() { - for { - select { - case event := <-ps.eventBus: - ps.mu.RLock() - subscriberCount := len(ps.subscribers) - - if subscriberCount == 0 { - // No subscribers, discard event - ps.mu.RUnlock() - continue - } - - // Forward to all subscribers (non-blocking) - for subID, ch := range ps.subscribers { - select { - case ch <- event: - // Sent successfully - default: - // Subscriber is slow or channel full, skip - log.Trace("[PUBSUB] Subscriber %s is slow, skipping event type=%s", subID, event.Type) - } - } - ps.mu.RUnlock() - - case <-ps.stopCh: - return - } - } -} - -// Publish sends an event to the event bus -// This is called by trace writers (e.g., manager.addUpdateAndBroadcast) -func (ps *PubSub) Publish(event *types.TraceUpdate) { - if ps.stopped { - return - } - - select { - case ps.eventBus <- event: - // Event published successfully - default: - // Event bus full, this shouldn't happen with large buffer (log as warning) - log.Warn("[PUBSUB] Event bus full, discarding event type=%s", event.Type) - } -} - -// Subscribe creates a new subscription and returns a channel for receiving updates -// The caller is responsible for reading from the channel and closing it when done -func (ps *PubSub) Subscribe(bufferSize int) (<-chan *types.TraceUpdate, string) { - // Generate unique subscriber ID - subID, _ := gonanoid.Generate("0123456789abcdefghijklmnopqrstuvwxyz", 12) - - // Create subscriber channel - ch := make(chan *types.TraceUpdate, bufferSize) - - // Register subscriber - ps.mu.Lock() - ps.subscribers[subID] = ch - ps.mu.Unlock() - - return ch, subID -} - -// Unsubscribe removes a subscriber and closes its channel -func (ps *PubSub) Unsubscribe(subID string) { - ps.mu.Lock() - defer ps.mu.Unlock() - - ch, exists := ps.subscribers[subID] - if !exists { - return - } - - // Remove from map - delete(ps.subscribers, subID) - - // Close channel - close(ch) -} - -// SubscriberCount returns the number of active subscribers -func (ps *PubSub) SubscriberCount() int { - ps.mu.RLock() - defer ps.mu.RUnlock() - return len(ps.subscribers) -} - -// Stop stops the forwarding service and closes all subscriber channels -func (ps *PubSub) Stop() { - if ps.stopped { - return - } - - ps.stopped = true - close(ps.stopCh) - - // Close all subscriber channels - ps.mu.Lock() - for _, ch := range ps.subscribers { - close(ch) - } - ps.subscribers = make(map[string]chan *types.TraceUpdate) - ps.mu.Unlock() -} diff --git a/trace/pubsub/subscriber.go b/trace/pubsub/subscriber.go deleted file mode 100644 index 26ada520..00000000 --- a/trace/pubsub/subscriber.go +++ /dev/null @@ -1,62 +0,0 @@ -package pubsub - -import ( - "time" - - "github.com/yaoapp/kun/log" - "github.com/yaoapp/yao/trace/types" -) - -// Subscriber represents a subscription to trace updates -type Subscriber struct { - ID string - Channel <-chan *types.TraceUpdate - pubsub *PubSub -} - -// Unsubscribe removes the subscription and closes the channel -func (s *Subscriber) Unsubscribe() { - s.pubsub.Unsubscribe(s.ID) -} - -// SubscribeWithHistory creates a subscription and replays historical updates first -// historicalUpdates: updates to replay before starting live stream -// bufferSize: size of the subscription channel buffer -func (ps *PubSub) SubscribeWithHistory(historicalUpdates []*types.TraceUpdate, bufferSize int) *Subscriber { - // Create subscription - ch, subID := ps.Subscribe(bufferSize) - - // Create subscriber - sub := &Subscriber{ - ID: subID, - Channel: ch, - pubsub: ps, - } - - // Replay historical updates in a goroutine - // This allows the subscription to start immediately - go func() { - // Get writable channel for replay - ps.mu.RLock() - writeCh, exists := ps.subscribers[subID] - ps.mu.RUnlock() - - if !exists { - return - } - - // Replay all historical updates (blocking send to ensure delivery) - for i, update := range historicalUpdates { - select { - case writeCh <- update: - // Sent successfully - case <-time.After(5 * time.Second): - // Timeout - subscriber is too slow or disconnected - log.Trace("[PUBSUB] Subscriber %s timed out during replay at update %d/%d", subID, i, len(historicalUpdates)) - return - } - } - }() - - return sub -} diff --git a/trace/space.go b/trace/space.go index 2806e7e4..1b3131cc 100644 --- a/trace/space.go +++ b/trace/space.go @@ -6,18 +6,18 @@ import ( "github.com/yaoapp/yao/trace/types" ) -// space implements the Space interface for custom space operations +// space implements the Space interface for custom space operations. +// Uses context.Background() for driver calls to decouple from caller context +// (fixes the context-fork bug where parent cancellation breaks child ops). type space struct { - ctx context.Context traceID string data *types.TraceSpace driver types.Driver } // NewSpace creates a new space instance -func NewSpace(ctx context.Context, traceID string, data *types.TraceSpace, driver types.Driver) types.Space { +func NewSpace(traceID string, data *types.TraceSpace, driver types.Driver) types.Space { return &space{ - ctx: ctx, traceID: traceID, data: data, driver: driver, @@ -31,32 +31,32 @@ func (s *space) ID() string { // Set stores a value by key func (s *space) Set(key string, value any) error { - return s.driver.SetSpaceKey(s.ctx, s.traceID, s.data.ID, key, value) + return s.driver.SetSpaceKey(context.Background(), s.traceID, s.data.ID, key, value) } // Get retrieves a value by key func (s *space) Get(key string) (any, error) { - return s.driver.GetSpaceKey(s.ctx, s.traceID, s.data.ID, key) + return s.driver.GetSpaceKey(context.Background(), s.traceID, s.data.ID, key) } // Has checks if a key exists func (s *space) Has(key string) bool { - return s.driver.HasSpaceKey(s.ctx, s.traceID, s.data.ID, key) + return s.driver.HasSpaceKey(context.Background(), s.traceID, s.data.ID, key) } // Delete removes a key-value pair func (s *space) Delete(key string) error { - return s.driver.DeleteSpaceKey(s.ctx, s.traceID, s.data.ID, key) + return s.driver.DeleteSpaceKey(context.Background(), s.traceID, s.data.ID, key) } // Clear removes all key-value pairs func (s *space) Clear() error { - return s.driver.ClearSpaceKeys(s.ctx, s.traceID, s.data.ID) + return s.driver.ClearSpaceKeys(context.Background(), s.traceID, s.data.ID) } // Keys returns all keys in the space func (s *space) Keys() []string { - keys, err := s.driver.ListSpaceKeys(s.ctx, s.traceID, s.data.ID) + keys, err := s.driver.ListSpaceKeys(context.Background(), s.traceID, s.data.ID) if err != nil { return nil } diff --git a/trace/state.go b/trace/state.go index 63e4915e..e49404e2 100644 --- a/trace/state.go +++ b/trace/state.go @@ -2,16 +2,13 @@ package trace import ( "fmt" - "sync/atomic" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/trace/types" ) -// State management using channel-based serialization (no locks needed) -// All state mutations go through a single worker goroutine - -// managerState holds all mutable state (accessed only by state worker) +// managerState holds all mutable state for a trace. +// Protected by manager.mu β€” all access goes through state* methods which acquire the lock. type managerState struct { rootNode *types.TraceNode currentNodes []*types.TraceNode @@ -19,344 +16,133 @@ type managerState struct { traceStatus types.TraceStatus completed bool updates []*types.TraceUpdate - // Note: subscribers moved to SubscriptionManager (no longer in state) -} - -// State command interface - all commands are processed serially -type stateCommand interface { - execute(s *managerState) -} - -// Commands with response channels for synchronous operations - -// --- Root Node Commands --- - -type cmdSetRoot struct { - node *types.TraceNode -} - -func (c *cmdSetRoot) execute(s *managerState) { - s.rootNode = c.node -} - -type cmdGetRoot struct { - resp chan *types.TraceNode -} - -func (c *cmdGetRoot) execute(s *managerState) { - c.resp <- s.rootNode -} - -// --- Current Nodes Commands --- - -type cmdSetCurrentNodes struct { - nodes []*types.TraceNode -} - -func (c *cmdSetCurrentNodes) execute(s *managerState) { - s.currentNodes = c.nodes -} - -type cmdGetCurrentNodes struct { - resp chan []*types.TraceNode -} - -func (c *cmdGetCurrentNodes) execute(s *managerState) { - // Return a copy to prevent external mutation - nodes := make([]*types.TraceNode, len(s.currentNodes)) - copy(nodes, s.currentNodes) - c.resp <- nodes -} - -type cmdUpdateRootAndCurrent struct { - root *types.TraceNode - current []*types.TraceNode -} - -func (c *cmdUpdateRootAndCurrent) execute(s *managerState) { - s.rootNode = c.root - s.currentNodes = c.current -} - -// --- Space Commands --- - -type cmdGetSpace struct { - id string - resp chan *types.TraceSpace -} - -func (c *cmdGetSpace) execute(s *managerState) { - c.resp <- s.spaces[c.id] -} - -type cmdSetSpace struct { - id string - space *types.TraceSpace -} - -func (c *cmdSetSpace) execute(s *managerState) { - s.spaces[c.id] = c.space -} - -type cmdDeleteSpace struct { - id string -} - -func (c *cmdDeleteSpace) execute(s *managerState) { - delete(s.spaces, c.id) -} - -type cmdGetAllSpaces struct { - resp chan []*types.TraceSpace -} - -func (c *cmdGetAllSpaces) execute(s *managerState) { - spaces := make([]*types.TraceSpace, 0, len(s.spaces)) - for _, space := range s.spaces { - spaces = append(spaces, space) - } - c.resp <- spaces -} - -// --- Trace Status Commands --- - -type cmdSetTraceStatus struct { - status types.TraceStatus -} - -func (c *cmdSetTraceStatus) execute(s *managerState) { - s.traceStatus = c.status -} - -type cmdGetTraceStatus struct { - resp chan types.TraceStatus -} - -func (c *cmdGetTraceStatus) execute(s *managerState) { - c.resp <- s.traceStatus -} - -// --- Completion Commands --- - -type cmdMarkCompleted struct { - resp chan bool // Returns true if marked, false if already completed -} - -func (c *cmdMarkCompleted) execute(s *managerState) { - if s.completed { - c.resp <- false - } else { - s.completed = true - c.resp <- true - } -} - -type cmdIsCompleted struct { - resp chan bool -} - -func (c *cmdIsCompleted) execute(s *managerState) { - c.resp <- s.completed -} - -// --- Update Commands --- - -type cmdAddUpdate struct { - update *types.TraceUpdate -} - -func (c *cmdAddUpdate) execute(s *managerState) { - s.updates = append(s.updates, c.update) -} - -type cmdGetUpdates struct { - since int64 - resp chan []*types.TraceUpdate -} - -func (c *cmdGetUpdates) execute(s *managerState) { - filtered := make([]*types.TraceUpdate, 0) - for _, update := range s.updates { - if update.Timestamp >= c.since { - filtered = append(filtered, update) - } - } - c.resp <- filtered -} - -type cmdSetUpdates struct { - updates []*types.TraceUpdate -} - -func (c *cmdSetUpdates) execute(s *managerState) { - s.updates = c.updates -} - -// --- Subscriber Commands (REMOVED - now handled by SubscriptionManager) --- -// Subscriber management has been decoupled from state machine for better separation of concerns - -// --- Space KV Commands (for concurrent safety) --- -// These ensure all operations on a space are serialized through state worker - -type cmdSpaceKVOp struct { - spaceID string - fn func() error - resp chan error -} - -func (c *cmdSpaceKVOp) execute(s *managerState) { - // Execute the operation (typically a driver call) - // The function is provided by caller and executed serially here - err := c.fn() - c.resp <- err -} - -// State worker - processes all commands serially in a single goroutine. -// Exits only when stateCmdChan is closed by Release(). The for-range loop -// automatically drains any buffered commands before returning (Go spec guarantee). -func (m *manager) startStateWorker() { - state := &managerState{ - rootNode: nil, - currentNodes: []*types.TraceNode{}, - spaces: make(map[string]*types.TraceSpace), - traceStatus: types.TraceStatusPending, - completed: false, - updates: make([]*types.TraceUpdate, 0, 100), - } - for cmd := range m.stateCmdChan { - cmd.execute(state) - } -} - -// Helper methods for manager to send commands - -// safeSend sends a command to the state worker channel. Returns false if the -// manager is closed, context is cancelled, or the channel was closed mid-send. -// The atomic closed flag provides a fast-path rejection before touching the channel, -// which is critical in CGO callback stacks where recover() may not work. -func (m *manager) safeSend(cmd stateCommand) (ok bool) { - if atomic.LoadInt32(&m.closed) == 1 { - return false - } - defer func() { - if r := recover(); r != nil { - ok = false - } - }() - select { - case <-m.ctx.Done(): - return false - case m.stateCmdChan <- cmd: - return true - } } func (m *manager) stateSetRoot(node *types.TraceNode) { - m.safeSend(&cmdSetRoot{node: node}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.rootNode = node } func (m *manager) stateGetRoot() *types.TraceNode { - resp := make(chan *types.TraceNode, 1) - if !m.safeSend(&cmdGetRoot{resp: resp}) { - return nil // Context cancelled - } - return <-resp + m.mu.Lock() + defer m.mu.Unlock() + return m.state.rootNode } func (m *manager) stateSetCurrentNodes(nodes []*types.TraceNode) { - m.safeSend(&cmdSetCurrentNodes{nodes: nodes}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.currentNodes = nodes } func (m *manager) stateGetCurrentNodes() []*types.TraceNode { - resp := make(chan []*types.TraceNode, 1) - if !m.safeSend(&cmdGetCurrentNodes{resp: resp}) { - return nil // Context cancelled + m.mu.Lock() + defer m.mu.Unlock() + if m.state.currentNodes == nil { + return nil } - return <-resp + nodes := make([]*types.TraceNode, len(m.state.currentNodes)) + copy(nodes, m.state.currentNodes) + return nodes } func (m *manager) stateUpdateRootAndCurrent(root *types.TraceNode, current []*types.TraceNode) { - m.safeSend(&cmdUpdateRootAndCurrent{root: root, current: current}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.rootNode = root + m.state.currentNodes = current } func (m *manager) stateGetSpace(id string) (*types.TraceSpace, bool) { - resp := make(chan *types.TraceSpace, 1) - if !m.safeSend(&cmdGetSpace{id: id, resp: resp}) { - return nil, false // Context cancelled - } - space := <-resp - return space, space != nil + m.mu.Lock() + defer m.mu.Unlock() + space, ok := m.state.spaces[id] + return space, ok } func (m *manager) stateSetSpace(id string, space *types.TraceSpace) { - m.safeSend(&cmdSetSpace{id: id, space: space}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.spaces[id] = space } func (m *manager) stateDeleteSpace(id string) { - m.safeSend(&cmdDeleteSpace{id: id}) + m.mu.Lock() + defer m.mu.Unlock() + delete(m.state.spaces, id) } func (m *manager) stateGetAllSpaces() []*types.TraceSpace { - resp := make(chan []*types.TraceSpace, 1) - if !m.safeSend(&cmdGetAllSpaces{resp: resp}) { - return nil // Context cancelled + m.mu.Lock() + defer m.mu.Unlock() + spaces := make([]*types.TraceSpace, 0, len(m.state.spaces)) + for _, space := range m.state.spaces { + spaces = append(spaces, space) } - return <-resp + return spaces } func (m *manager) stateSetTraceStatus(status types.TraceStatus) { - m.safeSend(&cmdSetTraceStatus{status: status}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.traceStatus = status } func (m *manager) stateGetTraceStatus() types.TraceStatus { - resp := make(chan types.TraceStatus, 1) - if !m.safeSend(&cmdGetTraceStatus{resp: resp}) { - return types.TraceStatusCancelled // Context cancelled - } - return <-resp + m.mu.Lock() + defer m.mu.Unlock() + return m.state.traceStatus } func (m *manager) stateMarkCompleted() bool { - resp := make(chan bool, 1) - if !m.safeSend(&cmdMarkCompleted{resp: resp}) { - return true // Context cancelled, treat as completed + m.mu.Lock() + defer m.mu.Unlock() + if m.state.completed { + return false } - return <-resp + m.state.completed = true + return true } func (m *manager) stateIsCompleted() bool { - resp := make(chan bool, 1) - if !m.safeSend(&cmdIsCompleted{resp: resp}) { - return true // Context cancelled, treat as completed - } - return <-resp + m.mu.Lock() + defer m.mu.Unlock() + return m.state.completed } func (m *manager) stateAddUpdate(update *types.TraceUpdate) { - m.safeSend(&cmdAddUpdate{update: update}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.updates = append(m.state.updates, update) } func (m *manager) stateGetUpdates(since int64) []*types.TraceUpdate { - resp := make(chan []*types.TraceUpdate, 1) - if !m.safeSend(&cmdGetUpdates{since: since, resp: resp}) { - return nil // Context cancelled + m.mu.Lock() + defer m.mu.Unlock() + filtered := make([]*types.TraceUpdate, 0) + for _, update := range m.state.updates { + if update.Timestamp >= since { + filtered = append(filtered, update) + } } - return <-resp + return filtered } func (m *manager) stateSetUpdates(updates []*types.TraceUpdate) { + m.mu.Lock() + defer m.mu.Unlock() log.Trace("[STATE] stateSetUpdates: setting %d updates for trace %s", len(updates), m.traceID) - m.safeSend(&cmdSetUpdates{updates: updates}) + m.state.updates = updates } -// Subscription management methods removed - now handled by SubscriptionManager -// See subscription_manager.go and subscription.go for the new implementation - -// stateExecuteSpaceOp executes a space operation serially through state worker +// stateExecuteSpaceOp executes a space operation while holding the lock. func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error { - resp := make(chan error, 1) - if !m.safeSend(&cmdSpaceKVOp{spaceID: spaceID, fn: fn, resp: resp}) { - return fmt.Errorf("trace %s: state worker stopped", m.traceID) + m.mu.Lock() + defer m.mu.Unlock() + err := fn() + if err != nil { + return fmt.Errorf("trace %s: space op failed: %w", m.traceID, err) } - return <-resp + return nil } diff --git a/trace/subscription.go b/trace/subscription.go index 7537e36b..ac76cf12 100644 --- a/trace/subscription.go +++ b/trace/subscription.go @@ -3,12 +3,18 @@ package trace import ( "fmt" + "github.com/yaoapp/yao/event" + eventTypes "github.com/yaoapp/yao/event/types" "github.com/yaoapp/yao/trace/types" ) +func dedupKey(u *types.TraceUpdate) string { + return fmt.Sprintf("%s:%s:%d", u.Type, u.NodeID, u.Timestamp) +} + // Subscribe creates a new subscription for trace updates (replays all historical events from the beginning) func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) { - return m.subscribe(0) // Subscribe from beginning to get all historical events + return m.subscribe(0) } // SubscribeFrom creates a subscription starting from a specific timestamp @@ -16,24 +22,60 @@ func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) return m.subscribe(since) } -// subscribe is the internal implementation for subscriptions +// subscribe creates a subscription channel that first replays historical +// updates, then streams live events via the event service's Subscriber. +// The subscriber is registered BEFORE reading historical state to prevent +// missing events that occur between the state snapshot and subscriber setup. func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) { - // Get historical updates - updates := m.stateGetUpdates(since) - - // Use manager's pubsub reference (always available) - if m.pubsub == nil { - return nil, fmt.Errorf("pubsub service not initialized for trace: %s", m.traceID) - } - - // Create subscription with historical replay - // Buffer size should be large enough to hold historical updates plus some live updates - // Using max of 1000 or len(updates)+100 to handle large traces bufferSize := 1000 - if len(updates)+100 > bufferSize { - bufferSize = len(updates) + 100 - } - sub := m.pubsub.SubscribeWithHistory(updates, bufferSize) - return sub.Channel, nil + out := make(chan *types.TraceUpdate, bufferSize) + + // Register live subscriber FIRST to avoid missing events between snapshot and subscribe. + liveCh := make(chan *eventTypes.Event, bufferSize) + traceID := m.traceID + subID := event.Subscribe("trace.*", liveCh, event.Filter(func(ev *eventTypes.Event) bool { + update, ok := ev.Payload.(*types.TraceUpdate) + if !ok { + return false + } + return update.TraceID == traceID + })) + + // THEN snapshot historical updates (may overlap with live events). + historical := m.stateGetUpdates(since) + + // Build a set of historical event identifiers for dedup. + // Key: "type:nodeID:timestamp" is unique enough for trace events. + histSeen := make(map[string]struct{}, len(historical)) + for _, u := range historical { + histSeen[dedupKey(u)] = struct{}{} + } + + go func() { + defer close(out) + defer event.Unsubscribe(subID) + + for _, update := range historical { + out <- update + } + + for ev := range liveCh { + update, ok := ev.Payload.(*types.TraceUpdate) + if !ok { + continue + } + key := dedupKey(update) + if _, dup := histSeen[key]; dup { + delete(histSeen, key) + continue + } + out <- update + if update.Type == types.UpdateTypeComplete { + return + } + } + }() + + return out, nil } diff --git a/trace/trace.go b/trace/trace.go index ebdb837b..3450567c 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -4,13 +4,11 @@ import ( "context" "fmt" "sync" - "sync/atomic" "time" gonanoid "github.com/matoous/go-nanoid/v2" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/trace/local" - "github.com/yaoapp/yao/trace/pubsub" "github.com/yaoapp/yao/trace/store" "github.com/yaoapp/yao/trace/types" ) @@ -21,14 +19,10 @@ const ( Store = "store" // Gou store storage ) -// Global trace registry and pubsub services +// Global trace registry var ( registry = make(map[string]*types.TraceInfo) registryMu sync.RWMutex - - // Each trace has its own independent pubsub service - pubsubRegistry = make(map[string]*pubsub.PubSub) - pubsubRegistryMu sync.RWMutex ) // getDriver creates a driver instance based on driver type and options @@ -145,23 +139,9 @@ func New(ctx context.Context, driver string, option *types.TraceOption, driverOp return LoadFromStorage(ctx, driver, traceID, driverOptions...) } - // Create independent PubSub service for this trace - pubsubService := pubsub.New() - - // Register pubsub service - pubsubRegistryMu.Lock() - pubsubRegistry[traceID] = pubsubService - pubsubRegistryMu.Unlock() - - // Create Manager instance with the driver and pubsub reference - // Manager uses pubsub only for publishing, doesn't manage its lifecycle - manager, err := NewManager(ctx, traceID, drv, pubsubService, option) + // Create Manager instance with the driver + manager, err := NewManager(ctx, traceID, drv, option) if err != nil { - // Clean up pubsub if manager creation fails - pubsubRegistryMu.Lock() - delete(pubsubRegistry, traceID) - pubsubRegistryMu.Unlock() - pubsubService.Stop() return "", nil, fmt.Errorf("failed to create manager: %w", err) } @@ -198,13 +178,6 @@ func New(ctx context.Context, driver string, option *types.TraceOption, driverOp return traceID, manager, nil } -// GetPubSub returns the pubsub service for a trace -func GetPubSub(traceID string) *pubsub.PubSub { - pubsubRegistryMu.RLock() - defer pubsubRegistryMu.RUnlock() - return pubsubRegistry[traceID] -} - // Load loads an existing trace by ID from the registry // Returns: manager, error // traceID: the trace ID to load @@ -253,29 +226,9 @@ func LoadFromStorage(ctx context.Context, driver string, traceID string, driverO return "", nil, fmt.Errorf("trace not found in storage: %s", traceID) } - // Create or reuse PubSub service for this trace - pubsubRegistryMu.Lock() - pubsubService, exists := pubsubRegistry[traceID] - if !exists { - pubsubService = pubsub.New() - pubsubRegistry[traceID] = pubsubService - } - pubsubRegistryMu.Unlock() - - // Create Manager instance with the driver and pubsub reference - // Note: We need to reconstruct the manager from stored data - // TODO: Implement proper restoration of manager state from storage - // For loaded traces, we don't have the original option, so pass nil - manager, err := NewManager(ctx, traceID, drv, pubsubService, nil) + manager, err := NewManager(ctx, traceID, drv, nil) if err != nil { drv.Close() - if !exists { - // Clean up pubsub if we just created it - pubsubRegistryMu.Lock() - delete(pubsubRegistry, traceID) - pubsubRegistryMu.Unlock() - pubsubService.Stop() - } return "", nil, fmt.Errorf("failed to create manager: %w", err) } @@ -356,12 +309,6 @@ func MarkCancelled(traceID string, reason string) error { log.Trace("[TRACE] MarkCancelled: starting to mark nodes and trace as cancelled") - // Get independent pubsub service - ps := GetPubSub(traceID) - if ps != nil { - log.Trace("[TRACE] MarkCancelled: current subscriber count: %d", ps.SubscriberCount()) - } - now := time.Now().UnixMilli() // Use background context since the original context is cancelled @@ -395,12 +342,7 @@ func MarkCancelled(traceID string, reason string) error { log.Trace("[TRACE] MarkCancelled: failed to save node %s: %v", node.ID, err) } - // Broadcast node failed event (also saves to disk) - subscriberCount := 0 - if ps := GetPubSub(traceID); ps != nil { - subscriberCount = ps.SubscriberCount() - } - log.Trace("[TRACE] MarkCancelled: publishing node failed event for node %s to %d subscribers", node.ID, subscriberCount) + log.Trace("[TRACE] MarkCancelled: broadcasting node failed event for node %s", node.ID) mgr.addUpdateAndBroadcast(&types.TraceUpdate{ Type: types.UpdateTypeNodeFailed, TraceID: traceID, @@ -445,12 +387,7 @@ func MarkCancelled(traceID string, reason string) error { mgr.stateSetTraceStatus(types.TraceStatusCancelled) mgr.stateMarkCompleted() - // Broadcast completion update (saves to disk and publishes to subscribers) - subscriberCount := 0 - if ps := GetPubSub(traceID); ps != nil { - subscriberCount = ps.SubscriberCount() - } - log.Trace("[TRACE] MarkCancelled: publishing completion update to %d subscribers", subscriberCount) + log.Trace("[TRACE] MarkCancelled: broadcasting completion update") totalDuration := int64(0) if rootNode.CreatedAt > 0 { totalDuration = now - rootNode.CreatedAt @@ -488,35 +425,7 @@ func Release(traceID string) error { return fmt.Errorf("trace not found in registry: %s", traceID) } - // Stop manager with safe three-step shutdown sequence. - // Order matters: flag blocks new writes -> cancel unblocks in-flight safeSend -> - // close terminates state worker (which drains remaining buffer first). - if mgr, ok := info.Manager.(*manager); ok { - // Step 1: Set closed flag β€” new safeSend calls return false immediately - atomic.StoreInt32(&mgr.closed, 1) - - // Step 2: Cancel context β€” unblocks any safeSend blocked in select on ctx.Done - if mgr.cancel != nil { - mgr.cancel() - } - - // Step 3: Close channel β€” state worker for-range exits after draining buffer - close(mgr.stateCmdChan) - } - - // Stop independent PubSub service - pubsubRegistryMu.Lock() - ps, psExists := pubsubRegistry[traceID] - if psExists { - delete(pubsubRegistry, traceID) - } - pubsubRegistryMu.Unlock() - - if psExists && ps != nil { - subscriberCount := ps.SubscriberCount() - log.Trace("[TRACE] Release: stopping pubsub service with %d active subscribers", subscriberCount) - ps.Stop() - } + _ = info.Manager log.Trace("[TRACE] Release: completed") return nil diff --git a/trace/trace_basic_test.go b/trace/trace_basic_test.go index 77564a51..e4a2512b 100644 --- a/trace/trace_basic_test.go +++ b/trace/trace_basic_test.go @@ -13,11 +13,8 @@ import ( ) func TestMain(m *testing.M) { - // Prepare test environment (initializes stores, models, etc.) test.Prepare(&testing.T{}, config.Conf) defer test.Clean() - - // Run tests os.Exit(m.Run()) } @@ -234,15 +231,17 @@ func TestContextCancellation(t *testing.T) { defer trace.Release(traceID) defer trace.Remove(context.Background(), d.DriverType, traceID, d.DriverOptions...) - // Cancel context + // Cancel the creation context β€” trace operations should still work. + // This is the core context-fork fix: trace managers no longer bind + // to the caller's context, so parent cancellation cannot break child ops. cancel() - // Operations should fail with context error - _, err = manager.Add("test", types.TraceNodeOption{ + node, err := manager.Add("test", types.TraceNodeOption{ Label: "Test", Type: "test", }) - assert.Error(t, err) + assert.NoError(t, err) + assert.NotNil(t, node) }) } } diff --git a/trace/trace_lifecycle_test.go b/trace/trace_lifecycle_test.go index c28f4452..88fa1a24 100644 --- a/trace/trace_lifecycle_test.go +++ b/trace/trace_lifecycle_test.go @@ -186,9 +186,10 @@ func TestConcurrentReleaseAndMarkCancelled(t *testing.T) { } } -// TestSafeSendAfterClosed verifies that operations using safeSend after -// Release return gracefully instead of panicking. -func TestSafeSendAfterClosed(t *testing.T) { +// TestOperationsAfterRelease verifies that using a manager reference after +// Release does not panic. The manager is removed from registry but its +// in-memory state remains valid (no channel close or context cancel). +func TestOperationsAfterRelease(t *testing.T) { drivers := trace.GetTestDrivers() for _, d := range drivers { @@ -205,21 +206,13 @@ func TestSafeSendAfterClosed(t *testing.T) { err = trace.Release(traceID) assert.NoError(t, err) - // All of these internally use safeSend. After Release they should - // return nil/error/zero-value, never panic. - manager.Info("post-close info") - manager.Debug("post-close debug") - manager.Error("post-close error") - manager.Warn("post-close warn") + // After Release, the manager object is still usable (state in memory). + // These calls should not panic. + manager.Info("post-release info") + manager.Debug("post-release debug") root, _ := manager.GetRootNode() - assert.Nil(t, root) - - nodes, _ := manager.GetCurrentNodes() - assert.Nil(t, nodes) - - status := manager.IsComplete() - assert.True(t, status) + assert.NotNil(t, root) }) } } From 7fa54c50f3e5033ab89dd9e7f834376547c010e2 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 23 Feb 2026 19:37:17 +0800 Subject: [PATCH 3/3] Fix event service startup error handling to ignore already started state - Update the error handling in the event service startup to allow for the case where the service is already running, preventing unnecessary test failures. --- test/utils.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/utils.go b/test/utils.go index 8b3efb80..f3ae31f7 100644 --- a/test/utils.go +++ b/test/utils.go @@ -495,7 +495,7 @@ func Prepare(t *testing.T, cfg config.Config, opts ...interface{}) { startRuntime(t, cfg) // Start event service (trace handler registered via blank import above) - if err := event.Start(); err != nil { + if err := event.Start(); err != nil && err != event.ErrAlreadyStart { t.Fatalf("Failed to start event service: %v", err) } }