Expose HTTP transient retry control and preserve shared capacity

This commit is contained in:
2026-09-26 15:31:31 +00:00
parent acf511fc0d
commit a4ef91fb7e
7 changed files with 167 additions and 10 deletions

View File

@@ -13,6 +13,7 @@ type runRequestDTO struct {
Vars map[string]string `json:"vars,omitempty"`
Model *modelOverrideRequestDTO `json:"model,omitempty"`
IncludeRawOutput bool `json:"include_raw_output,omitempty"`
TransientRetries *int `json:"transient_retries,omitempty"`
}
type inputRefDTO struct {

View File

@@ -76,6 +76,10 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusBadRequest, "invalid_request", "prompt_id is required")
return
}
if req.TransientRetries != nil && *req.TransientRetries < 0 {
writeError(w, http.StatusBadRequest, "invalid_request", "transient_retries must be greater than or equal to 0")
return
}
var mappedInputs map[string]promptkit.ArtifactRef
if len(req.Inputs) > 0 {
mappedInputs = make(map[string]promptkit.ArtifactRef, len(req.Inputs))
@@ -94,13 +98,14 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}
res, err := h.runner.Run(r.Context(), promptkit.RunRequest{
PromptID: req.PromptID,
PromptVersion: req.PromptVersion,
ProfileID: req.ProfileID,
SessionID: req.SessionID,
Inputs: mappedInputs,
Vars: req.Vars,
Execution: model,
PromptID: req.PromptID,
PromptVersion: req.PromptVersion,
ProfileID: req.ProfileID,
SessionID: req.SessionID,
Inputs: mappedInputs,
Vars: req.Vars,
Execution: model,
TransientRetries: req.TransientRetries,
})
if err != nil {
status, code, message := mapRunError(err)
@@ -224,7 +229,7 @@ func mapRunError(err error) (int, string, string) {
return http.StatusBadRequest, "prompt_render_failed", "failed to render prompt"
case errors.Is(err, promptkit.ErrCapacityExceeded):
return http.StatusServiceUnavailable, "capacity_exceeded", "model backend capacity is exhausted"
case errors.Is(err, promptkit.ErrLLMGenerate):
case errors.Is(err, promptkit.ErrLLMGenerate), errors.Is(err, promptkit.ErrTransport), errors.Is(err, promptkit.ErrMalformedResponse), errors.Is(err, promptkit.ErrRetriesExhausted):
return http.StatusBadGateway, "llm_failed", "model generation request failed"
case errors.Is(err, promptkit.ErrValidation):
return http.StatusInternalServerError, "validation_runtime_failed", "validation runtime failed"

View File

@@ -64,6 +64,42 @@ func TestMaintainedHTTPRunExampleMatchesRequestContract(t *testing.T) {
assertHTTPErrorCode(t, invalidW, http.StatusBadRequest, "invalid_json")
}
func TestHandlerTransientRetryOverridePresenceAndStrictness(t *testing.T) {
for _, tc := range []struct {
name string
body string
want *int
status int
code string
}{
{name: "omitted", body: `{"prompt_id":"p"}`},
{name: "null", body: `{"prompt_id":"p","transient_retries":null}`},
{name: "zero", body: `{"prompt_id":"p","transient_retries":0}`, want: intPointer(0)},
{name: "positive", body: `{"prompt_id":"p","transient_retries":3}`, want: intPointer(3)},
{name: "negative", body: `{"prompt_id":"p","transient_retries":-1}`, status: http.StatusBadRequest, code: "invalid_request"},
{name: "string", body: `{"prompt_id":"p","transient_retries":"2"}`, status: http.StatusBadRequest, code: "invalid_json"},
{name: "fractional", body: `{"prompt_id":"p","transient_retries":1.5}`, status: http.StatusBadRequest, code: "invalid_json"},
} {
t.Run(tc.name, func(t *testing.T) {
runner := &fakeRunner{result: &promptkit.RunResult{}}
w := httptest.NewRecorder()
NewHandler(runner).ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/v1/runs", strings.NewReader(tc.body)))
if tc.status != 0 {
assertHTTPErrorCode(t, w, tc.status, tc.code)
if runner.last.PromptID != "" {
t.Fatalf("runner invoked for invalid request: %+v", runner.last)
}
return
}
if w.Code != http.StatusOK || !reflect.DeepEqual(runner.last.TransientRetries, tc.want) {
t.Fatalf("retry mapping: status=%d got=%v want=%v body=%s", w.Code, runner.last.TransientRetries, tc.want, w.Body.String())
}
})
}
}
func intPointer(value int) *int { return &value }
type handlerLLMClient struct{}
func (handlerLLMClient) Generate(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
@@ -774,6 +810,107 @@ model: test
}
}
func TestHandlerBuiltInTransientRetriesInheritAndAllowExplicitZero(t *testing.T) {
var hits int32
provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if atomic.AddInt32(&hits, 1) == 1 {
http.Error(w, "private provider detail", http.StatusServiceUnavailable)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"recovered"}}]}`))
}))
defer provider.Close()
h := newRetryHandler(t, provider.URL+"/v1", intPointer(1), 0, nil)
w := httptest.NewRecorder()
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/v1/runs", strings.NewReader(`{"prompt_id":"p"}`)))
if w.Code != http.StatusOK || atomic.LoadInt32(&hits) != 2 {
t.Fatalf("inherited retry: status=%d hits=%d body=%s", w.Code, atomic.LoadInt32(&hits), w.Body.String())
}
atomic.StoreInt32(&hits, 0)
w = httptest.NewRecorder()
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/v1/runs", strings.NewReader(`{"prompt_id":"p","transient_retries":0}`)))
assertHTTPErrorCode(t, w, http.StatusBadGateway, "llm_failed")
if atomic.LoadInt32(&hits) != 1 || strings.Contains(w.Body.String(), "private provider detail") {
t.Fatalf("explicit zero or safe error failed: hits=%d body=%s", atomic.LoadInt32(&hits), w.Body.String())
}
}
func TestHandlerRetainsSharedBackendLimitDuringRetry(t *testing.T) {
started := make(chan struct{}, 1)
release := make(chan struct{})
defer func() {
select {
case <-release:
default:
close(release)
}
}()
var hits, active, peak int32
provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
current := atomic.AddInt32(&active, 1)
defer atomic.AddInt32(&active, -1)
for {
old := atomic.LoadInt32(&peak)
if current <= old || atomic.CompareAndSwapInt32(&peak, old, current) {
break
}
}
if atomic.AddInt32(&hits, 1) == 1 {
started <- struct{}{}
<-release
http.Error(w, "temporary", http.StatusServiceUnavailable)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"ok"}}]}`))
}))
defer provider.Close()
queueCapacity := 0
h := newRetryHandler(t, provider.URL+"/v1", intPointer(1), 1, &queueCapacity)
first := httptest.NewRecorder()
done := make(chan struct{})
go func() {
defer close(done)
h.ServeHTTP(first, httptest.NewRequest(http.MethodPost, "/v1/runs", strings.NewReader(`{"prompt_id":"p"}`)))
}()
select {
case <-started:
case <-time.After(3 * time.Second):
t.Fatal("first provider call did not start")
}
second := httptest.NewRecorder()
h.ServeHTTP(second, httptest.NewRequest(http.MethodPost, "/v1/runs", strings.NewReader(`{"prompt_id":"p","transient_retries":0}`)))
assertHTTPErrorCode(t, second, http.StatusServiceUnavailable, "capacity_exceeded")
close(release)
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("retrying request did not complete")
}
if first.Code != http.StatusOK || atomic.LoadInt32(&hits) != 2 || atomic.LoadInt32(&peak) > 1 {
t.Fatalf("shared backend limit failed: first=%d hits=%d peak=%d body=%s", first.Code, atomic.LoadInt32(&hits), atomic.LoadInt32(&peak), first.Body.String())
}
}
func newRetryHandler(t *testing.T, endpoint string, retries *int, concurrency int, queue *int) *Handler {
t.Helper()
promptDir := t.TempDir()
profileDir := t.TempDir()
if err := os.WriteFile(filepath.Join(promptDir, "prompt.yaml"), []byte("id: p\nversion: \"1\"\ndefault_profile: limited\nmessages:\n - role: user\n content: hi\noutput:\n format: text\n validation_mode: none\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(profileDir, "profile.yaml"), []byte("id: limited\nbackend: limited\nmodel: test\n"), 0o644); err != nil {
t.Fatal(err)
}
engine, err := promptkit.NewEngine(promptkit.Config{PromptDir: promptDir, ProfileDir: profileDir, TransientRetries: retries}, promptkit.WithBackend(promptkit.Backend{ID: "limited", Endpoint: endpoint, ConcurrencyLimit: concurrency, QueueCapacity: queue}))
if err != nil {
t.Fatal(err)
}
return NewHandler(engine)
}
func TestHandlerRejectsOverlongSessionIDAndNonStringReasoningEffort(t *testing.T) {
engine := newHandlerEngine(t)
for _, body := range []string{
@@ -1126,6 +1263,10 @@ func TestHandlerPublicErrorMapping(t *testing.T) {
{name: "prompt render", err: wrap(promptkit.ErrPromptRender, fmt.Errorf("render failed")), status: http.StatusBadRequest, code: "prompt_render_failed", message: "failed to render prompt", avoidCause: "render failed"},
{name: "capacity", err: &promptkit.CapacityError{BackendID: "private-backend"}, status: http.StatusServiceUnavailable, code: "capacity_exceeded", message: "model backend capacity is exhausted", avoidCause: "private-backend"},
{name: "llm", err: wrap(promptkit.ErrLLMGenerate, fmt.Errorf("llm failed")), status: http.StatusBadGateway, code: "llm_failed", message: "model generation request failed", avoidCause: "llm failed"},
{name: "transport", err: wrap(promptkit.ErrTransport, fmt.Errorf("private transport detail")), status: http.StatusBadGateway, code: "llm_failed", message: "model generation request failed", avoidCause: "private transport detail"},
{name: "malformed response", err: wrap(promptkit.ErrMalformedResponse, fmt.Errorf("private response detail")), status: http.StatusBadGateway, code: "llm_failed", message: "model generation request failed", avoidCause: "private response detail"},
{name: "retry exhaustion", err: wrap(promptkit.ErrRetriesExhausted, fmt.Errorf("private provider detail")), status: http.StatusBadGateway, code: "llm_failed", message: "model generation request failed", avoidCause: "private provider detail"},
{name: "retry exhaustion value", err: fmt.Errorf("generation: %w", &promptkit.RetryExhaustedError{}), status: http.StatusBadGateway, code: "llm_failed", message: "model generation request failed"},
{name: "validation runtime", err: wrap(promptkit.ErrValidation, fmt.Errorf("validator broke")), status: http.StatusInternalServerError, code: "validation_runtime_failed", message: "validation runtime failed", avoidCause: "validator broke"},
}