diff --git a/internal/proxy/managedfailover_test.go b/internal/proxy/managedfailover_test.go index ba2fd9c..905eefc 100644 --- a/internal/proxy/managedfailover_test.go +++ b/internal/proxy/managedfailover_test.go @@ -130,6 +130,77 @@ func TestManagedDispatchAdvancesToFallbackOnResponseBodyReadError(t *testing.T) } } +// A managed turn may need several model rounds. After the primary fails and a +// fallback asks for a tool, the next round must stay on that fallback instead +// of paying the primary failure timeout again. +func TestManagedTurnKeepsFallbackStickyAcrossToolRounds(t *testing.T) { + toolSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"balance":5000}`)) + })) + defer toolSrv.Close() + + var primaryCalls, fallbackCalls int + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("decode request: %v", err) + } + w.Header().Set("Content-Type", "application/json") + switch payload["model"] { + case "primary-model": + primaryCalls++ + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":{"message":"primary unavailable"}}`)) + case "fallback-model": + fallbackCalls++ + if fallbackCalls == 1 { + _, _ = w.Write([]byte(`{"id":"chatcmpl-tool","choices":[{"message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"trading-api.get_market_context","arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`)) + return + } + _, _ = w.Write([]byte(`{"id":"chatcmpl-final","choices":[{"message":{"role":"assistant","content":"fallback finished"},"finish_reason":"stop"}]}`)) + default: + t.Fatalf("unexpected model %#v", payload["model"]) + } + })) + defer backend.Close() + + reg := provider.NewRegistry("") + reg.Set("vercel", &provider.Provider{ + Name: "vercel", BaseURL: backend.URL + "/v1", APIKey: "gateway-key", Auth: "bearer", + }) + policy := &agentctx.ModelPolicy{ + Mode: "clamp", + Allowed: []agentctx.AllowedModel{ + {Slot: "primary", Ref: "vercel/primary-model"}, + {Slot: "fallback", Ref: "vercel/fallback-model"}, + }, + } + h := NewHandler(reg, stubContextLoaderWithToolsAndPolicy("tiverton", "tiverton:dummy123", + managedToolManifestForURL(toolSrv.URL, http.MethodGet, "/api/v1/market_context/{claw_id}", ""), policy), + logging.New(io.Discard)) + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + bytes.NewBufferString(`{"model":"vercel/primary-model","messages":[{"role":"user","content":"check"}]}`)) + req.Header.Set("Authorization", "Bearer tiverton:dummy123") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected fallback turn to complete, got %d: %s", w.Code, w.Body.String()) + } + if primaryCalls != 1 { + t.Fatalf("primary calls = %d; want 1 across the whole managed turn", primaryCalls) + } + if fallbackCalls != 2 { + t.Fatalf("fallback calls = %d; want tool round plus final round", fallbackCalls) + } + if !strings.Contains(w.Body.String(), "fallback finished") { + t.Fatalf("unexpected downstream body: %s", w.Body.String()) + } +} + // With no declared fallback the terminal 502 and its message are preserved. func TestManagedDispatchKeepsTerminal502WhenNoFallbackRemains(t *testing.T) { toolSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/internal/proxy/toolmediation.go b/internal/proxy/toolmediation.go index 24753c3..f339960 100644 --- a/internal/proxy/toolmediation.go +++ b/internal/proxy/toolmediation.go @@ -280,6 +280,7 @@ func (h *Handler) handleManagedOpenAI(w http.ResponseWriter, r *http.Request, ag lastProvider = resp.ProviderName lastUpstreamModel = resp.UpstreamModel + candidates = stickyManagedCandidates(candidates, resp) usage, _ := cost.ExtractUsage(resp.Body) usageAgg.AddRound(agentID, usage, resp.ProviderName, resp.UpstreamModel, h) @@ -615,6 +616,7 @@ func (h *Handler) handleManagedAnthropic(w http.ResponseWriter, r *http.Request, lastProvider = resp.ProviderName lastUpstreamModel = resp.UpstreamModel + candidates = stickyManagedCandidates(candidates, resp) usage, _ := cost.ExtractUsage(resp.Body) usageAgg.AddRound(agentID, usage, resp.ProviderName, resp.UpstreamModel, h) @@ -883,6 +885,24 @@ func (h *Handler) dispatchCandidatesJSON(ctx context.Context, r *http.Request, a return h.dispatchCandidatesJSONWithReasoning(ctx, r, agentID, requestedModel, payload, candidates, requestInfo, nil) } +// Once a managed turn falls through to a declared fallback, keep subsequent +// model rounds on that candidate (or a later one). Retrying the failed primary +// after every tool result turns one upstream outage into N full candidate +// timeouts and can consume the managed turn deadline before it can answer. +// Match provider and upstream model because several policy slots may share a +// gateway provider while selecting different models. +func stickyManagedCandidates(candidates []dispatchCandidate, resp *capturedResponse) []dispatchCandidate { + if resp == nil { + return candidates + } + for i, candidate := range candidates { + if candidate.ProviderName == resp.ProviderName && candidate.UpstreamModel == resp.UpstreamModel { + return candidates[i:] + } + } + return candidates +} + func (h *Handler) dispatchCandidatesJSONWithReasoning(ctx context.Context, r *http.Request, agentID string, requestedModel string, payload map[string]any, candidates []dispatchCandidate, requestInfo *logging.RequestInfo, reasoningReplay []responsesReasoningReplay) (*capturedResponse, int, string, error) { sawCooldown := false start := time.Now()