package extensions import ( "context" "encoding/json" "os" "path/filepath" "strings" "testing" "time" "unicode/utf8" ) func TestBoundInterceptContextPreservesUTF8AndLimit(t *testing.T) { got := boundInterceptContext(strings.Repeat("界", maxInterceptContext)) if len(got) > maxInterceptContext { t.Fatalf("bounded context bytes = %d, want <= %d", len(got), maxInterceptContext) } if !utf8.ValidString(got) { t.Fatal("bounded context is not valid UTF-8") } if !strings.HasSuffix(got, interceptContextTruncatedMarker) { t.Fatalf("bounded context missing truncation marker") } } // TestInterceptAllThreeEvents exercises tool_call / turn_start / // assistant_message interception end-to-end via a bash extension that // blocks `rm -rf`, rewrites any other bash command to `echo GUARDED`, // suppresses nothing at turn_start, and rewrites SECRET → [redacted] // in assistant messages. // // Validates: block, modified_args, pass-through, and replace_text. func TestInterceptAllThreeEvents(t *testing.T) { if _, err := os.Stat("/bin/bash"); err != nil { t.Skip("no /bin/bash") } extDir := t.TempDir() script := `#!/bin/bash emit() { printf '%s\n' "$1"; } emit '{"type":"hello","protocol_version":2,"name":"itest","version":"0.1.0","capabilities":["events"]}' emit '{"type":"subscribe","events":[],"intercept":["tool_call","turn_start","assistant_message"]}' emit '{"type":"ready"}' while IFS= read -r line; do t=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("type",""))') id=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("id",""))') ev=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("event",""))') if [[ "$t" == "shutdown" ]]; then emit '{"type":"shutdown_ack"}'; exit 0; fi if [[ "$t" != "event_intercept" ]]; then continue; fi case "$ev" in tool_call) name=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("tool_name",""))') args=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.dumps(json.load(sys.stdin).get("tool_args",{})))') cmd=$(printf '%s' "$args" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("command",""))') if [[ "$name" == "bash" && "$cmd" == *"rm -rf"* ]]; then emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"block\":true,\"reason\":\"refused: rm -rf\"}" elif [[ "$name" == "bash" && -n "$cmd" ]]; then timeout=$(printf '%s' "$args" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("timeout",0))') new=$(python3 -c "import json,sys;print(json.dumps({'command':'echo GUARDED: '+sys.argv[1], 'timeout': int(sys.argv[2])}))" "$cmd" "$timeout") emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"modified_args\":$new}" else emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}" fi ;; turn_start) emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}" ;; assistant_message) text=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("text",""))') if [[ "$text" == *"SECRET"* ]]; then new=$(python3 -c "import sys;print(sys.argv[1].replace('SECRET','[redacted]'))" "$text") rt=$(python3 -c 'import json,sys;print(json.dumps(sys.argv[1]))' "$new") emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\",\"replace_text\":$rt}" else emit "{\"type\":\"event_intercept_response\",\"id\":\"$id\"}" fi ;; esac done ` if err := os.WriteFile(filepath.Join(extDir, "ext.sh"), []byte(script), 0o755); err != nil { t.Fatal(err) } manifest := `{"name":"itest","version":"0.1.0","exec":"./ext.sh","enabled":true}` if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil { t.Fatal(err) } // ZutHome is unused here; we load the extension explicitly. m := New(t.TempDir(), "", "0.0.0-test", "anthropic", "claude-test", nil) t.Cleanup(func() { m.Stop(2 * time.Second) }) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() if errs := m.LoadExplicit(ctx, []string{extDir}); len(errs) > 0 { t.Fatalf("LoadExplicit: %v", errs) } m.WaitForReady(3 * time.Second) // tool_call: rm -rf is blocked res := m.InterceptToolCall(ctx, "T1", "bash", json.RawMessage(`{"command":"rm -rf /tmp/foo","timeout":30}`)) if !res.Block || !strings.Contains(res.Reason, "refused") { t.Errorf("rm -rf: want block+reason, got %+v", res) } // tool_call: non-dangerous bash gets args rewritten res = m.InterceptToolCall(ctx, "T2", "bash", json.RawMessage(`{"command":"ls -la","timeout":30}`)) if res.Block { t.Errorf("ls -la: want allow, got block %q", res.Reason) } if len(res.ModifiedArgs) == 0 { t.Errorf("ls -la: want modified_args, got nothing") } else { var obj struct { Command string `json:"command"` Timeout int `json:"timeout"` } if err := json.Unmarshal(res.ModifiedArgs, &obj); err != nil { t.Errorf("modified_args: %v", err) } else { if !strings.HasPrefix(obj.Command, "echo GUARDED") { t.Errorf("modified_args command=%q, want echo GUARDED prefix", obj.Command) } if obj.Timeout != 30 { t.Errorf("modified_args timeout=%d, want 30", obj.Timeout) } } } // tool_call: non-bash tool untouched res = m.InterceptToolCall(ctx, "T3", "read", json.RawMessage(`{"path":"/etc/hosts"}`)) if res.Block { t.Errorf("read: want allow, got block %q", res.Reason) } if len(res.ModifiedArgs) != 0 { t.Errorf("read: want no modified_args, got %s", string(res.ModifiedArgs)) } // turn_start: always allowed if r := m.InterceptTurnStart(ctx, 1); r.Block { t.Errorf("turn_start: want allow, got block %q", r.Reason) } // assistant_message: SECRET gets redacted r := m.InterceptAssistantMessage(ctx, "your password is SECRET123") if r.Block { t.Errorf("msg: want allow, got block %q", r.Reason) } if !strings.Contains(r.ReplaceText, "[redacted]") { t.Errorf("msg: want [redacted] in replacement, got %q", r.ReplaceText) } // assistant_message: clean text unchanged r = m.InterceptAssistantMessage(ctx, "hello world") if r.Block { t.Errorf("clean: want allow") } if r.ReplaceText != "" { t.Errorf("clean: want no replace, got %q", r.ReplaceText) } } // TestReloadRespawnsExtensions verifies Reload tears down and brings // back the same extension, and the onReload callback fires. func TestReloadRespawnsExtensions(t *testing.T) { if _, err := os.Stat("/bin/bash"); err != nil { t.Skip("no /bin/bash") } extDir := t.TempDir() script := `#!/bin/bash emit() { printf '%s\n' "$1"; } emit '{"type":"hello","protocol_version":2,"name":"rtest","version":"0.1.0","capabilities":["commands"]}' emit '{"type":"register_command","name":"rtest","description":"test"}' emit '{"type":"ready"}' while IFS= read -r line; do t=$(printf '%s' "$line" | python3 -c 'import sys,json;print(json.load(sys.stdin).get("type",""))') if [[ "$t" == "shutdown" ]]; then emit '{"type":"shutdown_ack"}'; exit 0; fi done ` if err := os.WriteFile(filepath.Join(extDir, "ext.sh"), []byte(script), 0o755); err != nil { t.Fatal(err) } manifest := `{"name":"rtest","version":"0.1.0","exec":"./ext.sh","enabled":true}` if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil { t.Fatal(err) } hooks := &stubHooks{} m := New(t.TempDir(), "", "0.0.0-test", "anthropic", "claude-test", hooks) t.Cleanup(func() { m.Stop(2 * time.Second) }) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() if errs := m.LoadExplicit(ctx, []string{extDir}); len(errs) > 0 { t.Fatalf("LoadExplicit: %v", errs) } m.WaitForReady(3 * time.Second) before := len(m.All()) if before == 0 { t.Fatal("extension didn't load") } reloadFired := false m.SetOnReload(func() { reloadFired = true }) stats := m.Reload(ctx, 2*time.Second) if stats.Stopped != before { t.Errorf("Stopped: want %d, got %d", before, stats.Stopped) } if stats.Loaded != before { t.Errorf("Loaded: want %d, got %d", before, stats.Loaded) } if !reloadFired { t.Error("onReload callback didn't fire") } hooks.mu.Lock() clears := append([]string(nil), hooks.chromeClears...) hooks.mu.Unlock() if len(clears) == 0 || clears[0] != "rtest" { t.Fatalf("extension chrome cleanup callbacks = %v", clears) } // Registered command should still be there after reload. if !m.HasCommand("rtest") { t.Error("rtest command not re-registered after reload") } }