From f8720e16ed558455e33cc8bf3d7dbbda807ccbef Mon Sep 17 00:00:00 2001
From: pikaxinge <2392811793@qq.com>
Date: Sat, 28 Feb 2026 06:52:54 +0000
Subject: [PATCH] fix(retry): address review edge cases in notices and deadline
 handling

---
 pkg/tools/toolloop.go       |  2 +-
 pkg/tools/toolloop_test.go  | 44 +++++++++++++++++++++++++++++++++++++
 pkg/utils/llm_retry.go      |  5 ++++-
 pkg/utils/llm_retry_test.go | 35 +++++++++++++++++++++++++++++
 4 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go
index 53f2c52e8..e524916f6 100644
--- a/pkg/tools/toolloop.go
+++ b/pkg/tools/toolloop.go
@@ -74,7 +74,7 @@ func RunToolLoop(
 			if existingNotify != nil {
 				existingNotify(notice)
 			}
-			if config.RetryNotice != nil && channel != "" && chatID != "" {
+			if config.RetryNotice != nil {
 				config.RetryNotice(utils.FormatLLMRetryNotice(notice))
 			}
 		}
diff --git a/pkg/tools/toolloop_test.go b/pkg/tools/toolloop_test.go
index e156f02f5..e707955a6 100644
--- a/pkg/tools/toolloop_test.go
+++ b/pkg/tools/toolloop_test.go
@@ -81,6 +81,50 @@ func TestRunToolLoop_TransientRetry(t *testing.T) {
 	}
 }
 
+func TestRunToolLoop_TransientRetry_NoticeWithoutChannelContext(t *testing.T) {
+	provider := &flakyToolLoopProvider{
+		errors: []error{
+			fmt.Errorf("API request failed: status: 502 body: bad gateway"),
+		},
+	}
+
+	notices := make([]string, 0, 1)
+	cfg := ToolLoopConfig{
+		Provider:      provider,
+		Model:         "test-model",
+		MaxIterations: 1,
+		RetryPolicy: &utils.RetryPolicy{
+			AttemptTimeouts: []time.Duration{time.Second, time.Second},
+		},
+		RetryNotice: func(content string) {
+			notices = append(notices, content)
+		},
+	}
+
+	result, err := RunToolLoop(
+		context.Background(),
+		cfg,
+		[]providers.Message{{Role: "user", Content: "hello"}},
+		"",
+		"",
+	)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result == nil || result.Content != "ok" {
+		t.Fatalf("unexpected result: %+v", result)
+	}
+	if provider.calls != 2 {
+		t.Fatalf("provider.calls = %d, want 2", provider.calls)
+	}
+	if len(notices) != 1 {
+		t.Fatalf("notices = %d, want 1", len(notices))
+	}
+	if !strings.Contains(strings.ToLower(notices[0]), "retry") {
+		t.Fatalf("notice = %q, want retry hint", notices[0])
+	}
+}
+
 func TestRunToolLoop_NonRetryableError_NoRetry(t *testing.T) {
 	provider := &flakyToolLoopProvider{
 		errors: []error{
diff --git a/pkg/utils/llm_retry.go b/pkg/utils/llm_retry.go
index dabf2a15b..a53725667 100644
--- a/pkg/utils/llm_retry.go
+++ b/pkg/utils/llm_retry.go
@@ -120,7 +120,10 @@ func DoWithRetry[T any](ctx context.Context, policy RetryPolicy, fn func(context
 		if attemptTimeout := policy.AttemptTimeouts[attempt]; attemptTimeout > 0 {
 			timeout, ok := boundedAttemptTimeout(runCtx, attemptTimeout)
 			if !ok {
-				return zero, runCtx.Err()
+				if err := runCtx.Err(); err != nil {
+					return zero, err
+				}
+				return zero, context.DeadlineExceeded
 			}
 			attemptCtx, cancelAttempt = context.WithTimeout(runCtx, timeout)
 		}
diff --git a/pkg/utils/llm_retry_test.go b/pkg/utils/llm_retry_test.go
index c51729598..e83e50065 100644
--- a/pkg/utils/llm_retry_test.go
+++ b/pkg/utils/llm_retry_test.go
@@ -9,6 +9,19 @@ import (
 	"github.com/sipeed/picoclaw/pkg/providers"
 )
 
+type staleDeadlineContext struct {
+	context.Context
+	deadline time.Time
+}
+
+func (c staleDeadlineContext) Deadline() (time.Time, bool) {
+	return c.deadline, true
+}
+
+func (c staleDeadlineContext) Err() error {
+	return nil
+}
+
 func TestLLMRetry_ClassifyRetryDecision_429WithRetryAfter(t *testing.T) {
 	err := errors.New("API request failed:\n  Status: 429\n  Retry-After: 7")
 	decision := ClassifyRetryDecision(err)
@@ -163,3 +176,25 @@ func TestLLMRetry_ExtractRetryAfter_HTTPDate(t *testing.T) {
 		t.Fatalf("delay = %v, want 1m", delay)
 	}
 }
+
+func TestLLMRetry_DoWithRetry_ExpiredDeadlineReturnsDeadlineExceeded(t *testing.T) {
+	ctx := staleDeadlineContext{
+		Context:  context.Background(),
+		deadline: time.Now().Add(-time.Second),
+	}
+
+	calls := 0
+	_, err := DoWithRetry(ctx, RetryPolicy{
+		AttemptTimeouts: []time.Duration{time.Second},
+	}, func(context.Context) (string, error) {
+		calls++
+		return "ok", nil
+	})
+
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("err = %v, want context deadline exceeded", err)
+	}
+	if calls != 0 {
+		t.Fatalf("calls = %d, want 0", calls)
+	}
+}