review: enforce single HALF-OPEN probe, injectable clock, preserve backend errors

Address all 5 existing review comments plus 2 additional findings:

1. HALF-OPEN probe tracking: add probeInFlight flag so only one probe
   passes in HALF-OPEN state; concurrent requests get ErrCircuitOpen.
   Add TestCircuitBreaker_HalfOpenBlocksConcurrentProbes.

2. Nil map guard: getCircuitBreaker initializes circuitBreakers map
   when nil, preventing panic in test constructors that bypass NewUnified.

3. Preserve backend error text: rate-limit detection now returns the
   original upstream error message via extractRateLimitErrorText instead
   of wrapping in ErrCircuitOpen. ErrCircuitOpen is only returned when
   cb.Allow() rejects the call before contacting the backend.

4. TOML-only doc comments: clarify that rate_limit_threshold and
   rate_limit_cooldown are not wired for stdin JSON config.

5. Injectable clock: replace time.Now() in circuit breaker with nowFunc
   field (default time.Now). Tests use deterministic fake time instead
   of flaky time.Sleep-based assertions.

6. Operator precedence: add explicit parentheses in isRateLimitText
   for the compound rate limit + 403 condition.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: lpcox

internal/config/config_core.go

@@ -219,11 +219,13 @@ type ServerConfig struct {
 	// RateLimitThreshold is the number of consecutive rate-limit errors from this backend
 	// that will trip the circuit breaker (transition CLOSED → OPEN). When OPEN, requests
 	// are immediately rejected until the cooldown period elapses. Default: 3.
+	// Supported in file-based config (TOML/JSON); stdin JSON config does not currently accept this field.
 	RateLimitThreshold int `toml:"rate_limit_threshold" json:"rate_limit_threshold,omitempty"`
 
 	// RateLimitCooldown is the number of seconds the circuit breaker stays OPEN before
 	// allowing a single probe request (transition OPEN → HALF-OPEN). If the probe
 	// succeeds the circuit closes; if rate-limited again it re-opens. Default: 60.
+	// Supported in file-based config (TOML/JSON); stdin JSON config does not currently accept this field.
 	RateLimitCooldown int `toml:"rate_limit_cooldown" json:"rate_limit_cooldown,omitempty"`
 }
 

internal/server/circuit_breaker.go

@@ -61,11 +61,16 @@ type circuitBreaker struct {
 	openedAt          time.Time
 	// resetAt is the time when the upstream rate limit resets, parsed from
 	// the X-RateLimit-Reset header or the tool response message.
-	resetAt  time.Time
-	serverID string
+	resetAt       time.Time
+	probeInFlight bool
+	serverID      string
 
 	threshold int
 	cooldown  time.Duration
+
+	// nowFunc returns the current time. Defaults to time.Now; overridden in tests
+	// to avoid flaky time.Sleep-based assertions.
+	nowFunc func() time.Time
 }
 
 // newCircuitBreaker creates a circuit breaker for the given server ID.
@@ -83,6 +88,7 @@ func newCircuitBreaker(serverID string, threshold int, cooldown time.Duration) *
 		state:     circuitClosed,
 		threshold: threshold,
 		cooldown:  cooldown,
+		nowFunc:   time.Now,
 	}
 }
 
@@ -114,7 +120,7 @@ func (cb *circuitBreaker) Allow() error {
 	case circuitOpen:
 		// Check whether we should transition to HALF-OPEN.
 		// We use the upstream reset time when available, otherwise the cooldown.
-		now := time.Now()
+		now := cb.nowFunc()
 		var openUntil time.Time
 		if !cb.resetAt.IsZero() && cb.resetAt.After(cb.openedAt) {
 			openUntil = cb.resetAt
@@ -125,12 +131,18 @@ func (cb *circuitBreaker) Allow() error {
 			logCircuitBreaker.Printf("server %q circuit breaker OPEN → HALF-OPEN after cooldown", cb.serverID)
 			logger.LogInfo("backend", "circuit breaker for server %q transitioning OPEN → HALF-OPEN", cb.serverID)
 			cb.state = circuitHalfOpen
-			return nil // allow the probe
+			cb.probeInFlight = true
+			return nil // allow the single probe
 		}
 		return &ErrCircuitOpen{ServerID: cb.serverID, ResetAt: cb.resetAt}
 
 	case circuitHalfOpen:
-		// One probe is allowed through; further requests are blocked.
+		// Only one probe is allowed; further requests are blocked until the probe resolves.
+		if cb.probeInFlight {
+			return &ErrCircuitOpen{ServerID: cb.serverID, ResetAt: cb.resetAt}
+		}
+		// This shouldn't normally happen (probe resolved but state wasn't updated),
+		// but allow through defensively.
 		return nil
 	}
 
@@ -145,6 +157,7 @@ func (cb *circuitBreaker) RecordSuccess() {
 
 	prev := cb.state
 	cb.consecutiveErrors = 0
+	cb.probeInFlight = false
 	if cb.state == circuitHalfOpen {
 		cb.state = circuitClosed
 		cb.resetAt = time.Time{}
@@ -163,6 +176,7 @@ func (cb *circuitBreaker) RecordRateLimit(resetAt time.Time) {
 	defer cb.mu.Unlock()
 
 	cb.consecutiveErrors++
+	cb.probeInFlight = false
 	if !resetAt.IsZero() {
 		cb.resetAt = resetAt
 	}
@@ -171,7 +185,7 @@ func (cb *circuitBreaker) RecordRateLimit(resetAt time.Time) {
 	case circuitClosed:
 		if cb.consecutiveErrors >= cb.threshold {
 			cb.state = circuitOpen
-			cb.openedAt = time.Now()
+			cb.openedAt = cb.nowFunc()
 			logger.LogError("backend",
 				"circuit breaker for server %q OPENED after %d consecutive rate-limit errors; resets at %s",
 				cb.serverID, cb.consecutiveErrors, formatResetAt(cb.resetAt))
@@ -185,7 +199,7 @@ func (cb *circuitBreaker) RecordRateLimit(resetAt time.Time) {
 	case circuitHalfOpen:
 		// Probe failed — re-open the circuit.
 		cb.state = circuitOpen
-		cb.openedAt = time.Now()
+		cb.openedAt = cb.nowFunc()
 		logger.LogError("backend",
 			"circuit breaker for server %q re-OPENED after probe was rate-limited; resets at %s",
 			cb.serverID, formatResetAt(cb.resetAt))
@@ -213,6 +227,27 @@ func formatResetAt(t time.Time) string {
 	return fmt.Sprintf("%s (in %s)", t.UTC().Format(time.RFC3339), time.Until(t).Round(time.Second))
 }
 
+// extractRateLimitErrorText extracts the text content from a raw tool result
+// that has been identified as a rate-limit error. Returns the original backend
+// message so agents see the actual upstream error rather than a synthetic one.
+func extractRateLimitErrorText(result interface{}) string {
+	m, ok := result.(map[string]interface{})
+	if !ok {
+		return "rate limit exceeded"
+	}
+	contents, _ := m["content"].([]interface{})
+	for _, c := range contents {
+		cm, ok := c.(map[string]interface{})
+		if !ok {
+			continue
+		}
+		if text, ok := cm["text"].(string); ok && text != "" {
+			return text
+		}
+	}
+	return "rate limit exceeded"
+}
+
 // isRateLimitToolResult reports whether a raw tool call result indicates
 // a rate-limit error from the GitHub MCP server. It inspects the `isError`
 // flag and the text content for well-known rate-limit phrases.
@@ -251,7 +286,7 @@ func isRateLimitToolResult(result interface{}) (bool, time.Time) {
 func isRateLimitText(text string) bool {
 	lower := strings.ToLower(text)
 	return strings.Contains(lower, "rate limit exceeded") ||
-		strings.Contains(lower, "rate limit") && strings.Contains(lower, "403") ||
+		(strings.Contains(lower, "rate limit") && strings.Contains(lower, "403")) ||
 		strings.Contains(lower, "api rate limit") ||
 		strings.Contains(lower, "secondary rate limit") ||
 		strings.Contains(lower, "too many requests")

internal/server/circuit_breaker_test.go

@@ -61,15 +61,19 @@ func TestCircuitBreaker_SuccessResetsCounter(t *testing.T) {
 // TestCircuitBreaker_HalfOpenAfterCooldown verifies OPEN → HALF-OPEN transition.
 func TestCircuitBreaker_HalfOpenAfterCooldown(t *testing.T) {
 	t.Parallel()
-	// Use a very short cooldown so the test doesn't sleep long.
-	cb := newCircuitBreaker("test", 1, 10*time.Millisecond)
+	fakeNow := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+	cb := newCircuitBreaker("test", 1, time.Minute)
+	cb.nowFunc = func() time.Time { return fakeNow }
 
 	cb.RecordRateLimit(time.Time{})
 	require.Equal(t, circuitOpen, cb.State(), "should be OPEN after 1 error")
 
-	// Wait for cooldown.
-	time.Sleep(20 * time.Millisecond)
+	// Before cooldown: still OPEN.
+	fakeNow = fakeNow.Add(30 * time.Second)
+	require.Error(t, cb.Allow(), "should reject before cooldown elapses")
 
+	// After cooldown: transitions to HALF-OPEN.
+	fakeNow = fakeNow.Add(31 * time.Second)
 	err := cb.Allow()
 	assert.NoError(t, err, "should allow probe after cooldown")
 	assert.Equal(t, circuitHalfOpen, cb.State(), "should be HALF-OPEN after cooldown")
@@ -78,12 +82,14 @@ func TestCircuitBreaker_HalfOpenAfterCooldown(t *testing.T) {
 // TestCircuitBreaker_HalfOpenClosesOnSuccess verifies HALF-OPEN → CLOSED on probe success.
 func TestCircuitBreaker_HalfOpenClosesOnSuccess(t *testing.T) {
 	t.Parallel()
-	cb := newCircuitBreaker("test", 1, 10*time.Millisecond)
+	fakeNow := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+	cb := newCircuitBreaker("test", 1, time.Minute)
+	cb.nowFunc = func() time.Time { return fakeNow }
 
 	cb.RecordRateLimit(time.Time{})
 	require.Equal(t, circuitOpen, cb.State())
 
-	time.Sleep(20 * time.Millisecond)
+	fakeNow = fakeNow.Add(2 * time.Minute)
 	require.NoError(t, cb.Allow()) // probe allowed
 
 	cb.RecordSuccess()
@@ -94,12 +100,14 @@ func TestCircuitBreaker_HalfOpenClosesOnSuccess(t *testing.T) {
 // TestCircuitBreaker_HalfOpenReOpensOnRateLimit verifies HALF-OPEN → OPEN on probe failure.
 func TestCircuitBreaker_HalfOpenReOpensOnRateLimit(t *testing.T) {
 	t.Parallel()
-	cb := newCircuitBreaker("test", 1, 10*time.Millisecond)
+	fakeNow := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+	cb := newCircuitBreaker("test", 1, time.Minute)
+	cb.nowFunc = func() time.Time { return fakeNow }
 
 	cb.RecordRateLimit(time.Time{})
 	require.Equal(t, circuitOpen, cb.State())
 
-	time.Sleep(20 * time.Millisecond)
+	fakeNow = fakeNow.Add(2 * time.Minute)
 	require.NoError(t, cb.Allow()) // probe allowed
 
 	cb.RecordRateLimit(time.Time{})
@@ -114,21 +122,54 @@ func TestCircuitBreaker_HalfOpenReOpensOnRateLimit(t *testing.T) {
 // TestCircuitBreaker_ResetAtFromHeader verifies the reset time from upstream is used.
 func TestCircuitBreaker_ResetAtFromHeader(t *testing.T) {
 	t.Parallel()
-	cb := newCircuitBreaker("test", 1, 60*time.Second)
-	future := time.Now().Add(5 * time.Millisecond)
-	cb.RecordRateLimit(future)
+	fakeNow := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+	cb := newCircuitBreaker("test", 1, time.Hour)
+	cb.nowFunc = func() time.Time { return fakeNow }
+
+	resetAt := fakeNow.Add(30 * time.Second)
+	cb.RecordRateLimit(resetAt)
 	require.Equal(t, circuitOpen, cb.State())
 
 	// Before the reset time: still OPEN.
+	fakeNow = fakeNow.Add(15 * time.Second)
 	require.Error(t, cb.Allow())
 
-	// After the reset time: transitions to HALF-OPEN.
-	time.Sleep(10 * time.Millisecond)
+	// After the reset time: transitions to HALF-OPEN (before cooldown would elapse).
+	fakeNow = fakeNow.Add(20 * time.Second)
 	err := cb.Allow()
 	assert.NoError(t, err, "should allow probe after reset time")
 	assert.Equal(t, circuitHalfOpen, cb.State())
 }
 
+// TestCircuitBreaker_HalfOpenBlocksConcurrentProbes verifies that only one probe is allowed in HALF-OPEN.
+func TestCircuitBreaker_HalfOpenBlocksConcurrentProbes(t *testing.T) {
+	t.Parallel()
+	fakeNow := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+	cb := newCircuitBreaker("test", 1, time.Minute)
+	cb.nowFunc = func() time.Time { return fakeNow }
+
+	cb.RecordRateLimit(time.Time{})
+	require.Equal(t, circuitOpen, cb.State())
+
+	// Advance past cooldown to trigger HALF-OPEN.
+	fakeNow = fakeNow.Add(2 * time.Minute)
+
+	// First Allow() should succeed (the probe).
+	require.NoError(t, cb.Allow())
+	assert.Equal(t, circuitHalfOpen, cb.State())
+
+	// Second Allow() should be rejected — probe is already in flight.
+	err := cb.Allow()
+	require.Error(t, err, "concurrent requests in HALF-OPEN should be rejected")
+	var openErr *ErrCircuitOpen
+	require.ErrorAs(t, err, &openErr)
+
+	// After the probe succeeds, requests should be allowed again.
+	cb.RecordSuccess()
+	assert.Equal(t, circuitClosed, cb.State())
+	assert.NoError(t, cb.Allow())
+}
+
 // TestCircuitBreaker_DefaultsApplied verifies zero-value config gets sensible defaults.
 func TestCircuitBreaker_DefaultsApplied(t *testing.T) {
 	t.Parallel()
@@ -362,3 +403,33 @@ func TestParseRateLimitResetHeader(t *testing.T) {
 		})
 	}
 }
+
+// TestExtractRateLimitErrorText verifies extraction of error text from backend results.
+func TestExtractRateLimitErrorText(t *testing.T) {
+	t.Parallel()
+
+	t.Run("extracts text from standard rate-limit result", func(t *testing.T) {
+		t.Parallel()
+		result := map[string]interface{}{
+			"isError": true,
+			"content": []interface{}{
+				map[string]interface{}{
+					"type": "text",
+					"text": "failed to search: 403 API rate limit exceeded [rate reset in 42s]",
+				},
+			},
+		}
+		assert.Equal(t, "failed to search: 403 API rate limit exceeded [rate reset in 42s]", extractRateLimitErrorText(result))
+	})
+
+	t.Run("returns fallback for nil result", func(t *testing.T) {
+		t.Parallel()
+		assert.Equal(t, "rate limit exceeded", extractRateLimitErrorText(nil))
+	})
+
+	t.Run("returns fallback for empty content", func(t *testing.T) {
+		t.Parallel()
+		result := map[string]interface{}{"isError": true, "content": []interface{}{}}
+		assert.Equal(t, "rate limit exceeded", extractRateLimitErrorText(result))
+	})
+}

internal/server/unified.go

@@ -386,6 +386,9 @@ func buildCircuitBreakers(cfg *config.Config) map[string]*circuitBreaker {
 // getCircuitBreaker returns the circuit breaker for serverID, creating one with
 // defaults if none exists (e.g., when called from tests that bypass NewUnified).
 func (us *UnifiedServer) getCircuitBreaker(serverID string) *circuitBreaker {
+	if us.circuitBreakers == nil {
+		us.circuitBreakers = make(map[string]*circuitBreaker)
+	}
 	if cb, ok := us.circuitBreakers[serverID]; ok {
 		return cb
 	}
@@ -588,9 +591,14 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName
 		cb.RecordRateLimit(resetAt)
 		execSpan.SetAttributes(attribute.Bool("rate_limit.hit", true))
 		httpStatusCode = 429
-		// Return the original error message so the agent can see it.
-		return newErrorCallToolResult(fmt.Errorf("backend server %q rate-limited: %w",
-			serverID, &ErrCircuitOpen{ServerID: serverID, ResetAt: resetAt}))
+		// Preserve the original backend error text so the agent sees the actual upstream
+		// rate-limit details. ErrCircuitOpen is only returned when cb.Allow() rejects
+		// the call before contacting the backend.
+		errText := extractRateLimitErrorText(backendResult)
+		return &sdk.CallToolResult{
+			Content: []sdk.Content{&sdk.TextContent{Text: errText}},
+			IsError: true,
+		}, backendResult, nil
 	}
 	cb.RecordSuccess()