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