fix(router): improve endpoint selection logic with retries

pkg/embedclient/router.go

@@ -67,7 +67,19 @@ type RouterSlot struct {
 }
 
 func (r *TargetRouter) Embed(ctx context.Context, req Request) (Response, error) {
-	slot := r.pick()
+	tried := make(map[*endpointSlot]bool, len(r.slots))
+	var lastErr error
+
+	for range r.slots {
+		if ctx.Err() != nil {
+			return Response{}, ctx.Err()
+		}
+
+		slot := r.pickExcluding(tried)
+		if slot == nil {
+			break
+		}
+		tried[slot] = true
 
 		slot.mu.Lock()
 		slot.inflight++
@@ -76,31 +88,38 @@ func (r *TargetRouter) Embed(ctx context.Context, req Request) (Response, error)
 		}
 		slot.mu.Unlock()
 
-	defer func() {
+		timeout := time.Duration(r.cfg.TimeoutSecs) * time.Second
+		reqCtx, cancel := context.WithTimeout(ctx, timeout)
+		resp, err := slot.client.Embed(reqCtx, req)
+		cancel()
+
 		slot.mu.Lock()
 		slot.inflight--
 		if r.metrics != nil {
 			r.metrics.SetEndpointInflight(r.cfg.TargetName, slot.url, float64(slot.inflight))
 		}
 		slot.mu.Unlock()
-	}()
 
-	timeout := time.Duration(r.cfg.TimeoutSecs) * time.Second
-	ctx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
-
-	resp, err := slot.client.Embed(ctx, req)
 		if err != nil {
 			r.onFailure(slot, err)
-		return Response{}, fmt.Errorf("router embed [%s]: %w", slot.url, err)
+			lastErr = fmt.Errorf("router embed [%s]: %w", slot.url, err)
+			continue
 		}
 
 		r.onSuccess(slot)
 		return resp, nil
+	}
+
+	if lastErr != nil {
+		return Response{}, lastErr
+	}
+	return Response{}, fmt.Errorf("router embed: no endpoints available")
 }
 
-// pick selects the best available slot.
-func (r *TargetRouter) pick() *endpointSlot {
+// pickExcluding selects the best available slot not in the excluded set.
+// On the first call (empty excluded map) it behaves like the original pick().
+// During retry loops, already-tried slots are excluded so each attempt uses a fresh endpoint.
+func (r *TargetRouter) pickExcluding(excluded map[*endpointSlot]bool) *endpointSlot {
 	cooldown := time.Duration(r.cfg.CooldownSecs) * time.Second
 	now := time.Now()
 
@@ -108,10 +127,12 @@ func (r *TargetRouter) pick() *endpointSlot {
 	bestScore := -1 << 30
 
 	for _, s := range r.slots {
+		if excluded[s] {
+			continue
+		}
 		s.mu.Lock()
 		inCooldown := !s.lastFail.IsZero() && now.Sub(s.lastFail) < cooldown
 		score := s.priority - s.inflight
-		lastFail := s.lastFail
 		s.mu.Unlock()
 
 		if inCooldown {
@@ -120,14 +141,16 @@ func (r *TargetRouter) pick() *endpointSlot {
 		if best == nil || score > bestScore {
 			best = s
 			bestScore = score
-			_ = lastFail
 		}
 	}
 
-	// All in cooldown — fall back to oldest failure
+	// All non-excluded slots are in cooldown — fall back to the one with the oldest failure.
 	if best == nil {
 		var oldest time.Time
 		for _, s := range r.slots {
+			if excluded[s] {
+				continue
+			}
 			s.mu.Lock()
 			lf := s.lastFail
 			s.mu.Unlock()