From 61add0ae46cbc6ec20bd3ed8dc7600157fd449c1 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Fri, 7 Aug 2026 13:50:20 +0800 Subject: [PATCH 1/5] feat: add optional shell command bridge --- cmd/aiscan/cli_test.go | 17 + cmd/aiscan/main.go | 12 + cmd/aiscan/setup.go | 1 + core/config/loader.go | 3 + core/config/options.go | 1 + go.mod | 1 + go.sum | 2 + pkg/commands/bash.go | 83 +++- pkg/commands/command_bridge.go | 512 +++++++++++++++++++++++++ pkg/commands/command_bridge_test.go | 256 +++++++++++++ pkg/commands/command_bridge_unix.go | 57 +++ pkg/commands/command_bridge_windows.go | 69 ++++ pkg/commands/factory.go | 16 +- pkg/commands/register.go | 5 + pkg/commands/register_test.go | 22 ++ pkg/runner/app.go | 28 +- pkg/runner/app_test.go | 23 ++ pkg/runner/application_builder.go | 2 + pkg/runner/application_config.go | 1 + pkg/runner/provider_config_test.go | 3 +- 20 files changed, 1103 insertions(+), 11 deletions(-) create mode 100644 pkg/commands/command_bridge.go create mode 100644 pkg/commands/command_bridge_test.go create mode 100644 pkg/commands/command_bridge_unix.go create mode 100644 pkg/commands/command_bridge_windows.go diff --git a/cmd/aiscan/cli_test.go b/cmd/aiscan/cli_test.go index 37d0fe97..f942502b 100644 --- a/cmd/aiscan/cli_test.go +++ b/cmd/aiscan/cli_test.go @@ -181,6 +181,23 @@ func TestParseCLIRejectsResumeWithExplicitFile(t *testing.T) { } } +func TestAgentCommandBridgeFlagIsOptIn(t *testing.T) { + parsed, err := parseCLI([]string{"agent", "--command-bridge", "-p", "test"}) + if err != nil { + t.Fatalf("parseCLI() error = %v", err) + } + if !parsed.Option.CommandBridge { + t.Fatal("--command-bridge was not propagated to agent options") + } + parsed, err = parseCLI([]string{"agent", "-p", "test"}) + if err != nil { + t.Fatalf("parseCLI() default error = %v", err) + } + if parsed.Option.CommandBridge { + t.Fatal("command bridge must remain disabled by default") + } +} + func TestParseCLIRootTimeoutAppliesToAgent(t *testing.T) { parsed, err := parseCLI([]string{"--timeout", "45", "agent", "-p", "test"}) if err != nil { diff --git a/cmd/aiscan/main.go b/cmd/aiscan/main.go index d787af77..8c8cc63b 100644 --- a/cmd/aiscan/main.go +++ b/cmd/aiscan/main.go @@ -1,5 +1,17 @@ package main +import ( + "os" + + "github.com/chainreactors/aiscan/pkg/commands" +) + func main() { + if code, ok := commands.RunCommandBridgeProxyIfRequested(); ok { + if code != 0 { + os.Exit(code) + } + return + } aiscan() } diff --git a/cmd/aiscan/setup.go b/cmd/aiscan/setup.go index c62adae6..bfd0c783 100644 --- a/cmd/aiscan/setup.go +++ b/cmd/aiscan/setup.go @@ -99,6 +99,7 @@ func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine deps := &commands.Deps{ WorkDir: workDir, BashTimeout: toolCfg.BashTimeout, + CommandBridge: toolCfg.CommandBridge, SkillStore: skillStore, ScannerProxy: scanCfg.Proxy, Logger: logger, diff --git a/core/config/loader.go b/core/config/loader.go index 067e370b..0b2b3e3b 100644 --- a/core/config/loader.go +++ b/core/config/loader.go @@ -134,6 +134,9 @@ func mergeOption(dst, src *Option) { if !dst.SaveSession && src.SaveSession { dst.SaveSession = true } + if !dst.CommandBridge && src.CommandBridge { + dst.CommandBridge = true + } mergeOutputOptions(&dst.OutputOptions, &src.OutputOptions) dst.DataDir = ResolveString(dst.DataDir, src.DataDir) } diff --git a/core/config/options.go b/core/config/options.go index fcf88abc..0b0de5f8 100644 --- a/core/config/options.go +++ b/core/config/options.go @@ -89,6 +89,7 @@ type AgentOptions struct { Resume string `short:"r" long:"resume" description:"Resume agent context from an AOP JSONL session file"` SaveSession bool `long:"save-session" config:"save_session" description:"Auto-select a .aiscan/sessions/*.jsonl recording path"` CaptureProviderFrames bool `long:"capture-provider-frames" config:"capture_provider_frames" description:"Emit exact provider request/response frames as sensitive AOP events"` + CommandBridge bool `long:"command-bridge" config:"command_bridge" description:"Expose in-memory commands to shell composition through a process-local bridge"` } type AgentTransport string diff --git a/go.mod b/go.mod index ee66dac2..aade70d6 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ tool ( require ( connectrpc.com/connect v1.20.0 + github.com/Microsoft/go-winio v0.6.2 github.com/alecthomas/chroma/v2 v2.14.0 github.com/carapace-sh/carapace v1.11.6 github.com/chainreactors/crtm v0.0.3-0.20260618163257-073207497076 diff --git a/go.sum b/go.sum index 19defe2a..0ba0f39b 100644 --- a/go.sum +++ b/go.sum @@ -70,6 +70,8 @@ github.com/M09ic/go-ntlmssp v0.0.0-20230312133735-dcccd454dfe0 h1:9Y+BdzDIHfpKy0 github.com/M09ic/go-ntlmssp v0.0.0-20230312133735-dcccd454dfe0/go.mod h1:yMNEF6ulbFipt3CakMhcmcNVACshPRG4Ap4l00V+mMs= github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0= github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM= +github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= +github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= github.com/Mzack9999/gcache v0.0.0-20230410081825-519e28eab057 h1:KFac3SiGbId8ub47e7kd2PLZeACxc1LkiiNoDOFRClE= github.com/Mzack9999/gcache v0.0.0-20230410081825-519e28eab057/go.mod h1:iLB2pivrPICvLOuROKmlqURtFIEsoJZaMidQfCG1+D4= github.com/Mzack9999/go-http-digest-auth-client v0.6.1-0.20220414142836-eb8883508809 h1:ZbFL+BDfBqegi+/Ssh7im5+aQfBRx6it+kHnC7jaDU8= diff --git a/pkg/commands/bash.go b/pkg/commands/bash.go index a99e6153..c0541d2e 100644 --- a/pkg/commands/bash.go +++ b/pkg/commands/bash.go @@ -8,6 +8,7 @@ import ( "os/exec" "sort" "strings" + "sync" "time" "github.com/chainreactors/aiscan/agent/inbox" @@ -44,6 +45,9 @@ type BashTool struct { tasks *tmux.Manager commandNames func() []string resolveCommand func(string) (Command, bool) + bridge *commandBridge + bridgeErr error + closeOnce sync.Once } func NewBashTool(workDir string, timeout int) *BashTool { @@ -60,7 +64,59 @@ func (t *BashTool) SetCommandResolver(fn func(string) (Command, bool)) { t.resolveCommand = fn } func (t *BashTool) Name() string { return "bash" } -func (t *BashTool) Close() { t.tasks.Shutdown() } +func (t *BashTool) Close() { + t.closeOnce.Do(func() { + if t.bridge != nil { + t.bridge.shutdown() + } + t.tasks.Shutdown() + if t.bridge != nil { + t.bridge.cleanup() + } + }) +} + +func (t *BashTool) EnableCommandBridge(registry *CommandRegistry) error { + if t.bridge != nil { + return t.SyncCommandBridgeAliases() + } + bridge, err := newCommandBridge(registry) + if err != nil { + t.bridgeErr = err + return err + } + t.bridge = bridge + if err := t.SyncCommandBridgeAliases(); err != nil { + t.bridgeErr = err + bridge.close() + t.bridge = nil + return err + } + t.bridgeErr = nil + return nil +} + +func (t *BashTool) SyncCommandBridgeAliases() error { + if t.bridge == nil || t.commandNames == nil { + return t.bridgeErr + } + if err := t.bridge.syncAliases(t.commandNames()); err != nil { + t.bridgeErr = err + return err + } + return nil +} + +func (t *BashTool) CommandBridgeError() error { return t.bridgeErr } + +func (t *BashTool) CommandBridgeEnabled() bool { return t.bridge != nil } + +func (t *BashTool) CommandBridgeRuntimeDir() string { + if t.bridge == nil { + return "" + } + return t.bridge.runtimeDir +} func (t *BashTool) WithScannerProxy(proxy string) *BashTool { t.scannerProxy = proxy @@ -208,7 +264,16 @@ func (t *BashTool) Start(ctx context.Context, command string, options BashExecOp if workDir == "" { workDir = t.workDir } - env := t.runEnv(options.Env) + env := t.runEnv(ctx, options.Env) + if t.bridge != nil { + execution := newExecution(t.tasks, command, nil, workDir, env) + info, err := t.tasks.Create(workDir, command, options.Name, timeout, env, "") + if err != nil { + return nil, err + } + execution.bind(info) + return execution, nil + } left, right, hasPipe := splitPipeline(command) leftToken := firstCommandToken(left) if cmd, ok := t.resolve(leftToken); ok { @@ -447,7 +512,7 @@ func (t *BashTool) collectResult(execution *Execution) *coretool.Result { return result } -func (t *BashTool) runEnv(overrides map[string]string) []string { +func (t *BashTool) runEnv(ctx context.Context, overrides map[string]string) []string { values := make(map[string]string) for _, item := range t.proxyEnv() { if key, value, ok := strings.Cut(item, "="); ok { @@ -457,6 +522,18 @@ func (t *BashTool) runEnv(overrides map[string]string) []string { for key, value := range overrides { values[key] = value } + if t.bridge != nil { + for _, item := range t.bridge.environment(coretool.InvocationFromContext(ctx)) { + if key, value, ok := strings.Cut(item, "="); ok { + values[key] = value + } + } + path := values["PATH"] + if path == "" { + path = os.Getenv("PATH") + } + values["PATH"] = t.bridge.runtimeDir + string(os.PathListSeparator) + path + } keys := make([]string, 0, len(values)) for key := range values { keys = append(keys, key) diff --git a/pkg/commands/command_bridge.go b/pkg/commands/command_bridge.go new file mode 100644 index 00000000..f5644a3f --- /dev/null +++ b/pkg/commands/command_bridge.go @@ -0,0 +1,512 @@ +package commands + +import ( + "bufio" + "context" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +const ( + commandBridgeMarkerEnv = "AISCAN_COMMAND_BRIDGE" + commandBridgeEndpointEnv = "AISCAN_COMMAND_BRIDGE_ENDPOINT" + commandBridgeCommandEnv = "AISCAN_COMMAND_BRIDGE_COMMAND" + commandBridgeExecutableEnv = "AISCAN_COMMAND_BRIDGE_EXECUTABLE" + commandBridgeCallIDEnv = "AISCAN_COMMAND_BRIDGE_CALL_ID" + commandBridgeSessionEnv = "AISCAN_COMMAND_BRIDGE_SESSION_ID" + commandBridgeTurnEnv = "AISCAN_COMMAND_BRIDGE_TURN_ID" + commandBridgeEmitterEnv = "AISCAN_COMMAND_BRIDGE_EMITTER" + + commandBridgeProtocolVersion = 1 + commandBridgeChunkSize = 32 << 10 + commandBridgeMaxFrameSize = 1 << 20 + commandBridgeDialTimeout = 5 * time.Second +) + +type commandBridgeFrame struct { + Type string `json:"type"` + Version int `json:"version,omitempty"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Dir string `json:"dir,omitempty"` + Invocation coretool.Invocation `json:"invocation,omitempty"` + Data []byte `json:"data,omitempty"` + ExitCode int `json:"exit_code,omitempty"` +} + +type commandBridge struct { + registry *CommandRegistry + executable string + runtimeDir string + endpoint string + listener net.Listener + cancel context.CancelFunc + + mu sync.Mutex + aliases map[string]string + connections map[net.Conn]struct{} + wg sync.WaitGroup + shutdownOnce sync.Once + cleanupOnce sync.Once +} + +func newCommandBridge(registry *CommandRegistry) (*commandBridge, error) { + if registry == nil { + return nil, fmt.Errorf("command bridge requires a registry") + } + executable, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("resolve command bridge executable: %w", err) + } + executable, err = filepath.Abs(executable) + if err != nil { + return nil, fmt.Errorf("resolve command bridge executable path: %w", err) + } + if err := cleanupStaleCommandBridgeRuntime(); err != nil { + return nil, err + } + + root := commandBridgeRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + return nil, fmt.Errorf("create command bridge runtime root: %w", err) + } + _ = os.Chmod(root, 0o700) + runtimeDir, err := os.MkdirTemp(root, strconv.Itoa(os.Getpid())+"-") + if err != nil { + return nil, fmt.Errorf("create command bridge runtime: %w", err) + } + cleanup := func() { _ = os.RemoveAll(runtimeDir) } + _ = os.Chmod(runtimeDir, 0o700) + endpoint := commandBridgeEndpoint(runtimeDir) + listener, err := listenCommandBridge(endpoint) + if err != nil { + cleanup() + return nil, fmt.Errorf("listen command bridge: %w", err) + } + ctx, cancel := context.WithCancel(context.Background()) + bridge := &commandBridge{ + registry: registry, executable: executable, runtimeDir: runtimeDir, + endpoint: endpoint, listener: listener, cancel: cancel, + aliases: make(map[string]string), connections: make(map[net.Conn]struct{}), + } + bridge.wg.Add(1) + go bridge.accept(ctx) + return bridge, nil +} + +func commandBridgeRuntimeRoot() string { + return filepath.Join(os.TempDir(), "aiscan-command-bridge") +} + +func cleanupStaleCommandBridgeRuntime() error { + root := commandBridgeRuntimeRoot() + entries, err := os.ReadDir(root) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return fmt.Errorf("read command bridge runtime root: %w", err) + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + pidText, _, ok := strings.Cut(entry.Name(), "-") + pid, parseErr := strconv.Atoi(pidText) + if !ok || parseErr != nil || pid <= 0 { + continue + } + if !commandBridgeProcessAlive(pid) { + _ = os.RemoveAll(filepath.Join(root, entry.Name())) + } + } + return nil +} + +func (b *commandBridge) accept(ctx context.Context) { + defer b.wg.Done() + for { + conn, err := b.listener.Accept() + if err != nil { + if ctx.Err() != nil { + return + } + continue + } + b.mu.Lock() + b.connections[conn] = struct{}{} + b.wg.Add(1) + b.mu.Unlock() + go b.handle(ctx, conn) + } +} + +func (b *commandBridge) handle(parent context.Context, conn net.Conn) { + defer b.wg.Done() + defer func() { + _ = conn.Close() + b.mu.Lock() + delete(b.connections, conn) + b.mu.Unlock() + }() + + reader := bufio.NewReader(conn) + header, err := readCommandBridgeFrame(reader) + if err != nil { + return + } + writer := &commandBridgeFrameWriter{writer: conn} + if header.Type != "request" || header.Version != commandBridgeProtocolVersion { + _, _ = io.WriteString(&commandBridgeStreamWriter{writer: writer, frameType: "stderr"}, "unsupported command bridge protocol\n") + _ = writer.write(commandBridgeFrame{Type: "final", ExitCode: 125}) + return + } + command, ok := b.registry.Get(header.Command) + if !ok || command.Run == nil { + message := "unknown in-memory command: " + header.Command + _, _ = io.WriteString(&commandBridgeStreamWriter{writer: writer, frameType: "stderr"}, message+"\n") + _ = writer.write(commandBridgeFrame{Type: "final", ExitCode: 127}) + return + } + + ctx, cancel := context.WithCancel(parent) + defer cancel() + ctx = coretool.ContextWithInvocation(ctx, header.Invocation) + stdinReader, stdinWriter := io.Pipe() + defer stdinReader.Close() + readDone := make(chan error, 1) + go func() { + readDone <- readCommandBridgeInput(reader, stdinWriter, cancel) + }() + + stdout := &commandBridgeStreamWriter{writer: writer, frameType: "stdout"} + stderr := &commandBridgeStreamWriter{writer: writer, frameType: "stderr"} + execution := newExecution(nil, command.Name, normalizeNoColor(command.Name, header.Args), header.Dir, nil) + execution.ID = header.Invocation.CallID + execution.StartedAt = time.Now() + execution.setIO(stdinReader, stdout, stderr) + _, runErr := command.Run(ctx, execution) + execution.EndedAt = time.Now() + _ = stdinReader.Close() + cancel() + + exitCode := commandBridgeExitCode(runErr) + if runErr != nil && !errors.Is(runErr, context.Canceled) && !errors.Is(runErr, context.DeadlineExceeded) { + _, _ = io.WriteString(stderr, runErr.Error()+"\n") + } + _ = writer.write(commandBridgeFrame{Type: "final", ExitCode: exitCode}) + _ = conn.Close() + select { + case <-readDone: + case <-time.After(time.Second): + } +} + +func readCommandBridgeInput(reader *bufio.Reader, stdin *io.PipeWriter, cancel context.CancelFunc) error { + stdinOpen := true + closeStdin := func() { + if stdinOpen { + stdinOpen = false + _ = stdin.Close() + } + } + defer closeStdin() + for { + frame, err := readCommandBridgeFrame(reader) + if err != nil { + cancel() + return err + } + switch frame.Type { + case "stdin": + if stdinOpen && len(frame.Data) > 0 { + if _, err := stdin.Write(frame.Data); err != nil { + return err + } + } + case "stdin_eof": + closeStdin() + case "cancel": + cancel() + return context.Canceled + default: + cancel() + return fmt.Errorf("unexpected command bridge frame %q", frame.Type) + } + } +} + +type commandBridgeFrameWriter struct { + mu sync.Mutex + writer io.Writer +} + +func (w *commandBridgeFrameWriter) write(frame commandBridgeFrame) error { + w.mu.Lock() + defer w.mu.Unlock() + return writeCommandBridgeFrame(w.writer, frame) +} + +type commandBridgeStreamWriter struct { + writer *commandBridgeFrameWriter + frameType string +} + +func (w *commandBridgeStreamWriter) Write(data []byte) (int, error) { + for offset := 0; offset < len(data); { + end := offset + commandBridgeChunkSize + if end > len(data) { + end = len(data) + } + if err := w.writer.write(commandBridgeFrame{Type: w.frameType, Data: data[offset:end]}); err != nil { + return offset, err + } + offset = end + } + return len(data), nil +} + +func commandBridgeExitCode(err error) int { + if err == nil { + return 0 + } + var exitCoder interface{ ExitCode() int } + if errors.As(err, &exitCoder) { + code := exitCoder.ExitCode() + if code >= 0 && code <= 255 { + return code + } + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return 130 + } + return 1 +} + +func writeCommandBridgeFrame(writer io.Writer, frame commandBridgeFrame) error { + data, err := json.Marshal(frame) + if err != nil { + return err + } + if len(data) > commandBridgeMaxFrameSize { + return fmt.Errorf("command bridge frame too large: %d", len(data)) + } + var size [4]byte + binary.BigEndian.PutUint32(size[:], uint32(len(data))) + if err := writeCommandBridgeBytes(writer, size[:]); err != nil { + return err + } + return writeCommandBridgeBytes(writer, data) +} + +func writeCommandBridgeBytes(writer io.Writer, data []byte) error { + for len(data) > 0 { + n, err := writer.Write(data) + if err != nil { + return err + } + if n <= 0 { + return io.ErrShortWrite + } + data = data[n:] + } + return nil +} + +func readCommandBridgeFrame(reader io.Reader) (commandBridgeFrame, error) { + var size [4]byte + if _, err := io.ReadFull(reader, size[:]); err != nil { + return commandBridgeFrame{}, err + } + length := binary.BigEndian.Uint32(size[:]) + if length == 0 || length > commandBridgeMaxFrameSize { + return commandBridgeFrame{}, fmt.Errorf("invalid command bridge frame size: %d", length) + } + data := make([]byte, int(length)) + if _, err := io.ReadFull(reader, data); err != nil { + return commandBridgeFrame{}, err + } + var frame commandBridgeFrame + if err := json.Unmarshal(data, &frame); err != nil { + return commandBridgeFrame{}, err + } + return frame, nil +} + +func (b *commandBridge) syncAliases(names []string) error { + b.mu.Lock() + defer b.mu.Unlock() + sorted := append([]string(nil), names...) + sort.Strings(sorted) + for _, name := range sorted { + if !validCommandBridgeName(name) { + continue + } + if _, ok := b.aliases[name]; ok { + continue + } + path, err := createCommandBridgeAlias(b.executable, b.runtimeDir, name) + if err != nil { + return fmt.Errorf("create command bridge alias %s: %w", name, err) + } + b.aliases[name] = path + } + return nil +} + +func validCommandBridgeName(name string) bool { + if name == "" || name == "." || name == ".." { + return false + } + for _, r := range name { + if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' { + continue + } + return false + } + return true +} + +func (b *commandBridge) environment(invocation coretool.Invocation) []string { + return []string{ + commandBridgeMarkerEnv + "=1", + commandBridgeEndpointEnv + "=" + b.endpoint, + commandBridgeExecutableEnv + "=" + b.executable, + commandBridgeCallIDEnv + "=" + invocation.CallID, + commandBridgeSessionEnv + "=" + invocation.SessionID, + commandBridgeTurnEnv + "=" + invocation.TurnID, + commandBridgeEmitterEnv + "=" + invocation.Emitter, + } +} + +func (b *commandBridge) shutdown() { + b.shutdownOnce.Do(func() { + b.cancel() + _ = b.listener.Close() + b.mu.Lock() + for conn := range b.connections { + _ = conn.Close() + } + b.mu.Unlock() + b.wg.Wait() + }) +} + +func (b *commandBridge) cleanup() { + b.cleanupOnce.Do(func() { + for attempt := 0; attempt < 5; attempt++ { + if err := os.RemoveAll(b.runtimeDir); err == nil { + break + } + time.Sleep(time.Duration(attempt+1) * 20 * time.Millisecond) + } + _ = os.Remove(commandBridgeRuntimeRoot()) + }) +} + +func (b *commandBridge) close() { + b.shutdown() + b.cleanup() +} + +// RunCommandBridgeProxyIfRequested detects a bridge wrapper invocation before +// normal CLI parsing. Callers should return normally for code 0 and os.Exit for +// a non-zero code when ok is true. +func RunCommandBridgeProxyIfRequested() (code int, ok bool) { + command := strings.TrimSpace(os.Getenv(commandBridgeCommandEnv)) + if os.Getenv(commandBridgeMarkerEnv) != "1" || command == "" { + return 0, false + } + if !validCommandBridgeName(command) { + fmt.Fprintln(os.Stderr, "invalid command bridge alias") + return 126, true + } + return runCommandBridgeProxy(command, os.Args[1:]), true +} + +func runCommandBridgeProxy(command string, args []string) int { + endpoint := os.Getenv(commandBridgeEndpointEnv) + if endpoint == "" { + fmt.Fprintln(os.Stderr, "AIScan command bridge environment is incomplete") + return 126 + } + ctx, cancel := context.WithTimeout(context.Background(), commandBridgeDialTimeout) + conn, err := dialCommandBridge(ctx, endpoint) + cancel() + if err != nil { + fmt.Fprintf(os.Stderr, "connect AIScan command bridge: %s\n", err) + return 125 + } + defer conn.Close() + cwd, _ := os.Getwd() + header := commandBridgeFrame{ + Type: "request", Version: commandBridgeProtocolVersion, + Command: command, Args: append([]string(nil), args...), Dir: cwd, + Invocation: coretool.Invocation{ + WorkDir: cwd, CallID: os.Getenv(commandBridgeCallIDEnv), + SessionID: os.Getenv(commandBridgeSessionEnv), TurnID: os.Getenv(commandBridgeTurnEnv), + Emitter: os.Getenv(commandBridgeEmitterEnv), + }, + } + writer := &commandBridgeFrameWriter{writer: conn} + if err := writer.write(header); err != nil { + fmt.Fprintf(os.Stderr, "start AIScan command bridge request: %s\n", err) + return 125 + } + go streamCommandBridgeStdin(writer) + + reader := bufio.NewReader(conn) + for { + frame, err := readCommandBridgeFrame(reader) + if err != nil { + fmt.Fprintf(os.Stderr, "read AIScan command bridge response: %s\n", err) + return 125 + } + switch frame.Type { + case "stdout": + _, _ = os.Stdout.Write(frame.Data) + case "stderr": + _, _ = os.Stderr.Write(frame.Data) + case "final": + flushCommandBridgeProxyOutput() + return frame.ExitCode + default: + fmt.Fprintf(os.Stderr, "unexpected AIScan command bridge response %q\n", frame.Type) + return 125 + } + } +} + +func streamCommandBridgeStdin(writer *commandBridgeFrameWriter) { + info, err := os.Stdin.Stat() + if err != nil || info.Mode()&os.ModeCharDevice != 0 { + _ = writer.write(commandBridgeFrame{Type: "stdin_eof"}) + return + } + buffer := make([]byte, commandBridgeChunkSize) + for { + n, readErr := os.Stdin.Read(buffer) + if n > 0 { + if err := writer.write(commandBridgeFrame{Type: "stdin", Data: buffer[:n]}); err != nil { + return + } + } + if readErr != nil { + _ = writer.write(commandBridgeFrame{Type: "stdin_eof"}) + return + } + } +} diff --git a/pkg/commands/command_bridge_test.go b/pkg/commands/command_bridge_test.go new file mode 100644 index 00000000..e92a1aea --- /dev/null +++ b/pkg/commands/command_bridge_test.go @@ -0,0 +1,256 @@ +package commands + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +func TestMain(m *testing.M) { + if code, ok := RunCommandBridgeProxyIfRequested(); ok { + if code != 0 { + os.Exit(code) + } + return + } + os.Exit(m.Run()) +} + +type bridgeTestExitError struct{ code int } + +func (e bridgeTestExitError) Error() string { return fmt.Sprintf("bridge test exit %d", e.code) } +func (e bridgeTestExitError) ExitCode() int { return e.code } + +type bridgeTestCommands struct { + started chan struct{} + canceled chan struct{} + once sync.Once +} + +func newBridgeTestBash(t *testing.T) (*BashTool, *CommandRegistry, *bridgeTestCommands) { + t.Helper() + state := &bridgeTestCommands{started: make(chan struct{}), canceled: make(chan struct{})} + registry := NewRegistry() + registry.Register(Command{Name: "bridge_echo", Run: func(_ context.Context, execution *Execution) (any, error) { + fmt.Fprintln(execution.Stdout, strings.Join(execution.Args, " ")) + return nil, nil + }}, "test") + registry.Register(Command{Name: "bridge_upper", Run: func(_ context.Context, execution *Execution) (any, error) { + data, err := io.ReadAll(execution.Stdin) + if err != nil { + return nil, err + } + _, err = execution.Stdout.Write(bytes.ToUpper(data)) + return nil, err + }}, "test") + registry.Register(Command{Name: "bridge_fail", Run: func(context.Context, *Execution) (any, error) { + return nil, bridgeTestExitError{code: 7} + }}, "test") + registry.Register(Command{Name: "bridge_context", Run: func(ctx context.Context, execution *Execution) (any, error) { + invocation := coretool.InvocationFromContext(ctx) + fmt.Fprintf(execution.Stdout, "dir=%s call=%s session=%s turn=%s emitter=%s\n", + execution.Dir, invocation.CallID, invocation.SessionID, invocation.TurnID, invocation.Emitter) + return nil, nil + }}, "test") + registry.Register(Command{Name: "bridge_wait", Run: func(ctx context.Context, _ *Execution) (any, error) { + state.once.Do(func() { close(state.started) }) + <-ctx.Done() + close(state.canceled) + return nil, ctx.Err() + }}, "test") + + bash := NewBashTool(t.TempDir(), 10) + bash.SetCommandNames(registry.Names) + bash.SetCommandResolver(registry.Get) + if err := bash.EnableCommandBridge(registry); err != nil { + t.Fatalf("EnableCommandBridge: %v", err) + } + t.Cleanup(bash.Close) + return bash, registry, state +} + +func runBridgeCommand(t *testing.T, bash *BashTool, ctx context.Context, command string, workDir string) (*Execution, string) { + t.Helper() + var output strings.Builder + execution, err := bash.RunForeground(ctx, command, BashExecOptions{ + WorkDir: workDir, + OnOutput: func(data []byte) { + _, _ = output.Write(data) + }, + }) + if err != nil { + t.Fatalf("RunForeground(%q): %v", command, err) + } + return execution, output.String() +} + +func TestCommandBridgeDisabledByDefault(t *testing.T) { + bash := NewBashTool(t.TempDir(), 5) + defer bash.Close() + if bash.CommandBridgeEnabled() { + t.Fatal("command bridge must be disabled by default") + } + if bash.CommandBridgeRuntimeDir() != "" { + t.Fatalf("disabled bridge runtime dir = %q", bash.CommandBridgeRuntimeDir()) + } +} + +func TestCommandBridgeMarkerDoesNotHijackNormalChildProcess(t *testing.T) { + t.Setenv(commandBridgeMarkerEnv, "1") + t.Setenv(commandBridgeCommandEnv, "") + if code, ok := RunCommandBridgeProxyIfRequested(); ok { + t.Fatalf("marker-only child was treated as proxy with code %d", code) + } +} + +func TestCommandBridgeShellComposition(t *testing.T) { + bash, _, _ := newBridgeTestBash(t) + workDir := t.TempDir() + + execution, output := runBridgeCommand(t, bash, context.Background(), "bridge_echo one && bridge_echo two", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "one") || !strings.Contains(output, "two") { + t.Fatalf("and composition exit=%d output=%q", execution.ExitCode, output) + } + + execution, output = runBridgeCommand(t, bash, context.Background(), "bridge_fail || bridge_echo recovered", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "recovered") { + t.Fatalf("or composition exit=%d output=%q", execution.ExitCode, output) + } + + execution, output = runBridgeCommand(t, bash, context.Background(), "bridge_echo hello | bridge_upper", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "HELLO") { + t.Fatalf("pipeline exit=%d output=%q", execution.ExitCode, output) + } + + execution, _ = runBridgeCommand(t, bash, context.Background(), "bridge_fail && bridge_echo unreachable", workDir) + if execution.ExitCode != 7 { + t.Fatalf("short-circuit exit code = %d, want 7", execution.ExitCode) + } +} + +func TestCommandBridgeRedirectionAndInvocationContext(t *testing.T) { + bash, _, _ := newBridgeTestBash(t) + workDir := t.TempDir() + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{ + WorkDir: workDir, CallID: "call-1", SessionID: "session-1", TurnID: "turn-1", Emitter: "runner", + }) + + execution, _ := runBridgeCommand(t, bash, ctx, "bridge_context > bridge-context.txt", workDir) + if execution.ExitCode != 0 { + t.Fatalf("redirection exit code = %d", execution.ExitCode) + } + data, err := os.ReadFile(filepath.Join(workDir, "bridge-context.txt")) + if err != nil { + t.Fatalf("read redirected output: %v", err) + } + text := string(data) + for _, expected := range []string{"dir=" + workDir, "call=call-1", "session=session-1", "turn=turn-1", "emitter=runner"} { + if !strings.Contains(text, expected) { + t.Fatalf("redirected context missing %q: %q", expected, text) + } + } +} + +func TestCommandBridgeSyncsCommandsRegisteredLater(t *testing.T) { + bash, registry, _ := newBridgeTestBash(t) + registry.Register(Command{Name: "bridge_late", Run: func(_ context.Context, execution *Execution) (any, error) { + fmt.Fprintln(execution.Stdout, "late") + return nil, nil + }}, "test") + if err := bash.SyncCommandBridgeAliases(); err != nil { + t.Fatalf("SyncCommandBridgeAliases: %v", err) + } + execution, output := runBridgeCommand(t, bash, context.Background(), "bridge_late && bridge_echo ready", t.TempDir()) + if execution.ExitCode != 0 || !strings.Contains(output, "late") || !strings.Contains(output, "ready") { + t.Fatalf("late alias exit=%d output=%q", execution.ExitCode, output) + } +} + +func TestCommandBridgeCloseCancelsCallsAndRemovesRuntime(t *testing.T) { + bash, _, state := newBridgeTestBash(t) + runtimeDir := bash.CommandBridgeRuntimeDir() + if _, err := os.Stat(runtimeDir); err != nil { + t.Fatalf("runtime directory before close: %v", err) + } + + execution, err := bash.Start(context.Background(), "bridge_wait", BashExecOptions{WorkDir: t.TempDir()}) + if err != nil { + t.Fatalf("Start bridge_wait: %v", err) + } + select { + case <-state.started: + case <-time.After(5 * time.Second): + t.Fatal("bridge command did not start") + } + bash.Close() + bash.Close() + select { + case <-state.canceled: + case <-time.After(5 * time.Second): + t.Fatal("bridge command was not canceled") + } + if err := execution.Wait(context.Background()); err != nil && !errors.Is(err, context.Canceled) { + t.Fatalf("wait after close: %v", err) + } + if _, err := os.Stat(runtimeDir); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("runtime directory still exists after close: %v", err) + } +} + +func TestCommandBridgeDisconnectCancelsRunningCommand(t *testing.T) { + bash, _, state := newBridgeTestBash(t) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + conn, err := dialCommandBridge(ctx, bash.bridge.endpoint) + if err != nil { + t.Fatalf("dial command bridge: %v", err) + } + writer := &commandBridgeFrameWriter{writer: conn} + if err := writer.write(commandBridgeFrame{ + Type: "request", Version: commandBridgeProtocolVersion, + Command: "bridge_wait", Dir: t.TempDir(), + }); err != nil { + t.Fatalf("write request: %v", err) + } + if err := writer.write(commandBridgeFrame{Type: "stdin_eof"}); err != nil { + t.Fatalf("write stdin eof: %v", err) + } + select { + case <-state.started: + case <-ctx.Done(): + t.Fatal("bridge command did not start") + } + _ = conn.Close() + select { + case <-state.canceled: + case <-ctx.Done(): + t.Fatal("disconnect did not cancel bridge command") + } +} + +func TestCommandBridgeStartupReclaimsOwnedStaleRuntime(t *testing.T) { + root := commandBridgeRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + t.Fatal(err) + } + stale, err := os.MkdirTemp(root, "1073741824-stale-") + if err != nil { + t.Fatal(err) + } + if err := cleanupStaleCommandBridgeRuntime(); err != nil { + t.Fatalf("cleanup stale runtime: %v", err) + } + if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("stale runtime still exists: %v", err) + } +} diff --git a/pkg/commands/command_bridge_unix.go b/pkg/commands/command_bridge_unix.go new file mode 100644 index 00000000..fdc87cec --- /dev/null +++ b/pkg/commands/command_bridge_unix.go @@ -0,0 +1,57 @@ +//go:build !windows + +package commands + +import ( + "context" + "errors" + "net" + "os" + "path/filepath" + "syscall" +) + +func commandBridgeEndpoint(runtimeDir string) string { + return filepath.Join(runtimeDir, "bridge.sock") +} + +func listenCommandBridge(endpoint string) (net.Listener, error) { + _ = os.Remove(endpoint) + listener, err := net.Listen("unix", endpoint) + if err != nil { + return nil, err + } + if err := os.Chmod(endpoint, 0o600); err != nil { + _ = listener.Close() + return nil, err + } + return listener, nil +} + +func dialCommandBridge(ctx context.Context, endpoint string) (net.Conn, error) { + var dialer net.Dialer + return dialer.DialContext(ctx, "unix", endpoint) +} + +func createCommandBridgeAlias(executable, runtimeDir, name string) (string, error) { + _ = executable + path := filepath.Join(runtimeDir, name) + if _, err := os.Stat(path); err == nil { + return path, nil + } + content := "#!/bin/sh\n" + commandBridgeCommandEnv + "='" + name + "' exec \"$" + commandBridgeExecutableEnv + "\" \"$@\"\n" + if err := os.WriteFile(path, []byte(content), 0o700); err != nil { + return "", err + } + return path, nil +} + +func commandBridgeProcessAlive(pid int) bool { + err := syscall.Kill(pid, 0) + return err == nil || errors.Is(err, syscall.EPERM) +} + +func flushCommandBridgeProxyOutput() { + _ = os.Stdout.Sync() + _ = os.Stderr.Sync() +} diff --git a/pkg/commands/command_bridge_windows.go b/pkg/commands/command_bridge_windows.go new file mode 100644 index 00000000..1ecce4c3 --- /dev/null +++ b/pkg/commands/command_bridge_windows.go @@ -0,0 +1,69 @@ +//go:build windows + +package commands + +import ( + "context" + "fmt" + "net" + "os" + "path/filepath" + "time" + + "github.com/Microsoft/go-winio" + "golang.org/x/sys/windows" +) + +func commandBridgeEndpoint(runtimeDir string) string { + return `\\.\pipe\aiscan-command-bridge-` + fmt.Sprintf("%d-%s", os.Getpid(), filepath.Base(runtimeDir)) +} + +func listenCommandBridge(endpoint string) (net.Listener, error) { + user, err := windows.GetCurrentProcessToken().GetTokenUser() + if err != nil { + return nil, err + } + return winio.ListenPipe(endpoint, &winio.PipeConfig{ + SecurityDescriptor: "D:P(A;;GA;;;" + user.User.Sid.String() + ")", + InputBufferSize: commandBridgeChunkSize, + OutputBufferSize: commandBridgeChunkSize, + }) +} + +func dialCommandBridge(ctx context.Context, endpoint string) (net.Conn, error) { + return winio.DialPipeContext(ctx, endpoint) +} + +func createCommandBridgeAlias(executable, runtimeDir, name string) (string, error) { + _ = executable + path := filepath.Join(runtimeDir, name+".cmd") + if _, err := os.Stat(path); err == nil { + return path, nil + } + content := "@echo off\r\n" + + "set \"" + commandBridgeCommandEnv + "=" + name + "\"\r\n" + + "\"%" + commandBridgeExecutableEnv + "%\" %*\r\n" + + "exit /b %errorlevel%\r\n" + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + return "", err + } + return path, nil +} + +func commandBridgeProcessAlive(pid int) bool { + handle, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(pid)) + if err != nil { + return err == windows.ERROR_ACCESS_DENIED + } + _ = windows.CloseHandle(handle) + return true +} + +func flushCommandBridgeProxyOutput() { + _ = os.Stdout.Sync() + _ = os.Stderr.Sync() + // ConPTY can report the proxy process exit before consuming its final + // console write. A short drain window prevents the last command in a chain + // from losing output when cmd.exe exits immediately afterwards. + time.Sleep(15 * time.Millisecond) +} diff --git a/pkg/commands/factory.go b/pkg/commands/factory.go index ef1df105..b068258e 100644 --- a/pkg/commands/factory.go +++ b/pkg/commands/factory.go @@ -29,10 +29,11 @@ type SkillSource interface { type Deps struct { *deps.Bag - WorkDir string - BashTimeout int - SkillStore SkillSource - RunnerMode bool + WorkDir string + BashTimeout int + SkillStore SkillSource + RunnerMode bool + CommandBridge bool Provider provider.Provider ScannerProxy string @@ -96,4 +97,11 @@ func BuildPlan(plan capability.Plan, deps *Deps, reg *CommandRegistry) { } f.Build(deps, reg) } + if tool, ok := reg.GetTool("bash"); ok { + if bash, ok := tool.(*BashTool); ok && bash.CommandBridgeEnabled() { + if err := bash.SyncCommandBridgeAliases(); err != nil { + deps.GetLogger().Warnf("command bridge alias sync failed: %s", err) + } + } + } } diff --git a/pkg/commands/register.go b/pkg/commands/register.go index 367fddc2..137b0aad 100644 --- a/pkg/commands/register.go +++ b/pkg/commands/register.go @@ -33,6 +33,11 @@ func init() { bash.SetCommandNames(reg.Names) bash.SetCommandResolver(reg.Get) reg.RegisterTool(bash) + if deps.CommandBridge { + if err := bash.EnableCommandBridge(reg); err != nil { + deps.GetLogger().Warnf("command bridge disabled after startup failure: %s", err) + } + } tmuxCmd := NewTmuxCommand(bash) reg.Register(tmuxCmd, "core") diff --git a/pkg/commands/register_test.go b/pkg/commands/register_test.go index 1b625953..bf7eaf8b 100644 --- a/pkg/commands/register_test.go +++ b/pkg/commands/register_test.go @@ -29,3 +29,25 @@ func TestNativeListToolIsRunnerOnly(t *testing.T) { t.Fatal("runner mode must expose the native ls tool") } } + +func TestCommandBridgeFactoryIsOptIn(t *testing.T) { + disabled := NewRegistry() + BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir()}, disabled) + defer closeRegistryTools(disabled) + disabledTool, ok := disabled.GetTool("bash") + if !ok || disabledTool.(*BashTool).CommandBridgeEnabled() { + t.Fatal("regular factory must leave the command bridge disabled") + } + + enabled := NewRegistry() + BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir(), CommandBridge: true}, enabled) + defer closeRegistryTools(enabled) + enabledTool, ok := enabled.GetTool("bash") + if !ok { + t.Fatal("enabled factory did not register bash") + } + bash := enabledTool.(*BashTool) + if !bash.CommandBridgeEnabled() || bash.CommandBridgeError() != nil { + t.Fatalf("command bridge enabled=%v error=%v", bash.CommandBridgeEnabled(), bash.CommandBridgeError()) + } +} diff --git a/pkg/runner/app.go b/pkg/runner/app.go index b4c9cee1..29539752 100644 --- a/pkg/runner/app.go +++ b/pkg/runner/app.go @@ -123,7 +123,12 @@ func NewApp(ctx context.Context, rc ApplicationConfig) (*App, error) { a.setLLMHealth(LLMHealth{State: LLMHealthNotConfigured}) } - a.Commands = initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, a.Events, logger) + commandRegistry, err := initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, a.Events, logger) + if err != nil { + a.Close() + return nil, err + } + a.Commands = commandRegistry if rc.RecordFile != "" { if err := a.StartRecording(rc.RecordFile); err != nil { a.Close() @@ -360,12 +365,13 @@ func llmConfigLabel(providerName, model string) string { return providerName + "/" + model } -func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, events aop.EventEmitter, logger telemetry.Logger) *commands.CommandRegistry { +func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, events aop.EventEmitter, logger telemetry.Logger) (*commands.CommandRegistry, error) { cmdReg := commands.NewRegistry() workDir, _ := os.Getwd() deps := &commands.Deps{ WorkDir: workDir, BashTimeout: rc.Tools.BashTimeout, + CommandBridge: rc.Tools.CommandBridge, SkillStore: skillStore, Provider: llmProvider, Logger: logger, @@ -379,7 +385,23 @@ func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillSto OptionalTools: rc.Tools.OptionalTools, }) commands.BuildPlan(plan, deps, cmdReg) - return cmdReg + if rc.Tools.CommandBridge { + tool, ok := cmdReg.GetTool("bash") + if !ok { + return nil, fmt.Errorf("command bridge requires the bash tool") + } + bash, ok := tool.(*commands.BashTool) + if !ok { + return nil, fmt.Errorf("command bridge requires the built-in bash tool") + } + if err := bash.CommandBridgeError(); err != nil { + return nil, fmt.Errorf("start command bridge: %w", err) + } + if !bash.CommandBridgeEnabled() { + return nil, fmt.Errorf("command bridge was requested but is not active") + } + } + return cmdReg, nil } func executeRegistryCommand(ctx context.Context, reg *commands.CommandRegistry, commandLine string, timeout time.Duration) (string, error) { diff --git a/pkg/runner/app_test.go b/pkg/runner/app_test.go index 20c51d3b..2d5e45a5 100644 --- a/pkg/runner/app_test.go +++ b/pkg/runner/app_test.go @@ -16,11 +16,34 @@ import ( aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/utils/parsers" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/types/known/anypb" ) +func TestAppCloseRemovesCommandBridgeRuntime(t *testing.T) { + app, err := NewApp(context.Background(), ApplicationConfig{ + Tools: ToolConfig{CommandBridge: true}, SkipEngines: true, Logger: telemetry.NopLogger(), + }) + if err != nil { + t.Fatalf("NewApp: %v", err) + } + tool, ok := app.Commands.GetTool("bash") + if !ok { + t.Fatal("bash tool is missing") + } + runtimeDir := tool.(*commands.BashTool).CommandBridgeRuntimeDir() + if _, err := os.Stat(runtimeDir); err != nil { + t.Fatalf("command bridge runtime before close: %v", err) + } + app.Close() + app.Close() + if _, err := os.Stat(runtimeDir); !os.IsNotExist(err) { + t.Fatalf("command bridge runtime after close: %v", err) + } +} + func TestLogLLMProbeStatusReady(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/chat/completions" { diff --git a/pkg/runner/application_builder.go b/pkg/runner/application_builder.go index d3b31e4a..8d37b0f7 100644 --- a/pkg/runner/application_builder.go +++ b/pkg/runner/application_builder.go @@ -61,6 +61,7 @@ func MergeOptionExtras(rc ApplicationConfig, option *cfg.Option) ApplicationConf } rc.Scanner.UncoverCredentials = cloneStringMap(option.UncoverCredentials) rc.Tools.PlaywrightSession = option.PlaywrightSession + rc.Tools.CommandBridge = option.CommandBridge rc.CLISkillPaths = skillPathsFromOptions(option) rc.RecordFile = option.OutputFile return rc @@ -92,6 +93,7 @@ func AppConfig(option *cfg.Option, features RuntimeFeatures, logger telemetry.Lo Tools: ToolConfig{ Enabled: features.ToolsEnabled, BashTimeout: 300, + CommandBridge: option.CommandBridge, TavilyKeys: resolveTavilyKeys(option.TavilyKey, option.SearchConfig.TavilyKeys, cfg.DefaultTavilyKeys), PlaywrightSession: option.PlaywrightSession, OptionalTools: option.Tools, diff --git a/pkg/runner/application_config.go b/pkg/runner/application_config.go index ff324d80..92afaaf9 100644 --- a/pkg/runner/application_config.go +++ b/pkg/runner/application_config.go @@ -43,6 +43,7 @@ type ScannerConfig struct { type ToolConfig struct { Enabled bool BashTimeout int + CommandBridge bool TavilyKeys string PlaywrightSession string OptionalTools []string // optional tool groups to enable (e.g. "search", "browser") diff --git a/pkg/runner/provider_config_test.go b/pkg/runner/provider_config_test.go index 2c3f0aca..5fd5cb91 100644 --- a/pkg/runner/provider_config_test.go +++ b/pkg/runner/provider_config_test.go @@ -110,9 +110,10 @@ func TestMergeOptionExtrasLayersNonProtoFields(t *testing.T) { PlaywrightSession: "browser-1", UncoverCredentials: map[string]string{"SHODAN_API_KEY": "shodan-key"}, MiscOptions: cfg.MiscOptions{OutputFile: "session.jsonl"}, + AgentOptions: cfg.AgentOptions{CommandBridge: true}, } rc = MergeOptionExtras(rc, option) - if rc.Tools.PlaywrightSession != "browser-1" || rc.Scanner.UncoverCredentials["SHODAN_API_KEY"] != "shodan-key" || rc.RecordFile != "session.jsonl" { + if rc.Tools.PlaywrightSession != "browser-1" || !rc.Tools.CommandBridge || rc.Scanner.UncoverCredentials["SHODAN_API_KEY"] != "shodan-key" || rc.RecordFile != "session.jsonl" { t.Fatalf("extras = %+v", rc) } } From 7ca507c5fcdcaa205d37fb998aff5b660b727542 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Fri, 7 Aug 2026 14:11:40 +0800 Subject: [PATCH 2/5] refactor: keep command bridge as a PTY adapter --- pkg/commands/bash.go | 23 ++--------------------- pkg/commands/command_bridge_test.go | 9 +++------ pkg/commands/factory.go | 2 +- pkg/commands/register_test.go | 6 +++--- pkg/runner/app.go | 27 +++------------------------ pkg/runner/app_test.go | 23 ----------------------- 6 files changed, 12 insertions(+), 78 deletions(-) diff --git a/pkg/commands/bash.go b/pkg/commands/bash.go index c0541d2e..6d2946a4 100644 --- a/pkg/commands/bash.go +++ b/pkg/commands/bash.go @@ -46,7 +46,6 @@ type BashTool struct { commandNames func() []string resolveCommand func(string) (Command, bool) bridge *commandBridge - bridgeErr error closeOnce sync.Once } @@ -82,40 +81,22 @@ func (t *BashTool) EnableCommandBridge(registry *CommandRegistry) error { } bridge, err := newCommandBridge(registry) if err != nil { - t.bridgeErr = err return err } t.bridge = bridge if err := t.SyncCommandBridgeAliases(); err != nil { - t.bridgeErr = err bridge.close() t.bridge = nil return err } - t.bridgeErr = nil return nil } func (t *BashTool) SyncCommandBridgeAliases() error { if t.bridge == nil || t.commandNames == nil { - return t.bridgeErr - } - if err := t.bridge.syncAliases(t.commandNames()); err != nil { - t.bridgeErr = err - return err - } - return nil -} - -func (t *BashTool) CommandBridgeError() error { return t.bridgeErr } - -func (t *BashTool) CommandBridgeEnabled() bool { return t.bridge != nil } - -func (t *BashTool) CommandBridgeRuntimeDir() string { - if t.bridge == nil { - return "" + return nil } - return t.bridge.runtimeDir + return t.bridge.syncAliases(t.commandNames()) } func (t *BashTool) WithScannerProxy(proxy string) *BashTool { diff --git a/pkg/commands/command_bridge_test.go b/pkg/commands/command_bridge_test.go index e92a1aea..57db7166 100644 --- a/pkg/commands/command_bridge_test.go +++ b/pkg/commands/command_bridge_test.go @@ -97,11 +97,8 @@ func runBridgeCommand(t *testing.T, bash *BashTool, ctx context.Context, command func TestCommandBridgeDisabledByDefault(t *testing.T) { bash := NewBashTool(t.TempDir(), 5) defer bash.Close() - if bash.CommandBridgeEnabled() { - t.Fatal("command bridge must be disabled by default") - } - if bash.CommandBridgeRuntimeDir() != "" { - t.Fatalf("disabled bridge runtime dir = %q", bash.CommandBridgeRuntimeDir()) + if bash.bridge != nil { + t.Fatal("disabled bridge allocated runtime state") } } @@ -178,7 +175,7 @@ func TestCommandBridgeSyncsCommandsRegisteredLater(t *testing.T) { func TestCommandBridgeCloseCancelsCallsAndRemovesRuntime(t *testing.T) { bash, _, state := newBridgeTestBash(t) - runtimeDir := bash.CommandBridgeRuntimeDir() + runtimeDir := bash.bridge.runtimeDir if _, err := os.Stat(runtimeDir); err != nil { t.Fatalf("runtime directory before close: %v", err) } diff --git a/pkg/commands/factory.go b/pkg/commands/factory.go index b068258e..dc234104 100644 --- a/pkg/commands/factory.go +++ b/pkg/commands/factory.go @@ -98,7 +98,7 @@ func BuildPlan(plan capability.Plan, deps *Deps, reg *CommandRegistry) { f.Build(deps, reg) } if tool, ok := reg.GetTool("bash"); ok { - if bash, ok := tool.(*BashTool); ok && bash.CommandBridgeEnabled() { + if bash, ok := tool.(*BashTool); ok && bash.bridge != nil { if err := bash.SyncCommandBridgeAliases(); err != nil { deps.GetLogger().Warnf("command bridge alias sync failed: %s", err) } diff --git a/pkg/commands/register_test.go b/pkg/commands/register_test.go index bf7eaf8b..17dfe574 100644 --- a/pkg/commands/register_test.go +++ b/pkg/commands/register_test.go @@ -35,7 +35,7 @@ func TestCommandBridgeFactoryIsOptIn(t *testing.T) { BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir()}, disabled) defer closeRegistryTools(disabled) disabledTool, ok := disabled.GetTool("bash") - if !ok || disabledTool.(*BashTool).CommandBridgeEnabled() { + if !ok || disabledTool.(*BashTool).bridge != nil { t.Fatal("regular factory must leave the command bridge disabled") } @@ -47,7 +47,7 @@ func TestCommandBridgeFactoryIsOptIn(t *testing.T) { t.Fatal("enabled factory did not register bash") } bash := enabledTool.(*BashTool) - if !bash.CommandBridgeEnabled() || bash.CommandBridgeError() != nil { - t.Fatalf("command bridge enabled=%v error=%v", bash.CommandBridgeEnabled(), bash.CommandBridgeError()) + if bash.bridge == nil { + t.Fatal("command bridge was not enabled") } } diff --git a/pkg/runner/app.go b/pkg/runner/app.go index 29539752..c325e421 100644 --- a/pkg/runner/app.go +++ b/pkg/runner/app.go @@ -123,12 +123,7 @@ func NewApp(ctx context.Context, rc ApplicationConfig) (*App, error) { a.setLLMHealth(LLMHealth{State: LLMHealthNotConfigured}) } - commandRegistry, err := initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, a.Events, logger) - if err != nil { - a.Close() - return nil, err - } - a.Commands = commandRegistry + a.Commands = initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, a.Events, logger) if rc.RecordFile != "" { if err := a.StartRecording(rc.RecordFile); err != nil { a.Close() @@ -365,7 +360,7 @@ func llmConfigLabel(providerName, model string) string { return providerName + "/" + model } -func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, events aop.EventEmitter, logger telemetry.Logger) (*commands.CommandRegistry, error) { +func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, events aop.EventEmitter, logger telemetry.Logger) *commands.CommandRegistry { cmdReg := commands.NewRegistry() workDir, _ := os.Getwd() deps := &commands.Deps{ @@ -385,23 +380,7 @@ func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillSto OptionalTools: rc.Tools.OptionalTools, }) commands.BuildPlan(plan, deps, cmdReg) - if rc.Tools.CommandBridge { - tool, ok := cmdReg.GetTool("bash") - if !ok { - return nil, fmt.Errorf("command bridge requires the bash tool") - } - bash, ok := tool.(*commands.BashTool) - if !ok { - return nil, fmt.Errorf("command bridge requires the built-in bash tool") - } - if err := bash.CommandBridgeError(); err != nil { - return nil, fmt.Errorf("start command bridge: %w", err) - } - if !bash.CommandBridgeEnabled() { - return nil, fmt.Errorf("command bridge was requested but is not active") - } - } - return cmdReg, nil + return cmdReg } func executeRegistryCommand(ctx context.Context, reg *commands.CommandRegistry, commandLine string, timeout time.Duration) (string, error) { diff --git a/pkg/runner/app_test.go b/pkg/runner/app_test.go index 2d5e45a5..20c51d3b 100644 --- a/pkg/runner/app_test.go +++ b/pkg/runner/app_test.go @@ -16,34 +16,11 @@ import ( aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" - "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/utils/parsers" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/types/known/anypb" ) -func TestAppCloseRemovesCommandBridgeRuntime(t *testing.T) { - app, err := NewApp(context.Background(), ApplicationConfig{ - Tools: ToolConfig{CommandBridge: true}, SkipEngines: true, Logger: telemetry.NopLogger(), - }) - if err != nil { - t.Fatalf("NewApp: %v", err) - } - tool, ok := app.Commands.GetTool("bash") - if !ok { - t.Fatal("bash tool is missing") - } - runtimeDir := tool.(*commands.BashTool).CommandBridgeRuntimeDir() - if _, err := os.Stat(runtimeDir); err != nil { - t.Fatalf("command bridge runtime before close: %v", err) - } - app.Close() - app.Close() - if _, err := os.Stat(runtimeDir); !os.IsNotExist(err) { - t.Fatalf("command bridge runtime after close: %v", err) - } -} - func TestLogLLMProbeStatusReady(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/chat/completions" { From 5f5c093fa20c5cb6701f927da210969cda34bf4b Mon Sep 17 00:00:00 2001 From: M09Ic Date: Fri, 7 Aug 2026 14:38:50 +0800 Subject: [PATCH 3/5] fix: reclaim exited Windows bridge runtimes --- pkg/commands/command_bridge.go | 8 ++++++-- pkg/commands/command_bridge_windows.go | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pkg/commands/command_bridge.go b/pkg/commands/command_bridge.go index f5644a3f..e9feff15 100644 --- a/pkg/commands/command_bridge.go +++ b/pkg/commands/command_bridge.go @@ -407,11 +407,15 @@ func (b *commandBridge) shutdown() { func (b *commandBridge) cleanup() { b.cleanupOnce.Do(func() { - for attempt := 0; attempt < 5; attempt++ { + deadline := time.Now().Add(2 * time.Second) + for { if err := os.RemoveAll(b.runtimeDir); err == nil { break } - time.Sleep(time.Duration(attempt+1) * 20 * time.Millisecond) + if time.Now().After(deadline) { + break + } + time.Sleep(50 * time.Millisecond) } _ = os.Remove(commandBridgeRuntimeRoot()) }) diff --git a/pkg/commands/command_bridge_windows.go b/pkg/commands/command_bridge_windows.go index 1ecce4c3..4ca7ed8f 100644 --- a/pkg/commands/command_bridge_windows.go +++ b/pkg/commands/command_bridge_windows.go @@ -55,8 +55,12 @@ func commandBridgeProcessAlive(pid int) bool { if err != nil { return err == windows.ERROR_ACCESS_DENIED } - _ = windows.CloseHandle(handle) - return true + defer windows.CloseHandle(handle) + var exitCode uint32 + if err := windows.GetExitCodeProcess(handle, &exitCode); err != nil { + return true + } + return exitCode == 259 // STILL_ACTIVE } func flushCommandBridgeProxyOutput() { From 38b0893d0a2cf5acead55f62069d7fb86f4bd920 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Fri, 7 Aug 2026 15:53:39 +0800 Subject: [PATCH 4/5] refactor: make shell command adapter transparent --- cmd/aiscan/cli_test.go | 17 - cmd/aiscan/main.go | 12 - cmd/aiscan/setup.go | 1 - core/config/loader.go | 3 - core/config/options.go | 1 - pkg/commands/bash.go | 93 ++- pkg/commands/command_bridge.go | 516 ----------------- pkg/commands/command_bridge_test.go | 253 -------- pkg/commands/factory.go | 16 +- pkg/commands/register.go | 6 +- pkg/commands/register_test.go | 29 +- pkg/commands/shell_command_adapter.go | 547 ++++++++++++++++++ pkg/commands/shell_command_adapter_test.go | 247 ++++++++ ..._unix.go => shell_command_adapter_unix.go} | 16 +- ...ws.go => shell_command_adapter_windows.go} | 22 +- pkg/runner/app.go | 1 - pkg/runner/application_builder.go | 2 - pkg/runner/application_config.go | 1 - pkg/runner/provider_config_test.go | 3 +- 19 files changed, 892 insertions(+), 894 deletions(-) delete mode 100644 pkg/commands/command_bridge.go delete mode 100644 pkg/commands/command_bridge_test.go create mode 100644 pkg/commands/shell_command_adapter.go create mode 100644 pkg/commands/shell_command_adapter_test.go rename pkg/commands/{command_bridge_unix.go => shell_command_adapter_unix.go} (56%) rename pkg/commands/{command_bridge_windows.go => shell_command_adapter_windows.go} (66%) diff --git a/cmd/aiscan/cli_test.go b/cmd/aiscan/cli_test.go index f942502b..37d0fe97 100644 --- a/cmd/aiscan/cli_test.go +++ b/cmd/aiscan/cli_test.go @@ -181,23 +181,6 @@ func TestParseCLIRejectsResumeWithExplicitFile(t *testing.T) { } } -func TestAgentCommandBridgeFlagIsOptIn(t *testing.T) { - parsed, err := parseCLI([]string{"agent", "--command-bridge", "-p", "test"}) - if err != nil { - t.Fatalf("parseCLI() error = %v", err) - } - if !parsed.Option.CommandBridge { - t.Fatal("--command-bridge was not propagated to agent options") - } - parsed, err = parseCLI([]string{"agent", "-p", "test"}) - if err != nil { - t.Fatalf("parseCLI() default error = %v", err) - } - if parsed.Option.CommandBridge { - t.Fatal("command bridge must remain disabled by default") - } -} - func TestParseCLIRootTimeoutAppliesToAgent(t *testing.T) { parsed, err := parseCLI([]string{"--timeout", "45", "agent", "-p", "test"}) if err != nil { diff --git a/cmd/aiscan/main.go b/cmd/aiscan/main.go index 8c8cc63b..d787af77 100644 --- a/cmd/aiscan/main.go +++ b/cmd/aiscan/main.go @@ -1,17 +1,5 @@ package main -import ( - "os" - - "github.com/chainreactors/aiscan/pkg/commands" -) - func main() { - if code, ok := commands.RunCommandBridgeProxyIfRequested(); ok { - if code != 0 { - os.Exit(code) - } - return - } aiscan() } diff --git a/cmd/aiscan/setup.go b/cmd/aiscan/setup.go index bfd0c783..c62adae6 100644 --- a/cmd/aiscan/setup.go +++ b/cmd/aiscan/setup.go @@ -99,7 +99,6 @@ func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine deps := &commands.Deps{ WorkDir: workDir, BashTimeout: toolCfg.BashTimeout, - CommandBridge: toolCfg.CommandBridge, SkillStore: skillStore, ScannerProxy: scanCfg.Proxy, Logger: logger, diff --git a/core/config/loader.go b/core/config/loader.go index 0b2b3e3b..067e370b 100644 --- a/core/config/loader.go +++ b/core/config/loader.go @@ -134,9 +134,6 @@ func mergeOption(dst, src *Option) { if !dst.SaveSession && src.SaveSession { dst.SaveSession = true } - if !dst.CommandBridge && src.CommandBridge { - dst.CommandBridge = true - } mergeOutputOptions(&dst.OutputOptions, &src.OutputOptions) dst.DataDir = ResolveString(dst.DataDir, src.DataDir) } diff --git a/core/config/options.go b/core/config/options.go index 0b0de5f8..fcf88abc 100644 --- a/core/config/options.go +++ b/core/config/options.go @@ -89,7 +89,6 @@ type AgentOptions struct { Resume string `short:"r" long:"resume" description:"Resume agent context from an AOP JSONL session file"` SaveSession bool `long:"save-session" config:"save_session" description:"Auto-select a .aiscan/sessions/*.jsonl recording path"` CaptureProviderFrames bool `long:"capture-provider-frames" config:"capture_provider_frames" description:"Emit exact provider request/response frames as sensitive AOP events"` - CommandBridge bool `long:"command-bridge" config:"command_bridge" description:"Expose in-memory commands to shell composition through a process-local bridge"` } type AgentTransport string diff --git a/pkg/commands/bash.go b/pkg/commands/bash.go index 6d2946a4..959e8362 100644 --- a/pkg/commands/bash.go +++ b/pkg/commands/bash.go @@ -45,7 +45,9 @@ type BashTool struct { tasks *tmux.Manager commandNames func() []string resolveCommand func(string) (Command, bool) - bridge *commandBridge + shellRegistry *CommandRegistry + adapterMu sync.Mutex + shellAdapter *shellCommandAdapter closeOnce sync.Once } @@ -65,38 +67,44 @@ func (t *BashTool) SetCommandResolver(fn func(string) (Command, bool)) { func (t *BashTool) Name() string { return "bash" } func (t *BashTool) Close() { t.closeOnce.Do(func() { - if t.bridge != nil { - t.bridge.shutdown() + t.adapterMu.Lock() + adapter := t.shellAdapter + if adapter != nil { + adapter.shutdown() } + t.adapterMu.Unlock() t.tasks.Shutdown() - if t.bridge != nil { - t.bridge.cleanup() + if adapter != nil { + adapter.cleanup() } }) } -func (t *BashTool) EnableCommandBridge(registry *CommandRegistry) error { - if t.bridge != nil { - return t.SyncCommandBridgeAliases() - } - bridge, err := newCommandBridge(registry) - if err != nil { - return err - } - t.bridge = bridge - if err := t.SyncCommandBridgeAliases(); err != nil { - bridge.close() - t.bridge = nil - return err - } - return nil +func (t *BashTool) attachShellCommands(registry *CommandRegistry) { + t.shellRegistry = registry } -func (t *BashTool) SyncCommandBridgeAliases() error { - if t.bridge == nil || t.commandNames == nil { - return nil +func (t *BashTool) ensureShellCommands() (*shellCommandAdapter, error) { + if t.shellRegistry == nil { + return nil, nil + } + t.adapterMu.Lock() + defer t.adapterMu.Unlock() + if t.shellAdapter == nil { + adapter, err := newShellCommandAdapter(t.shellRegistry) + if err != nil { + return nil, err + } + t.shellAdapter = adapter + } + if t.commandNames != nil { + if err := t.shellAdapter.syncAliases(t.commandNames()); err != nil { + t.shellAdapter.close() + t.shellAdapter = nil + return nil, err + } } - return t.bridge.syncAliases(t.commandNames()) + return t.shellAdapter, nil } func (t *BashTool) WithScannerProxy(proxy string) *BashTool { @@ -245,18 +253,39 @@ func (t *BashTool) Start(ctx context.Context, command string, options BashExecOp if workDir == "" { workDir = t.workDir } - env := t.runEnv(ctx, options.Env) - if t.bridge != nil { + left, right, hasPipe := splitPipeline(command) + leftToken := firstCommandToken(left) + if !hasPipe { + if cmd, ok := t.resolve(leftToken); ok { + if tokens, err := SplitCommandLine(left); err == nil { + if args, syntaxErr := stripShellSyntax(tokens[1:]); syntaxErr == nil { + args = normalizeNoColor(cmd.Name, args) + return t.startBuiltin(ctx, cmd, args, timeout, workDir, t.runEnv(options.Env, nil, ""), options) + } + } + } + } + adapter, err := t.ensureShellCommands() + if err != nil { + return nil, err + } + if adapter != nil { + contextID := adapter.retainContext(ctx) + env := t.runEnv(options.Env, adapter, contextID) execution := newExecution(t.tasks, command, nil, workDir, env) info, err := t.tasks.Create(workDir, command, options.Name, timeout, env, "") if err != nil { + adapter.releaseContext(contextID) return nil, err } execution.bind(info) + go func() { + <-t.tasks.Done(execution.ID) + adapter.releaseContext(contextID) + }() return execution, nil } - left, right, hasPipe := splitPipeline(command) - leftToken := firstCommandToken(left) + env := t.runEnv(options.Env, nil, "") if cmd, ok := t.resolve(leftToken); ok { tokens, err := SplitCommandLine(left) if err != nil { @@ -493,7 +522,7 @@ func (t *BashTool) collectResult(execution *Execution) *coretool.Result { return result } -func (t *BashTool) runEnv(ctx context.Context, overrides map[string]string) []string { +func (t *BashTool) runEnv(overrides map[string]string, adapter *shellCommandAdapter, shellContextID string) []string { values := make(map[string]string) for _, item := range t.proxyEnv() { if key, value, ok := strings.Cut(item, "="); ok { @@ -503,8 +532,8 @@ func (t *BashTool) runEnv(ctx context.Context, overrides map[string]string) []st for key, value := range overrides { values[key] = value } - if t.bridge != nil { - for _, item := range t.bridge.environment(coretool.InvocationFromContext(ctx)) { + if adapter != nil && shellContextID != "" { + for _, item := range adapter.environment(shellContextID) { if key, value, ok := strings.Cut(item, "="); ok { values[key] = value } @@ -513,7 +542,7 @@ func (t *BashTool) runEnv(ctx context.Context, overrides map[string]string) []st if path == "" { path = os.Getenv("PATH") } - values["PATH"] = t.bridge.runtimeDir + string(os.PathListSeparator) + path + values["PATH"] = adapter.runtimeDir + string(os.PathListSeparator) + path } keys := make([]string, 0, len(values)) for key := range values { diff --git a/pkg/commands/command_bridge.go b/pkg/commands/command_bridge.go deleted file mode 100644 index e9feff15..00000000 --- a/pkg/commands/command_bridge.go +++ /dev/null @@ -1,516 +0,0 @@ -package commands - -import ( - "bufio" - "context" - "encoding/binary" - "encoding/json" - "errors" - "fmt" - "io" - "net" - "os" - "path/filepath" - "sort" - "strconv" - "strings" - "sync" - "time" - - coretool "github.com/chainreactors/aiscan/core/tool" -) - -const ( - commandBridgeMarkerEnv = "AISCAN_COMMAND_BRIDGE" - commandBridgeEndpointEnv = "AISCAN_COMMAND_BRIDGE_ENDPOINT" - commandBridgeCommandEnv = "AISCAN_COMMAND_BRIDGE_COMMAND" - commandBridgeExecutableEnv = "AISCAN_COMMAND_BRIDGE_EXECUTABLE" - commandBridgeCallIDEnv = "AISCAN_COMMAND_BRIDGE_CALL_ID" - commandBridgeSessionEnv = "AISCAN_COMMAND_BRIDGE_SESSION_ID" - commandBridgeTurnEnv = "AISCAN_COMMAND_BRIDGE_TURN_ID" - commandBridgeEmitterEnv = "AISCAN_COMMAND_BRIDGE_EMITTER" - - commandBridgeProtocolVersion = 1 - commandBridgeChunkSize = 32 << 10 - commandBridgeMaxFrameSize = 1 << 20 - commandBridgeDialTimeout = 5 * time.Second -) - -type commandBridgeFrame struct { - Type string `json:"type"` - Version int `json:"version,omitempty"` - Command string `json:"command,omitempty"` - Args []string `json:"args,omitempty"` - Dir string `json:"dir,omitempty"` - Invocation coretool.Invocation `json:"invocation,omitempty"` - Data []byte `json:"data,omitempty"` - ExitCode int `json:"exit_code,omitempty"` -} - -type commandBridge struct { - registry *CommandRegistry - executable string - runtimeDir string - endpoint string - listener net.Listener - cancel context.CancelFunc - - mu sync.Mutex - aliases map[string]string - connections map[net.Conn]struct{} - wg sync.WaitGroup - shutdownOnce sync.Once - cleanupOnce sync.Once -} - -func newCommandBridge(registry *CommandRegistry) (*commandBridge, error) { - if registry == nil { - return nil, fmt.Errorf("command bridge requires a registry") - } - executable, err := os.Executable() - if err != nil { - return nil, fmt.Errorf("resolve command bridge executable: %w", err) - } - executable, err = filepath.Abs(executable) - if err != nil { - return nil, fmt.Errorf("resolve command bridge executable path: %w", err) - } - if err := cleanupStaleCommandBridgeRuntime(); err != nil { - return nil, err - } - - root := commandBridgeRuntimeRoot() - if err := os.MkdirAll(root, 0o700); err != nil { - return nil, fmt.Errorf("create command bridge runtime root: %w", err) - } - _ = os.Chmod(root, 0o700) - runtimeDir, err := os.MkdirTemp(root, strconv.Itoa(os.Getpid())+"-") - if err != nil { - return nil, fmt.Errorf("create command bridge runtime: %w", err) - } - cleanup := func() { _ = os.RemoveAll(runtimeDir) } - _ = os.Chmod(runtimeDir, 0o700) - endpoint := commandBridgeEndpoint(runtimeDir) - listener, err := listenCommandBridge(endpoint) - if err != nil { - cleanup() - return nil, fmt.Errorf("listen command bridge: %w", err) - } - ctx, cancel := context.WithCancel(context.Background()) - bridge := &commandBridge{ - registry: registry, executable: executable, runtimeDir: runtimeDir, - endpoint: endpoint, listener: listener, cancel: cancel, - aliases: make(map[string]string), connections: make(map[net.Conn]struct{}), - } - bridge.wg.Add(1) - go bridge.accept(ctx) - return bridge, nil -} - -func commandBridgeRuntimeRoot() string { - return filepath.Join(os.TempDir(), "aiscan-command-bridge") -} - -func cleanupStaleCommandBridgeRuntime() error { - root := commandBridgeRuntimeRoot() - entries, err := os.ReadDir(root) - if errors.Is(err, os.ErrNotExist) { - return nil - } - if err != nil { - return fmt.Errorf("read command bridge runtime root: %w", err) - } - for _, entry := range entries { - if !entry.IsDir() { - continue - } - pidText, _, ok := strings.Cut(entry.Name(), "-") - pid, parseErr := strconv.Atoi(pidText) - if !ok || parseErr != nil || pid <= 0 { - continue - } - if !commandBridgeProcessAlive(pid) { - _ = os.RemoveAll(filepath.Join(root, entry.Name())) - } - } - return nil -} - -func (b *commandBridge) accept(ctx context.Context) { - defer b.wg.Done() - for { - conn, err := b.listener.Accept() - if err != nil { - if ctx.Err() != nil { - return - } - continue - } - b.mu.Lock() - b.connections[conn] = struct{}{} - b.wg.Add(1) - b.mu.Unlock() - go b.handle(ctx, conn) - } -} - -func (b *commandBridge) handle(parent context.Context, conn net.Conn) { - defer b.wg.Done() - defer func() { - _ = conn.Close() - b.mu.Lock() - delete(b.connections, conn) - b.mu.Unlock() - }() - - reader := bufio.NewReader(conn) - header, err := readCommandBridgeFrame(reader) - if err != nil { - return - } - writer := &commandBridgeFrameWriter{writer: conn} - if header.Type != "request" || header.Version != commandBridgeProtocolVersion { - _, _ = io.WriteString(&commandBridgeStreamWriter{writer: writer, frameType: "stderr"}, "unsupported command bridge protocol\n") - _ = writer.write(commandBridgeFrame{Type: "final", ExitCode: 125}) - return - } - command, ok := b.registry.Get(header.Command) - if !ok || command.Run == nil { - message := "unknown in-memory command: " + header.Command - _, _ = io.WriteString(&commandBridgeStreamWriter{writer: writer, frameType: "stderr"}, message+"\n") - _ = writer.write(commandBridgeFrame{Type: "final", ExitCode: 127}) - return - } - - ctx, cancel := context.WithCancel(parent) - defer cancel() - ctx = coretool.ContextWithInvocation(ctx, header.Invocation) - stdinReader, stdinWriter := io.Pipe() - defer stdinReader.Close() - readDone := make(chan error, 1) - go func() { - readDone <- readCommandBridgeInput(reader, stdinWriter, cancel) - }() - - stdout := &commandBridgeStreamWriter{writer: writer, frameType: "stdout"} - stderr := &commandBridgeStreamWriter{writer: writer, frameType: "stderr"} - execution := newExecution(nil, command.Name, normalizeNoColor(command.Name, header.Args), header.Dir, nil) - execution.ID = header.Invocation.CallID - execution.StartedAt = time.Now() - execution.setIO(stdinReader, stdout, stderr) - _, runErr := command.Run(ctx, execution) - execution.EndedAt = time.Now() - _ = stdinReader.Close() - cancel() - - exitCode := commandBridgeExitCode(runErr) - if runErr != nil && !errors.Is(runErr, context.Canceled) && !errors.Is(runErr, context.DeadlineExceeded) { - _, _ = io.WriteString(stderr, runErr.Error()+"\n") - } - _ = writer.write(commandBridgeFrame{Type: "final", ExitCode: exitCode}) - _ = conn.Close() - select { - case <-readDone: - case <-time.After(time.Second): - } -} - -func readCommandBridgeInput(reader *bufio.Reader, stdin *io.PipeWriter, cancel context.CancelFunc) error { - stdinOpen := true - closeStdin := func() { - if stdinOpen { - stdinOpen = false - _ = stdin.Close() - } - } - defer closeStdin() - for { - frame, err := readCommandBridgeFrame(reader) - if err != nil { - cancel() - return err - } - switch frame.Type { - case "stdin": - if stdinOpen && len(frame.Data) > 0 { - if _, err := stdin.Write(frame.Data); err != nil { - return err - } - } - case "stdin_eof": - closeStdin() - case "cancel": - cancel() - return context.Canceled - default: - cancel() - return fmt.Errorf("unexpected command bridge frame %q", frame.Type) - } - } -} - -type commandBridgeFrameWriter struct { - mu sync.Mutex - writer io.Writer -} - -func (w *commandBridgeFrameWriter) write(frame commandBridgeFrame) error { - w.mu.Lock() - defer w.mu.Unlock() - return writeCommandBridgeFrame(w.writer, frame) -} - -type commandBridgeStreamWriter struct { - writer *commandBridgeFrameWriter - frameType string -} - -func (w *commandBridgeStreamWriter) Write(data []byte) (int, error) { - for offset := 0; offset < len(data); { - end := offset + commandBridgeChunkSize - if end > len(data) { - end = len(data) - } - if err := w.writer.write(commandBridgeFrame{Type: w.frameType, Data: data[offset:end]}); err != nil { - return offset, err - } - offset = end - } - return len(data), nil -} - -func commandBridgeExitCode(err error) int { - if err == nil { - return 0 - } - var exitCoder interface{ ExitCode() int } - if errors.As(err, &exitCoder) { - code := exitCoder.ExitCode() - if code >= 0 && code <= 255 { - return code - } - } - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - return 130 - } - return 1 -} - -func writeCommandBridgeFrame(writer io.Writer, frame commandBridgeFrame) error { - data, err := json.Marshal(frame) - if err != nil { - return err - } - if len(data) > commandBridgeMaxFrameSize { - return fmt.Errorf("command bridge frame too large: %d", len(data)) - } - var size [4]byte - binary.BigEndian.PutUint32(size[:], uint32(len(data))) - if err := writeCommandBridgeBytes(writer, size[:]); err != nil { - return err - } - return writeCommandBridgeBytes(writer, data) -} - -func writeCommandBridgeBytes(writer io.Writer, data []byte) error { - for len(data) > 0 { - n, err := writer.Write(data) - if err != nil { - return err - } - if n <= 0 { - return io.ErrShortWrite - } - data = data[n:] - } - return nil -} - -func readCommandBridgeFrame(reader io.Reader) (commandBridgeFrame, error) { - var size [4]byte - if _, err := io.ReadFull(reader, size[:]); err != nil { - return commandBridgeFrame{}, err - } - length := binary.BigEndian.Uint32(size[:]) - if length == 0 || length > commandBridgeMaxFrameSize { - return commandBridgeFrame{}, fmt.Errorf("invalid command bridge frame size: %d", length) - } - data := make([]byte, int(length)) - if _, err := io.ReadFull(reader, data); err != nil { - return commandBridgeFrame{}, err - } - var frame commandBridgeFrame - if err := json.Unmarshal(data, &frame); err != nil { - return commandBridgeFrame{}, err - } - return frame, nil -} - -func (b *commandBridge) syncAliases(names []string) error { - b.mu.Lock() - defer b.mu.Unlock() - sorted := append([]string(nil), names...) - sort.Strings(sorted) - for _, name := range sorted { - if !validCommandBridgeName(name) { - continue - } - if _, ok := b.aliases[name]; ok { - continue - } - path, err := createCommandBridgeAlias(b.executable, b.runtimeDir, name) - if err != nil { - return fmt.Errorf("create command bridge alias %s: %w", name, err) - } - b.aliases[name] = path - } - return nil -} - -func validCommandBridgeName(name string) bool { - if name == "" || name == "." || name == ".." { - return false - } - for _, r := range name { - if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' { - continue - } - return false - } - return true -} - -func (b *commandBridge) environment(invocation coretool.Invocation) []string { - return []string{ - commandBridgeMarkerEnv + "=1", - commandBridgeEndpointEnv + "=" + b.endpoint, - commandBridgeExecutableEnv + "=" + b.executable, - commandBridgeCallIDEnv + "=" + invocation.CallID, - commandBridgeSessionEnv + "=" + invocation.SessionID, - commandBridgeTurnEnv + "=" + invocation.TurnID, - commandBridgeEmitterEnv + "=" + invocation.Emitter, - } -} - -func (b *commandBridge) shutdown() { - b.shutdownOnce.Do(func() { - b.cancel() - _ = b.listener.Close() - b.mu.Lock() - for conn := range b.connections { - _ = conn.Close() - } - b.mu.Unlock() - b.wg.Wait() - }) -} - -func (b *commandBridge) cleanup() { - b.cleanupOnce.Do(func() { - deadline := time.Now().Add(2 * time.Second) - for { - if err := os.RemoveAll(b.runtimeDir); err == nil { - break - } - if time.Now().After(deadline) { - break - } - time.Sleep(50 * time.Millisecond) - } - _ = os.Remove(commandBridgeRuntimeRoot()) - }) -} - -func (b *commandBridge) close() { - b.shutdown() - b.cleanup() -} - -// RunCommandBridgeProxyIfRequested detects a bridge wrapper invocation before -// normal CLI parsing. Callers should return normally for code 0 and os.Exit for -// a non-zero code when ok is true. -func RunCommandBridgeProxyIfRequested() (code int, ok bool) { - command := strings.TrimSpace(os.Getenv(commandBridgeCommandEnv)) - if os.Getenv(commandBridgeMarkerEnv) != "1" || command == "" { - return 0, false - } - if !validCommandBridgeName(command) { - fmt.Fprintln(os.Stderr, "invalid command bridge alias") - return 126, true - } - return runCommandBridgeProxy(command, os.Args[1:]), true -} - -func runCommandBridgeProxy(command string, args []string) int { - endpoint := os.Getenv(commandBridgeEndpointEnv) - if endpoint == "" { - fmt.Fprintln(os.Stderr, "AIScan command bridge environment is incomplete") - return 126 - } - ctx, cancel := context.WithTimeout(context.Background(), commandBridgeDialTimeout) - conn, err := dialCommandBridge(ctx, endpoint) - cancel() - if err != nil { - fmt.Fprintf(os.Stderr, "connect AIScan command bridge: %s\n", err) - return 125 - } - defer conn.Close() - cwd, _ := os.Getwd() - header := commandBridgeFrame{ - Type: "request", Version: commandBridgeProtocolVersion, - Command: command, Args: append([]string(nil), args...), Dir: cwd, - Invocation: coretool.Invocation{ - WorkDir: cwd, CallID: os.Getenv(commandBridgeCallIDEnv), - SessionID: os.Getenv(commandBridgeSessionEnv), TurnID: os.Getenv(commandBridgeTurnEnv), - Emitter: os.Getenv(commandBridgeEmitterEnv), - }, - } - writer := &commandBridgeFrameWriter{writer: conn} - if err := writer.write(header); err != nil { - fmt.Fprintf(os.Stderr, "start AIScan command bridge request: %s\n", err) - return 125 - } - go streamCommandBridgeStdin(writer) - - reader := bufio.NewReader(conn) - for { - frame, err := readCommandBridgeFrame(reader) - if err != nil { - fmt.Fprintf(os.Stderr, "read AIScan command bridge response: %s\n", err) - return 125 - } - switch frame.Type { - case "stdout": - _, _ = os.Stdout.Write(frame.Data) - case "stderr": - _, _ = os.Stderr.Write(frame.Data) - case "final": - flushCommandBridgeProxyOutput() - return frame.ExitCode - default: - fmt.Fprintf(os.Stderr, "unexpected AIScan command bridge response %q\n", frame.Type) - return 125 - } - } -} - -func streamCommandBridgeStdin(writer *commandBridgeFrameWriter) { - info, err := os.Stdin.Stat() - if err != nil || info.Mode()&os.ModeCharDevice != 0 { - _ = writer.write(commandBridgeFrame{Type: "stdin_eof"}) - return - } - buffer := make([]byte, commandBridgeChunkSize) - for { - n, readErr := os.Stdin.Read(buffer) - if n > 0 { - if err := writer.write(commandBridgeFrame{Type: "stdin", Data: buffer[:n]}); err != nil { - return - } - } - if readErr != nil { - _ = writer.write(commandBridgeFrame{Type: "stdin_eof"}) - return - } - } -} diff --git a/pkg/commands/command_bridge_test.go b/pkg/commands/command_bridge_test.go deleted file mode 100644 index 57db7166..00000000 --- a/pkg/commands/command_bridge_test.go +++ /dev/null @@ -1,253 +0,0 @@ -package commands - -import ( - "bytes" - "context" - "errors" - "fmt" - "io" - "os" - "path/filepath" - "strings" - "sync" - "testing" - "time" - - coretool "github.com/chainreactors/aiscan/core/tool" -) - -func TestMain(m *testing.M) { - if code, ok := RunCommandBridgeProxyIfRequested(); ok { - if code != 0 { - os.Exit(code) - } - return - } - os.Exit(m.Run()) -} - -type bridgeTestExitError struct{ code int } - -func (e bridgeTestExitError) Error() string { return fmt.Sprintf("bridge test exit %d", e.code) } -func (e bridgeTestExitError) ExitCode() int { return e.code } - -type bridgeTestCommands struct { - started chan struct{} - canceled chan struct{} - once sync.Once -} - -func newBridgeTestBash(t *testing.T) (*BashTool, *CommandRegistry, *bridgeTestCommands) { - t.Helper() - state := &bridgeTestCommands{started: make(chan struct{}), canceled: make(chan struct{})} - registry := NewRegistry() - registry.Register(Command{Name: "bridge_echo", Run: func(_ context.Context, execution *Execution) (any, error) { - fmt.Fprintln(execution.Stdout, strings.Join(execution.Args, " ")) - return nil, nil - }}, "test") - registry.Register(Command{Name: "bridge_upper", Run: func(_ context.Context, execution *Execution) (any, error) { - data, err := io.ReadAll(execution.Stdin) - if err != nil { - return nil, err - } - _, err = execution.Stdout.Write(bytes.ToUpper(data)) - return nil, err - }}, "test") - registry.Register(Command{Name: "bridge_fail", Run: func(context.Context, *Execution) (any, error) { - return nil, bridgeTestExitError{code: 7} - }}, "test") - registry.Register(Command{Name: "bridge_context", Run: func(ctx context.Context, execution *Execution) (any, error) { - invocation := coretool.InvocationFromContext(ctx) - fmt.Fprintf(execution.Stdout, "dir=%s call=%s session=%s turn=%s emitter=%s\n", - execution.Dir, invocation.CallID, invocation.SessionID, invocation.TurnID, invocation.Emitter) - return nil, nil - }}, "test") - registry.Register(Command{Name: "bridge_wait", Run: func(ctx context.Context, _ *Execution) (any, error) { - state.once.Do(func() { close(state.started) }) - <-ctx.Done() - close(state.canceled) - return nil, ctx.Err() - }}, "test") - - bash := NewBashTool(t.TempDir(), 10) - bash.SetCommandNames(registry.Names) - bash.SetCommandResolver(registry.Get) - if err := bash.EnableCommandBridge(registry); err != nil { - t.Fatalf("EnableCommandBridge: %v", err) - } - t.Cleanup(bash.Close) - return bash, registry, state -} - -func runBridgeCommand(t *testing.T, bash *BashTool, ctx context.Context, command string, workDir string) (*Execution, string) { - t.Helper() - var output strings.Builder - execution, err := bash.RunForeground(ctx, command, BashExecOptions{ - WorkDir: workDir, - OnOutput: func(data []byte) { - _, _ = output.Write(data) - }, - }) - if err != nil { - t.Fatalf("RunForeground(%q): %v", command, err) - } - return execution, output.String() -} - -func TestCommandBridgeDisabledByDefault(t *testing.T) { - bash := NewBashTool(t.TempDir(), 5) - defer bash.Close() - if bash.bridge != nil { - t.Fatal("disabled bridge allocated runtime state") - } -} - -func TestCommandBridgeMarkerDoesNotHijackNormalChildProcess(t *testing.T) { - t.Setenv(commandBridgeMarkerEnv, "1") - t.Setenv(commandBridgeCommandEnv, "") - if code, ok := RunCommandBridgeProxyIfRequested(); ok { - t.Fatalf("marker-only child was treated as proxy with code %d", code) - } -} - -func TestCommandBridgeShellComposition(t *testing.T) { - bash, _, _ := newBridgeTestBash(t) - workDir := t.TempDir() - - execution, output := runBridgeCommand(t, bash, context.Background(), "bridge_echo one && bridge_echo two", workDir) - if execution.ExitCode != 0 || !strings.Contains(output, "one") || !strings.Contains(output, "two") { - t.Fatalf("and composition exit=%d output=%q", execution.ExitCode, output) - } - - execution, output = runBridgeCommand(t, bash, context.Background(), "bridge_fail || bridge_echo recovered", workDir) - if execution.ExitCode != 0 || !strings.Contains(output, "recovered") { - t.Fatalf("or composition exit=%d output=%q", execution.ExitCode, output) - } - - execution, output = runBridgeCommand(t, bash, context.Background(), "bridge_echo hello | bridge_upper", workDir) - if execution.ExitCode != 0 || !strings.Contains(output, "HELLO") { - t.Fatalf("pipeline exit=%d output=%q", execution.ExitCode, output) - } - - execution, _ = runBridgeCommand(t, bash, context.Background(), "bridge_fail && bridge_echo unreachable", workDir) - if execution.ExitCode != 7 { - t.Fatalf("short-circuit exit code = %d, want 7", execution.ExitCode) - } -} - -func TestCommandBridgeRedirectionAndInvocationContext(t *testing.T) { - bash, _, _ := newBridgeTestBash(t) - workDir := t.TempDir() - ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{ - WorkDir: workDir, CallID: "call-1", SessionID: "session-1", TurnID: "turn-1", Emitter: "runner", - }) - - execution, _ := runBridgeCommand(t, bash, ctx, "bridge_context > bridge-context.txt", workDir) - if execution.ExitCode != 0 { - t.Fatalf("redirection exit code = %d", execution.ExitCode) - } - data, err := os.ReadFile(filepath.Join(workDir, "bridge-context.txt")) - if err != nil { - t.Fatalf("read redirected output: %v", err) - } - text := string(data) - for _, expected := range []string{"dir=" + workDir, "call=call-1", "session=session-1", "turn=turn-1", "emitter=runner"} { - if !strings.Contains(text, expected) { - t.Fatalf("redirected context missing %q: %q", expected, text) - } - } -} - -func TestCommandBridgeSyncsCommandsRegisteredLater(t *testing.T) { - bash, registry, _ := newBridgeTestBash(t) - registry.Register(Command{Name: "bridge_late", Run: func(_ context.Context, execution *Execution) (any, error) { - fmt.Fprintln(execution.Stdout, "late") - return nil, nil - }}, "test") - if err := bash.SyncCommandBridgeAliases(); err != nil { - t.Fatalf("SyncCommandBridgeAliases: %v", err) - } - execution, output := runBridgeCommand(t, bash, context.Background(), "bridge_late && bridge_echo ready", t.TempDir()) - if execution.ExitCode != 0 || !strings.Contains(output, "late") || !strings.Contains(output, "ready") { - t.Fatalf("late alias exit=%d output=%q", execution.ExitCode, output) - } -} - -func TestCommandBridgeCloseCancelsCallsAndRemovesRuntime(t *testing.T) { - bash, _, state := newBridgeTestBash(t) - runtimeDir := bash.bridge.runtimeDir - if _, err := os.Stat(runtimeDir); err != nil { - t.Fatalf("runtime directory before close: %v", err) - } - - execution, err := bash.Start(context.Background(), "bridge_wait", BashExecOptions{WorkDir: t.TempDir()}) - if err != nil { - t.Fatalf("Start bridge_wait: %v", err) - } - select { - case <-state.started: - case <-time.After(5 * time.Second): - t.Fatal("bridge command did not start") - } - bash.Close() - bash.Close() - select { - case <-state.canceled: - case <-time.After(5 * time.Second): - t.Fatal("bridge command was not canceled") - } - if err := execution.Wait(context.Background()); err != nil && !errors.Is(err, context.Canceled) { - t.Fatalf("wait after close: %v", err) - } - if _, err := os.Stat(runtimeDir); !errors.Is(err, os.ErrNotExist) { - t.Fatalf("runtime directory still exists after close: %v", err) - } -} - -func TestCommandBridgeDisconnectCancelsRunningCommand(t *testing.T) { - bash, _, state := newBridgeTestBash(t) - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - conn, err := dialCommandBridge(ctx, bash.bridge.endpoint) - if err != nil { - t.Fatalf("dial command bridge: %v", err) - } - writer := &commandBridgeFrameWriter{writer: conn} - if err := writer.write(commandBridgeFrame{ - Type: "request", Version: commandBridgeProtocolVersion, - Command: "bridge_wait", Dir: t.TempDir(), - }); err != nil { - t.Fatalf("write request: %v", err) - } - if err := writer.write(commandBridgeFrame{Type: "stdin_eof"}); err != nil { - t.Fatalf("write stdin eof: %v", err) - } - select { - case <-state.started: - case <-ctx.Done(): - t.Fatal("bridge command did not start") - } - _ = conn.Close() - select { - case <-state.canceled: - case <-ctx.Done(): - t.Fatal("disconnect did not cancel bridge command") - } -} - -func TestCommandBridgeStartupReclaimsOwnedStaleRuntime(t *testing.T) { - root := commandBridgeRuntimeRoot() - if err := os.MkdirAll(root, 0o700); err != nil { - t.Fatal(err) - } - stale, err := os.MkdirTemp(root, "1073741824-stale-") - if err != nil { - t.Fatal(err) - } - if err := cleanupStaleCommandBridgeRuntime(); err != nil { - t.Fatalf("cleanup stale runtime: %v", err) - } - if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) { - t.Fatalf("stale runtime still exists: %v", err) - } -} diff --git a/pkg/commands/factory.go b/pkg/commands/factory.go index dc234104..ef1df105 100644 --- a/pkg/commands/factory.go +++ b/pkg/commands/factory.go @@ -29,11 +29,10 @@ type SkillSource interface { type Deps struct { *deps.Bag - WorkDir string - BashTimeout int - SkillStore SkillSource - RunnerMode bool - CommandBridge bool + WorkDir string + BashTimeout int + SkillStore SkillSource + RunnerMode bool Provider provider.Provider ScannerProxy string @@ -97,11 +96,4 @@ func BuildPlan(plan capability.Plan, deps *Deps, reg *CommandRegistry) { } f.Build(deps, reg) } - if tool, ok := reg.GetTool("bash"); ok { - if bash, ok := tool.(*BashTool); ok && bash.bridge != nil { - if err := bash.SyncCommandBridgeAliases(); err != nil { - deps.GetLogger().Warnf("command bridge alias sync failed: %s", err) - } - } - } } diff --git a/pkg/commands/register.go b/pkg/commands/register.go index 137b0aad..4158edc4 100644 --- a/pkg/commands/register.go +++ b/pkg/commands/register.go @@ -33,11 +33,7 @@ func init() { bash.SetCommandNames(reg.Names) bash.SetCommandResolver(reg.Get) reg.RegisterTool(bash) - if deps.CommandBridge { - if err := bash.EnableCommandBridge(reg); err != nil { - deps.GetLogger().Warnf("command bridge disabled after startup failure: %s", err) - } - } + bash.attachShellCommands(reg) tmuxCmd := NewTmuxCommand(bash) reg.Register(tmuxCmd, "core") diff --git a/pkg/commands/register_test.go b/pkg/commands/register_test.go index 17dfe574..d8de8900 100644 --- a/pkg/commands/register_test.go +++ b/pkg/commands/register_test.go @@ -30,24 +30,19 @@ func TestNativeListToolIsRunnerOnly(t *testing.T) { } } -func TestCommandBridgeFactoryIsOptIn(t *testing.T) { - disabled := NewRegistry() - BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir()}, disabled) - defer closeRegistryTools(disabled) - disabledTool, ok := disabled.GetTool("bash") - if !ok || disabledTool.(*BashTool).bridge != nil { - t.Fatal("regular factory must leave the command bridge disabled") - } - - enabled := NewRegistry() - BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir(), CommandBridge: true}, enabled) - defer closeRegistryTools(enabled) - enabledTool, ok := enabled.GetTool("bash") +func TestCoreFactoryAttachesRegisteredCommandsToShell(t *testing.T) { + registry := NewRegistry() + BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir()}, registry) + defer closeRegistryTools(registry) + tool, ok := registry.GetTool("bash") if !ok { - t.Fatal("enabled factory did not register bash") + t.Fatal("core factory did not register bash") + } + bash := tool.(*BashTool) + if bash.shellRegistry != registry { + t.Fatal("registered commands were not attached to bash") } - bash := enabledTool.(*BashTool) - if bash.bridge == nil { - t.Fatal("command bridge was not enabled") + if bash.shellAdapter != nil { + t.Fatal("shell adapter must remain lazy until a real shell is needed") } } diff --git a/pkg/commands/shell_command_adapter.go b/pkg/commands/shell_command_adapter.go new file mode 100644 index 00000000..c7d0c754 --- /dev/null +++ b/pkg/commands/shell_command_adapter.go @@ -0,0 +1,547 @@ +package commands + +import ( + "bufio" + "context" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +const ( + shellCommandAdapterMarkerEnv = "AISCAN_SHELL_COMMAND" + shellCommandAdapterEndpointEnv = "AISCAN_SHELL_COMMAND_ENDPOINT" + shellCommandAdapterCommandEnv = "AISCAN_SHELL_COMMAND_COMMAND" + shellCommandAdapterExecutableEnv = "AISCAN_SHELL_COMMAND_EXECUTABLE" + shellCommandAdapterContextEnv = "AISCAN_SHELL_COMMAND_CONTEXT" + + shellCommandAdapterProtocolVersion = 1 + shellCommandAdapterChunkSize = 32 << 10 + shellCommandAdapterMaxFrameSize = 1 << 20 + shellCommandAdapterDialTimeout = 5 * time.Second +) + +func init() { + if code, ok := runShellCommandProxyIfRequested(); ok { + os.Exit(code) + } +} + +type shellCommandAdapterFrame struct { + Type string `json:"type"` + Version int `json:"version,omitempty"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Dir string `json:"dir,omitempty"` + ContextID string `json:"context_id,omitempty"` + Data []byte `json:"data,omitempty"` + ExitCode int `json:"exit_code,omitempty"` +} + +type shellCommandAdapter struct { + registry *CommandRegistry + executable string + runtimeDir string + endpoint string + listener net.Listener + cancel context.CancelFunc + + mu sync.Mutex + aliases map[string]string + contexts map[string]context.Context + nextContext uint64 + connections map[net.Conn]struct{} + wg sync.WaitGroup + shutdownOnce sync.Once + cleanupOnce sync.Once +} + +func newShellCommandAdapter(registry *CommandRegistry) (*shellCommandAdapter, error) { + if registry == nil { + return nil, fmt.Errorf("shell command adapter requires a registry") + } + executable, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("resolve shell command executable: %w", err) + } + executable, err = filepath.Abs(executable) + if err != nil { + return nil, fmt.Errorf("resolve shell command executable path: %w", err) + } + if err := cleanupStaleShellCommandAdapterRuntime(); err != nil { + return nil, err + } + + root := shellCommandAdapterRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + return nil, fmt.Errorf("create shell command runtime root: %w", err) + } + _ = os.Chmod(root, 0o700) + runtimeDir, err := os.MkdirTemp(root, strconv.Itoa(os.Getpid())+"-") + if err != nil { + return nil, fmt.Errorf("create shell command runtime: %w", err) + } + cleanup := func() { _ = os.RemoveAll(runtimeDir) } + _ = os.Chmod(runtimeDir, 0o700) + endpoint := shellCommandAdapterEndpoint(runtimeDir) + listener, err := listenShellCommandAdapter(endpoint) + if err != nil { + cleanup() + return nil, fmt.Errorf("listen for shell commands: %w", err) + } + ctx, cancel := context.WithCancel(context.Background()) + adapter := &shellCommandAdapter{ + registry: registry, executable: executable, runtimeDir: runtimeDir, + endpoint: endpoint, listener: listener, cancel: cancel, + aliases: make(map[string]string), contexts: make(map[string]context.Context), + connections: make(map[net.Conn]struct{}), + } + adapter.wg.Add(1) + go adapter.accept(ctx) + return adapter, nil +} + +func shellCommandAdapterRuntimeRoot() string { + return filepath.Join(os.TempDir(), "aiscan-shell-commands") +} + +func cleanupStaleShellCommandAdapterRuntime() error { + root := shellCommandAdapterRuntimeRoot() + entries, err := os.ReadDir(root) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return fmt.Errorf("read shell command runtime root: %w", err) + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + pidText, _, ok := strings.Cut(entry.Name(), "-") + pid, parseErr := strconv.Atoi(pidText) + if !ok || parseErr != nil || pid <= 0 { + continue + } + if !shellCommandAdapterProcessAlive(pid) { + _ = os.RemoveAll(filepath.Join(root, entry.Name())) + } + } + return nil +} + +func (b *shellCommandAdapter) accept(ctx context.Context) { + defer b.wg.Done() + for { + conn, err := b.listener.Accept() + if err != nil { + if ctx.Err() != nil { + return + } + continue + } + b.mu.Lock() + b.connections[conn] = struct{}{} + b.wg.Add(1) + b.mu.Unlock() + go b.handle(ctx, conn) + } +} + +func (b *shellCommandAdapter) handle(parent context.Context, conn net.Conn) { + defer b.wg.Done() + defer func() { + _ = conn.Close() + b.mu.Lock() + delete(b.connections, conn) + b.mu.Unlock() + }() + + reader := bufio.NewReader(conn) + header, err := readShellCommandAdapterFrame(reader) + if err != nil { + return + } + writer := &shellCommandAdapterFrameWriter{writer: conn} + if header.Type != "request" || header.Version != shellCommandAdapterProtocolVersion { + _, _ = io.WriteString(&shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"}, "unsupported shell command protocol\n") + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: 125}) + return + } + command, ok := b.registry.Get(header.Command) + if !ok || command.Run == nil { + message := "unknown in-memory command: " + header.Command + _, _ = io.WriteString(&shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"}, message+"\n") + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: 127}) + return + } + baseCtx, ok := b.context(header.ContextID) + if !ok { + _, _ = io.WriteString(&shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"}, "shell command context expired\n") + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: 125}) + return + } + + ctx, cancel := context.WithCancel(baseCtx) + defer cancel() + stopParent := context.AfterFunc(parent, cancel) + defer stopParent() + stdinReader, stdinWriter := io.Pipe() + defer stdinReader.Close() + readDone := make(chan error, 1) + go func() { + readDone <- readShellCommandAdapterInput(reader, stdinWriter, cancel) + }() + + stdout := &shellCommandAdapterStreamWriter{writer: writer, frameType: "stdout"} + stderr := &shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"} + execution := newExecution(nil, command.Name, normalizeNoColor(command.Name, header.Args), header.Dir, nil) + execution.ID = coretool.InvocationFromContext(ctx).CallID + execution.StartedAt = time.Now() + execution.setIO(stdinReader, stdout, stderr) + _, runErr := command.Run(ctx, execution) + execution.EndedAt = time.Now() + _ = stdinReader.Close() + cancel() + + exitCode := shellCommandAdapterExitCode(runErr) + if runErr != nil && !errors.Is(runErr, context.Canceled) && !errors.Is(runErr, context.DeadlineExceeded) { + _, _ = io.WriteString(stderr, runErr.Error()+"\n") + } + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: exitCode}) + _ = conn.Close() + select { + case <-readDone: + case <-time.After(time.Second): + } +} + +func readShellCommandAdapterInput(reader *bufio.Reader, stdin *io.PipeWriter, cancel context.CancelFunc) error { + stdinOpen := true + closeStdin := func() { + if stdinOpen { + stdinOpen = false + _ = stdin.Close() + } + } + defer closeStdin() + for { + frame, err := readShellCommandAdapterFrame(reader) + if err != nil { + cancel() + return err + } + switch frame.Type { + case "stdin": + if stdinOpen && len(frame.Data) > 0 { + if _, err := stdin.Write(frame.Data); err != nil { + return err + } + } + case "stdin_eof": + closeStdin() + case "cancel": + cancel() + return context.Canceled + default: + cancel() + return fmt.Errorf("unexpected shell command frame %q", frame.Type) + } + } +} + +type shellCommandAdapterFrameWriter struct { + mu sync.Mutex + writer io.Writer +} + +func (w *shellCommandAdapterFrameWriter) write(frame shellCommandAdapterFrame) error { + w.mu.Lock() + defer w.mu.Unlock() + return writeShellCommandAdapterFrame(w.writer, frame) +} + +type shellCommandAdapterStreamWriter struct { + writer *shellCommandAdapterFrameWriter + frameType string +} + +func (w *shellCommandAdapterStreamWriter) Write(data []byte) (int, error) { + for offset := 0; offset < len(data); { + end := offset + shellCommandAdapterChunkSize + if end > len(data) { + end = len(data) + } + if err := w.writer.write(shellCommandAdapterFrame{Type: w.frameType, Data: data[offset:end]}); err != nil { + return offset, err + } + offset = end + } + return len(data), nil +} + +func shellCommandAdapterExitCode(err error) int { + if err == nil { + return 0 + } + var exitCoder interface{ ExitCode() int } + if errors.As(err, &exitCoder) { + code := exitCoder.ExitCode() + if code >= 0 && code <= 255 { + return code + } + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return 130 + } + return 1 +} + +func writeShellCommandAdapterFrame(writer io.Writer, frame shellCommandAdapterFrame) error { + data, err := json.Marshal(frame) + if err != nil { + return err + } + if len(data) > shellCommandAdapterMaxFrameSize { + return fmt.Errorf("shell command frame too large: %d", len(data)) + } + var size [4]byte + binary.BigEndian.PutUint32(size[:], uint32(len(data))) + if err := writeShellCommandAdapterBytes(writer, size[:]); err != nil { + return err + } + return writeShellCommandAdapterBytes(writer, data) +} + +func writeShellCommandAdapterBytes(writer io.Writer, data []byte) error { + for len(data) > 0 { + n, err := writer.Write(data) + if err != nil { + return err + } + if n <= 0 { + return io.ErrShortWrite + } + data = data[n:] + } + return nil +} + +func readShellCommandAdapterFrame(reader io.Reader) (shellCommandAdapterFrame, error) { + var size [4]byte + if _, err := io.ReadFull(reader, size[:]); err != nil { + return shellCommandAdapterFrame{}, err + } + length := binary.BigEndian.Uint32(size[:]) + if length == 0 || length > shellCommandAdapterMaxFrameSize { + return shellCommandAdapterFrame{}, fmt.Errorf("invalid shell command frame size: %d", length) + } + data := make([]byte, int(length)) + if _, err := io.ReadFull(reader, data); err != nil { + return shellCommandAdapterFrame{}, err + } + var frame shellCommandAdapterFrame + if err := json.Unmarshal(data, &frame); err != nil { + return shellCommandAdapterFrame{}, err + } + return frame, nil +} + +func (b *shellCommandAdapter) syncAliases(names []string) error { + b.mu.Lock() + defer b.mu.Unlock() + sorted := append([]string(nil), names...) + sort.Strings(sorted) + for _, name := range sorted { + if !validShellCommandAdapterName(name) { + continue + } + if _, ok := b.aliases[name]; ok { + continue + } + path, err := createShellCommandAdapterAlias(b.executable, b.runtimeDir, name) + if err != nil { + return fmt.Errorf("create shell command alias %s: %w", name, err) + } + b.aliases[name] = path + } + return nil +} + +func validShellCommandAdapterName(name string) bool { + if name == "" || name == "." || name == ".." { + return false + } + for _, r := range name { + if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' { + continue + } + return false + } + return true +} + +func (b *shellCommandAdapter) retainContext(ctx context.Context) string { + if ctx == nil { + ctx = context.Background() + } + b.mu.Lock() + defer b.mu.Unlock() + b.nextContext++ + id := strconv.FormatUint(b.nextContext, 10) + b.contexts[id] = ctx + return id +} + +func (b *shellCommandAdapter) context(id string) (context.Context, bool) { + b.mu.Lock() + defer b.mu.Unlock() + ctx, ok := b.contexts[id] + return ctx, ok +} + +func (b *shellCommandAdapter) releaseContext(id string) { + b.mu.Lock() + delete(b.contexts, id) + b.mu.Unlock() +} + +func (b *shellCommandAdapter) environment(contextID string) []string { + return []string{ + shellCommandAdapterMarkerEnv + "=1", + shellCommandAdapterEndpointEnv + "=" + b.endpoint, + shellCommandAdapterExecutableEnv + "=" + b.executable, + shellCommandAdapterContextEnv + "=" + contextID, + } +} + +func (b *shellCommandAdapter) shutdown() { + b.shutdownOnce.Do(func() { + b.cancel() + _ = b.listener.Close() + b.mu.Lock() + for conn := range b.connections { + _ = conn.Close() + } + b.mu.Unlock() + b.wg.Wait() + }) +} + +func (b *shellCommandAdapter) cleanup() { + b.cleanupOnce.Do(func() { + deadline := time.Now().Add(2 * time.Second) + for { + if err := os.RemoveAll(b.runtimeDir); err == nil { + break + } + if time.Now().After(deadline) { + break + } + time.Sleep(50 * time.Millisecond) + } + _ = os.Remove(shellCommandAdapterRuntimeRoot()) + }) +} + +func (b *shellCommandAdapter) close() { + b.shutdown() + b.cleanup() +} + +// runShellCommandProxyIfRequested dispatches the process-local PATH shim. The +// marker and command variables are only injected into shim children, so normal +// AIScan and embedding-host startup is unchanged. +func runShellCommandProxyIfRequested() (code int, ok bool) { + command := strings.TrimSpace(os.Getenv(shellCommandAdapterCommandEnv)) + if os.Getenv(shellCommandAdapterMarkerEnv) != "1" || command == "" { + return 0, false + } + if !validShellCommandAdapterName(command) { + fmt.Fprintln(os.Stderr, "invalid shell command alias") + return 126, true + } + return runShellCommandAdapterProxy(command, os.Args[1:]), true +} + +func runShellCommandAdapterProxy(command string, args []string) int { + endpoint := os.Getenv(shellCommandAdapterEndpointEnv) + if endpoint == "" { + fmt.Fprintln(os.Stderr, "AIScan shell command environment is incomplete") + return 126 + } + ctx, cancel := context.WithTimeout(context.Background(), shellCommandAdapterDialTimeout) + conn, err := dialShellCommandAdapter(ctx, endpoint) + cancel() + if err != nil { + fmt.Fprintf(os.Stderr, "connect AIScan shell command adapter: %s\n", err) + return 125 + } + defer conn.Close() + cwd, _ := os.Getwd() + header := shellCommandAdapterFrame{ + Type: "request", Version: shellCommandAdapterProtocolVersion, + Command: command, Args: append([]string(nil), args...), Dir: cwd, + ContextID: os.Getenv(shellCommandAdapterContextEnv), + } + writer := &shellCommandAdapterFrameWriter{writer: conn} + if err := writer.write(header); err != nil { + fmt.Fprintf(os.Stderr, "start AIScan shell command request: %s\n", err) + return 125 + } + go streamShellCommandAdapterStdin(writer) + + reader := bufio.NewReader(conn) + for { + frame, err := readShellCommandAdapterFrame(reader) + if err != nil { + fmt.Fprintf(os.Stderr, "read AIScan shell command response: %s\n", err) + return 125 + } + switch frame.Type { + case "stdout": + _, _ = os.Stdout.Write(frame.Data) + case "stderr": + _, _ = os.Stderr.Write(frame.Data) + case "final": + flushShellCommandAdapterProxyOutput() + return frame.ExitCode + default: + fmt.Fprintf(os.Stderr, "unexpected AIScan shell command response %q\n", frame.Type) + return 125 + } + } +} + +func streamShellCommandAdapterStdin(writer *shellCommandAdapterFrameWriter) { + info, err := os.Stdin.Stat() + if err != nil || info.Mode()&os.ModeCharDevice != 0 { + _ = writer.write(shellCommandAdapterFrame{Type: "stdin_eof"}) + return + } + buffer := make([]byte, shellCommandAdapterChunkSize) + for { + n, readErr := os.Stdin.Read(buffer) + if n > 0 { + if err := writer.write(shellCommandAdapterFrame{Type: "stdin", Data: buffer[:n]}); err != nil { + return + } + } + if readErr != nil { + _ = writer.write(shellCommandAdapterFrame{Type: "stdin_eof"}) + return + } + } +} diff --git a/pkg/commands/shell_command_adapter_test.go b/pkg/commands/shell_command_adapter_test.go new file mode 100644 index 00000000..e2499bda --- /dev/null +++ b/pkg/commands/shell_command_adapter_test.go @@ -0,0 +1,247 @@ +package commands + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +type adapterTestExitError struct{ code int } + +func (e adapterTestExitError) Error() string { return fmt.Sprintf("adapter test exit %d", e.code) } +func (e adapterTestExitError) ExitCode() int { return e.code } + +type adapterTestCommands struct { + started chan struct{} + canceled chan struct{} + once sync.Once +} + +func newAdapterTestBash(t *testing.T) (*BashTool, *CommandRegistry, *adapterTestCommands) { + t.Helper() + state := &adapterTestCommands{started: make(chan struct{}), canceled: make(chan struct{})} + registry := NewRegistry() + registry.Register(Command{Name: "memory_echo", Run: func(_ context.Context, execution *Execution) (any, error) { + fmt.Fprintln(execution.Stdout, strings.Join(execution.Args, " ")) + return nil, nil + }}, "test") + registry.Register(Command{Name: "memory_upper", Run: func(_ context.Context, execution *Execution) (any, error) { + data, err := io.ReadAll(execution.Stdin) + if err != nil { + return nil, err + } + _, err = execution.Stdout.Write(bytes.ToUpper(data)) + return nil, err + }}, "test") + registry.Register(Command{Name: "memory_fail", Run: func(context.Context, *Execution) (any, error) { + return nil, adapterTestExitError{code: 7} + }}, "test") + registry.Register(Command{Name: "memory_context", Run: func(ctx context.Context, execution *Execution) (any, error) { + invocation := coretool.InvocationFromContext(ctx) + fmt.Fprintf(execution.Stdout, "dir=%s call=%s session=%s turn=%s emitter=%s\n", + execution.Dir, invocation.CallID, invocation.SessionID, invocation.TurnID, invocation.Emitter) + return nil, nil + }}, "test") + registry.Register(Command{Name: "memory_wait", Run: func(ctx context.Context, _ *Execution) (any, error) { + state.once.Do(func() { close(state.started) }) + <-ctx.Done() + close(state.canceled) + return nil, ctx.Err() + }}, "test") + + bash := NewBashTool(t.TempDir(), 10) + bash.SetCommandNames(registry.Names) + bash.SetCommandResolver(registry.Get) + bash.attachShellCommands(registry) + t.Cleanup(bash.Close) + return bash, registry, state +} + +func runAdapterCommand(t *testing.T, bash *BashTool, ctx context.Context, command string, workDir string) (*Execution, string) { + t.Helper() + var output strings.Builder + execution, err := bash.RunForeground(ctx, command, BashExecOptions{ + WorkDir: workDir, + OnOutput: func(data []byte) { + _, _ = output.Write(data) + }, + }) + if err != nil { + t.Fatalf("RunForeground(%q): %v", command, err) + } + return execution, output.String() +} + +func TestShellCommandAdapterIsLazy(t *testing.T) { + bash, _, _ := newAdapterTestBash(t) + execution, output := runAdapterCommand(t, bash, context.Background(), "memory_echo direct", t.TempDir()) + if execution.ExitCode != 0 || !strings.Contains(output, "direct") { + t.Fatalf("direct command exit=%d output=%q", execution.ExitCode, output) + } + if bash.shellAdapter != nil { + t.Fatal("simple registered command allocated shell runtime state") + } +} + +func TestShellCommandMarkerDoesNotHijackNormalChildProcess(t *testing.T) { + t.Setenv(shellCommandAdapterMarkerEnv, "1") + t.Setenv(shellCommandAdapterCommandEnv, "") + if code, ok := runShellCommandProxyIfRequested(); ok { + t.Fatalf("marker-only child was treated as proxy with code %d", code) + } +} + +func TestShellCommandComposition(t *testing.T) { + bash, _, _ := newAdapterTestBash(t) + workDir := t.TempDir() + + execution, output := runAdapterCommand(t, bash, context.Background(), "memory_echo one && memory_echo two", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "one") || !strings.Contains(output, "two") { + t.Fatalf("and composition exit=%d output=%q", execution.ExitCode, output) + } + + execution, output = runAdapterCommand(t, bash, context.Background(), "memory_fail || memory_echo recovered", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "recovered") { + t.Fatalf("or composition exit=%d output=%q", execution.ExitCode, output) + } + + execution, output = runAdapterCommand(t, bash, context.Background(), "memory_echo hello | memory_upper", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "HELLO") { + t.Fatalf("pipeline exit=%d output=%q", execution.ExitCode, output) + } + + execution, _ = runAdapterCommand(t, bash, context.Background(), "memory_fail && memory_echo unreachable", workDir) + if execution.ExitCode != 7 { + t.Fatalf("short-circuit exit code = %d, want 7", execution.ExitCode) + } +} + +func TestShellCommandRedirectionAndInvocationContext(t *testing.T) { + bash, _, _ := newAdapterTestBash(t) + workDir := t.TempDir() + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{ + WorkDir: workDir, CallID: "call-1", SessionID: "session-1", TurnID: "turn-1", Emitter: "runner", + }) + + execution, _ := runAdapterCommand(t, bash, ctx, "memory_context > adapter-context.txt", workDir) + if execution.ExitCode != 0 { + t.Fatalf("redirection exit code = %d", execution.ExitCode) + } + data, err := os.ReadFile(filepath.Join(workDir, "adapter-context.txt")) + if err != nil { + t.Fatalf("read redirected output: %v", err) + } + text := string(data) + for _, expected := range []string{"dir=" + workDir, "call=call-1", "session=session-1", "turn=turn-1", "emitter=runner"} { + if !strings.Contains(text, expected) { + t.Fatalf("redirected context missing %q: %q", expected, text) + } + } +} + +func TestShellCommandAdapterFindsCommandsRegisteredLater(t *testing.T) { + bash, registry, _ := newAdapterTestBash(t) + registry.Register(Command{Name: "scan", Run: func(_ context.Context, execution *Execution) (any, error) { + fmt.Fprintln(execution.Stdout, strings.Join(execution.Args, " ")) + return nil, nil + }}, "test") + command := "scan -i http://127.0.0.1:1 --timeout 1 --no-color && memory_echo ready" + execution, output := runAdapterCommand(t, bash, context.Background(), command, t.TempDir()) + if execution.ExitCode != 0 || !strings.Contains(output, "http://127.0.0.1:1") || !strings.Contains(output, "ready") { + t.Fatalf("late alias exit=%d output=%q", execution.ExitCode, output) + } +} + +func TestShellCommandCloseCancelsCallsAndRemovesRuntime(t *testing.T) { + bash, _, state := newAdapterTestBash(t) + execution, err := bash.Start(context.Background(), "memory_wait && memory_echo unreachable", BashExecOptions{WorkDir: t.TempDir()}) + if err != nil { + t.Fatalf("Start memory_wait: %v", err) + } + runtimeDir := bash.shellAdapter.runtimeDir + if _, err := os.Stat(runtimeDir); err != nil { + t.Fatalf("runtime directory before close: %v", err) + } + select { + case <-state.started: + case <-time.After(5 * time.Second): + t.Fatal("shell command did not start") + } + bash.Close() + bash.Close() + select { + case <-state.canceled: + case <-time.After(5 * time.Second): + t.Fatal("shell command was not canceled") + } + if err := execution.Wait(context.Background()); err != nil && !errors.Is(err, context.Canceled) { + t.Fatalf("wait after close: %v", err) + } + if _, err := os.Stat(runtimeDir); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("runtime directory still exists after close: %v", err) + } +} + +func TestShellCommandDisconnectCancelsRunningCommand(t *testing.T) { + bash, _, state := newAdapterTestBash(t) + adapter, err := bash.ensureShellCommands() + if err != nil { + t.Fatalf("ensure shell commands: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + conn, err := dialShellCommandAdapter(ctx, adapter.endpoint) + if err != nil { + t.Fatalf("dial shell command adapter: %v", err) + } + contextID := adapter.retainContext(context.Background()) + defer adapter.releaseContext(contextID) + writer := &shellCommandAdapterFrameWriter{writer: conn} + if err := writer.write(shellCommandAdapterFrame{ + Type: "request", Version: shellCommandAdapterProtocolVersion, + Command: "memory_wait", Dir: t.TempDir(), ContextID: contextID, + }); err != nil { + t.Fatalf("write request: %v", err) + } + if err := writer.write(shellCommandAdapterFrame{Type: "stdin_eof"}); err != nil { + t.Fatalf("write stdin eof: %v", err) + } + select { + case <-state.started: + case <-ctx.Done(): + t.Fatal("shell command did not start") + } + _ = conn.Close() + select { + case <-state.canceled: + case <-ctx.Done(): + t.Fatal("disconnect did not cancel shell command") + } +} + +func TestShellCommandStartupReclaimsOwnedStaleRuntime(t *testing.T) { + root := shellCommandAdapterRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + t.Fatal(err) + } + stale, err := os.MkdirTemp(root, "1073741824-stale-") + if err != nil { + t.Fatal(err) + } + if err := cleanupStaleShellCommandAdapterRuntime(); err != nil { + t.Fatalf("cleanup stale runtime: %v", err) + } + if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("stale runtime still exists: %v", err) + } +} diff --git a/pkg/commands/command_bridge_unix.go b/pkg/commands/shell_command_adapter_unix.go similarity index 56% rename from pkg/commands/command_bridge_unix.go rename to pkg/commands/shell_command_adapter_unix.go index fdc87cec..9b61a9c7 100644 --- a/pkg/commands/command_bridge_unix.go +++ b/pkg/commands/shell_command_adapter_unix.go @@ -11,11 +11,11 @@ import ( "syscall" ) -func commandBridgeEndpoint(runtimeDir string) string { - return filepath.Join(runtimeDir, "bridge.sock") +func shellCommandAdapterEndpoint(runtimeDir string) string { + return filepath.Join(runtimeDir, "commands.sock") } -func listenCommandBridge(endpoint string) (net.Listener, error) { +func listenShellCommandAdapter(endpoint string) (net.Listener, error) { _ = os.Remove(endpoint) listener, err := net.Listen("unix", endpoint) if err != nil { @@ -28,30 +28,30 @@ func listenCommandBridge(endpoint string) (net.Listener, error) { return listener, nil } -func dialCommandBridge(ctx context.Context, endpoint string) (net.Conn, error) { +func dialShellCommandAdapter(ctx context.Context, endpoint string) (net.Conn, error) { var dialer net.Dialer return dialer.DialContext(ctx, "unix", endpoint) } -func createCommandBridgeAlias(executable, runtimeDir, name string) (string, error) { +func createShellCommandAdapterAlias(executable, runtimeDir, name string) (string, error) { _ = executable path := filepath.Join(runtimeDir, name) if _, err := os.Stat(path); err == nil { return path, nil } - content := "#!/bin/sh\n" + commandBridgeCommandEnv + "='" + name + "' exec \"$" + commandBridgeExecutableEnv + "\" \"$@\"\n" + content := "#!/bin/sh\n" + shellCommandAdapterCommandEnv + "='" + name + "' exec \"$" + shellCommandAdapterExecutableEnv + "\" \"$@\"\n" if err := os.WriteFile(path, []byte(content), 0o700); err != nil { return "", err } return path, nil } -func commandBridgeProcessAlive(pid int) bool { +func shellCommandAdapterProcessAlive(pid int) bool { err := syscall.Kill(pid, 0) return err == nil || errors.Is(err, syscall.EPERM) } -func flushCommandBridgeProxyOutput() { +func flushShellCommandAdapterProxyOutput() { _ = os.Stdout.Sync() _ = os.Stderr.Sync() } diff --git a/pkg/commands/command_bridge_windows.go b/pkg/commands/shell_command_adapter_windows.go similarity index 66% rename from pkg/commands/command_bridge_windows.go rename to pkg/commands/shell_command_adapter_windows.go index 4ca7ed8f..6b2dce86 100644 --- a/pkg/commands/command_bridge_windows.go +++ b/pkg/commands/shell_command_adapter_windows.go @@ -14,35 +14,35 @@ import ( "golang.org/x/sys/windows" ) -func commandBridgeEndpoint(runtimeDir string) string { - return `\\.\pipe\aiscan-command-bridge-` + fmt.Sprintf("%d-%s", os.Getpid(), filepath.Base(runtimeDir)) +func shellCommandAdapterEndpoint(runtimeDir string) string { + return `\\.\pipe\aiscan-shell-commands-` + fmt.Sprintf("%d-%s", os.Getpid(), filepath.Base(runtimeDir)) } -func listenCommandBridge(endpoint string) (net.Listener, error) { +func listenShellCommandAdapter(endpoint string) (net.Listener, error) { user, err := windows.GetCurrentProcessToken().GetTokenUser() if err != nil { return nil, err } return winio.ListenPipe(endpoint, &winio.PipeConfig{ SecurityDescriptor: "D:P(A;;GA;;;" + user.User.Sid.String() + ")", - InputBufferSize: commandBridgeChunkSize, - OutputBufferSize: commandBridgeChunkSize, + InputBufferSize: shellCommandAdapterChunkSize, + OutputBufferSize: shellCommandAdapterChunkSize, }) } -func dialCommandBridge(ctx context.Context, endpoint string) (net.Conn, error) { +func dialShellCommandAdapter(ctx context.Context, endpoint string) (net.Conn, error) { return winio.DialPipeContext(ctx, endpoint) } -func createCommandBridgeAlias(executable, runtimeDir, name string) (string, error) { +func createShellCommandAdapterAlias(executable, runtimeDir, name string) (string, error) { _ = executable path := filepath.Join(runtimeDir, name+".cmd") if _, err := os.Stat(path); err == nil { return path, nil } content := "@echo off\r\n" + - "set \"" + commandBridgeCommandEnv + "=" + name + "\"\r\n" + - "\"%" + commandBridgeExecutableEnv + "%\" %*\r\n" + + "set \"" + shellCommandAdapterCommandEnv + "=" + name + "\"\r\n" + + "\"%" + shellCommandAdapterExecutableEnv + "%\" %*\r\n" + "exit /b %errorlevel%\r\n" if err := os.WriteFile(path, []byte(content), 0o600); err != nil { return "", err @@ -50,7 +50,7 @@ func createCommandBridgeAlias(executable, runtimeDir, name string) (string, erro return path, nil } -func commandBridgeProcessAlive(pid int) bool { +func shellCommandAdapterProcessAlive(pid int) bool { handle, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(pid)) if err != nil { return err == windows.ERROR_ACCESS_DENIED @@ -63,7 +63,7 @@ func commandBridgeProcessAlive(pid int) bool { return exitCode == 259 // STILL_ACTIVE } -func flushCommandBridgeProxyOutput() { +func flushShellCommandAdapterProxyOutput() { _ = os.Stdout.Sync() _ = os.Stderr.Sync() // ConPTY can report the proxy process exit before consuming its final diff --git a/pkg/runner/app.go b/pkg/runner/app.go index c325e421..b4c9cee1 100644 --- a/pkg/runner/app.go +++ b/pkg/runner/app.go @@ -366,7 +366,6 @@ func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillSto deps := &commands.Deps{ WorkDir: workDir, BashTimeout: rc.Tools.BashTimeout, - CommandBridge: rc.Tools.CommandBridge, SkillStore: skillStore, Provider: llmProvider, Logger: logger, diff --git a/pkg/runner/application_builder.go b/pkg/runner/application_builder.go index 8d37b0f7..d3b31e4a 100644 --- a/pkg/runner/application_builder.go +++ b/pkg/runner/application_builder.go @@ -61,7 +61,6 @@ func MergeOptionExtras(rc ApplicationConfig, option *cfg.Option) ApplicationConf } rc.Scanner.UncoverCredentials = cloneStringMap(option.UncoverCredentials) rc.Tools.PlaywrightSession = option.PlaywrightSession - rc.Tools.CommandBridge = option.CommandBridge rc.CLISkillPaths = skillPathsFromOptions(option) rc.RecordFile = option.OutputFile return rc @@ -93,7 +92,6 @@ func AppConfig(option *cfg.Option, features RuntimeFeatures, logger telemetry.Lo Tools: ToolConfig{ Enabled: features.ToolsEnabled, BashTimeout: 300, - CommandBridge: option.CommandBridge, TavilyKeys: resolveTavilyKeys(option.TavilyKey, option.SearchConfig.TavilyKeys, cfg.DefaultTavilyKeys), PlaywrightSession: option.PlaywrightSession, OptionalTools: option.Tools, diff --git a/pkg/runner/application_config.go b/pkg/runner/application_config.go index 92afaaf9..ff324d80 100644 --- a/pkg/runner/application_config.go +++ b/pkg/runner/application_config.go @@ -43,7 +43,6 @@ type ScannerConfig struct { type ToolConfig struct { Enabled bool BashTimeout int - CommandBridge bool TavilyKeys string PlaywrightSession string OptionalTools []string // optional tool groups to enable (e.g. "search", "browser") diff --git a/pkg/runner/provider_config_test.go b/pkg/runner/provider_config_test.go index 5fd5cb91..2c3f0aca 100644 --- a/pkg/runner/provider_config_test.go +++ b/pkg/runner/provider_config_test.go @@ -110,10 +110,9 @@ func TestMergeOptionExtrasLayersNonProtoFields(t *testing.T) { PlaywrightSession: "browser-1", UncoverCredentials: map[string]string{"SHODAN_API_KEY": "shodan-key"}, MiscOptions: cfg.MiscOptions{OutputFile: "session.jsonl"}, - AgentOptions: cfg.AgentOptions{CommandBridge: true}, } rc = MergeOptionExtras(rc, option) - if rc.Tools.PlaywrightSession != "browser-1" || !rc.Tools.CommandBridge || rc.Scanner.UncoverCredentials["SHODAN_API_KEY"] != "shodan-key" || rc.RecordFile != "session.jsonl" { + if rc.Tools.PlaywrightSession != "browser-1" || rc.Scanner.UncoverCredentials["SHODAN_API_KEY"] != "shodan-key" || rc.RecordFile != "session.jsonl" { t.Fatalf("extras = %+v", rc) } } From 5753f2bf3023469c88acec8df5edafca718f8eeb Mon Sep 17 00:00:00 2001 From: M09Ic Date: Fri, 7 Aug 2026 16:19:58 +0800 Subject: [PATCH 5/5] test: stabilize concurrent PTY option isolation --- pkg/commands/bash_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pkg/commands/bash_test.go b/pkg/commands/bash_test.go index a85ba5d1..0a990f72 100644 --- a/pkg/commands/bash_test.go +++ b/pkg/commands/bash_test.go @@ -479,7 +479,9 @@ func TestBashExecOptionsAreIsolatedAcrossConcurrentCalls(t *testing.T) { wg.Add(1) go func(i int) { defer wg.Done() - results[i], errs[i] = bash.RunForeground(context.Background(), `printf '%s\n' "$AISCAN_RUN_VALUE"; pwd`, BashExecOptions{ + // Keep the short-lived shell alive until the PTY reader is scheduled; + // this test exercises concurrent option isolation, not PTY drain timing. + results[i], errs[i] = bash.RunForeground(context.Background(), `printf '%s\n' "$AISCAN_RUN_VALUE"; pwd; sleep 0.05`, BashExecOptions{ WorkDir: dirs[i], Env: map[string]string{"AISCAN_RUN_VALUE": fmt.Sprintf("value-%d", i)}, OnOutput: func(data []byte) {