| 1 | package guardian |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | "reflect" |
| 10 | "slices" |
| 11 | "strings" |
| 12 | "sync" |
| 13 | "testing" |
| 14 | |
| 15 | "reasonix/internal/agent" |
| 16 | "reasonix/internal/event" |
| 17 | "reasonix/internal/provider" |
| 18 | "reasonix/internal/tool" |
| 19 | ) |
| 20 | |
| 21 | type scriptedProvider struct { |
| 22 | mu sync.Mutex |
| 23 | responses []scriptedResponse |
| 24 | requests []provider.Request |
| 25 | defaultUsage *provider.Usage |
| 26 | } |
| 27 | |
| 28 | type scriptedResponse struct { |
| 29 | text string |
| 30 | reasoning string |
| 31 | usage *provider.Usage |
| 32 | err error |
| 33 | } |
| 34 | |
| 35 | type reasoningScriptedProvider struct{ *scriptedProvider } |
| 36 | |
| 37 | func (*reasoningScriptedProvider) RequiresToolCallReasoning() bool { return true } |
| 38 | |
| 39 | func (p *scriptedProvider) Name() string { return "guardian-test" } |
| 40 | |
| 41 | func (p *scriptedProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 42 | p.mu.Lock() |
| 43 | p.requests = append(p.requests, req) |
| 44 | resp := scriptedResponse{text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`} |
| 45 | if len(p.responses) > 0 { |
| 46 | resp = p.responses[0] |
| 47 | p.responses = p.responses[1:] |
| 48 | } |
| 49 | p.mu.Unlock() |
| 50 | if resp.usage == nil && p.defaultUsage != nil { |
| 51 | usage := *p.defaultUsage |
| 52 | resp.usage = &usage |
| 53 | } |
| 54 | |
| 55 | ch := make(chan provider.Chunk, 4) |
| 56 | if resp.err != nil { |
| 57 | close(ch) |
| 58 | return nil, resp.err |
| 59 | } |
| 60 | if resp.reasoning != "" { |
| 61 | ch <- provider.Chunk{Type: provider.ChunkReasoning, Text: resp.reasoning} |
| 62 | } |
| 63 | if resp.text != "" { |
| 64 | ch <- provider.Chunk{Type: provider.ChunkText, Text: resp.text} |
| 65 | } |
| 66 | if resp.usage != nil { |
| 67 | ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: resp.usage} |
| 68 | } |
| 69 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 70 | close(ch) |
| 71 | return ch, nil |
| 72 | } |
| 73 | |
| 74 | func (p *scriptedProvider) requestsSnapshot() []provider.Request { |
| 75 | p.mu.Lock() |
| 76 | defer p.mu.Unlock() |
| 77 | return append([]provider.Request(nil), p.requests...) |
| 78 | } |
| 79 | |
| 80 | type captureSink struct { |
| 81 | mu sync.Mutex |
| 82 | events []event.Event |
| 83 | } |
| 84 | |
| 85 | func (s *captureSink) Emit(e event.Event) { |
| 86 | s.mu.Lock() |
| 87 | defer s.mu.Unlock() |
| 88 | s.events = append(s.events, e) |
| 89 | } |
| 90 | |
| 91 | func (s *captureSink) guardianEvents() []event.Event { |
| 92 | s.mu.Lock() |
| 93 | defer s.mu.Unlock() |
| 94 | var out []event.Event |
| 95 | for _, e := range s.events { |
| 96 | if e.Kind == event.GuardianAssessment { |
| 97 | out = append(out, e) |
| 98 | } |
| 99 | } |
| 100 | return out |
| 101 | } |
| 102 | |
| 103 | func TestParseAssessmentEnforcesCriticalDeny(t *testing.T) { |
| 104 | a, err := ParseAssessment(`{"risk_level":"critical","user_authorization":"high","outcome":"allow","rationale":"delete prod secrets"}`) |
| 105 | if err != nil { |
| 106 | t.Fatalf("ParseAssessment error: %v", err) |
| 107 | } |
| 108 | if a.Outcome != "deny" { |
| 109 | t.Fatalf("critical outcome = %q, want deny", a.Outcome) |
| 110 | } |
| 111 | } |
| 112 | |
| 113 | func TestParseAssessmentRejectsUnknownEnum(t *testing.T) { |
| 114 | if _, err := ParseAssessment(`{"risk_level":"spicy","user_authorization":"high","outcome":"allow","rationale":"x"}`); err == nil { |
| 115 | t.Fatal("ParseAssessment accepted unknown risk_level") |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | func TestGuardianReasoningOnlyStopRetriesInsteadOfReusingPriorVerdict(t *testing.T) { |
| 120 | base := &scriptedProvider{responses: []scriptedResponse{ |
| 121 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first action is safe"}`}, |
| 122 | {reasoning: "The second action is dangerous and must be denied.", usage: &provider.Usage{FinishReason: "stop"}}, |
| 123 | {text: `{"risk_level":"high","user_authorization":"low","outcome":"deny","rationale":"current action is unsafe"}`}, |
| 124 | }} |
| 125 | prov := &reasoningScriptedProvider{scriptedProvider: base} |
| 126 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 127 | parent := agent.NewSession("sys") |
| 128 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "review two different actions"}) |
| 129 | |
| 130 | if allow, _, err := gs.ReviewVerdict(context.Background(), "read_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 131 | t.Fatalf("first ReviewVerdict = allow %v err %v, want allow nil", allow, err) |
| 132 | } |
| 133 | allow, reason, err := gs.ReviewVerdict(context.Background(), "bash", json.RawMessage(`{"command":"dangerous command"}`), parent) |
| 134 | if err != nil || allow { |
| 135 | t.Fatalf("second ReviewVerdict = allow %v reason %q err %v, want deny with authentic verdict", allow, reason, err) |
| 136 | } |
| 137 | if !strings.Contains(reason, "current action is unsafe") { |
| 138 | t.Fatalf("second denial reason = %q, want current verdict rationale", reason) |
| 139 | } |
| 140 | if got := len(base.requestsSnapshot()); got != 3 { |
| 141 | t.Fatalf("provider requests = %d, want 3 (allow, reasoning-only, visible retry)", got) |
| 142 | } |
| 143 | } |
| 144 | |
| 145 | func TestGuardianRepeatedReasoningOnlyStopsFailClosedWithoutReusingPriorAllow(t *testing.T) { |
| 146 | base := &scriptedProvider{responses: []scriptedResponse{ |
| 147 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first action is safe"}`}, |
| 148 | {reasoning: "dangerous review 1", usage: &provider.Usage{FinishReason: "stop"}}, |
| 149 | {reasoning: "dangerous review 2", usage: &provider.Usage{FinishReason: "stop"}}, |
| 150 | {reasoning: "dangerous review 3", usage: &provider.Usage{FinishReason: "stop"}}, |
| 151 | }} |
| 152 | prov := &reasoningScriptedProvider{scriptedProvider: base} |
| 153 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 154 | parent := agent.NewSession("sys") |
| 155 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "review two different actions"}) |
| 156 | |
| 157 | if allow, _, err := gs.ReviewVerdict(context.Background(), "read_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 158 | t.Fatalf("first ReviewVerdict = allow %v err %v, want allow nil", allow, err) |
| 159 | } |
| 160 | before := gs.sess.Snapshot() |
| 161 | allow, _, err := gs.ReviewVerdict(context.Background(), "bash", json.RawMessage(`{"command":"dangerous command"}`), parent) |
| 162 | if allow || err == nil { |
| 163 | t.Fatalf("second ReviewVerdict = allow %v err %v, want fail-closed error", allow, err) |
| 164 | } |
| 165 | if got := len(base.requestsSnapshot()); got != 4 { |
| 166 | t.Fatalf("provider requests = %d, want 4 (prior allow plus three bounded empty finals)", got) |
| 167 | } |
| 168 | if after := gs.sess.Snapshot(); !reflect.DeepEqual(after, before) { |
| 169 | t.Fatalf("failed reasoning-only review did not roll back session:\nbefore=%+v\nafter=%+v", before, after) |
| 170 | } |
| 171 | |
| 172 | // A fresh review after the bounded failure must start from the prior clean |
| 173 | // verdict, not from any of the discarded retry turns. The scripted |
| 174 | // provider's default response is a visible allow verdict. |
| 175 | if allow, _, err := gs.ReviewVerdict(context.Background(), "read_file", json.RawMessage(`{"file_path":"c.txt"}`), parent); err != nil || !allow { |
| 176 | t.Fatalf("review after rollback = allow %v err %v, want allow nil", allow, err) |
| 177 | } |
| 178 | reqs := base.requestsSnapshot() |
| 179 | if got := len(reqs); got != 5 { |
| 180 | t.Fatalf("provider requests after recovery = %d, want 5", got) |
| 181 | } |
| 182 | last := reqs[len(reqs)-1] |
| 183 | for i := 1; i < len(last.Messages); i++ { |
| 184 | if last.Messages[i].Role == provider.RoleUser && last.Messages[i-1].Role == provider.RoleUser { |
| 185 | t.Fatalf("request after reasoning-only rollback carries consecutive user messages at index %d", i) |
| 186 | } |
| 187 | } |
| 188 | } |
| 189 | |
| 190 | func TestGuardianRollbackAfterRewriteDropsReasoningOnlyRetryTail(t *testing.T) { |
| 191 | gs := NewSession(&scriptedProvider{}, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 192 | before := gs.sess.Snapshot() |
| 193 | rewriteBefore := gs.sess.RewriteVersion() |
| 194 | |
| 195 | compactedWithEvidence := []provider.Message{ |
| 196 | {Role: provider.RoleSystem, Content: PolicyPrompt()}, |
| 197 | {Role: provider.RoleUser, Content: "<compaction-summary>\ncompacted reviews\n</compaction-summary>"}, |
| 198 | {Role: provider.RoleAssistant, Content: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"preserved verdict"}`}, |
| 199 | {Role: provider.RoleUser, Content: "review the current action"}, |
| 200 | {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "read-1", Name: "read_file", Arguments: `{"path":"policy.md"}`}}}, |
| 201 | {Role: provider.RoleTool, ToolCallID: "read-1", Name: "read_file", Content: "preserved read-only evidence"}, |
| 202 | } |
| 203 | gs.sess.Replace(compactedWithEvidence) |
| 204 | gs.sess.IncrementRewrite() |
| 205 | gs.sess.Add(provider.Message{Role: provider.RoleAssistant, ReasoningContent: "first hidden-only verdict"}) |
| 206 | gs.sess.Add(provider.Message{Role: provider.RoleUser, Content: "provide a visible verdict"}) |
| 207 | gs.sess.Add(provider.Message{Role: provider.RoleAssistant, ReasoningContent: "second hidden-only verdict"}) |
| 208 | gs.sess.Add(provider.Message{Role: provider.RoleUser, Content: "do not call tools; provide a visible verdict"}) |
| 209 | gs.sess.Add(provider.Message{Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "unpaired-1", Name: "read_file", Arguments: `{"path":"more.md"}`}}}) |
| 210 | |
| 211 | gs.rollbackReview(before, rewriteBefore) |
| 212 | |
| 213 | got := gs.sess.Snapshot() |
| 214 | if !reflect.DeepEqual(got, compactedWithEvidence) { |
| 215 | t.Fatalf("rewrite-aware rollback changed compacted or completed tool evidence:\n got=%+v\nwant=%+v", got, compactedWithEvidence) |
| 216 | } |
| 217 | if normalized := provider.NormalizeMessages(got); !reflect.DeepEqual(normalized, got) { |
| 218 | t.Fatalf("preserved guardian history is not provider-coherent:\n got=%+v\nnormalized=%+v", got, normalized) |
| 219 | } |
| 220 | } |
| 221 | |
| 222 | func TestTranscriptRenderKeepsFirstAndLastUserAnchors(t *testing.T) { |
| 223 | entries := []TranscriptEntry{{Kind: "user", Text: "first task"}} |
| 224 | for range maxRecentEntries + 5 { |
| 225 | entries = append(entries, TranscriptEntry{Kind: "assistant", Text: "assistant detail"}) |
| 226 | } |
| 227 | entries = append(entries, TranscriptEntry{Kind: "user", Text: "latest instruction"}) |
| 228 | rendered := FormatTranscript(entries) |
| 229 | if !strings.Contains(rendered, "first task") || !strings.Contains(rendered, "latest instruction") { |
| 230 | t.Fatalf("rendered transcript lost user anchors:\n%s", rendered) |
| 231 | } |
| 232 | if !strings.Contains(rendered, "Some conversation entries were omitted.") { |
| 233 | t.Fatalf("rendered transcript should mention omissions:\n%s", rendered) |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | func TestGuardianSaveLoadRestoresCursorForDeltaTranscript(t *testing.T) { |
| 238 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 239 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first ok"}`}, |
| 240 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"second ok"}`}, |
| 241 | }} |
| 242 | sink := &captureSink{} |
| 243 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 244 | parent := agent.NewSession("sys") |
| 245 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "first user request"}) |
| 246 | |
| 247 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 248 | t.Fatalf("first Review = allow %v err %v, want allow nil", allow, err) |
| 249 | } |
| 250 | path := filepath.Join(t.TempDir(), "session.guardian.jsonl") |
| 251 | if err := gs.Save(path); err != nil { |
| 252 | t.Fatalf("Save error: %v", err) |
| 253 | } |
| 254 | if data, err := os.ReadFile(cursorPathForGuardianPath(path)); err != nil || !strings.Contains(string(data), `"EntryCount":1`) { |
| 255 | t.Fatalf("cursor sidecar = %q err %v, want EntryCount 1", data, err) |
| 256 | } |
| 257 | |
| 258 | loaded := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 259 | if err := loaded.Load(path); err != nil { |
| 260 | t.Fatalf("Load error: %v", err) |
| 261 | } |
| 262 | if loaded.cursor.EntryCount != 1 { |
| 263 | t.Fatalf("loaded cursor = %+v, want EntryCount 1", loaded.cursor) |
| 264 | } |
| 265 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "second user request"}) |
| 266 | if allow, _, err := loaded.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"b.txt"}`), parent); err != nil || !allow { |
| 267 | t.Fatalf("second Review = allow %v err %v, want allow nil", allow, err) |
| 268 | } |
| 269 | |
| 270 | reqs := prov.requestsSnapshot() |
| 271 | if len(reqs) < 2 { |
| 272 | t.Fatalf("requests = %d, want >= 2", len(reqs)) |
| 273 | } |
| 274 | var delta string |
| 275 | for _, req := range reqs { |
| 276 | for _, m := range req.Messages { |
| 277 | if strings.Contains(m.Content, "TRANSCRIPT DELTA") && strings.Contains(m.Content, "second user request") { |
| 278 | delta = m.Content |
| 279 | break |
| 280 | } |
| 281 | } |
| 282 | if delta != "" { |
| 283 | break |
| 284 | } |
| 285 | } |
| 286 | if delta == "" { |
| 287 | t.Fatalf("second request did not include a delta transcript") |
| 288 | } |
| 289 | if !strings.Contains(delta, "second user request") { |
| 290 | t.Fatalf("delta transcript missing new parent entry:\n%s", delta) |
| 291 | } |
| 292 | if strings.Contains(delta, "first user request") { |
| 293 | t.Fatalf("delta transcript repeated old parent entry:\n%s", delta) |
| 294 | } |
| 295 | } |
| 296 | |
| 297 | func TestGuardianUsageDoesNotLeakAcrossReviews(t *testing.T) { |
| 298 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 299 | { |
| 300 | text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first ok"}`, |
| 301 | usage: &provider.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}, |
| 302 | }, |
| 303 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"second ok"}`}, |
| 304 | }} |
| 305 | sink := &captureSink{} |
| 306 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 307 | parent := agent.NewSession("sys") |
| 308 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do it"}) |
| 309 | |
| 310 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 311 | t.Fatalf("first Review = allow %v err %v, want allow nil", allow, err) |
| 312 | } |
| 313 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"b.txt"}`), parent); err != nil || !allow { |
| 314 | t.Fatalf("second Review = allow %v err %v, want allow nil", allow, err) |
| 315 | } |
| 316 | |
| 317 | events := sink.guardianEvents() |
| 318 | if len(events) != 2 { |
| 319 | t.Fatalf("guardian events = %d, want 2", len(events)) |
| 320 | } |
| 321 | if events[0].Guardian.Usage == nil || events[0].Guardian.Usage.TotalTokens != 12 { |
| 322 | t.Fatalf("first usage = %+v, want total 12", events[0].Guardian.Usage) |
| 323 | } |
| 324 | if events[1].Guardian.Usage != nil { |
| 325 | t.Fatalf("second usage leaked from first review: %+v", events[1].Guardian.Usage) |
| 326 | } |
| 327 | } |
| 328 | |
| 329 | func TestGuardianUsageAggregatesEveryModelCall(t *testing.T) { |
| 330 | base := &scriptedProvider{responses: []scriptedResponse{ |
| 331 | {reasoning: "The action is safe; a visible verdict is still required.", usage: &provider.Usage{PromptTokens: 3, CompletionTokens: 1, TotalTokens: 4, FinishReason: "stop"}}, |
| 332 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`, usage: &provider.Usage{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7, RequestCount: 2}}, |
| 333 | }} |
| 334 | prov := &reasoningScriptedProvider{scriptedProvider: base} |
| 335 | sink := &captureSink{} |
| 336 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 337 | parent := agent.NewSession("sys") |
| 338 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do it"}) |
| 339 | |
| 340 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 341 | t.Fatalf("Review = allow %v err %v, want allow nil", allow, err) |
| 342 | } |
| 343 | events := sink.guardianEvents() |
| 344 | if len(events) != 1 || events[0].Guardian.Usage == nil { |
| 345 | t.Fatalf("guardian events = %+v, want one usage-bearing event", events) |
| 346 | } |
| 347 | usage := events[0].Guardian.Usage |
| 348 | if usage.PromptTokens != 8 || usage.CompletionTokens != 3 || usage.TotalTokens != 11 || usage.RequestCount != 3 { |
| 349 | t.Fatalf("aggregated usage = %+v, want prompt=8 completion=3 total=11 requests=3", usage) |
| 350 | } |
| 351 | } |
| 352 | |
| 353 | func TestGuardianEmptyResponseFailsClosedUntilExplicitReview(t *testing.T) { |
| 354 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 355 | {usage: &provider.Usage{PromptTokens: 3, CompletionTokens: 1, TotalTokens: 4}}, |
| 356 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`, usage: &provider.Usage{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7}}, |
| 357 | }} |
| 358 | sink := &captureSink{} |
| 359 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 360 | parent := agent.NewSession("sys") |
| 361 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do it"}) |
| 362 | args := json.RawMessage(`{"file_path":"a.txt"}`) |
| 363 | |
| 364 | if allow, _, err := gs.ReviewVerdict(context.Background(), "write_file", args, parent); err == nil || allow { |
| 365 | t.Fatalf("first ReviewVerdict = allow %v err %v, want fail-closed error", allow, err) |
| 366 | } |
| 367 | if got := len(prov.requestsSnapshot()); got != 1 { |
| 368 | t.Fatalf("requests after empty response = %d, want 1", got) |
| 369 | } |
| 370 | events := sink.guardianEvents() |
| 371 | if len(events) != 1 || events[0].Guardian.Usage == nil { |
| 372 | t.Fatalf("guardian events = %+v, want one usage-bearing failure", events) |
| 373 | } |
| 374 | usage := events[0].Guardian.Usage |
| 375 | if usage.PromptTokens != 3 || usage.CompletionTokens != 1 || usage.TotalTokens != 4 || usage.RequestCount != 1 { |
| 376 | t.Fatalf("failed review usage = %+v, want prompt=3 completion=1 total=4 requests=1", usage) |
| 377 | } |
| 378 | |
| 379 | if allow, _, err := gs.ReviewVerdict(context.Background(), "write_file", args, parent); err != nil || !allow { |
| 380 | t.Fatalf("explicit second ReviewVerdict = allow %v err %v, want allow nil", allow, err) |
| 381 | } |
| 382 | if got := len(prov.requestsSnapshot()); got != 2 { |
| 383 | t.Fatalf("requests after explicit second review = %d, want 2", got) |
| 384 | } |
| 385 | events = sink.guardianEvents() |
| 386 | if len(events) != 2 || events[1].Guardian.Usage == nil { |
| 387 | t.Fatalf("guardian events = %+v, want two usage-bearing events", events) |
| 388 | } |
| 389 | usage = events[1].Guardian.Usage |
| 390 | if usage.PromptTokens != 5 || usage.CompletionTokens != 2 || usage.TotalTokens != 7 || usage.RequestCount != 1 { |
| 391 | t.Fatalf("second review usage = %+v, want prompt=5 completion=2 total=7 requests=1", usage) |
| 392 | } |
| 393 | } |
| 394 | |
| 395 | // TestGuardianReviewTurnsAlternateRoles pins the review request shape: the |
| 396 | // transcript evidence and action request ride in one combined user message per |
| 397 | // review, so the guardian session alternates user/assistant strictly and |
| 398 | // providers that reject consecutive same-role messages can run the guardian. |
| 399 | func TestGuardianReviewTurnsAlternateRoles(t *testing.T) { |
| 400 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 401 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"first ok"}`}, |
| 402 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"second ok"}`}, |
| 403 | }} |
| 404 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 405 | parent := agent.NewSession("sys") |
| 406 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do the thing"}) |
| 407 | |
| 408 | for i := range 2 { |
| 409 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 410 | t.Fatalf("review %d = allow %v err %v, want allow nil", i+1, allow, err) |
| 411 | } |
| 412 | } |
| 413 | |
| 414 | reqs := prov.requestsSnapshot() |
| 415 | if len(reqs) < 2 { |
| 416 | t.Fatalf("requests = %d, want >= 2", len(reqs)) |
| 417 | } |
| 418 | for r, req := range reqs { |
| 419 | for i := 1; i < len(req.Messages); i++ { |
| 420 | if req.Messages[i].Role == provider.RoleUser && req.Messages[i-1].Role == provider.RoleUser { |
| 421 | t.Fatalf("request %d carries consecutive user messages at index %d", r, i) |
| 422 | } |
| 423 | } |
| 424 | } |
| 425 | |
| 426 | // The combined message must still carry the evidence boundary and the action. |
| 427 | msgs := reqs[len(reqs)-1].Messages |
| 428 | var review string |
| 429 | for _, v := range slices.Backward(msgs) { |
| 430 | if v.Role == provider.RoleUser { |
| 431 | review = v.Content |
| 432 | break |
| 433 | } |
| 434 | } |
| 435 | for _, want := range []string{"untrusted evidence", "The agent has requested the following action", "write_file"} { |
| 436 | if !strings.Contains(review, want) { |
| 437 | t.Fatalf("combined review message missing %q:\n%s", want, review) |
| 438 | } |
| 439 | } |
| 440 | } |
| 441 | |
| 442 | // TestGuardianFailedReviewRollsBackSession pins the error-path rollback: |
| 443 | // agent.Run appends the combined review user message before the provider is |
| 444 | // reached, so a failed review must not leave it dangling — the next review |
| 445 | // would otherwise append another user message and strict-alternation providers |
| 446 | // would reject every request from then on. |
| 447 | func TestGuardianFailedReviewRollsBackSession(t *testing.T) { |
| 448 | prov := &scriptedProvider{responses: []scriptedResponse{ |
| 449 | {err: fmt.Errorf("provider unavailable")}, |
| 450 | {text: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`}, |
| 451 | }} |
| 452 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 453 | parent := agent.NewSession("sys") |
| 454 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "do the thing"}) |
| 455 | |
| 456 | allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent) |
| 457 | if err == nil && allow { |
| 458 | t.Fatal("first review should fail closed") |
| 459 | } |
| 460 | if n := gs.sess.Len(); n != 1 { |
| 461 | t.Fatalf("guardian session messages = %d after failed review, want rollback to system only", n) |
| 462 | } |
| 463 | |
| 464 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 465 | t.Fatalf("second review = allow %v err %v, want allow nil", allow, err) |
| 466 | } |
| 467 | reqs := prov.requestsSnapshot() |
| 468 | last := reqs[len(reqs)-1] |
| 469 | for i := 1; i < len(last.Messages); i++ { |
| 470 | if last.Messages[i].Role == provider.RoleUser && last.Messages[i-1].Role == provider.RoleUser { |
| 471 | t.Fatalf("request after failed review carries consecutive user messages at index %d", i) |
| 472 | } |
| 473 | } |
| 474 | } |
| 475 | |
| 476 | // TestGuardianLoadResetsLegacyConsecutiveUserSessions pins the load-time |
| 477 | // normalization: sessions saved by the old multi-message review shape carry |
| 478 | // consecutive user messages that would poison strict-alternation providers, so |
| 479 | // Load starts fresh instead of adopting them. |
| 480 | func TestGuardianLoadResetsLegacyConsecutiveUserSessions(t *testing.T) { |
| 481 | legacy := agent.NewSession(PolicyPrompt()) |
| 482 | legacy.Add(provider.Message{Role: provider.RoleUser, Content: "transcript evidence"}) |
| 483 | legacy.Add(provider.Message{Role: provider.RoleUser, Content: "action request"}) |
| 484 | legacy.Add(provider.Message{Role: provider.RoleAssistant, Content: `{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"ok"}`}) |
| 485 | path := filepath.Join(t.TempDir(), "session.guardian.jsonl") |
| 486 | if err := legacy.Save(path); err != nil { |
| 487 | t.Fatalf("Save legacy session: %v", err) |
| 488 | } |
| 489 | |
| 490 | gs := NewSession(&scriptedProvider{}, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, &captureSink{}) |
| 491 | if err := gs.Load(path); err != nil { |
| 492 | t.Fatalf("Load: %v", err) |
| 493 | } |
| 494 | if n := gs.sess.Len(); n != 1 { |
| 495 | t.Fatalf("loaded legacy session messages = %d, want reset to system only", n) |
| 496 | } |
| 497 | if gs.cursor.EntryCount != 0 { |
| 498 | t.Fatalf("cursor = %+v, want zeroed after reset", gs.cursor) |
| 499 | } |
| 500 | } |
| 501 | |
| 502 | // TestGuardianSessionAlternatesAfterCompaction reproduces the compaction seam: |
| 503 | // every compactEvery-th review runs CompactNow, and generic compaction inserts |
| 504 | // its digest as a RoleUser message that can land directly before a review's |
| 505 | // user turn — consecutive user roles again. The post-review normalization must |
| 506 | // keep the session strictly alternating across that fold. |
| 507 | func TestGuardianSessionAlternatesAfterCompaction(t *testing.T) { |
| 508 | prov := &scriptedProvider{defaultUsage: &provider.Usage{TotalTokens: 1}} // default allow verdict, also serves the summarizer |
| 509 | sink := &captureSink{} |
| 510 | gs := NewSession(prov, tool.NewRegistry(), PolicyPrompt(), "guardian-test", 0, nil, sink) |
| 511 | parent := agent.NewSession("sys") |
| 512 | |
| 513 | filler := strings.Repeat("parent transcript filler. ", 160) |
| 514 | for i := range compactEvery { |
| 515 | parent.Add(provider.Message{Role: provider.RoleUser, Content: fmt.Sprintf("turn %d: %s", i, filler)}) |
| 516 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 517 | t.Fatalf("review %d = allow %v err %v, want allow nil", i+1, allow, err) |
| 518 | } |
| 519 | } |
| 520 | |
| 521 | msgs := gs.sess.Snapshot() |
| 522 | events := sink.guardianEvents() |
| 523 | if len(events) != compactEvery { |
| 524 | t.Fatalf("guardian events = %d, want %d", len(events), compactEvery) |
| 525 | } |
| 526 | usage := events[len(events)-1].Guardian.Usage |
| 527 | if usage == nil || usage.TotalTokens != 2 || usage.RequestCount != 2 { |
| 528 | t.Fatalf("compacting review usage = %+v, want total=2 requests=2", usage) |
| 529 | } |
| 530 | for i := 1; i < len(msgs); i++ { |
| 531 | if msgs[i].Role == msgs[i-1].Role { |
| 532 | t.Fatalf("guardian session has consecutive %s messages at indexes %d/%d of %d", msgs[i].Role, i-1, i, len(msgs)) |
| 533 | } |
| 534 | } |
| 535 | |
| 536 | // Projection compaction leaves the canonical transcript untouched. The next |
| 537 | // review must nevertheless send the compacted view, with the digest and tail |
| 538 | // coalesced so strict-alternation providers do not receive adjacent users. |
| 539 | parent.Add(provider.Message{Role: provider.RoleUser, Content: "post-compaction review"}) |
| 540 | if allow, _, err := gs.Review(context.Background(), "write_file", json.RawMessage(`{"file_path":"a.txt"}`), parent); err != nil || !allow { |
| 541 | t.Fatalf("post-compaction review = allow %v err %v, want allow nil", allow, err) |
| 542 | } |
| 543 | reqs := prov.requestsSnapshot() |
| 544 | if len(reqs) <= compactEvery { |
| 545 | t.Fatalf("provider requests = %d, want a request after projection compaction", len(reqs)) |
| 546 | } |
| 547 | last := reqs[len(reqs)-1] |
| 548 | hasDigest := false |
| 549 | for i, m := range last.Messages { |
| 550 | if agent.IsCompactionSummary(m) { |
| 551 | hasDigest = true |
| 552 | } |
| 553 | if i > 0 && m.Role == provider.RoleUser && last.Messages[i-1].Role == provider.RoleUser { |
| 554 | t.Fatalf("post-compaction request carries consecutive user messages at index %d", i) |
| 555 | } |
| 556 | } |
| 557 | if !hasDigest { |
| 558 | t.Fatal("post-compaction request did not use the digest projection") |
| 559 | } |
| 560 | } |
| 561 |