Handle context cancellation in proxy logs (#3378)

Author: eleftherias

cmd/thv/app/logs.go

@@ -8,8 +8,10 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/adrg/xdg"
@@ -94,14 +96,20 @@ func logsCmdFunc(cmd *cobra.Command, args []string) error {
 	follow := viper.GetBool("follow")
 	proxy := viper.GetBool("proxy")
 
+	if follow {
+		var cancel context.CancelFunc
+		ctx, cancel = signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
+		defer cancel()
+	}
+
 	manager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to create workload manager: %w", err)
 	}
 
 	if proxy {
 		if follow {
-			return getProxyLogs(workloadName)
+			return getProxyLogs(ctx, workloadName)
 		}
 		// Use the shared manager method for non-follow proxy logs
 		// CLI gets all logs (0 = unlimited)
@@ -246,7 +254,7 @@ func reportPruneResults(prunedFiles, errs []string) {
 }
 
 // getProxyLogs reads and displays the proxy logs for a given workload in follow mode
-func getProxyLogs(workloadName string) error {
+func getProxyLogs(ctx context.Context, workloadName string) error {
 	// Get the proxy log file path
 	logFilePath, err := xdg.DataFile(fmt.Sprintf("toolhive/logs/%s.log", workloadName))
 	if err != nil {
@@ -262,11 +270,11 @@ func getProxyLogs(workloadName string) error {
 		return nil
 	}
 
-	return followProxyLogFile(cleanLogFilePath)
+	return followProxyLogFile(ctx, cleanLogFilePath)
 }
 
 // followProxyLogFile implements tail -f functionality for proxy logs
-func followProxyLogFile(logFilePath string) error {
+func followProxyLogFile(ctx context.Context, logFilePath string) error {
 	// Clean the file path to prevent path traversal
 	cleanLogFilePath := filepath.Clean(logFilePath)
 
@@ -295,6 +303,11 @@ func followProxyLogFile(logFilePath string) error {
 	}
 
 	// Follow the file for new content
+	contentCheckInterval := 100 * time.Millisecond
+
+	ticker := time.NewTicker(contentCheckInterval)
+	defer ticker.Stop()
+
 	for {
 		// Read any new content
 		buffer := make([]byte, 1024)
@@ -307,7 +320,12 @@ func followProxyLogFile(logFilePath string) error {
 			fmt.Print(string(buffer[:n]))
 		}
 
-		// Sleep briefly before checking for more content
-		time.Sleep(100 * time.Millisecond)
+		// Wait for next iteration or cancellation
+		select {
+		case <-ctx.Done():
+			return nil
+		case <-ticker.C:
+			// Continue to next iteration
+		}
 	}
 }