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