Merge pull request #239 from docker/implement-exit-behavior

Add context cancellation support for Ctrl+C during model response

Author: ericcurtin

cmd/cli/commands/run.go

@@ -2,11 +2,14 @@ package commands
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 
 	"github.com/charmbracelet/glamour"
 	"github.com/docker/model-runner/cmd/cli/commands/completion"
@@ -201,8 +204,36 @@ func generateInteractiveWithReadline(cmd *cobra.Command, desktopClient *desktop.
 		if sb.Len() > 0 && !multiline {
 			userInput := sb.String()
 
-			if err := chatWithMarkdown(cmd, desktopClient, backend, model, userInput, apiKey); err != nil {
-				cmd.PrintErr(handleClientError(err, "Failed to generate a response"))
+			// Create a cancellable context for the chat request
+			// This allows us to cancel the request if the user presses Ctrl+C during response generation
+			chatCtx, cancelChat := context.WithCancel(cmd.Context())
+			
+			// Set up signal handler to cancel the context on Ctrl+C
+			sigChan := make(chan os.Signal, 1)
+			signal.Notify(sigChan, syscall.SIGINT)
+			go func() {
+				select {
+				case <-sigChan:
+					cancelChat()
+				case <-chatCtx.Done():
+					// Context cancelled, exit goroutine
+				}
+			}()
+
+			err := chatWithMarkdownContext(chatCtx, cmd, desktopClient, backend, model, userInput, apiKey)
+			
+			// Clean up signal handler
+			signal.Stop(sigChan)
+			// Do not close sigChan to avoid race condition
+			cancelChat()
+
+			if err != nil {
+				// Check if the error is due to context cancellation (Ctrl+C during response)
+				if errors.Is(err, context.Canceled) {
+					cmd.Println()
+				} else {
+					cmd.PrintErr(handleClientError(err, "Failed to generate a response"))
+				}
 				sb.Reset()
 				continue
 			}
@@ -233,8 +264,36 @@ func generateInteractiveBasic(cmd *cobra.Command, desktopClient *desktop.Client,
 			continue
 		}
 
-		if err := chatWithMarkdown(cmd, desktopClient, backend, model, userInput, apiKey); err != nil {
-			cmd.PrintErr(handleClientError(err, "Failed to generate a response"))
+		// Create a cancellable context for the chat request
+		// This allows us to cancel the request if the user presses Ctrl+C during response generation
+		chatCtx, cancelChat := context.WithCancel(cmd.Context())
+		
+		// Set up signal handler to cancel the context on Ctrl+C
+		sigChan := make(chan os.Signal, 1)
+		signal.Notify(sigChan, syscall.SIGINT)
+		go func() {
+			select {
+			case <-sigChan:
+				cancelChat()
+			case <-chatCtx.Done():
+				// Context cancelled, exit goroutine
+				// Context cancelled, exit goroutine
+			}
+		}()
+
+		err = chatWithMarkdownContext(chatCtx, cmd, desktopClient, backend, model, userInput, apiKey)
+		
+		cancelChat()
+		signal.Stop(sigChan)
+		cancelChat()
+
+		if err != nil {
+			// Check if the error is due to context cancellation (Ctrl+C during response)
+			if errors.Is(err, context.Canceled) {
+				fmt.Println("\nUse Ctrl + d or /bye to exit.")
+			} else {
+				cmd.PrintErr(handleClientError(err, "Failed to generate a response"))
+			}
 			continue
 		}
 
@@ -425,21 +484,26 @@ func renderMarkdown(content string) (string, error) {
 
 // chatWithMarkdown performs chat and streams the response with selective markdown rendering.
 func chatWithMarkdown(cmd *cobra.Command, client *desktop.Client, backend, model, prompt, apiKey string) error {
+	return chatWithMarkdownContext(cmd.Context(), cmd, client, backend, model, prompt, apiKey)
+}
+
+// chatWithMarkdownContext performs chat with context support and streams the response with selective markdown rendering.
+func chatWithMarkdownContext(ctx context.Context, cmd *cobra.Command, client *desktop.Client, backend, model, prompt, apiKey string) error {
 	colorMode, _ := cmd.Flags().GetString("color")
 	useMarkdown := shouldUseMarkdown(colorMode)
 	debug, _ := cmd.Flags().GetBool("debug")
 
 	if !useMarkdown {
 		// Simple case: just stream as plain text
-		return client.Chat(backend, model, prompt, apiKey, func(content string) {
+		return client.ChatWithContext(ctx, backend, model, prompt, apiKey, func(content string) {
 			cmd.Print(content)
 		}, false)
 	}
 
 	// For markdown: use streaming buffer to render code blocks as they complete
 	markdownBuffer := NewStreamingMarkdownBuffer()
 
-	err := client.Chat(backend, model, prompt, apiKey, func(content string) {
+	err := client.ChatWithContext(ctx, backend, model, prompt, apiKey, func(content string) {
 		// Use the streaming markdown buffer to intelligently render content
 		rendered, err := markdownBuffer.AddContent(content, true)
 		if err != nil {

cmd/cli/desktop/desktop.go

@@ -366,6 +366,11 @@ func (c *Client) fullModelID(id string) (string, error) {
 
 // Chat performs a chat request and streams the response content with selective markdown rendering.
 func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(string), shouldUseMarkdown bool) error {
+	return c.ChatWithContext(context.Background(), backend, model, prompt, apiKey, outputFunc, shouldUseMarkdown)
+}
+
+// ChatWithContext performs a chat request with context support for cancellation and streams the response content with selective markdown rendering.
+func (c *Client) ChatWithContext(ctx context.Context, backend, model, prompt, apiKey string, outputFunc func(string), shouldUseMarkdown bool) error {
 	model = normalizeHuggingFaceModelName(model)
 	if !strings.Contains(strings.Trim(model, "/"), "/") {
 		// Do an extra API call to check if the model parameter isn't a model ID.
@@ -397,7 +402,8 @@ func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(str
 		completionsPath = inference.InferencePrefix + "/v1/chat/completions"
 	}
 
-	resp, err := c.doRequestWithAuth(
+	resp, err := c.doRequestWithAuthContext(
+		ctx,
 		http.MethodPost,
 		completionsPath,
 		bytes.NewReader(jsonData),
@@ -432,6 +438,13 @@ func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(str
 
 	scanner := bufio.NewScanner(resp.Body)
 	for scanner.Scan() {
+		// Check if context was cancelled
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		line := scanner.Text()
 		if line == "" {
 			continue
@@ -755,7 +768,11 @@ func (c *Client) doRequest(method, path string, body io.Reader) (*http.Response,
 
 // doRequestWithAuth is a helper function that performs HTTP requests with optional authentication
 func (c *Client) doRequestWithAuth(method, path string, body io.Reader, backend, apiKey string) (*http.Response, error) {
-	req, err := http.NewRequest(method, c.modelRunner.URL(path), body)
+	return c.doRequestWithAuthContext(context.Background(), method, path, body, backend, apiKey)
+}
+
+func (c *Client) doRequestWithAuthContext(ctx context.Context, method, path string, body io.Reader, backend, apiKey string) (*http.Response, error) {
+	req, err := http.NewRequestWithContext(ctx, method, c.modelRunner.URL(path), body)
 	if err != nil {
 		return nil, fmt.Errorf("error creating request: %w", err)
 	}