Expose HTTP transient retry control and preserve shared capacity
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"},
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user