| 1 | package agentrunner |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "bytes" |
| 6 | "context" |
| 7 | "encoding/json" |
| 8 | "errors" |
| 9 | "fmt" |
| 10 | "io" |
| 11 | "os/exec" |
| 12 | "strings" |
| 13 | "sync" |
| 14 | "time" |
| 15 | ) |
| 16 | |
| 17 | type PermissionOption struct { |
| 18 | ID string `json:"optionId"` |
| 19 | Name string `json:"name"` |
| 20 | Kind string `json:"kind"` |
| 21 | } |
| 22 | type Permission struct { |
| 23 | Title string |
| 24 | Details string |
| 25 | Options []PermissionOption |
| 26 | } |
| 27 | type ProcessOptions struct { |
| 28 | Command string |
| 29 | Args []string |
| 30 | Env []string |
| 31 | CWD string |
| 32 | Protocol string |
| 33 | Progress func(string) |
| 34 | Permission func(context.Context, Permission) (string, error) |
| 35 | } |
| 36 | type rpcFrame struct { |
| 37 | ID json.RawMessage `json:"id,omitempty"` |
| 38 | Method string `json:"method,omitempty"` |
| 39 | Params json.RawMessage `json:"params,omitempty"` |
| 40 | Result json.RawMessage `json:"result,omitempty"` |
| 41 | Error *rpcError `json:"error,omitempty"` |
| 42 | } |
| 43 | type rpcError struct { |
| 44 | Code int `json:"code"` |
| 45 | Message string `json:"message"` |
| 46 | } |
| 47 | type rpcProcess struct { |
| 48 | input io.WriteCloser |
| 49 | frames chan rpcFrame |
| 50 | finished chan struct{} |
| 51 | readErr error |
| 52 | writeMu sync.Mutex |
| 53 | cancelOnce sync.Once |
| 54 | sessionMu sync.RWMutex |
| 55 | opts ProcessOptions |
| 56 | session string |
| 57 | text strings.Builder |
| 58 | } |
| 59 | |
| 60 | func RunProcess(ctx context.Context, opts ProcessOptions, prompt string) (string, error) { |
| 61 | if opts.Command == "" { |
| 62 | return "" , errors.New("agent command is required" ) |
| 63 | } |
| 64 | if opts.Protocol == "command" { |
| 65 | return runCommand(ctx, opts, prompt) |
| 66 | } |
| 67 | if opts.Protocol != "acp" { |
| 68 | return "" , errors.New("agent protocol must be acp or command" ) |
| 69 | } |
| 70 | command := exec.CommandContext(ctx, opts.Command, opts.Args...) |
| 71 | configureProcessGroup(command) |
| 72 | var process *rpcProcess |
| 73 | command.Cancel = func() error { |
| 74 | if process != nil && process.hasSession() { |
| 75 | process.cancelSession() |
| 76 | } |
| 77 | terminateProcess(command) |
| 78 | return nil |
| 79 | } |
| 80 | command.Dir = opts.CWD |
| 81 | command.Env = opts.Env |
| 82 | command.WaitDelay = 2 * time.Second |
| 83 | input, err := command.StdinPipe() |
| 84 | if err != nil { |
| 85 | return "" , err |
| 86 | } |
| 87 | output, err := command.StdoutPipe() |
| 88 | if err != nil { |
| 89 | return "" , err |
| 90 | } |
| 91 | // Agent stderr belongs to the local runner, never to a room transcript. |
| 92 | command.Stderr = io.Discard |
| 93 | p := &rpcProcess{input: input, frames: make(chan rpcFrame, 1), finished: make(chan struct{}), opts: opts} |
| 94 | process = p |
| 95 | if err := command.Start(); err != nil { |
| 96 | return "" , fmt.Errorf("start agent: %w" , err) |
| 97 | } |
| 98 | readCtx, stopRead := context.WithCancel(ctx) |
| 99 | go p.read(readCtx, output) |
| 100 | defer func() { stopRead(); input.Close(); terminateProcess(command); command.Wait(); <-p.finished }() |
| 101 | init, err := p.call(ctx, 1, "initialize" , map[string]any{"protocolVersion" : 1, "clientCapabilities" : map[string]any{}, "clientInfo" : map[string]string{"name" : "tiny" , "version" : "1" }}) |
| 102 | if err != nil { |
| 103 | return "" , err |
| 104 | } |
| 105 | var initialized struct { |
| 106 | Version int `json:"protocolVersion"` |
| 107 | } |
| 108 | if json.Unmarshal(init, &initialized) != nil || initialized.Version != 1 { |
| 109 | return "" , errors.New("agent does not support ACP version 1" ) |
| 110 | } |
| 111 | raw, err := p.call(ctx, 2, "session/new" , map[string]any{"cwd" : opts.CWD, "mcpServers" : []any{}}) |
| 112 | if err != nil { |
| 113 | return "" , err |
| 114 | } |
| 115 | var session struct { |
| 116 | ID string `json:"sessionId"` |
| 117 | } |
| 118 | if json.Unmarshal(raw, &session) != nil || session.ID == "" { |
| 119 | return "" , errors.New("agent returned no session id" ) |
| 120 | } |
| 121 | p.setSession(session.ID) |
| 122 | raw, err = p.call(ctx, 3, "session/prompt" , map[string]any{"sessionId" : session.ID, "prompt" : []any{map[string]string{"type" : "text" , "text" : prompt}}}) |
| 123 | if err != nil { |
| 124 | if ctx.Err() != nil { |
| 125 | p.cancelSession() |
| 126 | } |
| 127 | return "" , err |
| 128 | } |
| 129 | var done struct { |
| 130 | Reason string `json:"stopReason"` |
| 131 | } |
| 132 | if json.Unmarshal(raw, &done) != nil { |
| 133 | return "" , errors.New("invalid ACP prompt result" ) |
| 134 | } |
| 135 | if done.Reason == "cancelled" { |
| 136 | return "" , context.Canceled |
| 137 | } |
| 138 | return strings.TrimSpace(p.text.String()), nil |
| 139 | } |
| 140 | func (p *rpcProcess) read(ctx context.Context, out io.Reader) { |
| 141 | defer close(p.finished) |
| 142 | defer close(p.frames) |
| 143 | scanner := bufio.NewScanner(out) |
| 144 | scanner.Buffer(make([]byte, 4096), 1<<20) |
| 145 | for scanner.Scan() { |
| 146 | var frame rpcFrame |
| 147 | if err := json.Unmarshal(scanner.Bytes(), &frame); err != nil { |
| 148 | p.readErr = fmt.Errorf("invalid ACP response: %w" , err) |
| 149 | return |
| 150 | } |
| 151 | select { |
| 152 | case p.frames <- frame: |
| 153 | case <-ctx.Done(): |
| 154 | return |
| 155 | } |
| 156 | } |
| 157 | p.readErr = scanner.Err() |
| 158 | } |
| 159 | func (p *rpcProcess) write(value any) error { |
| 160 | p.writeMu.Lock() |
| 161 | defer p.writeMu.Unlock() |
| 162 | return json.NewEncoder(p.input).Encode(value) |
| 163 | } |
| 164 | func (p *rpcProcess) call(ctx context.Context, id int, method string, params any) (json.RawMessage, error) { |
| 165 | if err := p.write(map[string]any{"jsonrpc" : "2.0" , "id" : id, "method" : method, "params" : params}); err != nil { |
| 166 | return nil, err |
| 167 | } |
| 168 | for { |
| 169 | select { |
| 170 | case <-ctx.Done(): |
| 171 | if method == "session/prompt" && p.session != "" { |
| 172 | p.cancelSession() |
| 173 | } |
| 174 | return nil, ctx.Err() |
| 175 | case frame, ok := <-p.frames: |
| 176 | if !ok { |
| 177 | if p.readErr != nil { |
| 178 | return nil, p.readErr |
| 179 | } |
| 180 | return nil, errors.New("agent process exited before replying" ) |
| 181 | } |
| 182 | if frame.Method != "" { |
| 183 | if err := p.receive(ctx, frame); err != nil { |
| 184 | return nil, err |
| 185 | } |
| 186 | continue |
| 187 | } |
| 188 | if string(frame.ID) != fmt.Sprint(id) { |
| 189 | continue |
| 190 | } |
| 191 | if frame.Error != nil { |
| 192 | return nil, fmt.Errorf("ACP %s: %s" , method, frame.Error.Message) |
| 193 | } |
| 194 | return frame.Result, nil |
| 195 | } |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | func (p *rpcProcess) cancelSession() { |
| 200 | p.cancelOnce.Do(func() { |
| 201 | session := p.getSession() |
| 202 | if session == "" { |
| 203 | return |
| 204 | } |
| 205 | done := make(chan struct{}) |
| 206 | go func() { |
| 207 | defer close(done) |
| 208 | _ = p.write(map[string]any{"jsonrpc" : "2.0" , "method" : "session/cancel" , "params" : map[string]string{"sessionId" : session}}) |
| 209 | }() |
| 210 | timer := time.NewTimer(100 * time.Millisecond) |
| 211 | defer timer.Stop() |
| 212 | select { |
| 213 | case <-done: |
| 214 | <-timer.C |
| 215 | case <-timer.C: |
| 216 | // Closing the pipe releases a blocked write before we join its goroutine. |
| 217 | _ = p.input.Close() |
| 218 | <-done |
| 219 | } |
| 220 | }) |
| 221 | } |
| 222 | |
| 223 | func (p *rpcProcess) setSession(session string) { |
| 224 | p.sessionMu.Lock() |
| 225 | p.session = session |
| 226 | p.sessionMu.Unlock() |
| 227 | } |
| 228 | func (p *rpcProcess) getSession() string { |
| 229 | p.sessionMu.RLock() |
| 230 | defer p.sessionMu.RUnlock() |
| 231 | return p.session |
| 232 | } |
| 233 | func (p *rpcProcess) hasSession() bool { return p.getSession() != "" } |
| 234 | func (p *rpcProcess) receive(ctx context.Context, frame rpcFrame) error { |
| 235 | var params struct { |
| 236 | Session string `json:"sessionId"` |
| 237 | Update struct { |
| 238 | Type string `json:"sessionUpdate"` |
| 239 | Title string `json:"title"` |
| 240 | Status string `json:"status"` |
| 241 | Content struct { |
| 242 | Type string `json:"type"` |
| 243 | Text string `json:"text"` |
| 244 | } `json:"content"` |
| 245 | } `json:"update"` |
| 246 | Tool struct { |
| 247 | Title string `json:"title"` |
| 248 | Input json.RawMessage `json:"rawInput"` |
| 249 | } `json:"toolCall"` |
| 250 | Options []PermissionOption `json:"options"` |
| 251 | } |
| 252 | if json.Unmarshal(frame.Params, ¶ms) != nil { |
| 253 | return errors.New("invalid ACP parameters" ) |
| 254 | } |
| 255 | if params.Session != "" && params.Session != p.session { |
| 256 | return errors.New("ACP update belongs to another session" ) |
| 257 | } |
| 258 | if frame.Method == "session/update" { |
| 259 | if params.Update.Type == "agent_message_chunk" && params.Update.Content.Type == "text" { |
| 260 | if p.text.Len()+len(params.Update.Content.Text) > 64000 { |
| 261 | return errors.New("agent reply exceeds 64 KiB" ) |
| 262 | } |
| 263 | p.text.WriteString(params.Update.Content.Text) |
| 264 | } |
| 265 | if (params.Update.Type == "tool_call" || params.Update.Type == "tool_call_update" ) && p.opts.Progress != nil { |
| 266 | p.opts.Progress(strings.TrimSpace(params.Update.Title + " " + params.Update.Status)) |
| 267 | } |
| 268 | return nil |
| 269 | } |
| 270 | if len(frame.ID) == 0 { |
| 271 | return nil |
| 272 | } |
| 273 | if frame.Method != "session/request_permission" { |
| 274 | return p.write(map[string]any{"jsonrpc" : "2.0" , "id" : frame.ID, "error" : rpcError{Code: -32601, Message: "Client capability unavailable" }}) |
| 275 | } |
| 276 | outcome := map[string]string{"outcome" : "cancelled" } |
| 277 | if p.opts.Permission != nil { |
| 278 | if len(params.Tool.Input) > 16000 { |
| 279 | return errors.New("permission request is too large to review" ) |
| 280 | } |
| 281 | choice, err := p.opts.Permission(ctx, Permission{Title: params.Tool.Title, Details: string(params.Tool.Input), Options: params.Options}) |
| 282 | if err != nil { |
| 283 | return err |
| 284 | } |
| 285 | for _, option := range params.Options { |
| 286 | if option.ID == choice { |
| 287 | outcome = map[string]string{"outcome" : "selected" , "optionId" : choice} |
| 288 | break |
| 289 | } |
| 290 | } |
| 291 | } |
| 292 | return p.write(map[string]any{"jsonrpc" : "2.0" , "id" : frame.ID, "result" : map[string]any{"outcome" : outcome}}) |
| 293 | } |
| 294 | |
| 295 | type boundedOutput struct { |
| 296 | bytes.Buffer |
| 297 | limit int |
| 298 | } |
| 299 | |
| 300 | func (b *boundedOutput) Write(p []byte) (int, error) { |
| 301 | if b.Len()+len(p) > b.limit { |
| 302 | return 0, errors.New("agent output exceeds 64 KiB" ) |
| 303 | } |
| 304 | return b.Buffer.Write(p) |
| 305 | } |
| 306 | func runCommand(ctx context.Context, opts ProcessOptions, prompt string) (string, error) { |
| 307 | command := exec.CommandContext(ctx, opts.Command, opts.Args...) |
| 308 | configureProcessGroup(command) |
| 309 | command.Cancel = func() error { |
| 310 | terminateProcess(command) |
| 311 | return nil |
| 312 | } |
| 313 | command.Dir = opts.CWD |
| 314 | command.Env = opts.Env |
| 315 | command.WaitDelay = 2 * time.Second |
| 316 | command.Stdin = strings.NewReader(prompt) |
| 317 | output := &boundedOutput{limit: 64000} |
| 318 | command.Stdout = output |
| 319 | command.Stderr = io.Discard |
| 320 | defer terminateProcess(command) |
| 321 | if err := command.Run(); err != nil { |
| 322 | return "" , fmt.Errorf("agent command failed: %w" , err) |
| 323 | } |
| 324 | return strings.TrimSpace(output.String()), nil |
| 325 | } |