refactor: add context support to sleep functions for improved cancellation handling

internal/version/VERSION

@@ -1 +1 @@
-v5.0.0-beta-2506201738
+v5.0.0-beta-2506201912

uixt/driver_utils.go

@@ -1,6 +1,7 @@
 package uixt
 
 import (
+	"context"
 	"crypto/md5"
 	"fmt"
 	"io"
@@ -274,7 +275,8 @@ func getSimulationDuration(params []float64) (milliseconds int64) {
 
 // sleepStrict sleeps strict duration with given params
 // startTime is used to correct sleep duration caused by process time
-func sleepStrict(startTime time.Time, strictMilliseconds int64) {
+// ctx allows for cancellation during sleep
+func sleepStrict(ctx context.Context, startTime time.Time, strictMilliseconds int64) {
 	var elapsed int64
 	if !startTime.IsZero() {
 		elapsed = time.Since(startTime).Milliseconds()
@@ -294,7 +296,18 @@ func sleepStrict(startTime time.Time, strictMilliseconds int64) {
 		Int64("elapsed(ms)", elapsed).
 		Int64("strictSleep(ms)", strictMilliseconds).
 		Msg("sleep remaining duration time")
-	time.Sleep(time.Duration(dur) * time.Millisecond)
+
+	// Use context-aware sleep instead of blocking time.Sleep
+	select {
+	case <-time.After(time.Duration(dur) * time.Millisecond):
+		// Normal completion
+		log.Debug().Int64("duration_ms", dur).Msg("strict sleep completed normally")
+	case <-ctx.Done():
+		// Interrupted by context cancellation (e.g., CTRL+C)
+		log.Info().Int64("planned_duration_ms", dur).
+			Msg("strict sleep interrupted by context cancellation")
+		return
+	}
 }
 
 // global file lock

uixt/driver_utils_test.go

@@ -1,6 +1,7 @@
 package uixt
 
 import (
+	"context"
 	"strings"
 	"testing"
 	"time"
@@ -30,8 +31,9 @@ func TestGetSimulationDuration(t *testing.T) {
 }
 
 func TestSleepStrict(t *testing.T) {
+	ctx := context.Background()
 	startTime := time.Now()
-	sleepStrict(startTime, 1230)
+	sleepStrict(ctx, startTime, 1230)
 	dur := time.Since(startTime).Milliseconds()
 	t.Log(dur)
 	if dur < 1230 || dur > 1300 {

uixt/mcp_tools_utility.go

@@ -68,7 +68,15 @@ func (t *ToolSleep) Implement() server.ToolHandlerFunc {
 			return nil, fmt.Errorf("unsupported sleep duration type: %T", v)
 		}
 
-		time.Sleep(duration)
+		// Use context-aware sleep instead of blocking time.Sleep
+		select {
+		case <-time.After(duration):
+			// Normal completion
+		case <-ctx.Done():
+			// Interrupted by context cancellation (e.g., CTRL+C)
+			log.Warn().Msg("sleep interrupted by cancellation")
+			return nil, fmt.Errorf("sleep interrupted: %w", ctx.Err())
+		}
 
 		message := fmt.Sprintf("Successfully slept for %v seconds", actualSeconds)
 		returnData := ToolSleep{
@@ -120,7 +128,18 @@ func (t *ToolSleepMS) Implement() server.ToolHandlerFunc {
 
 		// Sleep MS action logic
 		log.Info().Int64("milliseconds", unifiedReq.Milliseconds).Msg("sleeping in milliseconds")
-		time.Sleep(time.Duration(unifiedReq.Milliseconds) * time.Millisecond)
+
+		duration := time.Duration(unifiedReq.Milliseconds) * time.Millisecond
+
+		// Use context-aware sleep instead of blocking time.Sleep
+		select {
+		case <-time.After(duration):
+			// Normal completion
+		case <-ctx.Done():
+			// Interrupted by context cancellation (e.g., CTRL+C)
+			log.Warn().Msg("sleep interrupted by cancellation")
+			return nil, fmt.Errorf("sleep interrupted: %w", ctx.Err())
+		}
 
 		message := fmt.Sprintf("Successfully slept for %d milliseconds", unifiedReq.Milliseconds)
 		returnData := ToolSleepMS{Milliseconds: unifiedReq.Milliseconds}
@@ -170,8 +189,8 @@ func (t *ToolSleepRandom) Implement() server.ToolHandlerFunc {
 			return nil, err
 		}
 
-		// Sleep random action logic
-		sleepStrict(time.Now(), getSimulationDuration(unifiedReq.Params))
+		// Sleep random action logic with context support
+		sleepStrict(ctx, time.Now(), getSimulationDuration(unifiedReq.Params))
 
 		message := fmt.Sprintf("Successfully slept for random duration with params: %v", unifiedReq.Params)
 		returnData := ToolSleepRandom{Params: unifiedReq.Params}