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 timeout=$(printf '%s' "$args" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("timeout",0))')
60 new=$(python3 -c "import json,sys;print(json.dumps({'command':'echo GUARDED: '+sys.argv[1], 'timeout': int(sys.argv[2])}))" "$cmd" "$timeout")
61 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"modified_args\":$new}"
62 else
63 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}"
64 fi ;;
65 turn_start)
66 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}" ;;
67 assistant_message)
68 text=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("text",""))')
69 if [[ "$text" == *"SECRET"* ]]; then
70 new=$(python3 -c "import sys;print(sys.argv[1].replace('SECRET','[redacted]'))" "$text")
71 rt=$(python3 -c 'import json,sys;print(json.dumps(sys.argv[1]))' "$new")
72 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"replace_text\":$rt}"
73 else
74 emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}"
75 fi ;;
76 esac
77done
78`
79 if err := os.WriteFile(filepath.Join(extDir, "ext.sh"), []byte(script), 0o755); err != nil {
80 t.Fatal(err)
81 }
82 manifest := `{"name":"itest","version":"0.1.0","exec":"./ext.sh","enabled":true}`
83 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil {
84 t.Fatal(err)
85 }
86
87 // ZutHome is unused here; we load the extension explicitly.
88 m := New(t.TempDir(), "", "0.0.0-test", "anthropic", "claude-test", nil)
89 t.Cleanup(func() { m.Stop(2 * time.Second) })
90
91 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
92 defer cancel()
93 if errs := m.LoadExplicit(ctx, []string{extDir}); len(errs) > 0 {
94 t.Fatalf("LoadExplicit: %v", errs)
95 }
96 m.WaitForReady(3 * time.Second)
97
98 // tool_call: rm -rf is blocked
99 res := m.InterceptToolCall(ctx, "T1", "bash", json.RawMessage(`{"command":"rm -rf /tmp/foo","timeout":30}`))
100 if !res.Block || !strings.Contains(res.Reason, "refused") {
101 t.Errorf("rm -rf: want block+reason, got %+v", res)
102 }
103
104 // tool_call: non-dangerous bash gets args rewritten
105 res = m.InterceptToolCall(ctx, "T2", "bash", json.RawMessage(`{"command":"ls -la","timeout":30}`))
106 if res.Block {
107 t.Errorf("ls -la: want allow, got block %q", res.Reason)
108 }
109 if len(res.ModifiedArgs) == 0 {
110 t.Errorf("ls -la: want modified_args, got nothing")
111 } else {
112 var obj struct {
113 Command string `json:"command"`
114 Timeout int `json:"timeout"`
115 }
116 if err := json.Unmarshal(res.ModifiedArgs, &obj); err != nil {
117 t.Errorf("modified_args: %v", err)
118 } else {
119 if !strings.HasPrefix(obj.Command, "echo GUARDED") {
120 t.Errorf("modified_args command=%q, want echo GUARDED prefix", obj.Command)
121 }
122 if obj.Timeout != 30 {
123 t.Errorf("modified_args timeout=%d, want 30", obj.Timeout)
124 }
125 }
126 }
127
128 // tool_call: non-bash tool untouched
129 res = m.InterceptToolCall(ctx, "T3", "read", json.RawMessage(`{"path":"/etc/hosts"}`))
130 if res.Block {
131 t.Errorf("read: want allow, got block %q", res.Reason)
132 }
133 if len(res.ModifiedArgs) != 0 {
134 t.Errorf("read: want no modified_args, got %s", string(res.ModifiedArgs))
135 }
136
137 // turn_start: always allowed
138 if r := m.InterceptTurnStart(ctx, 1); r.Block {
139 t.Errorf("turn_start: want allow, got block %q", r.Reason)
140 }
141
142 // assistant_message: SECRET gets redacted
143 r := m.InterceptAssistantMessage(ctx, "your password is SECRET123")
144 if r.Block {
145 t.Errorf("msg: want allow, got block %q", r.Reason)
146 }
147 if !strings.Contains(r.ReplaceText, "[redacted]") {
148 t.Errorf("msg: want [redacted] in replacement, got %q", r.ReplaceText)
149 }
150
151 // assistant_message: clean text unchanged
152 r = m.InterceptAssistantMessage(ctx, "hello world")
153 if r.Block {
154 t.Errorf("clean: want allow")
155 }
156 if r.ReplaceText != "" {
157 t.Errorf("clean: want no replace, got %q", r.ReplaceText)
158 }
159}
160
161// TestReloadRespawnsExtensions verifies Reload tears down and brings
162// back the same extension, and the onReload callback fires.
163func TestReloadRespawnsExtensions(t *testing.T) {
164 if _, err := os.Stat("/bin/bash"); err != nil {
165 t.Skip("no /bin/bash")
166 }
167
168 extDir := t.TempDir()
169 script := `#!/bin/bash
170emit() { printf '%s\n' "$1"; }
171emit '{"type":"hello","protocol_version":2,"name":"rtest","version":"0.1.0","capabilities":["commands"]}'
172emit '{"type":"register_command","name":"rtest","description":"test"}'
173emit '{"type":"ready"}'
174while IFS= read -r line; do
175 t=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("type",""))')
176 if [[ "$t" == "shutdown" ]]; then emit '{"type":"shutdown_ack"}'; exit 0; fi
177done
178`
179 if err := os.WriteFile(filepath.Join(extDir, "ext.sh"), []byte(script), 0o755); err != nil {
180 t.Fatal(err)
181 }
182 manifest := `{"name":"rtest","version":"0.1.0","exec":"./ext.sh","enabled":true}`
183 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil {
184 t.Fatal(err)
185 }
186
187 hooks := &stubHooks{}
188 m := New(t.TempDir(), "", "0.0.0-test", "anthropic", "claude-test", hooks)
189 t.Cleanup(func() { m.Stop(2 * time.Second) })
190
191 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
192 defer cancel()
193 if errs := m.LoadExplicit(ctx, []string{extDir}); len(errs) > 0 {
194 t.Fatalf("LoadExplicit: %v", errs)
195 }
196 m.WaitForReady(3 * time.Second)
197 before := len(m.All())
198 if before == 0 {
199 t.Fatal("extension didn't load")
200 }
201
202 reloadFired := false
203 m.SetOnReload(func() { reloadFired = true })
204
205 stats := m.Reload(ctx, 2*time.Second)
206 if stats.Stopped != before {
207 t.Errorf("Stopped: want %d, got %d", before, stats.Stopped)
208 }
209 if stats.Loaded != before {
210 t.Errorf("Loaded: want %d, got %d", before, stats.Loaded)
211 }
212 if !reloadFired {
213 t.Error("onReload callback didn't fire")
214 }
215 hooks.mu.Lock()
216 clears := append([]string(nil), hooks.chromeClears...)
217 hooks.mu.Unlock()
218 if len(clears) == 0 || clears[0] != "rtest" {
219 t.Fatalf("extension chrome cleanup callbacks = %v", clears)
220 }
221
222 // Registered command should still be there after reload.
223 if !m.HasCommand("rtest") {
224 t.Error("rtest command not re-registered after reload")
225 }
226}