1package core
2
3import (
4 "encoding/json"
5 "sync"
6 "testing"
7)
8
9type recordingConfirmer struct {
10 mu sync.Mutex
11 calls []string
12 replies []ConfirmDecision
13 idx int
14}
15
16func (r *recordingConfirmer) Confirm(toolName, preview string) ConfirmDecision {
17 r.mu.Lock()
18 defer r.mu.Unlock()
19 r.calls = append(r.calls, toolName+"/"+preview)
20 if r.idx >= len(r.replies) {
21 return ConfirmDecision{Allow: false, Reason: "no reply queued"}
22 }
23 d := r.replies[r.idx]
24 r.idx++
25 return d
26}
27
28func TestConfirmGateNilAllowsEverything(t *testing.T) {
29 var g *ConfirmGate
30 allow, reason, args := g.Check("bash", "rm -rf /")
31 if !allow || reason != "" || args != nil {
32 t.Fatalf("nil gate should allow, got allow=%v reason=%q args=%s", allow, reason, args)
33 }
34}
35
36func TestConfirmGateNilInnerRefuses(t *testing.T) {
37 g := NewConfirmGate(nil)
38 allow, reason, _ := g.Check("bash", "ls")
39 if allow {
40 t.Fatal("gate with nil inner should refuse")
41 }
42 if reason == "" {
43 t.Fatal("refusal must carry a reason for the model to learn from")
44 }
45}
46
47func TestConfirmGateAllowOnce(t *testing.T) {
48 rc := &recordingConfirmer{replies: []ConfirmDecision{
49 {Allow: true},
50 {Allow: true},
51 }}
52 g := NewConfirmGate(rc)
53 if allow, _, _ := g.Check("bash", "ls"); !allow {
54 t.Fatal("call 1 should allow")
55 }
56 if allow, _, _ := g.Check("bash", "ls"); !allow {
57 t.Fatal("call 2 should allow")
58 }
59 // Two calls, two confirmer invocations (no remember).
60 if len(rc.calls) != 2 {
61 t.Errorf("want 2 confirmer calls, got %d", len(rc.calls))
62 }
63}
64
65func TestConfirmGateRememberTool(t *testing.T) {
66 rc := &recordingConfirmer{replies: []ConfirmDecision{
67 {Allow: true, RememberTool: true},
68 }}
69 g := NewConfirmGate(rc)
70
71 // First call prompts and remembers.
72 if allow, _, _ := g.Check("bash", "ls"); !allow {
73 t.Fatal("call 1 should allow")
74 }
75 // Second call short-circuits; confirmer must NOT be invoked.
76 if allow, _, _ := g.Check("bash", "pwd"); !allow {
77 t.Fatal("call 2 should allow from memory")
78 }
79 // Different tool still prompts.
80 rc.replies = append(rc.replies, ConfirmDecision{Allow: false, Reason: "no"})
81 if allow, reason, _ := g.Check("read", "foo.txt"); allow || reason != "no" {
82 t.Errorf("different tool should re-prompt; got allow=%v reason=%q", allow, reason)
83 }
84 if len(rc.calls) != 2 {
85 t.Errorf("want 2 confirmer calls (bash+read), got %d: %v", len(rc.calls), rc.calls)
86 }
87}
88
89func TestConfirmGateRememberAll(t *testing.T) {
90 rc := &recordingConfirmer{replies: []ConfirmDecision{
91 {Allow: true, RememberAll: true},
92 }}
93 g := NewConfirmGate(rc)
94 if allow, _, _ := g.Check("bash", "ls"); !allow {
95 t.Fatal("call 1 should allow")
96 }
97 // From now on everything short-circuits.
98 if allow, _, _ := g.Check("read", "foo.txt"); !allow {
99 t.Fatal("call 2 should allow")
100 }
101 if allow, _, _ := g.Check("write", "bar.txt"); !allow {
102 t.Fatal("call 3 should allow")
103 }
104 if len(rc.calls) != 1 {
105 t.Errorf("want 1 confirmer call (remember-all short-circuits the rest), got %d", len(rc.calls))
106 }
107}
108
109func TestConfirmGateRefuseSurfacesReason(t *testing.T) {
110 rc := &recordingConfirmer{replies: []ConfirmDecision{
111 {Allow: false, Reason: "do not run sudo"},
112 }}
113 g := NewConfirmGate(rc)
114 allow, reason, _ := g.Check("bash", "sudo rm -rf /")
115 if allow || reason != "do not run sudo" {
116 t.Errorf("want block + reason, got allow=%v reason=%q", allow, reason)
117 }
118}
119
120func TestConfirmGateRefuseEmptyReasonGetsDefault(t *testing.T) {
121 rc := &recordingConfirmer{replies: []ConfirmDecision{
122 {Allow: false},
123 }}
124 g := NewConfirmGate(rc)
125 _, reason, _ := g.Check("bash", "x")
126 if reason == "" {
127 t.Fatal("empty reason must be replaced with a sensible default")
128 }
129}
130
131func TestConfirmGateAllowAllRuntime(t *testing.T) {
132 rc := &recordingConfirmer{replies: []ConfirmDecision{
133 {Allow: false, Reason: "no"},
134 }}
135 g := NewConfirmGate(rc)
136 // Refuse first call
137 if allow, _, _ := g.Check("bash", "ls"); allow {
138 t.Fatal("call 1 should refuse")
139 }
140 // User types /yolo
141 g.AllowAll()
142 // All subsequent calls allowed without prompting.
143 if allow, _, _ := g.Check("bash", "rm -rf tmp"); !allow {
144 t.Fatal("after AllowAll, should allow")
145 }
146 if allow, _, _ := g.Check("read", "x"); !allow {
147 t.Fatal("after AllowAll, should allow any tool")
148 }
149 if len(rc.calls) != 1 {
150 t.Errorf("want 1 confirmer call (before AllowAll), got %d", len(rc.calls))
151 }
152}
153
154func TestBuildPreview(t *testing.T) {
155 cases := []struct {
156 name string
157 args string
158 want string
159 }{
160 {"bash command", `{"command":"ls -la"}`, "ls -la"},
161 {"path", `{"path":"/tmp/x.txt"}`, "/tmp/x.txt"},
162 {"file_path", `{"file_path":"a.go"}`, "a.go"},
163 {"url", `{"url":"https://example.com"}`, "https://example.com"},
164 {"truncation", `{"command":"` + string(make([]byte, 200)) + `"}`, ""},
165 {"unparseable", `{not json`, `{not json`},
166 {"empty", ``, ``},
167 }
168 for _, c := range cases {
169 t.Run(c.name, func(t *testing.T) {
170 got := BuildPreview(json.RawMessage(c.args), 50)
171 if c.name == "truncation" && !hasEllipsis(got) {
172 t.Errorf("%s: expected ellipsis truncation, got %q", c.name, got)
173 return
174 }
175 if c.want != "" && got != c.want {
176 t.Errorf("%s: want %q, got %q", c.name, c.want, got)
177 }
178 })
179 }
180}
181
182func hasEllipsis(s string) bool {
183 return len(s) > 0 && s[len(s)-len("\u2026"):] == "\u2026"
184}