1package core
2
3import (
4 "encoding/json"
5 "os"
6 "strings"
7 "testing"
8
9 "github.com/bnema/zut/packages/provider"
10)
11
12func TestSessionPersistsSanitizedRetryLifecycleOnUsage(t *testing.T) {
13 const privateDetail = "private provider response body"
14
15 session, err := NewSession(t.TempDir(), "/workspace", "provider", "model", "test")
16 if err != nil {
17 t.Fatal(err)
18 }
19 if err := session.AppendMessage(textMessage(provider.RoleUser, "hello")); err != nil {
20 t.Fatal(err)
21 }
22 session.RecordRetryLifecycle(RetryLifecycleRecord{
23 Event: RetryLifecycleRequestFailed, Scope: RetryScopeAgent,
24 Attempt: 1, MaxAttempts: 4, Reason: RetryReasonOverload,
25 })
26 session.RecordRetryLifecycle(RetryLifecycleRecord{
27 Event: RetryLifecycleRetryScheduled, Scope: RetryScopeAgent,
28 Attempt: 2, MaxAttempts: 4, Reason: RetryReasonOverload, DelayMS: 2000,
29 })
30 // RetryReason has a string representation for JSON, so the session layer
31 // clamps unknown values rather than trusting callers with persisted text.
32 session.RecordRetryLifecycle(RetryLifecycleRecord{
33 Event: RetryLifecycleRequestFailed, Scope: RetryScopeProvider,
34 Attempt: 1, MaxAttempts: 1, Reason: RetryReason(privateDetail), Terminal: true,
35 })
36 if err := session.AppendUsage(provider.Usage{InputTokens: 1}, provider.Usage{InputTokens: 1}); err != nil {
37 t.Fatal(err)
38 }
39 path := session.Path
40 if err := session.Close(); err != nil {
41 t.Fatal(err)
42 }
43
44 raw, err := os.ReadFile(path)
45 if err != nil {
46 t.Fatal(err)
47 }
48 if strings.Contains(string(raw), privateDetail) {
49 t.Fatalf("session leaked untrusted lifecycle text: %s", raw)
50 }
51
52 var lifecycle []RetryLifecycleRecord
53 for _, line := range strings.Split(strings.TrimSpace(string(raw)), "\n") {
54 var row struct {
55 Type string `json:"type"`
56 RetryLifecycle []RetryLifecycleRecord `json:"retry_lifecycle"`
57 }
58 if err := json.Unmarshal([]byte(line), &row); err != nil {
59 t.Fatal(err)
60 }
61 if row.Type == "usage" {
62 lifecycle = row.RetryLifecycle
63 }
64 }
65 if len(lifecycle) != 3 {
66 t.Fatalf("retry lifecycle = %#v, want 3 records", lifecycle)
67 }
68 if lifecycle[0].Reason != RetryReasonOverload || lifecycle[1].DelayMS != 2000 || lifecycle[2].Reason != RetryReasonUnknown || !lifecycle[2].Terminal {
69 t.Fatalf("retry lifecycle = %#v", lifecycle)
70 }
71
72 snapshot, err := ReadSessionSnapshot(path)
73 if err != nil {
74 t.Fatalf("ReadSessionSnapshot with retry metadata: %v", err)
75 }
76 if len(snapshot.Messages) != 1 || extractText(snapshot.Messages[0]) != "hello" {
77 t.Fatalf("snapshot messages changed by retry metadata: %#v", snapshot.Messages)
78 }
79}
80
81func TestSessionFlushRetryLifecycleWritesOnlyWhenPending(t *testing.T) {
82 session, err := NewSession(t.TempDir(), "/workspace", "provider", "model", "test")
83 if err != nil {
84 t.Fatal(err)
85 }
86 if err := session.AppendMessage(textMessage(provider.RoleUser, "hello")); err != nil {
87 t.Fatal(err)
88 }
89 before, err := os.Stat(session.Path)
90 if err != nil {
91 t.Fatal(err)
92 }
93 if err := session.FlushRetryLifecycle(provider.Usage{}, provider.Usage{}); err != nil {
94 t.Fatal(err)
95 }
96 after, err := os.Stat(session.Path)
97 if err != nil {
98 t.Fatal(err)
99 }
100 if after.Size() != before.Size() {
101 t.Fatalf("empty lifecycle flush changed session size: %d -> %d", before.Size(), after.Size())
102 }
103
104 session.RecordRetryLifecycle(RetryLifecycleRecord{
105 Event: RetryLifecycleRequestFailed, Scope: RetryScopeAgent,
106 Attempt: 1, MaxAttempts: 1, Reason: RetryReasonServer, Terminal: true,
107 })
108 if err := session.FlushRetryLifecycle(provider.Usage{}, provider.Usage{}); err != nil {
109 t.Fatal(err)
110 }
111 if err := session.Close(); err != nil {
112 t.Fatal(err)
113 }
114 raw, err := os.ReadFile(session.Path)
115 if err != nil {
116 t.Fatal(err)
117 }
118 if !strings.Contains(string(raw), `"retry_lifecycle"`) || !strings.Contains(string(raw), `"terminal":true`) {
119 t.Fatalf("terminal lifecycle was not flushed: %s", raw)
120 }
121}