Merge pull request #484 from dgageot/ctrl-c

Handle ctrl-c for all commands

Author: dgageot

cmd/root/new.go

@@ -1,7 +1,6 @@
 package root
 
 import (
-	"bufio"
 	"fmt"
 	"os"
 	"strings"
@@ -78,15 +77,13 @@ func NewNewCmd() *cobra.Command {
 			if len(args) > 0 {
 				prompt = strings.Join(args, " ")
 			} else {
-				reader := bufio.NewReader(os.Stdin)
-
 				fmt.Printf("%s\n", blue("------- Welcome to %s! -------", bold(AppName)))
 				fmt.Printf("%s\n\n", white("         (Ctrl+C to exit)"))
 				fmt.Printf("%s\n\n", blue("What should your agent/agent team do? (describe its purpose)"))
 				fmt.Print(blue("> "))
 
 				var err error
-				prompt, err = reader.ReadString('\n')
+				prompt, err = readLine(ctx)
 				if err != nil {
 					return fmt.Errorf("failed to read purpose: %w", err)
 				}
@@ -133,7 +130,7 @@ func NewNewCmd() *cobra.Command {
 						llmIsTyping = false
 					}
 
-					result := promptMaxIterationsContinue(e.MaxIterations)
+					result := promptMaxIterationsContinue(ctx, e.MaxIterations)
 					switch result {
 					case ConfirmationApprove:
 						rt.Resume(ctx, string(runtime.ResumeTypeApprove))

cmd/root/root.go

@@ -1,13 +1,16 @@
 package root
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"syscall"
 
 	"github.com/spf13/cobra"
 
@@ -124,11 +127,10 @@ func NewRootCmd() *cobra.Command {
 	return cmd
 }
 
-func Run() {
-	Execute()
-}
+func Execute() error {
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
 
-func Execute() {
 	// Set the version for automatic telemetry initialization
 	telemetry.SetGlobalTelemetryVersion(version.Version)
 
@@ -147,11 +149,13 @@ We collect anonymous usage data to help improve cagent. To disable:
 	}
 
 	rootCmd := NewRootCmd()
-	if err := rootCmd.Execute(); err != nil {
+	if err := rootCmd.ExecuteContext(ctx); err != nil {
 		envErr := &environment.RequiredEnvError{}
 		runtimeErr := RuntimeError{}
 
 		switch {
+		case ctx.Err() != nil:
+			return ctx.Err()
 		case errors.As(err, &envErr):
 			fmt.Fprintln(os.Stderr, "The following environment variables must be set:")
 			for _, v := range envErr.Missing {
@@ -170,8 +174,10 @@ We collect anonymous usage data to help improve cagent. To disable:
 			}
 		}
 
-		os.Exit(1)
+		return err
 	}
+
+	return nil
 }
 
 // setupLogging configures slog logging behavior.

cmd/root/run.go

@@ -506,7 +506,7 @@ func runWithoutTUI(ctx context.Context, agentFilename string, rt runtime.Runtime
 					llmIsTyping = false
 				}
 
-				result := promptMaxIterationsContinue(e.MaxIterations)
+				result := promptMaxIterationsContinue(ctx, e.MaxIterations)
 				switch result {
 				case ConfirmationApprove:
 					rt.Resume(ctx, string(runtime.ResumeTypeApprove))
@@ -524,11 +524,13 @@ func runWithoutTUI(ctx context.Context, agentFilename string, rt runtime.Runtime
 				}
 
 				serverURL := e.Meta["cagent/server_url"].(string)
-				result := promptOAuthAuthorization(serverURL)
-				switch result {
-				case ConfirmationApprove:
+				result := promptOAuthAuthorization(ctx, serverURL)
+				switch {
+				case ctx.Err() != nil:
+					return ctx.Err()
+				case result == ConfirmationApprove:
 					_ = rt.ResumeElicitation(ctx, "accept", nil)
-				case ConfirmationReject:
+				case result == ConfirmationReject:
 					_ = rt.ResumeElicitation(ctx, "decline", nil)
 					return fmt.Errorf("OAuth authorization rejected by user")
 				}

cmd/root/run_text_utils.go

@@ -2,6 +2,7 @@ package root
 
 import (
 	"bufio"
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -117,13 +118,12 @@ func printToolCallResponse(toolCall tools.ToolCall, response string) {
 	fmt.Printf("\n%s\n", white("%s response%s", bold(toolCall.Function.Name), formatToolCallResponse(response)))
 }
 
-func promptMaxIterationsContinue(maxIterations int) ConfirmationResult {
+func promptMaxIterationsContinue(ctx context.Context, maxIterations int) ConfirmationResult {
 	fmt.Printf("\n%s\n", yellow("⚠️  Maximum iterations (%d) reached. The agent may be stuck in a loop.", maxIterations))
 	fmt.Printf("%s\n", white("This can happen with smaller or less capable models."))
 	fmt.Printf("\n%s (y/n): ", blue("Do you want to continue for 10 more iterations?"))
 
-	reader := bufio.NewReader(os.Stdin)
-	response, err := reader.ReadString('\n')
+	response, err := readLine(ctx)
 	if err != nil {
 		fmt.Printf("\n%s\n", red("Failed to read input, exiting..."))
 		return ConfirmationAbort
@@ -139,15 +139,14 @@ func promptMaxIterationsContinue(maxIterations int) ConfirmationResult {
 	}
 }
 
-func promptOAuthAuthorization(serverURL string) ConfirmationResult {
+func promptOAuthAuthorization(ctx context.Context, serverURL string) ConfirmationResult {
 	fmt.Printf("\n%s\n", yellow("🔐 OAuth Authorization Required"))
 	fmt.Printf("%s %s (remote)\n", white("Server:"), blue(serverURL))
 	fmt.Printf("%s\n", white("This server requires OAuth authentication to access its tools."))
 	fmt.Printf("%s\n", white("Your browser will open automatically to complete the authorization."))
 	fmt.Printf("\n%s (y/n): ", blue("Do you want to authorize access?"))
 
-	reader := bufio.NewReader(os.Stdin)
-	response, err := reader.ReadString('\n')
+	response, err := readLine(ctx)
 	if err != nil {
 		fmt.Printf("\n%s\n", red("Failed to read input, aborting authorization..."))
 		return ConfirmationAbort
@@ -290,3 +289,30 @@ func formatJSONValue(key string, value any) string {
 		return fmt.Sprintf("%s: %s", bold(key), string(jsonBytes))
 	}
 }
+
+func readLine(ctx context.Context) (string, error) {
+	lines := make(chan string)
+	errs := make(chan error)
+
+	go func() {
+		defer close(lines)
+		defer close(errs)
+
+		reader := bufio.NewReader(os.Stdin)
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			errs <- err
+		} else {
+			lines <- line
+		}
+	}()
+
+	select {
+	case <-ctx.Done():
+		return "", ctx.Err()
+	case err := <-errs:
+		return "", err
+	case line := <-lines:
+		return line, nil
+	}
+}

main.go

@@ -1,9 +1,13 @@
 package main
 
 import (
+	"os"
+
 	"github.com/docker/cagent/cmd/root"
 )
 
 func main() {
-	root.Run()
+	if err := root.Execute(); err != nil {
+		os.Exit(1)
+	}
 }

pkg/agent/agent.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"log/slog"
 	"math/rand"
-	"sync"
 
 	"github.com/docker/cagent/pkg/memorymanager"
 	"github.com/docker/cagent/pkg/model/provider"
@@ -17,9 +16,7 @@ type Agent struct {
 	name               string
 	description        string
 	instruction        string
-	toolsets           []tools.ToolSet
-	startedToolsets    map[tools.ToolSet]bool
-	toolsetsMutex      sync.RWMutex
+	toolsets           []*StartableToolSet
 	models             []provider.Provider
 	subAgents          []*Agent
 	parents            []*Agent
@@ -36,9 +33,8 @@ type Agent struct {
 // New creates a new agent
 func New(name, prompt string, opts ...Opt) *Agent {
 	agent := &Agent{
-		name:            name,
-		instruction:     prompt,
-		startedToolsets: make(map[tools.ToolSet]bool),
+		name:        name,
+		instruction: prompt,
 	}
 
 	for _, opt := range opts {
@@ -143,7 +139,13 @@ func (a *Agent) ToolDisplayName(ctx context.Context, toolName string) string {
 }
 
 func (a *Agent) ToolSets() []tools.ToolSet {
-	return a.toolsets
+	var toolSets []tools.ToolSet
+
+	for _, ts := range a.toolsets {
+		toolSets = append(toolSets, ts)
+	}
+
+	return toolSets
 }
 
 // Commands returns the named commands configured for this agent.
@@ -152,12 +154,9 @@ func (a *Agent) Commands() map[string]string {
 }
 
 func (a *Agent) ensureToolSetsAreStarted() error {
-	a.toolsetsMutex.Lock()
-	defer a.toolsetsMutex.Unlock()
-
 	for _, toolSet := range a.toolsets {
 		// Skip if toolset is already started
-		if a.startedToolsets[toolSet] {
+		if toolSet.started.Load() {
 			continue
 		}
 
@@ -172,19 +171,16 @@ func (a *Agent) ensureToolSetsAreStarted() error {
 		}
 
 		// Mark toolset as started
-		a.startedToolsets[toolSet] = true
+		toolSet.started.Store(true)
 	}
 
 	return nil
 }
 
 func (a *Agent) StopToolSets() error {
-	a.toolsetsMutex.Lock()
-	defer a.toolsetsMutex.Unlock()
-
 	for _, toolSet := range a.toolsets {
 		// Only stop toolsets that are marked as started
-		if !a.startedToolsets[toolSet] {
+		if !toolSet.started.Load() {
 			continue
 		}
 
@@ -193,7 +189,7 @@ func (a *Agent) StopToolSets() error {
 		}
 
 		// Mark toolset as stopped
-		a.startedToolsets[toolSet] = false
+		toolSet.started.Store(false)
 	}
 
 	return nil

pkg/agent/opts.go

@@ -1,6 +1,8 @@
 package agent
 
 import (
+	"sync/atomic"
+
 	"github.com/docker/cagent/pkg/memorymanager"
 	"github.com/docker/cagent/pkg/model/provider"
 	"github.com/docker/cagent/pkg/tools"
@@ -15,8 +17,15 @@ func WithInstruction(prompt string) Opt {
 }
 
 func WithToolSets(toolSet ...tools.ToolSet) Opt {
+	var startableToolSet []*StartableToolSet
+	for _, ts := range toolSet {
+		startableToolSet = append(startableToolSet, &StartableToolSet{
+			ToolSet: ts,
+		})
+	}
+
 	return func(a *Agent) {
-		a.toolsets = toolSet
+		a.toolsets = startableToolSet
 	}
 }
 
@@ -96,3 +105,8 @@ func WithCommands(commands map[string]string) Opt {
 		a.commands = commands
 	}
 }
+
+type StartableToolSet struct {
+	tools.ToolSet
+	started atomic.Bool
+}