| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "sync" |
| 8 | "testing" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| 12 | // Sampling Err before signalling makes the lock-wait interleaving deterministic: |
| 13 | // the first check returns nil even if cancellation arrives before it returns. |
| 14 | type prepareEntryContext struct { |
| 15 | context.Context |
| 16 | entered chan struct{} |
| 17 | once sync.Once |
| 18 | } |
| 19 | |
| 20 | func (c *prepareEntryContext) Err() error { |
| 21 | err := c.Context.Err() |
| 22 | c.once.Do(func() { close(c.entered) }) |
| 23 | return err |
| 24 | } |
| 25 | |
| 26 | func TestPrepareRejectsCancelledContextBeforeMaintenance(t *testing.T) { |
| 27 | for _, tc := range []struct { |
| 28 | name string |
| 29 | turns int |
| 30 | deadline bool |
| 31 | }{ |
| 32 | {"below threshold", 0, false}, |
| 33 | {"above ceiling", 6, false}, |
| 34 | {"expired deadline", 6, true}, |
| 35 | } { |
| 36 | t.Run(tc.name, func(t *testing.T) { |
| 37 | prov := &failingSummaryProvider{} |
| 38 | a := agentOverForce(t, prov, foldableSessionOverForce(tc.turns)) |
| 39 | before, err := json.Marshal(a.modelVisibleMessages()) |
| 40 | if err != nil { |
| 41 | t.Fatal(err) |
| 42 | } |
| 43 | version := a.currentProjectionVersion() |
| 44 | canonical, canonicalVersion := a.sess.conversation.snapshotMessagesVersion() |
| 45 | canonicalBefore, err := json.Marshal(canonical) |
| 46 | if err != nil { |
| 47 | t.Fatal(err) |
| 48 | } |
| 49 | ctx, cancel := context.WithCancel(context.Background()) |
| 50 | want := error(context.Canceled) |
| 51 | if tc.deadline { |
| 52 | cancel() |
| 53 | ctx, cancel = context.WithDeadline(context.Background(), time.Unix(1, 0)) |
| 54 | want = context.DeadlineExceeded |
| 55 | } else { |
| 56 | cancel() |
| 57 | } |
| 58 | defer cancel() |
| 59 | trigger := CompactionTriggerOverflow |
| 60 | if tc.turns == 0 { |
| 61 | trigger = CompactionTriggerPressure |
| 62 | } |
| 63 | _, err = a.contextManager().Prepare(ctx, ContextPreparePolicy{Trigger: trigger}) |
| 64 | if !errors.Is(err, want) { |
| 65 | t.Fatalf("Prepare error = %v, want %v", err, want) |
| 66 | } |
| 67 | after, err := json.Marshal(a.modelVisibleMessages()) |
| 68 | if err != nil { |
| 69 | t.Fatal(err) |
| 70 | } |
| 71 | current, currentVersion := a.sess.conversation.snapshotMessagesVersion() |
| 72 | canonicalAfter, err := json.Marshal(current) |
| 73 | if err != nil { |
| 74 | t.Fatal(err) |
| 75 | } |
| 76 | if string(before) != string(after) || version != a.currentProjectionVersion() || a.sess.compactionState.LastReceipt != nil { |
| 77 | t.Fatal("cancelled Prepare changed projection or receipt") |
| 78 | } |
| 79 | if string(canonicalBefore) != string(canonicalAfter) || canonicalVersion != currentVersion { |
| 80 | t.Fatal("cancelled Prepare changed canonical history") |
| 81 | } |
| 82 | if prov.calls != 0 { |
| 83 | t.Fatalf("cancelled Prepare called provider %d times", prov.calls) |
| 84 | } |
| 85 | }) |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | func TestPrepareRejectsCancellationWhileWaitingForMaintenance(t *testing.T) { |
| 90 | prov := &failingSummaryProvider{} |
| 91 | a := agentOverForce(t, prov, foldableSessionOverForce(6)) |
| 92 | ctx, cancel := context.WithCancel(context.Background()) |
| 93 | defer cancel() |
| 94 | entry := &prepareEntryContext{Context: ctx, entered: make(chan struct{})} |
| 95 | a.sess.compactionRunMu.Lock() |
| 96 | result := make(chan error, 1) |
| 97 | go func() { |
| 98 | _, err := a.contextManager().Prepare(entry, ContextPreparePolicy{Trigger: CompactionTriggerOverflow}) |
| 99 | result <- err |
| 100 | }() |
| 101 | select { |
| 102 | case <-entry.entered: |
| 103 | case <-time.After(5 * time.Second): |
| 104 | cancel() |
| 105 | a.sess.compactionRunMu.Unlock() |
| 106 | <-result |
| 107 | t.Fatal("Prepare did not check cancellation before waiting for maintenance") |
| 108 | } |
| 109 | cancel() |
| 110 | a.sess.compactionRunMu.Unlock() |
| 111 | if err := <-result; !errors.Is(err, context.Canceled) { |
| 112 | t.Fatalf("Prepare error = %v, want cancellation", err) |
| 113 | } |
| 114 | if prov.calls != 0 || a.currentProjectionVersion() != 0 || a.sess.compactionState.LastReceipt != nil { |
| 115 | t.Fatal("cancelled waiter entered maintenance") |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | func TestPrepareAllowsLiveBelowThresholdContext(t *testing.T) { |
| 120 | prov := &failingSummaryProvider{} |
| 121 | a := agentOverForce(t, prov, foldableSessionOverForce(0)) |
| 122 | prepared, err := a.contextManager().Prepare(context.Background(), ContextPreparePolicy{}) |
| 123 | if err != nil { |
| 124 | t.Fatal(err) |
| 125 | } |
| 126 | if len(prepared.Messages) == 0 || prov.calls != 0 || a.currentProjectionVersion() != 0 { |
| 127 | t.Fatal("live fast path changed") |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | func TestPreparePreservesLegacyNilContext(t *testing.T) { |
| 132 | for _, manual := range []bool{false, true} { |
| 133 | turns := 0 |
| 134 | if manual { |
| 135 | turns = 6 |
| 136 | } |
| 137 | a := agentOverForce(t, &fakeProvider{reply: "digest"}, foldableSessionOverForce(turns)) |
| 138 | var err error |
| 139 | if manual { |
| 140 | err = a.CompactNow(nil, "") //nolint:staticcheck // Exercise the legacy nil-context compatibility boundary. |
| 141 | } else { |
| 142 | err = a.PrepareContext(nil) //nolint:staticcheck // Exercise the legacy nil-context compatibility boundary. |
| 143 | } |
| 144 | if err != nil { |
| 145 | t.Fatalf("manual=%v: %v", manual, err) |
| 146 | } |
| 147 | } |
| 148 | } |
| 149 |