clone

a self-hosted multitenant nostr relay with sites, file storage, git hosting, and inbox and outbox delivery

feat/standalone-relay
tinyrelay / internal / agentrunner / process.go
process.go 9.1K
1package agentrunner23import (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)1617type PermissionOption struct {18 ID string `json:"optionId"`19 Name string `json:"name"`20 Kind string `json:"kind"`21}22type Permission struct {23 Title string24 Details string25 Options []PermissionOption26}27type ProcessOptions struct {28 Command string29 Args []string30 Env []string31 CWD string32 Protocol string33 Progress func(string)34 Permission func(context.Context, Permission) (string, error)35}36type 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}43type rpcError struct {44 Code int `json:"code"`45 Message string `json:"message"`46}47type rpcProcess struct {48 input io.WriteCloser49 frames chan rpcFrame50 finished chan struct{}51 readErr error52 writeMu sync.Mutex53 cancelOnce sync.Once54 sessionMu sync.RWMutex55 opts ProcessOptions56 session string57 text strings.Builder58}5960func 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 *rpcProcess73 command.Cancel = func() error {74 if process != nil && process.hasSession() {75 process.cancelSession()76 }77 terminateProcess(command)78 return nil79 }80 command.Dir = opts.CWD81 command.Env = opts.Env82 command.WaitDelay = 2 * time.Second83 input, err := command.StdinPipe()84 if err != nil {85 return "", err86 }87 output, err := command.StdoutPipe()88 if err != nil {89 return "", err90 }91 // Agent stderr belongs to the local runner, never to a room transcript.92 command.Stderr = io.Discard93 p := &rpcProcess{input: input, frames: make(chan rpcFrame, 1), finished: make(chan struct{}), opts: opts}94 process = p95 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 "", err104 }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 "", err114 }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 "", err128 }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.Canceled137 }138 return strings.TrimSpace(p.text.String()), nil139}140func (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 rpcFrame147 if err := json.Unmarshal(scanner.Bytes(), &frame); err != nil {148 p.readErr = fmt.Errorf("invalid ACP response: %w", err)149 return150 }151 select {152 case p.frames <- frame:153 case <-ctx.Done():154 return155 }156 }157 p.readErr = scanner.Err()158}159func (p *rpcProcess) write(value any) error {160 p.writeMu.Lock()161 defer p.writeMu.Unlock()162 return json.NewEncoder(p.input).Encode(value)163}164func (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, err167 }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.readErr179 }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, err185 }186 continue187 }188 if string(frame.ID) != fmt.Sprint(id) {189 continue190 }191 if frame.Error != nil {192 return nil, fmt.Errorf("ACP %s: %s", method, frame.Error.Message)193 }194 return frame.Result, nil195 }196 }197}198199func (p *rpcProcess) cancelSession() {200 p.cancelOnce.Do(func() {201 session := p.getSession()202 if session == "" {203 return204 }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.C215 case <-timer.C:216 // Closing the pipe releases a blocked write before we join its goroutine.217 _ = p.input.Close()218 <-done219 }220 })221}222223func (p *rpcProcess) setSession(session string) {224 p.sessionMu.Lock()225 p.session = session226 p.sessionMu.Unlock()227}228func (p *rpcProcess) getSession() string {229 p.sessionMu.RLock()230 defer p.sessionMu.RUnlock()231 return p.session232}233func (p *rpcProcess) hasSession() bool { return p.getSession() != "" }234func (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, &params) != 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 nil269 }270 if len(frame.ID) == 0 {271 return nil272 }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 err284 }285 for _, option := range params.Options {286 if option.ID == choice {287 outcome = map[string]string{"outcome": "selected", "optionId": choice}288 break289 }290 }291 }292 return p.write(map[string]any{"jsonrpc": "2.0", "id": frame.ID, "result": map[string]any{"outcome": outcome}})293}294295type boundedOutput struct {296 bytes.Buffer297 limit int298}299300func (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}306func 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 nil312 }313 command.Dir = opts.CWD314 command.Env = opts.Env315 command.WaitDelay = 2 * time.Second316 command.Stdin = strings.NewReader(prompt)317 output := &boundedOutput{limit: 64000}318 command.Stdout = output319 command.Stderr = io.Discard320 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()), nil325}