exec tools

internal/claude/session.go

@@ -206,6 +206,11 @@ type Session struct {
 	pendingMu sync.Mutex
 	pending   map[string]chan controlResult
 
+	ctx        context.Context
+	stop       context.CancelFunc
+	inflightMu sync.Mutex
+	inflight   map[string]context.CancelFunc
+
 	hooks   map[string]HookFunc
 	servers map[string]*mcp.Server
 
@@ -236,12 +241,14 @@ func Start(ctx context.Context, opts Options) (*Session, error) {
 		done:       make(chan struct{}),
 		transcript: opts.Transcript,
 		pending:    map[string]chan controlResult{},
+		inflight:   map[string]context.CancelFunc{},
 		hooks:      map[string]HookFunc{},
 		servers:    map[string]*mcp.Server{},
 		id:         initialID(opts),
 		onBadLine:  opts.OnBadLine,
 	}
 	s.lastLine.Store(time.Now().UnixNano())
+	s.ctx, s.stop = context.WithCancel(context.Background())
 
 	args, err := s.buildArgs(opts)
 	if err != nil {
@@ -556,6 +563,7 @@ func (s *Session) record(dir string, line []byte) {
 func (s *Session) readLoop(stdout io.Reader) {
 	defer close(s.done)
 	defer close(s.events)
+	defer s.stop()
 
 	r := bufio.NewReader(stdout)
 	for {
@@ -606,10 +614,23 @@ func (s *Session) handleLine(line []byte) {
 		return
 
 	case "control_request":
-		go s.handleControlRequest(env.RequestID, env.Request)
+		ctx, cancel := context.WithCancel(s.ctx)
+		s.inflightMu.Lock()
+		s.inflight[env.RequestID] = cancel
+		s.inflightMu.Unlock()
+		go s.handleControlRequest(ctx, cancel, env.RequestID, env.Request)
 		return
 
-	case "control_cancel_request", "keep_alive":
+	case "control_cancel_request":
+		s.inflightMu.Lock()
+		cancel := s.inflight[env.RequestID]
+		s.inflightMu.Unlock()
+		if cancel != nil {
+			cancel()
+		}
+		return
+
+	case "keep_alive":
 		return
 	}
 
@@ -674,8 +695,15 @@ func (s *Session) handleControlResponse(raw json.RawMessage) {
 	ch <- controlResult{response: resp.Response}
 }
 
-func (s *Session) handleControlRequest(id string, raw json.RawMessage) {
-	result, err := s.dispatchControl(raw)
+func (s *Session) handleControlRequest(ctx context.Context, cancel context.CancelFunc, id string, raw json.RawMessage) {
+	defer func() {
+		s.inflightMu.Lock()
+		delete(s.inflight, id)
+		s.inflightMu.Unlock()
+		cancel()
+	}()
+
+	result, err := s.dispatchControl(ctx, raw)
 	if err != nil {
 		slog.Error("control request failed", "request_id", id, "err", err)
 		s.write(map[string]any{
@@ -699,7 +727,7 @@ func (s *Session) handleControlRequest(id string, raw json.RawMessage) {
 	})
 }
 
-func (s *Session) dispatchControl(raw json.RawMessage) (any, error) {
+func (s *Session) dispatchControl(ctx context.Context, raw json.RawMessage) (any, error) {
 	req := struct {
 		Subtype    string          `json:"subtype"`
 		CallbackID string          `json:"callback_id"`
@@ -712,8 +740,6 @@ func (s *Session) dispatchControl(raw json.RawMessage) (any, error) {
 		return nil, err
 	}
 
-	ctx := context.Background()
-
 	switch req.Subtype {
 	case "hook_callback":
 		fn := s.hooks[req.CallbackID]