From f26111d5bf5195dce0d3744de28eac4253a305db Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Thu, 22 Oct 2026 17:06:50 +0300 Subject: [PATCH] fix(storage): make StreamActivities cancel-safe at the source StreamActivities walked the activity bucket inside a BBolt read transaction and sent every match into a channel buffered at 100 with no way to stop. A consumer that stopped reading early (a cancelled 24h stats loop, an activity export whose client went away) left the producer blocked on the send for the life of the process, holding the read transaction and the mmap read lock: OpenTxN stayed at 1, writes could not remap, and db.Close hung. PR #1718 papered over it in its own consumer with a drain goroutine. The producer now takes a context, checks it before every send and selects on ctx.Done() around the send, so cancellation returns from db.View promptly and closes the channel. Every caller passes its request or loop context; the export path cancels on its write-error return; the 24h stats rollup no longer caches a truncated window and Suggest returns the context error instead of a partial result. drainActivities and its test are gone. New storage tests prove the leak (OpenTxN returns to 0, Close returns) and the full-stream path. --- internal/httpapi/activity.go | 24 +-- .../httpapi/activity_auth_identity_test.go | 3 +- internal/httpapi/activity_call_parity_test.go | 3 +- .../httpapi/activity_scope_params_test.go | 3 +- .../httpapi/activity_summary_scope_test.go | 3 +- internal/httpapi/contracts_test.go | 2 +- internal/httpapi/scope_round9_harness_test.go | 3 +- internal/httpapi/security_test.go | 2 +- internal/httpapi/server.go | 2 +- internal/runtime/profiles_service.go | 9 +- internal/runtime/profiles_service_test.go | 20 +++ internal/runtime/profiles_suggest.go | 23 +-- internal/runtime/profiles_suggest_test.go | 31 +--- internal/runtime/runtime.go | 8 +- internal/server/server.go | 4 +- internal/storage/activity.go | 32 +++- .../storage/activity_stream_cancel_test.go | 147 ++++++++++++++++++ 17 files changed, 238 insertions(+), 81 deletions(-) create mode 100644 internal/storage/activity_stream_cancel_test.go diff --git a/internal/httpapi/activity.go b/internal/httpapi/activity.go index f3f3b10e5..89e63bd75 100644 --- a/internal/httpapi/activity.go +++ b/internal/httpapi/activity.go @@ -747,8 +747,12 @@ func (s *Server) handleExportActivity(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Disposition", "attachment; filename=\""+filename+"\"") w.Header().Set("Cache-Control", "no-cache") - // Stream activities - activityCh := s.controller.StreamActivities(filter) + // Stream activities. The context is cancelled on return, so a write + // error (client gone) releases the producer's read transaction instead of + // leaving it parked on a full channel for the life of the process. + streamCtx, cancelStream := context.WithCancel(r.Context()) + defer cancelStream() + activityCh := s.controller.StreamActivities(streamCtx, filter) // Write CSV header if format is CSV if format == "csv" { @@ -939,9 +943,9 @@ func (s *Server) handleActivitySummary(w http.ResponseWriter, r *http.Request) { toolCounts := make(map[string]int) perServer := make(map[string]*perServerAccumulator) - // The stream holds a read transaction open until the channel is drained or - // closed, so this loop must always run to completion. - for a := range s.controller.StreamActivities(filter) { + // The producer releases its read transaction when the request context is + // cancelled, so a client that disconnects mid-scan does not pin the store. + for a := range s.controller.StreamActivities(r.Context(), filter) { totalCount++ // "How many rows" (totalCount) and "how many calls" (callCount) are @@ -1398,7 +1402,7 @@ func (s *Server) handleActivityUsage(w http.ResponseWriter, r *http.Request) { var resp *contracts.UsageAggregateResponse if params.scopeFiltered() { - resp = s.scopedUsageResponse(params, time.Now().UTC()) + resp = s.scopedUsageResponse(r.Context(), params, time.Now().UTC()) } else { snap := s.controller.UsageSnapshot() tokens, _ := s.controller.GetTokenSavings() @@ -1416,7 +1420,7 @@ func (s *Server) handleActivityUsage(w http.ResponseWriter, r *http.Request) { // admission rule instead. Per-tool metrics are therefore window-bounded here // (lifetime-cumulative on the unfiltered path), and the global tokens-saved // headline is omitted because it cannot be re-derived for a subset. -func (s *Server) scopedUsageResponse(p usageParams, now time.Time) *contracts.UsageAggregateResponse { +func (s *Server) scopedUsageResponse(ctx context.Context, p usageParams, now time.Time) *contracts.UsageAggregateResponse { filter := storage.ActivityFilter{ // Every record Apply may admit, including the call_tool_* mirrors it // drops itself: the population a cold-start rebuild scans. @@ -1433,9 +1437,9 @@ func (s *Server) scopedUsageResponse(p usageParams, now time.Time) *contracts.Us filter.AllowedServers = p.allowed } agg := internalRuntime.NewUsageAggregate() - // The stream holds a read transaction open until drained, so it always runs - // to completion. - for rec := range s.controller.StreamActivities(filter) { + // The producer releases its read transaction when the request context is + // cancelled, so a client that disconnects mid-scan does not pin the store. + for rec := range s.controller.StreamActivities(ctx, filter) { agg.Apply(rec) } agg.UpdatedAt = now diff --git a/internal/httpapi/activity_auth_identity_test.go b/internal/httpapi/activity_auth_identity_test.go index 6a87e8f00..af64b661f 100644 --- a/internal/httpapi/activity_auth_identity_test.go +++ b/internal/httpapi/activity_auth_identity_test.go @@ -1,6 +1,7 @@ package httpapi import ( + "context" "encoding/json" "net/http" "testing" @@ -166,7 +167,7 @@ type streamingActivityController struct { mockActivityController } -func (m *streamingActivityController) StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (m *streamingActivityController) StreamActivities(_ context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord { ch := make(chan *storage.ActivityRecord, len(m.activities)) for _, a := range m.activities { if filter.Matches(a) { diff --git a/internal/httpapi/activity_call_parity_test.go b/internal/httpapi/activity_call_parity_test.go index 3fd38eacf..149d2dfa6 100644 --- a/internal/httpapi/activity_call_parity_test.go +++ b/internal/httpapi/activity_call_parity_test.go @@ -1,6 +1,7 @@ package httpapi import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -49,7 +50,7 @@ func (m *callParityController) GetCurrentConfig() *config.Config { return &config.Config{APIKey: "test-key", Observability: config.DefaultObservabilityConfig()} } -func (m *callParityController) StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (m *callParityController) StreamActivities(_ context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord { ch := make(chan *storage.ActivityRecord) go func() { defer close(ch) diff --git a/internal/httpapi/activity_scope_params_test.go b/internal/httpapi/activity_scope_params_test.go index 06ff90d60..05602e4cb 100644 --- a/internal/httpapi/activity_scope_params_test.go +++ b/internal/httpapi/activity_scope_params_test.go @@ -1,6 +1,7 @@ package httpapi import ( + "context" "encoding/csv" "encoding/json" "net/http" @@ -52,7 +53,7 @@ func (c *scopeParamsController) ListActivities(f storage.ActivityFilter) ([]*sto return out, total, nil } -func (c *scopeParamsController) StreamActivities(f storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (c *scopeParamsController) StreamActivities(_ context.Context, f storage.ActivityFilter) <-chan *storage.ActivityRecord { c.lastFilter = f ch := make(chan *storage.ActivityRecord, len(c.records)) for _, r := range c.records { diff --git a/internal/httpapi/activity_summary_scope_test.go b/internal/httpapi/activity_summary_scope_test.go index 163f0b060..61e3cca29 100644 --- a/internal/httpapi/activity_summary_scope_test.go +++ b/internal/httpapi/activity_summary_scope_test.go @@ -1,6 +1,7 @@ package httpapi import ( + "context" "encoding/json" "fmt" "net/http" @@ -49,7 +50,7 @@ func (m *summaryScopeController) ListActivities(filter storage.ActivityFilter) ( return matched, total, nil } -func (m *summaryScopeController) StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (m *summaryScopeController) StreamActivities(_ context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord { ch := make(chan *storage.ActivityRecord) go func() { defer close(ch) diff --git a/internal/httpapi/contracts_test.go b/internal/httpapi/contracts_test.go index 80e9a49df..62eb66430 100644 --- a/internal/httpapi/contracts_test.go +++ b/internal/httpapi/contracts_test.go @@ -302,7 +302,7 @@ func (m *MockServerController) AggregateToolUsage(_ time.Time) (map[string]stora return map[string]storage.ToolUsageStat{}, nil } func (m *MockServerController) UsageSnapshot() *internalRuntime.UsageAggregate { return nil } -func (m *MockServerController) StreamActivities(_ storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (m *MockServerController) StreamActivities(_ context.Context, _ storage.ActivityFilter) <-chan *storage.ActivityRecord { ch := make(chan *storage.ActivityRecord) close(ch) return ch diff --git a/internal/httpapi/scope_round9_harness_test.go b/internal/httpapi/scope_round9_harness_test.go index b0b8dd7c7..2fbff4862 100644 --- a/internal/httpapi/scope_round9_harness_test.go +++ b/internal/httpapi/scope_round9_harness_test.go @@ -1,6 +1,7 @@ package httpapi import ( + "context" "fmt" "time" @@ -93,7 +94,7 @@ func (c *scopeController) ListActivities(filter storage.ActivityFilter) ([]*stor return matched, total, nil } -func (c *scopeController) StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (c *scopeController) StreamActivities(_ context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord { ch := make(chan *storage.ActivityRecord, len(scopeActivities())) for _, rec := range scopeActivities() { if filter.Matches(rec) { diff --git a/internal/httpapi/security_test.go b/internal/httpapi/security_test.go index 43324a847..4fb3e97f4 100644 --- a/internal/httpapi/security_test.go +++ b/internal/httpapi/security_test.go @@ -352,7 +352,7 @@ func (m *baseController) AggregateToolUsage(_ time.Time) (map[string]storage.Too return map[string]storage.ToolUsageStat{}, nil } func (m *baseController) UsageSnapshot() *runtime.UsageAggregate { return nil } -func (m *baseController) StreamActivities(_ storage.ActivityFilter) <-chan *storage.ActivityRecord { +func (m *baseController) StreamActivities(_ context.Context, _ storage.ActivityFilter) <-chan *storage.ActivityRecord { ch := make(chan *storage.ActivityRecord) close(ch) return ch diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index e23e7193e..0831b1cdc 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -264,7 +264,7 @@ type ServerController interface { // Activity logging (RFC-003) ListActivities(filter storage.ActivityFilter) ([]*storage.ActivityRecord, int, error) GetActivity(id string) (*storage.ActivityRecord, error) - StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord + StreamActivities(ctx context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord // AggregateToolUsage rolls up tool_call activity per (server,tool) since the // given time. Backs the global tools page usage columns (spec 050). AggregateToolUsage(since time.Time) (map[string]storage.ToolUsageStat, error) diff --git a/internal/runtime/profiles_service.go b/internal/runtime/profiles_service.go index ab4681bb6..5ca4784fb 100644 --- a/internal/runtime/profiles_service.go +++ b/internal/runtime/profiles_service.go @@ -316,10 +316,7 @@ func (s *ProfilesService) Stats24h(ctx context.Context, viewer ViewerScope) Acti StartTime: now.Add(-24 * time.Hour), AllowedServers: filterServers, } - for rec := range s.rt.StreamActivities(filter) { - if ctx.Err() != nil { - break - } + for rec := range s.rt.StreamActivities(ctx, filter) { blocked := rec.Status == "blocked" isCall := rec.Type == storage.ActivityTypeToolCall if !isCall && !blocked { @@ -341,6 +338,10 @@ func (s *ProfilesService) Stats24h(ctx context.Context, viewer ViewerScope) Acti add(stats.ByProfile, rec.EffectiveProfile()) add(stats.ByClient, rec.ClientID) } + if ctx.Err() != nil { + // The stream closed early: a truncated rollup is not worth caching. + return stats + } s.statsMu.Lock() s.statsCache[key] = cachedStats{at: now, stats: stats} diff --git a/internal/runtime/profiles_service_test.go b/internal/runtime/profiles_service_test.go index 65f77a840..23740acc1 100644 --- a/internal/runtime/profiles_service_test.go +++ b/internal/runtime/profiles_service_test.go @@ -551,6 +551,26 @@ func TestProfilesService_Stats24hCountsByProfileAndClient(t *testing.T) { assert.Equal(t, Counter{Calls: 2, Blocked: 2}, stats.ByClient["cursor"]) } +// A cancelled caller gets whatever was counted so far, but that truncated +// rollup must not be cached for the next 30 s of callers (the stream closes +// early on cancellation, so the loop ends as if it had seen everything). +func TestProfilesService_Stats24hDoesNotCacheTruncatedRollup(t *testing.T) { + h := newProfilesHarness(t) + sm := h.rt.StorageManager() + now := time.Now() + for i := 0; i < 250; i++ { // more than the stream's channel buffer + require.NoError(t, sm.SaveActivity(&storage.ActivityRecord{Type: storage.ActivityTypeToolCall, Status: "success", Profile: "ro", ClientID: "cursor", Timestamp: now})) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + truncated := h.svc.Stats24h(ctx, ViewerScope{}) + assert.Less(t, truncated.ByProfile["ro"].Calls, 250, "a cancelled context ends the stream early") + + stats := h.svc.Stats24h(context.Background(), ViewerScope{}) + assert.Equal(t, 250, stats.ByProfile["ro"].Calls, "the truncated rollup must not have been cached") +} + // A save surfaces one warning per stale tools.classify entry, with its reason // (#1446): the editors previously only learned of staleness from a later // effective-tools read. diff --git a/internal/runtime/profiles_suggest.go b/internal/runtime/profiles_suggest.go index 681f065cb..4ff7f3ad0 100644 --- a/internal/runtime/profiles_suggest.go +++ b/internal/runtime/profiles_suggest.go @@ -198,12 +198,7 @@ func (s *ProfilesService) Suggest(ctx context.Context, req SuggestRequest) (*Acc Profile: req.Profile, TokenName: req.Token, } - stream := s.rt.StreamActivities(filter) - for rec := range stream { - if ctx.Err() != nil { - drainActivities(stream) - return nil, ctx.Err() - } + for rec := range s.rt.StreamActivities(ctx, filter) { reason := profile.BlockReason(rec.EffectiveBlockReason()) if !profileBlockReasons[reason] { out.NonProfileBlocked++ @@ -250,6 +245,10 @@ func (s *ProfilesService) Suggest(ctx context.Context, req SuggestRequest) (*Acc g.clients[rec.ClientID] = true } } + if err := ctx.Err(); err != nil { + // The stream closed early; a partial window would read as a complete one. + return nil, err + } for _, g := range groups { keep, err := s.finishSuggestion(ctx, ev, cfg, g, out) @@ -378,18 +377,6 @@ func deniedByRuleNow(exp *AccessExplanation) bool { return true } -// drainActivities lets a StreamActivities producer run to its end after the -// consumer stops early. The producer sends into a bounded channel from inside -// a BBolt read transaction, so a consumer that walks away with records left -// would leave that transaction open for the life of the process, and an open -// read transaction keeps the database from remapping on a later write. -func drainActivities(ch <-chan *storage.ActivityRecord) { - go func() { - for range ch { //nolint:revive // discard the remainder - } - }() -} - // buildPatches applies every apply-ready suggestion (allow, add_server) to a // copy of the stored profile. Classify needs a tier and is left out. func buildPatches(cfg *config.Config, sugg []AccessSuggestion) []ProfilePatch { diff --git a/internal/runtime/profiles_suggest_test.go b/internal/runtime/profiles_suggest_test.go index cc8a65f22..4af969f96 100644 --- a/internal/runtime/profiles_suggest_test.go +++ b/internal/runtime/profiles_suggest_test.go @@ -343,8 +343,8 @@ func TestSuggest_CancelledContextReturnsWithRecordsLeft(t *testing.T) { _, err := r.svc.Suggest(ctx, SuggestRequest{}) require.ErrorIs(t, err, context.Canceled) - // The producer was left behind with records to send; it must be drained so - // its read transaction ends and the store keeps accepting writes. + // The producer was left behind with records to send; cancellation must have + // released its read transaction so the store keeps accepting writes. for i := 0; i < 50; i++ { r.block("research", "github", "create_issue", profile.BlockReasonTier, "ci-bot", time.Minute, fmt.Sprintf("w%d", i)) } @@ -352,30 +352,3 @@ func TestSuggest_CancelledContextReturnsWithRecordsLeft(t *testing.T) { require.NoError(t, err) require.NotNil(t, r.find(res, "github:create_issue", profile.BlockReasonTier)) } - -func TestDrainActivities_LetsTheProducerFinish(t *testing.T) { - ch := make(chan *storage.ActivityRecord, 100) - finished := make(chan struct{}) - go func() { - defer close(finished) - defer close(ch) - for i := 0; i < 250; i++ { // more than the buffer: the producer blocks without a reader - ch <- &storage.ActivityRecord{RequestID: fmt.Sprintf("r%d", i)} - } - }() - // The consumer takes a few records and walks away, as Suggest does on cancel. - <-ch - <-ch - select { - case <-finished: - t.Fatal("the producer cannot have finished with no reader") - case <-time.After(50 * time.Millisecond): - } - - drainActivities(ch) - select { - case <-finished: - case <-time.After(5 * time.Second): - t.Fatal("drainActivities did not let the producer finish") - } -} diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 0f1a04405..38111504c 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -4349,14 +4349,16 @@ func (r *Runtime) UsageSnapshot() *UsageAggregate { return r.activityService.UsageSnapshot() } -// StreamActivities returns a channel that yields activity records matching the filter. -func (r *Runtime) StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord { +// StreamActivities returns a channel that yields activity records matching the +// filter. A consumer that stops reading early must cancel ctx so the producer +// releases its read transaction (see storage.Manager.StreamActivities). +func (r *Runtime) StreamActivities(ctx context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord { if r.storageManager == nil { ch := make(chan *storage.ActivityRecord) close(ch) return ch } - return r.storageManager.StreamActivities(filter) + return r.storageManager.StreamActivities(ctx, filter) } // ListToolApprovals returns all tool approval records for a server (Spec 032). diff --git a/internal/server/server.go b/internal/server/server.go index a3d6fb202..44e9d0e2c 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -4586,8 +4586,8 @@ func (s *Server) GetActivity(id string) (*storage.ActivityRecord, error) { } // StreamActivities returns a channel that yields activity records matching the filter. -func (s *Server) StreamActivities(filter storage.ActivityFilter) <-chan *storage.ActivityRecord { - return s.runtime.StreamActivities(filter) +func (s *Server) StreamActivities(ctx context.Context, filter storage.ActivityFilter) <-chan *storage.ActivityRecord { + return s.runtime.StreamActivities(ctx, filter) } // AggregateToolUsage rolls up tool_call activity per (server,tool) since the diff --git a/internal/storage/activity.go b/internal/storage/activity.go index 4f24aea13..c36316d2c 100644 --- a/internal/storage/activity.go +++ b/internal/storage/activity.go @@ -1,7 +1,9 @@ package storage import ( + "context" "encoding/json" + "errors" "fmt" "time" @@ -288,10 +290,17 @@ func (m *Manager) AggregateToolUsage(since time.Time) (map[string]ToolUsageStat, } // StreamActivities returns a channel that yields activity records matching the filter. -// The channel is closed when all matching records have been sent. -// This is useful for streaming large exports without loading all records into memory. -// Respects filter.Limit and filter.Offset for bounded streaming. -func (m *Manager) StreamActivities(filter ActivityFilter) <-chan *ActivityRecord { +// The channel is closed when all matching records have been sent, or as soon as +// ctx is done. This is useful for streaming large exports without loading all +// records into memory. Respects filter.Limit and filter.Offset for bounded streaming. +// +// The producer walks the activity bucket inside a BBolt read transaction and +// sends into a bounded channel. A consumer that stops reading early MUST cancel +// ctx: that is what lets the producer give up its pending send, return from +// db.View and release the read transaction (an open read transaction pins the +// mmap and blocks db.Close and remaps on later writes). Without cancellation an +// abandoned stream holds that transaction for the life of the process. +func (m *Manager) StreamActivities(ctx context.Context, filter ActivityFilter) <-chan *ActivityRecord { ch := make(chan *ActivityRecord, 100) go func() { @@ -328,14 +337,23 @@ func (m *Manager) StreamActivities(filter ActivityFilter) <-chan *ActivityRecord return nil } - ch <- &record - sent++ + // A done context wins even when the buffer has room: select + // picks a ready case at random, so check first. + if err := ctx.Err(); err != nil { + return err + } + select { + case ch <- &record: + sent++ + case <-ctx.Done(): + return ctx.Err() + } } return nil }) - if err != nil { + if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { m.logger.Errorw("Error streaming activities", "error", err) } }() diff --git a/internal/storage/activity_stream_cancel_test.go b/internal/storage/activity_stream_cancel_test.go new file mode 100644 index 000000000..01c357358 --- /dev/null +++ b/internal/storage/activity_stream_cancel_test.go @@ -0,0 +1,147 @@ +package storage + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// seedStreamRecords writes n tool-call records so the stream has more than +// the channel buffer (100) left to send after one is consumed. +func seedStreamRecords(t *testing.T, m *Manager, n int) { + t.Helper() + base := time.Now().UTC().Add(-time.Hour) + for i := 0; i < n; i++ { + require.NoError(t, m.SaveActivity(&ActivityRecord{ + ID: fmt.Sprintf("stream-cancel-%04d", i), + Type: ActivityTypeToolCall, + ServerName: "srv", + ToolName: "tool", + Status: ActivityStatusSuccess, + Timestamp: base.Add(time.Duration(i) * time.Millisecond), + })) + } +} + +// newStreamTestManager opens a manager in a temp dir without registering +// Close as cleanup: a failing assertion means a stream still pins a read +// transaction, and Close would then hang the whole package run. +func newStreamTestManager(t *testing.T) *Manager { + t.Helper() + m, err := NewManager(t.TempDir(), zap.NewNop().Sugar()) + require.NoError(t, err) + return m +} + +func waitForOpenTx(t *testing.T, m *Manager, want int) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for { + if m.GetDB().Stats().OpenTxN == want { + return + } + if time.Now().After(deadline) { + t.Fatalf("OpenTxN = %d, want %d after 5s", m.GetDB().Stats().OpenTxN, want) + } + time.Sleep(10 * time.Millisecond) + } +} + +func closeWithin(t *testing.T, m *Manager, d time.Duration) { + t.Helper() + done := make(chan error, 1) + go func() { done <- m.Close() }() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(d): + t.Fatalf("Manager.Close did not return within %s: a stream still holds a read transaction", d) + } +} + +// A consumer that stops reading after the first record and cancels its +// context must not pin the producer's BBolt read transaction: the producer +// returns from db.View, OpenTxN drops to 0 and Close does not hang. +func TestStreamActivities_CancelReleasesReadTransaction(t *testing.T) { + manager := newStreamTestManager(t) + seedStreamRecords(t, manager, 300) + + ctx, cancel := context.WithCancel(context.Background()) + ch := manager.StreamActivities(ctx, ActivityFilter{Types: []string{string(ActivityTypeToolCall)}}) + + first, ok := <-ch + require.True(t, ok) + require.NotNil(t, first) + cancel() // walk away with >100 records left to send + + waitForOpenTx(t, manager, 0) + // The channel is closed once the producer gives up, so a late reader + // cannot block on it. + deadline := time.After(5 * time.Second) + for { + select { + case _, open := <-ch: + if !open { + closeWithin(t, manager, 5*time.Second) + return + } + case <-deadline: + t.Fatal("stream channel was not closed after cancellation") + } + } +} + +// A consumer that walks away WITHOUT cancelling still leaks the transaction +// (that is the contract: the context is the release). This guards the test +// above against passing vacuously if the buffer ever grows past the seed. +func TestStreamActivities_AbandonedWithoutCancelHoldsTransaction(t *testing.T) { + manager := newStreamTestManager(t) + seedStreamRecords(t, manager, 300) + + ctx, cancel := context.WithCancel(context.Background()) + ch := manager.StreamActivities(ctx, ActivityFilter{Types: []string{string(ActivityTypeToolCall)}}) + _, ok := <-ch + require.True(t, ok) + + time.Sleep(200 * time.Millisecond) + require.Equal(t, 1, manager.GetDB().Stats().OpenTxN, "producer should be parked on a full channel inside db.View") + + cancel() + waitForOpenTx(t, manager, 0) + closeWithin(t, manager, 5*time.Second) +} + +// Cancelling before the first read yields no records and a closed channel. +func TestStreamActivities_CancelledContextYieldsNothing(t *testing.T) { + manager := newStreamTestManager(t) + seedStreamRecords(t, manager, 5) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + ch := manager.StreamActivities(ctx, ActivityFilter{}) + n := 0 + for range ch { + n++ + } + require.Equal(t, 0, n, "a done context is checked before every send") + waitForOpenTx(t, manager, 0) + closeWithin(t, manager, 5*time.Second) +} + +// The full stream still delivers everything when the context stays live. +func TestStreamActivities_LiveContextStreamsAll(t *testing.T) { + manager := newStreamTestManager(t) + seedStreamRecords(t, manager, 250) + + n := 0 + for range manager.StreamActivities(context.Background(), ActivityFilter{Types: []string{string(ActivityTypeToolCall)}}) { + n++ + } + require.Equal(t, 250, n) + waitForOpenTx(t, manager, 0) + closeWithin(t, manager, 5*time.Second) +}