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