1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "os"
7 "path/filepath"
8 "strings"
9 "testing"
10 "time"
11 "unicode/utf8"
12)
13
14func TestBoundInterceptContextPreservesUTF8AndLimit(t *testing.T) {
15 got := boundInterceptContext(strings.Repeat("界", maxInterceptContext))
16 if len(got) > maxInterceptContext {
17 t.Fatalf("bounded context bytes = %d, want <= %d", len(got), maxInterceptContext)
18 }
19 if !utf8.ValidString(got) {
20 t.Fatal("bounded context is not valid UTF-8")
21 }
22 if !strings.HasSuffix(got, interceptContextTruncatedMarker) {
23 t.Fatalf("bounded context missing truncation marker")
24 }
25}
26
27// TestInterceptAllThreeEvents exercises tool_call / turn_start /
28// assistant_message interception end-to-end via a bash extension that
29// blocks `rm -rf`, rewrites any other bash command to `echo GUARDED`,
30// suppresses nothing at turn_start, and rewrites SECRET → [redacted]
31// in assistant messages.
32//
33// Validates: block, modified_args, pass-through, and replace_text.
34func TestInterceptAllThreeEvents(t *testing.T) {
35 if _, err := os.Stat("/bin/bash"); err != nil {
36 t.Skip("no /bin/bash")
37 }
38
39 extDir := t.TempDir()
40 script := `#!/bin/bash
41emit() { printf '%s\n' "$1"; }
42emit '{"type":"hello","protocol_version":2,"name":"itest","version":"0.1.0","capabilities":["events"]}'
43emit '{"type":"subscribe","events":[],"intercept":["tool_call","turn_start","assistant_message"]}'
44emit '{"type":"ready"}'
45while IFS= read -r line; do
46 t=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("type",""))')
47 id=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("id",""))')
48 ev=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("event",""))')
49 if [[ "$t" == "shutdown" ]]; then emit '{"type":"shutdown_ack"}'; exit 0; fi
50 if [[ "$t" != "event_intercept" ]]; then continue; fi
51 case "$ev" in
52 tool_call)
53 name=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("tool_name",""))')
54 args=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.dumps(json.load(sys.stdin).get("tool_args",{})))')
55 cmd=$(printf '%s' "$args" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("command",""))')
56 if [[ "$name" == "bash" && "$cmd" == *"rm -rf"* ]]; then
57 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"block\":true,\"reason\":\"refused: rm -rf\"}"
58 elif [[ "$name" == "bash" && -n "$cmd" ]]; then
59 new=$(python3 -c "import json,sys;print(json.dumps({'command':'echo GUARDED: '+sys.argv[1]}))" "$cmd")
60 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"modified_args\":$new}"
61 else
62 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}"
63 fi ;;
64 turn_start)
65 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}" ;;
66 assistant_message)
67 text=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("text",""))')
68 if [[ "$text" == *"SECRET"* ]]; then
69 new=$(python3 -c "import sys;print(sys.argv[1].replace('SECRET','[redacted]'))" "$text")
70 rt=$(python3 -c 'import json,sys;print(json.dumps(sys.argv[1]))' "$new")
71 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"replace_text\":$rt}"
72 else
73 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}"
74 fi ;;
75 esac
76done
77`
78 if err := os.WriteFile(filepath.Join(extDir, "ext.sh"), []byte(script), 0o755); err != nil {
79 t.Fatal(err)
80 }
81 manifest := `{"name":"itest","version":"0.1.0","exec":"./ext.sh","enabled":true}`
82 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil {
83 t.Fatal(err)
84 }
85
86 // ZutHome is unused here; we load the extension explicitly.
87 m := New(t.TempDir(), "", "0.0.0-test", "anthropic", "claude-test", nil)
88 t.Cleanup(func() { m.Stop(2 * time.Second) })
89
90 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
91 defer cancel()
92 if errs := m.LoadExplicit(ctx, []string{extDir}); len(errs) > 0 {
93 t.Fatalf("LoadExplicit: %v", errs)
94 }
95 m.WaitForReady(3 * time.Second)
96
97 // tool_call: rm -rf is blocked
98 res := m.InterceptToolCall(ctx, "T1", "bash", json.RawMessage(`{"command":"rm -rf /tmp/foo"}`))
99 if !res.Block || !strings.Contains(res.Reason, "refused") {
100 t.Errorf("rm -rf: want block+reason, got %+v", res)
101 }
102
103 // tool_call: non-dangerous bash gets args rewritten
104 res = m.InterceptToolCall(ctx, "T2", "bash", json.RawMessage(`{"command":"ls -la"}`))
105 if res.Block {
106 t.Errorf("ls -la: want allow, got block %q", res.Reason)
107 }
108 if len(res.ModifiedArgs) == 0 {
109 t.Errorf("ls -la: want modified_args, got nothing")
110 } else {
111 var obj map[string]string
112 if err := json.Unmarshal(res.ModifiedArgs, &obj); err != nil {
113 t.Errorf("modified_args: %v", err)
114 } else if !strings.HasPrefix(obj["command"], "echo GUARDED") {
115 t.Errorf("modified_args command=%q, want echo GUARDED prefix", obj["command"])
116 }
117 }
118
119 // tool_call: non-bash tool untouched
120 res = m.InterceptToolCall(ctx, "T3", "read", json.RawMessage(`{"path":"/etc/hosts"}`))
121 if res.Block {
122 t.Errorf("read: want allow, got block %q", res.Reason)
123 }
124 if len(res.ModifiedArgs) != 0 {
125 t.Errorf("read: want no modified_args, got %s", string(res.ModifiedArgs))
126 }
127
128 // turn_start: always allowed
129 if r := m.InterceptTurnStart(ctx, 1); r.Block {
130 t.Errorf("turn_start: want allow, got block %q", r.Reason)
131 }
132
133 // assistant_message: SECRET gets redacted
134 r := m.InterceptAssistantMessage(ctx, "your password is SECRET123")
135 if r.Block {
136 t.Errorf("msg: want allow, got block %q", r.Reason)
137 }
138 if !strings.Contains(r.ReplaceText, "[redacted]") {
139 t.Errorf("msg: want [redacted] in replacement, got %q", r.ReplaceText)
140 }
141
142 // assistant_message: clean text unchanged
143 r = m.InterceptAssistantMessage(ctx, "hello world")
144 if r.Block {
145 t.Errorf("clean: want allow")
146 }
147 if r.ReplaceText != "" {
148 t.Errorf("clean: want no replace, got %q", r.ReplaceText)
149 }
150}
151
152// TestReloadRespawnsExtensions verifies Reload tears down and brings
153// back the same extension, and the onReload callback fires.
154func TestReloadRespawnsExtensions(t *testing.T) {
155 if _, err := os.Stat("/bin/bash"); err != nil {
156 t.Skip("no /bin/bash")
157 }
158
159 extDir := t.TempDir()
160 script := `#!/bin/bash
161emit() { printf '%s\n' "$1"; }
162emit '{"type":"hello","protocol_version":2,"name":"rtest","version":"0.1.0","capabilities":["commands"]}'
163emit '{"type":"register_command","name":"rtest","description":"test"}'
164emit '{"type":"ready"}'
165while IFS= read -r line; do
166 t=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("type",""))')
167 if [[ "$t" == "shutdown" ]]; then emit '{"type":"shutdown_ack"}'; exit 0; fi
168done
169`
170 if err := os.WriteFile(filepath.Join(extDir, "ext.sh"), []byte(script), 0o755); err != nil {
171 t.Fatal(err)
172 }
173 manifest := `{"name":"rtest","version":"0.1.0","exec":"./ext.sh","enabled":true}`
174 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil {
175 t.Fatal(err)
176 }
177
178 hooks := &stubHooks{}
179 m := New(t.TempDir(), "", "0.0.0-test", "anthropic", "claude-test", hooks)
180 t.Cleanup(func() { m.Stop(2 * time.Second) })
181
182 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
183 defer cancel()
184 if errs := m.LoadExplicit(ctx, []string{extDir}); len(errs) > 0 {
185 t.Fatalf("LoadExplicit: %v", errs)
186 }
187 m.WaitForReady(3 * time.Second)
188 before := len(m.All())
189 if before == 0 {
190 t.Fatal("extension didn't load")
191 }
192
193 reloadFired := false
194 m.SetOnReload(func() { reloadFired = true })
195
196 stats := m.Reload(ctx, 2*time.Second)
197 if stats.Stopped != before {
198 t.Errorf("Stopped: want %d, got %d", before, stats.Stopped)
199 }
200 if stats.Loaded != before {
201 t.Errorf("Loaded: want %d, got %d", before, stats.Loaded)
202 }
203 if !reloadFired {
204 t.Error("onReload callback didn't fire")
205 }
206 hooks.mu.Lock()
207 clears := append([]string(nil), hooks.chromeClears...)
208 hooks.mu.Unlock()
209 if len(clears) == 0 || clears[0] != "rtest" {
210 t.Fatalf("extension chrome cleanup callbacks = %v", clears)
211 }
212
213 // Registered command should still be there after reload.
214 if !m.HasCommand("rtest") {
215 t.Error("rtest command not re-registered after reload")
216 }
217}