Add a pause that holds an agent after its current step

internal/codex/session.go

@@ -4,12 +4,15 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/rand"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"log/slog"
+	"net"
+	"net/http"
 	"os"
 	"strconv"
 	"sync"
@@ -17,6 +20,7 @@ import (
 	"time"
 
 	"g.fc.run/ai/byre/internal/agentinput"
+	"g.fc.run/ai/byre/internal/hookcmd"
 	"g.fc.run/ai/byre/internal/mcp"
 	"g.fc.run/ai/byre/internal/procgroup"
 )
@@ -32,6 +36,7 @@ type Options struct {
 	Reminder   string
 	MCPServers []*mcp.Server
 	StdioMCP   []mcp.Stdio
+	Hooks      map[string]hookcmd.Func
 	OnBadLine  func(line []byte, reason string)
 }
 
@@ -83,6 +88,7 @@ type Session struct {
 
 	servers  map[string]*mcp.Server
 	reminder string
+	hookSrv  *http.Server
 
 	onBadLine func(line []byte, reason string)
 	lastLine  atomic.Int64
@@ -94,7 +100,10 @@ func Start(ctx context.Context, opts Options) (*Session, error) {
 		return nil, err
 	}
 
-	err = s.open(ctx, opts)
+	config, err := s.serveHooks(opts)
+	if err == nil {
+		err = s.open(ctx, opts, config)
+	}
 	if err != nil {
 		s.Close()
 		return nil, err
@@ -102,6 +111,38 @@ func Start(ctx context.Context, opts Options) (*Session, error) {
 	return s, nil
 }
 
+func (s *Session) serveHooks(opts Options) (map[string]any, error) {
+	config := map[string]any{}
+	if len(opts.StdioMCP) > 0 {
+		config["mcp_servers"] = stdioConfig(opts.StdioMCP)
+	}
+	if len(opts.Hooks) == 0 {
+		return config, nil
+	}
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		return nil, err
+	}
+	token := rand.Text()
+	mux := http.NewServeMux()
+	mux.HandleFunc("POST /"+token+"/hook/{event}", hookcmd.Handler(opts.Hooks))
+	s.hookSrv = &http.Server{Handler: mux}
+	go s.hookSrv.Serve(ln)
+
+	hooks := map[string]any{}
+	for event := range opts.Hooks {
+		hooks[event] = []map[string]any{{"hooks": []map[string]any{{
+			"type":    "command",
+			"command": hookcmd.Command("http://" + ln.Addr().String() + "/" + token + "/hook/" + event),
+			"timeout": hookcmd.Timeout,
+		}}}}
+	}
+	config["hooks"] = hooks
+	config["bypass_hook_trust"] = true
+	return config, nil
+}
+
 func ReadRateLimits(ctx context.Context, cwd string, env []string) (RateLimitSnapshot, error) {
 	s, err := spawn(Options{Cwd: cwd, Env: env})
 	if err != nil {
@@ -177,7 +218,7 @@ func newSession(opts Options, stdin *os.File, stdout io.Reader) *Session {
 	return s
 }
 
-func (s *Session) open(ctx context.Context, opts Options) error {
+func (s *Session) open(ctx context.Context, opts Options, config map[string]any) error {
 	err := s.initialize(ctx)
 	if err != nil {
 		return err
@@ -192,7 +233,7 @@ func (s *Session) open(ctx context.Context, opts Options) error {
 			ApprovalPolicy:        json.RawMessage(`"never"`),
 			Sandbox:               SandboxModeDangerFullAccess,
 			DeveloperInstructions: opts.Guidance,
-			Config:                stdioConfig(opts.StdioMCP),
+			Config:                config,
 		}, &resp)
 		if err != nil {
 			return fmt.Errorf("resume thread: %w", err)
@@ -211,7 +252,7 @@ func (s *Session) open(ctx context.Context, opts Options) error {
 			ApprovalPolicy:        json.RawMessage(`"never"`),
 			Sandbox:               SandboxModeDangerFullAccess,
 			DeveloperInstructions: opts.Guidance,
-			Config:                stdioConfig(opts.StdioMCP),
+			Config:                config,
 		}, &resp)
 		if err != nil {
 			return fmt.Errorf("fork thread: %w", err)
@@ -228,7 +269,7 @@ func (s *Session) open(ctx context.Context, opts Options) error {
 		Sandbox:               SandboxModeDangerFullAccess,
 		DeveloperInstructions: opts.Guidance,
 		DynamicTools:          dynamicTools(opts.MCPServers),
-		Config:                stdioConfig(opts.StdioMCP),
+		Config:                config,
 	}, &resp)
 	if err != nil {
 		return fmt.Errorf("start thread: %w", err)
@@ -251,14 +292,11 @@ func (s *Session) initialize(ctx context.Context) error {
 }
 
 func stdioConfig(servers []mcp.Stdio) map[string]any {
-	if len(servers) == 0 {
-		return nil
-	}
 	cfg := map[string]any{}
 	for _, srv := range servers {
 		cfg[srv.Name] = map[string]any{"command": srv.Command, "args": srv.Args}
 	}
-	return map[string]any{"mcp_servers": cfg}
+	return cfg
 }
 
 func dynamicTools(servers []*mcp.Server) []DynamicToolSpec {
@@ -466,7 +504,14 @@ func (s *Session) notify(method string, params any) error {
 	return s.write(m, time.Now().Add(rpcTimeout))
 }
 
+func (s *Session) closeHooks() {
+	if s.hookSrv != nil {
+		s.hookSrv.Close()
+	}
+}
+
 func (s *Session) Close() error {
+	defer s.closeHooks()
 	s.writeLock <- struct{}{}
 	s.stdin.Close()
 	<-s.writeLock
@@ -482,6 +527,7 @@ func (s *Session) Close() error {
 }
 
 func (s *Session) Kill() {
+	defer s.closeHooks()
 	s.cmd.Kill()
 	<-s.done
 	s.cmd.Wait()