| 1 | package guardian |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | "strings" |
| 10 | "sync" |
| 11 | "testing" |
| 12 | |
| 13 | "reasonix/internal/agent" |
| 14 | "reasonix/internal/event" |
| 15 | "reasonix/internal/provider" |
| 16 | "reasonix/internal/tool" |
| 17 | ) |
| 18 | |
| 19 | type scriptedProvider struct { |
| 20 | mu sync.Mutex |
| 21 | responses []scriptedResponse |
| 22 | requests []provider.Request |
| 23 | defaultUsage *provider.Usage |
| 24 | } |
| 25 | |
| 26 | type scriptedResponse struct { |
| 27 | text string |
| 28 | usage *provider.Usage |
| 29 | err error |
| 30 | } |
| 31 | |
| 32 | func (p *scriptedProvider) Name() string { return "guardian-test" } |
| 33 | |
| 34 | func (p *scriptedProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 35 | p.mu.Lock() |
| 36 | p.requests = append(p.requests, req) |
| 37 | resp := scriptedResponse{text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`} |
| 38 | if len(p.responses) > 0 { |
| 39 | resp = p.responses[0] |
| 40 | p.responses = p.responses[1:] |
| 41 | } |
| 42 | p.mu.Unlock() |
| 43 | if resp.usage == nil && p.defaultUsage != nil { |
| 44 | usage := *p.defaultUsage |
| 45 | resp.usage = &usage |
| 46 | } |
| 47 | |
| 48 | ch := make(chan provider.Chunk, 3) |
| 49 | if resp.err != nil { |
| 50 | close(ch) |
| 51 | return nil, resp.err |
| 52 | } |
| 53 | if resp.text != "" { |
| 54 | ch <- provider.Chunk{Type: provider.ChunkText, Text: resp.text} |
| 55 | } |
| 56 | if resp.usage != nil { |
| 57 | ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: resp.usage} |
| 58 | } |
| 59 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 60 | close(ch) |
| 61 | return ch, nil |
| 62 | } |
| 63 | |
| 64 | func (p *scriptedProvider) requestsSnapshot() []provider.Request { |
| 65 | p.mu.Lock() |
| 66 | defer p.mu.Unlock() |
| 67 | return append([]provider.Request(nil), p.requests...) |
| 68 | } |
| 69 | |
| 70 | type captureSink struct { |
| 71 | mu sync.Mutex |
| 72 | events []event.Event |
| 73 | } |
| 74 | |
| 75 | func (s *captureSink) Emit(e event.Event) { |
| 76 | s.mu.Lock() |
| 77 | defer s.mu.Unlock() |
| 78 | s.events = append(s.events, e) |
| 79 | } |
| 80 | |
| 81 | func (s *captureSink) guardianEvents() []event.Event { |
| 82 | s.mu.Lock() |
| 83 | defer s.mu.Unlock() |
| 84 | var out []event.Event |
| 85 | for _, e := range s.events { |
| 86 | if e.Kind == event.GuardianAssessment { |
| 87 | out = append(out, e) |
| 88 | } |
| 89 | } |
| 90 | return out |
| 91 | } |
| 92 | |
| 93 | func TestParseAssessmentEnforcesCriticalDeny(t *testing.T) { |
| 94 | a, err := ParseAssessment(`{"risk_level":"critical","user_authorization":"high","outcome":"allow","rationale":"delete prod secrets"}`) |
| 95 | if err != nil { |
| 96 | t.Fatalf("ParseAssessment error: %v", err) |
| 97 | } |
| 98 | if a.Outcome != "deny" { |
| 99 | t.Fatalf("critical outcome = %q, want deny", a.Outcome) |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | func TestParseAssessmentRejectsUnknownEnum(t *testing.T) { |
| 104 | if _, err := ParseAssessment(`{"risk_level":"spicy","user_authorization":"high","outcome":"allow","rationale":"x"}`); err == nil { |
| 105 | t.Fatal("ParseAssessment accepted unknown risk_level") |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | func TestTranscriptRenderKeepsFirstAndLastUserAnchors(t *testing.T) { |
| 110 | entries := []TranscriptEntry{{Kind: "user", Text: "first task"}} |
| 111 | for i := 0; i < maxRecentEntries+5; i++ { |
| 112 | entries = append(entries, TranscriptEntry{Kind: "assistant", Text: "assistant detail"}) |
| 113 | } |
| 114 | entries = append(entries, TranscriptEntry{Kind: "user", Text: "latest instruction"}) |
| 115 | rendered := FormatTranscript(entries) |
| 116 | if !strings.Contains(rendered, "first task") || !strings.Contains(rendered, "latest instruction") { |
| 117 | t.Fatalf("rendered transcript lost user anchors:\n%s", rendered) |
| 118 | } |
| 119 | if !strings.Contains(rendered, "Some conversation entries were omitted.") { |
| 120 | t.Fatalf("rendered transcript should mention omissions:\n%s", rendered) |
| 121 | } |
| 122 | } |
| 123 | |
| 124 | func TestGuardianSaveLoadRestoresCursorForDeltaTranscript(t *testing.T) { |
| 125 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 126 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first ok"}`}, |
| 127 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"second ok"}`}, |
| 128 | }} |
| 129 | sink := &captureSink{} |
| 130 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 131 | parent := agent.NewSession("sys") |
| 132 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "first user request"}) |
| 133 | |
| 134 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 135 | t.Fatalf("first Review = allow %v err %v, want allow nil", allow, err) |
| 136 | } |
| 137 | path := filepath.Join(t.TempDir(), "session.guardian.jsonl") |
| 138 | if err := gs.Save(path); err != nil { |
| 139 | t.Fatalf("Save error: %v", err) |
| 140 | } |
| 141 | if data, err := os.ReadFile(cursorPathForGuardianPath(path)); err != nil || !strings.Contains(string(data), `"EntryCount":1`) { |
| 142 | t.Fatalf("cursor sidecar = %q err %v, want EntryCount 1", data, err) |
| 143 | } |
| 144 | |
| 145 | loaded := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 146 | if err := loaded.Load(path); err != nil { |
| 147 | t.Fatalf("Load error: %v", err) |
| 148 | } |
| 149 | if loaded.cursor.EntryCount != 1 { |
| 150 | t.Fatalf("loaded cursor = %+v, want EntryCount 1", loaded.cursor) |
| 151 | } |
| 152 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "second user request"}) |
| 153 | if allow, _, err := loaded.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"b.txt"}`), parent); err != nil || !allow { |
| 154 | t.Fatalf("second Review = allow %v err %v, want allow nil", allow, err) |
| 155 | } |
| 156 | |
| 157 | reqs := prov.requestsSnapshot() |
| 158 | if len(reqs) < 2 { |
| 159 | t.Fatalf("requests = %d, want >= 2", len(reqs)) |
| 160 | } |
| 161 | var delta string |
| 162 | for _, req := range reqs { |
| 163 | for _, m := range req.Messages { |
| 164 | if strings.Contains(m.Content, "TRANSCRIPT DELTA") && strings.Contains(m.Content, "second user request") { |
| 165 | delta = m.Content |
| 166 | break |
| 167 | } |
| 168 | } |
| 169 | if delta != "" { |
| 170 | break |
| 171 | } |
| 172 | } |
| 173 | if delta == "" { |
| 174 | t.Fatalf("second request did not include a delta transcript") |
| 175 | } |
| 176 | if !strings.Contains(delta, "second user request") { |
| 177 | t.Fatalf("delta transcript missing new parent entry:\n%s", delta) |
| 178 | } |
| 179 | if strings.Contains(delta, "first user request") { |
| 180 | t.Fatalf("delta transcript repeated old parent entry:\n%s", delta) |
| 181 | } |
| 182 | } |
| 183 | |
| 184 | func TestGuardianUsageDoesNotLeakAcrossReviews(t *testing.T) { |
| 185 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 186 | { |
| 187 | text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first ok"}`, |
| 188 | usage: &provider.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}, |
| 189 | }, |
| 190 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"second ok"}`}, |
| 191 | }} |
| 192 | sink := &captureSink{} |
| 193 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 194 | parent := agent.NewSession("sys") |
| 195 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do it"}) |
| 196 | |
| 197 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 198 | t.Fatalf("first Review = allow %v err %v, want allow nil", allow, err) |
| 199 | } |
| 200 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"b.txt"}`), parent); err != nil || !allow { |
| 201 | t.Fatalf("second Review = allow %v err %v, want allow nil", allow, err) |
| 202 | } |
| 203 | |
| 204 | events := sink.guardianEvents() |
| 205 | if len(events) != 2 { |
| 206 | t.Fatalf("guardian events = %d, want 2", len(events)) |
| 207 | } |
| 208 | if events[0].Guardian.Usage == nil || events[0].Guardian.Usage.TotalTokens != 12 { |
| 209 | t.Fatalf("first usage = %+v, want total 12", events[0].Guardian.Usage) |
| 210 | } |
| 211 | if events[1].Guardian.Usage != nil { |
| 212 | t.Fatalf("second usage leaked from first review: %+v", events[1].Guardian.Usage) |
| 213 | } |
| 214 | } |
| 215 | |
| 216 | func TestGuardianUsageAggregatesEveryModelCall(t *testing.T) { |
| 217 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 218 | {text: "", usage: &provider.Usage{PromptTokens: 3, CompletionTokens: 1, TotalTokens: 4}}, |
| 219 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`, usage: &provider.Usage{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7, RequestCount: 2}}, |
| 220 | }} |
| 221 | sink := &captureSink{} |
| 222 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 223 | parent := agent.NewSession("sys") |
| 224 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do it"}) |
| 225 | |
| 226 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 227 | t.Fatalf("Review = allow %v err %v, want allow nil", allow, err) |
| 228 | } |
| 229 | events := sink.guardianEvents() |
| 230 | if len(events) != 1 || events[0].Guardian.Usage == nil { |
| 231 | t.Fatalf("guardian events = %+v, want one usage-bearing event", events) |
| 232 | } |
| 233 | usage := events[0].Guardian.Usage |
| 234 | if usage.PromptTokens != 8 || usage.CompletionTokens != 3 || usage.TotalTokens != 11 || usage.RequestCount != 3 { |
| 235 | t.Fatalf("aggregated usage = %+v, want prompt=8 completion=3 total=11 requests=3", usage) |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | // TestGuardianReviewTurnsAlternateRoles pins the review request shape: the |
| 240 | // transcript evidence and action request ride in one combined user message per |
| 241 | // review, so the guardian session alternates user/assistant strictly and |
| 242 | // providers that reject consecutive same-role messages can run the guardian. |
| 243 | func TestGuardianReviewTurnsAlternateRoles(t *testing.T) { |
| 244 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 245 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first ok"}`}, |
| 246 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"second ok"}`}, |
| 247 | }} |
| 248 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 249 | parent := agent.NewSession("sys") |
| 250 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do the thing"}) |
| 251 | |
| 252 | for i := 0; i < 2; i++ { |
| 253 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 254 | t.Fatalf("review %d = allow %v err %v, want allow nil", i+1, allow, err) |
| 255 | } |
| 256 | } |
| 257 | |
| 258 | reqs := prov.requestsSnapshot() |
| 259 | if len(reqs) < 2 { |
| 260 | t.Fatalf("requests = %d, want >= 2", len(reqs)) |
| 261 | } |
| 262 | for r, req := range reqs { |
| 263 | for i := 1; i < len(req.Messages); i++ { |
| 264 | if req.Messages[i].Role == provider.RoleUser && req.Messages[i-1].Role == provider.RoleUser { |
| 265 | t.Fatalf("request %d carries consecutive user messages at index %d", r, i) |
| 266 | } |
| 267 | } |
| 268 | } |
| 269 | |
| 270 | // The combined message must still carry the evidence boundary and the action. |
| 271 | msgs := reqs[len(reqs)-1].Messages |
| 272 | var review string |
| 273 | for i := len(msgs) - 1; i >= 0; i-- { |
| 274 | if msgs[i].Role == provider.RoleUser { |
| 275 | review = msgs[i].Content |
| 276 | break |
| 277 | } |
| 278 | } |
| 279 | for _, want := range []string{"untrusted evidence", "The agent has requested the following action", "write_file"} { |
| 280 | if !strings.Contains(review, want) { |
| 281 | t.Fatalf("combined review message missing %q:\n%s", want, review) |
| 282 | } |
| 283 | } |
| 284 | } |
| 285 | |
| 286 | // TestGuardianFailedReviewRollsBackSession pins the error-path rollback: |
| 287 | // agent.Run appends the combined review user message before the provider is |
| 288 | // reached, so a failed review must not leave it dangling — the next review |
| 289 | // would otherwise append another user message and strict-alternation providers |
| 290 | // would reject every request from then on. |
| 291 | func TestGuardianFailedReviewRollsBackSession(t *testing.T) { |
| 292 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 293 | {err: fmt.Errorf("provider unavailable")}, |
| 294 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`}, |
| 295 | }} |
| 296 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 297 | parent := agent.NewSession("sys") |
| 298 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do the thing"}) |
| 299 | |
| 300 | allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent) |
| 301 | if err == nil && allow { |
| 302 | t.Fatal("first review should fail closed") |
| 303 | } |
| 304 | if n := gs.sess.Len(); n != 1 { |
| 305 | t.Fatalf("guardian session messages = %d after failed review, want rollback to system only", n) |
| 306 | } |
| 307 | |
| 308 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 309 | t.Fatalf("second review = allow %v err %v, want allow nil", allow, err) |
| 310 | } |
| 311 | reqs := prov.requestsSnapshot() |
| 312 | last := reqs[len(reqs)-1] |
| 313 | for i := 1; i < len(last.Messages); i++ { |
| 314 | if last.Messages[i].Role == provider.RoleUser && last.Messages[i-1].Role == provider.RoleUser { |
| 315 | t.Fatalf("request after failed review carries consecutive user messages at index %d", i) |
| 316 | } |
| 317 | } |
| 318 | } |
| 319 | |
| 320 | // TestGuardianLoadResetsLegacyConsecutiveUserSessions pins the load-time |
| 321 | // normalization: sessions saved by the old multi-message review shape carry |
| 322 | // consecutive user messages that would poison strict-alternation providers, so |
| 323 | // Load starts fresh instead of adopting them. |
| 324 | func TestGuardianLoadResetsLegacyConsecutiveUserSessions(t *testing.T) { |
| 325 | legacy := agent.NewSession(PolicyPrompt()) |
| 326 | legacy.Add(provider.Message{Role: provider.RoleUser, Content: "transcript evidence"}) |
| 327 | legacy.Add(provider.Message{Role: provider.RoleUser, Content: "action request"}) |
| 328 | legacy.Add(provider.Message{Role: provider.RoleAssistant, Content: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`}) |
| 329 | path := filepath.Join(t.TempDir(), "session.guardian.jsonl") |
| 330 | if err := legacy.Save(path); err != nil { |
| 331 | t.Fatalf("Save legacy session: %v", err) |
| 332 | } |
| 333 | |
| 334 | gs := NewSession(&scriptedProvider{}, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 335 | if err := gs.Load(path); err != nil { |
| 336 | t.Fatalf("Load: %v", err) |
| 337 | } |
| 338 | if n := gs.sess.Len(); n != 1 { |
| 339 | t.Fatalf("loaded legacy session messages = %d, want reset to system only", n) |
| 340 | } |
| 341 | if gs.cursor.EntryCount != 0 { |
| 342 | t.Fatalf("cursor = %+v, want zeroed after reset", gs.cursor) |
| 343 | } |
| 344 | } |
| 345 | |
| 346 | // TestGuardianSessionAlternatesAfterCompaction reproduces the compaction seam: |
| 347 | // every compactEvery-th review runs CompactNow, and generic compaction inserts |
| 348 | // its digest as a RoleUser message that can land directly before a review's |
| 349 | // user turn — consecutive user roles again. The post-review normalization must |
| 350 | // keep the session strictly alternating across that fold. |
| 351 | func TestGuardianSessionAlternatesAfterCompaction(t *testing.T) { |
| 352 | prov := &scriptedProvider{defaultUsage: &provider.Usage{TotalTokens: 1}} // default allow verdict, also serves the summarizer |
| 353 | sink := &captureSink{} |
| 354 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 355 | parent := agent.NewSession("sys") |
| 356 | |
| 357 | filler := strings.Repeat("parent transcript filler. ", 160) |
| 358 | for i := 0; i < compactEvery; i++ { |
| 359 | parent.Add(provider.Message{Role: provider.RoleUser, Content: fmt.Sprintf("turn %d: %s", i, filler)}) |
| 360 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 361 | t.Fatalf("review %d = allow %v err %v, want allow nil", i+1, allow, err) |
| 362 | } |
| 363 | } |
| 364 | |
| 365 | msgs := gs.sess.Snapshot() |
| 366 | var hasDigest bool |
| 367 | for _, m := range msgs { |
| 368 | if agent.IsCompactionSummary(m) { |
| 369 | hasDigest = true |
| 370 | } |
| 371 | } |
| 372 | if !hasDigest { |
| 373 | t.Fatal("test setup: guardian compaction did not fold anything, the digest adjacency is not exercised") |
| 374 | } |
| 375 | events := sink.guardianEvents() |
| 376 | if len(events) != compactEvery { |
| 377 | t.Fatalf("guardian events = %d, want %d", len(events), compactEvery) |
| 378 | } |
| 379 | usage := events[len(events)-1].Guardian.Usage |
| 380 | if usage == nil || usage.TotalTokens != 2 || usage.RequestCount != 2 { |
| 381 | t.Fatalf("compacting review usage = %+v, want total=2 requests=2", usage) |
| 382 | } |
| 383 | for i := 1; i < len(msgs); i++ { |
| 384 | if msgs[i].Role == msgs[i-1].Role { |
| 385 | t.Fatalf("guardian session has consecutive %s messages at indexes %d/%d of %d", msgs[i].Role, i-1, i, len(msgs)) |
| 386 | } |
| 387 | } |
| 388 | } |
| 389 |