Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
518e78db54 | ||
|
|
cb925618bc | ||
|
|
7cc96121e4 | ||
|
|
bba9033ac5 |
@@ -0,0 +1,103 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
// ErrAuth means the vendor rejected the credentials.
|
||||||
|
ErrAuth = errors.New("llm: authentication failed")
|
||||||
|
// ErrRateLimited means the vendor throttled the request. The concrete
|
||||||
|
// error is a *RateLimitedError carrying the requested wait.
|
||||||
|
ErrRateLimited = errors.New("llm: rate limited")
|
||||||
|
// ErrTimeout means the request did not finish within its timeout.
|
||||||
|
ErrTimeout = errors.New("llm: request timed out")
|
||||||
|
// ErrBadOutput means the model replied but the reply is unusable, e.g.
|
||||||
|
// structured output that does not parse.
|
||||||
|
ErrBadOutput = errors.New("llm: unusable model output")
|
||||||
|
// ErrProvider means the vendor returned an unexpected HTTP status. The
|
||||||
|
// concrete error is a *ProviderError carrying the status and a redacted
|
||||||
|
// body excerpt.
|
||||||
|
ErrProvider = errors.New("llm: provider error")
|
||||||
|
)
|
||||||
|
|
||||||
|
// RateLimitedError is a throttled request. It matches ErrRateLimited.
|
||||||
|
type RateLimitedError struct {
|
||||||
|
// RetryAfter is how long the vendor asked callers to wait; 0 means it
|
||||||
|
// did not say.
|
||||||
|
RetryAfter time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *RateLimitedError) Error() string {
|
||||||
|
if e.RetryAfter == 0 {
|
||||||
|
return ErrRateLimited.Error()
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s: retry after %v", ErrRateLimited, e.RetryAfter)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Is reports whether target is ErrRateLimited.
|
||||||
|
func (e *RateLimitedError) Is(target error) bool { return target == ErrRateLimited }
|
||||||
|
|
||||||
|
// MaxBodyExcerpt is the longest response-body excerpt, in bytes, a
|
||||||
|
// ProviderError keeps.
|
||||||
|
const MaxBodyExcerpt = 512
|
||||||
|
|
||||||
|
// redacted replaces each secret in a ProviderError body.
|
||||||
|
const redacted = "[REDACTED]"
|
||||||
|
|
||||||
|
// ProviderError is an unexpected HTTP status from the vendor. It matches
|
||||||
|
// ErrProvider. Build it with NewProviderError, which redacts the body.
|
||||||
|
type ProviderError struct {
|
||||||
|
status int
|
||||||
|
body string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewProviderError returns the error for an unexpected HTTP status.
|
||||||
|
// secrets must hold the API key (and any other credential sent with the
|
||||||
|
// request). Every occurrence of each non-empty secret in body is replaced
|
||||||
|
// before the body is cut to MaxBodyExcerpt bytes, so no part of a secret
|
||||||
|
// survives. If secrets holds no non-empty value the body cannot be
|
||||||
|
// redacted, so no excerpt is kept at all.
|
||||||
|
func NewProviderError(status int, body []byte, secrets []string) *ProviderError {
|
||||||
|
s := string(body)
|
||||||
|
redactedAny := false
|
||||||
|
for _, secret := range secrets {
|
||||||
|
if secret != "" {
|
||||||
|
s = strings.ReplaceAll(s, secret, redacted)
|
||||||
|
redactedAny = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !redactedAny {
|
||||||
|
return &ProviderError{status: status}
|
||||||
|
}
|
||||||
|
s = strings.TrimSpace(s)
|
||||||
|
if len(s) > MaxBodyExcerpt {
|
||||||
|
cut := MaxBodyExcerpt
|
||||||
|
for cut > 0 && !utf8.RuneStart(s[cut]) {
|
||||||
|
cut--
|
||||||
|
}
|
||||||
|
s = s[:cut] + "…"
|
||||||
|
}
|
||||||
|
return &ProviderError{status: status, body: s}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Status returns the HTTP status code.
|
||||||
|
func (e *ProviderError) Status() int { return e.status }
|
||||||
|
|
||||||
|
// Body returns the redacted, truncated response-body excerpt; it is empty
|
||||||
|
// when the vendor sent no body or no secret was given to redact it with.
|
||||||
|
func (e *ProviderError) Body() string { return e.body }
|
||||||
|
|
||||||
|
func (e *ProviderError) Error() string {
|
||||||
|
if e.body == "" {
|
||||||
|
return fmt.Sprintf("%s: HTTP %d", ErrProvider, e.status)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s: HTTP %d: %s", ErrProvider, e.status, e.body)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Is reports whether target is ErrProvider.
|
||||||
|
func (e *ProviderError) Is(target error) bool { return target == ErrProvider }
|
||||||
@@ -0,0 +1,121 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSentinelErrorsMatchWhenWrapped(t *testing.T) {
|
||||||
|
for _, sentinel := range []error{ErrAuth, ErrTimeout, ErrBadOutput} {
|
||||||
|
wrapped := fmt.Errorf("openai: generate: %w", sentinel)
|
||||||
|
if !errors.Is(wrapped, sentinel) {
|
||||||
|
t.Errorf("errors.Is(%v, %v) = false", wrapped, sentinel)
|
||||||
|
}
|
||||||
|
for _, other := range []error{ErrAuth, ErrTimeout, ErrBadOutput, ErrRateLimited, ErrProvider} {
|
||||||
|
if other != sentinel && errors.Is(wrapped, other) {
|
||||||
|
t.Errorf("errors.Is(%v, %v) = true", wrapped, other)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRateLimitedError(t *testing.T) {
|
||||||
|
err := fmt.Errorf("anthropic: %w", &RateLimitedError{RetryAfter: 30 * time.Second})
|
||||||
|
if !errors.Is(err, ErrRateLimited) {
|
||||||
|
t.Fatal("errors.Is(err, ErrRateLimited) = false")
|
||||||
|
}
|
||||||
|
if errors.Is(err, ErrProvider) || errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatal("rate-limited error matches an unrelated sentinel")
|
||||||
|
}
|
||||||
|
var rl *RateLimitedError
|
||||||
|
if !errors.As(err, &rl) {
|
||||||
|
t.Fatal("errors.As(err, *RateLimitedError) = false")
|
||||||
|
}
|
||||||
|
if rl.RetryAfter != 30*time.Second {
|
||||||
|
t.Fatalf("RetryAfter = %v, want 30s", rl.RetryAfter)
|
||||||
|
}
|
||||||
|
if got := err.Error(); got != "anthropic: llm: rate limited: retry after 30s" {
|
||||||
|
t.Fatalf("Error() = %q", got)
|
||||||
|
}
|
||||||
|
if got := (&RateLimitedError{}).Error(); got != "llm: rate limited" {
|
||||||
|
t.Fatalf("Error() without RetryAfter = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderError(t *testing.T) {
|
||||||
|
err := fmt.Errorf("gemini: %w", NewProviderError(503, []byte(` {"error":"overloaded"} `), []string{"sk-test"}))
|
||||||
|
if !errors.Is(err, ErrProvider) {
|
||||||
|
t.Fatal("errors.Is(err, ErrProvider) = false")
|
||||||
|
}
|
||||||
|
if errors.Is(err, ErrRateLimited) || errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatal("provider error matches an unrelated sentinel")
|
||||||
|
}
|
||||||
|
var pe *ProviderError
|
||||||
|
if !errors.As(err, &pe) {
|
||||||
|
t.Fatal("errors.As(err, *ProviderError) = false")
|
||||||
|
}
|
||||||
|
if pe.Status() != 503 || pe.Body() != `{"error":"overloaded"}` {
|
||||||
|
t.Fatalf("Status, Body = %d, %q", pe.Status(), pe.Body())
|
||||||
|
}
|
||||||
|
if got := err.Error(); got != `gemini: llm: provider error: HTTP 503: {"error":"overloaded"}` {
|
||||||
|
t.Fatalf("Error() = %q", got)
|
||||||
|
}
|
||||||
|
if got := NewProviderError(500, nil, []string{"sk-test"}).Error(); got != "llm: provider error: HTTP 500" {
|
||||||
|
t.Fatalf("Error() without body = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorRedactsSecrets(t *testing.T) {
|
||||||
|
const key = "sk-live-0123456789abcdef"
|
||||||
|
body := []byte(`{"error":"invalid key sk-live-0123456789abcdef","echo":"Bearer sk-live-0123456789abcdef"}`)
|
||||||
|
pe := NewProviderError(400, body, []string{key, ""})
|
||||||
|
for _, s := range []string{pe.Error(), pe.Body()} {
|
||||||
|
if strings.Contains(s, key) {
|
||||||
|
t.Fatalf("secret leaked: %q", s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
want := `{"error":"invalid key [REDACTED]","echo":"Bearer [REDACTED]"}`
|
||||||
|
if pe.Body() != want {
|
||||||
|
t.Fatalf("Body() = %q, want %q", pe.Body(), want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorRedactsBeforeTruncating(t *testing.T) {
|
||||||
|
const key = "sk-secret-straddling-the-cut"
|
||||||
|
// The key starts just before the excerpt limit, so truncating first
|
||||||
|
// would leave its prefix in the excerpt.
|
||||||
|
body := strings.Repeat("x", MaxBodyExcerpt-5) + key
|
||||||
|
pe := NewProviderError(500, []byte(body), []string{key})
|
||||||
|
if strings.Contains(pe.Body(), "sk-se") {
|
||||||
|
t.Fatalf("secret prefix leaked: %q", pe.Body())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorTruncatesOnRuneBoundary(t *testing.T) {
|
||||||
|
body := strings.Repeat("a", MaxBodyExcerpt-1) + "é" + strings.Repeat("b", 100)
|
||||||
|
pe := NewProviderError(500, []byte(body), []string{"sk-test"})
|
||||||
|
want := strings.Repeat("a", MaxBodyExcerpt-1) + "…"
|
||||||
|
if pe.Body() != want {
|
||||||
|
t.Fatalf("Body() has %d bytes, want %d", len(pe.Body()), len(want))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorOmitsBodyWithoutSecret(t *testing.T) {
|
||||||
|
const key = "sk-live-0123456789abcdef"
|
||||||
|
body := []byte(`{"error":"invalid key sk-live-0123456789abcdef"}`)
|
||||||
|
for _, secrets := range [][]string{nil, {}, {""}} {
|
||||||
|
pe := NewProviderError(401, body, secrets)
|
||||||
|
if pe.Body() != "" {
|
||||||
|
t.Fatalf("secrets %q: Body() = %q, want empty", secrets, pe.Body())
|
||||||
|
}
|
||||||
|
if got := pe.Error(); got != "llm: provider error: HTTP 401" {
|
||||||
|
t.Fatalf("secrets %q: Error() = %q", secrets, got)
|
||||||
|
}
|
||||||
|
if strings.Contains(pe.Error(), key) {
|
||||||
|
t.Fatalf("secrets %q: secret leaked: %q", secrets, pe.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"slices"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Fake is a scripted Provider for tests. Each method returns its queued
|
||||||
|
// results in order, one per call, and records every Request Generate
|
||||||
|
// receives. A call with nothing left in its queue fails, so an unexpected
|
||||||
|
// call is never silently answered. It is safe for concurrent use.
|
||||||
|
type Fake struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
generate []GenerateResult
|
||||||
|
models []ListModelsResult
|
||||||
|
verify []error
|
||||||
|
requests []Request
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ Provider = (*Fake)(nil)
|
||||||
|
|
||||||
|
// GenerateResult is one scripted Generate outcome.
|
||||||
|
type GenerateResult struct {
|
||||||
|
Response Response
|
||||||
|
Err error
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListModelsResult is one scripted ListModels outcome.
|
||||||
|
type ListModelsResult struct {
|
||||||
|
Models []Model
|
||||||
|
Err error
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueueGenerate appends results to the Generate script.
|
||||||
|
func (f *Fake) QueueGenerate(results ...GenerateResult) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.generate = append(f.generate, results...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueueListModels appends results to the ListModels script.
|
||||||
|
func (f *Fake) QueueListModels(results ...ListModelsResult) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.models = append(f.models, results...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueueVerify appends results to the Verify script; a nil entry is success.
|
||||||
|
func (f *Fake) QueueVerify(results ...error) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.verify = append(f.verify, results...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate records req and returns the next scripted result.
|
||||||
|
func (f *Fake) Generate(_ context.Context, req Request) (Response, error) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.requests = append(f.requests, cloneRequest(req))
|
||||||
|
if len(f.generate) == 0 {
|
||||||
|
return Response{}, fmt.Errorf("llm: fake: unscripted Generate call %d", len(f.requests))
|
||||||
|
}
|
||||||
|
r := f.generate[0]
|
||||||
|
f.generate = f.generate[1:]
|
||||||
|
return r.Response, r.Err
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListModels returns the next scripted result.
|
||||||
|
func (f *Fake) ListModels(context.Context) ([]Model, error) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
if len(f.models) == 0 {
|
||||||
|
return nil, fmt.Errorf("llm: fake: unscripted ListModels call")
|
||||||
|
}
|
||||||
|
r := f.models[0]
|
||||||
|
f.models = f.models[1:]
|
||||||
|
return slices.Clone(r.Models), r.Err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify returns the next scripted result.
|
||||||
|
func (f *Fake) Verify(context.Context) error {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
if len(f.verify) == 0 {
|
||||||
|
return fmt.Errorf("llm: fake: unscripted Verify call")
|
||||||
|
}
|
||||||
|
err := f.verify[0]
|
||||||
|
f.verify = f.verify[1:]
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Requests returns a copy of every Request Generate received, in order.
|
||||||
|
func (f *Fake) Requests() []Request {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
out := make([]Request, len(f.requests))
|
||||||
|
for i, r := range f.requests {
|
||||||
|
out[i] = cloneRequest(r)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remaining reports how many scripted results of each kind are unused, so a
|
||||||
|
// test can assert its script was consumed.
|
||||||
|
func (f *Fake) Remaining() (generate, listModels, verify int) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
return len(f.generate), len(f.models), len(f.verify)
|
||||||
|
}
|
||||||
|
|
||||||
|
// cloneRequest deep-copies r so a recorded Request is unaffected by later
|
||||||
|
// changes the caller makes to its slices and pointers.
|
||||||
|
func cloneRequest(r Request) Request {
|
||||||
|
r.Messages = slices.Clone(r.Messages)
|
||||||
|
if r.Temperature != nil {
|
||||||
|
t := *r.Temperature
|
||||||
|
r.Temperature = &t
|
||||||
|
}
|
||||||
|
if r.Schema != nil {
|
||||||
|
s := JSONSchema{Name: r.Schema.Name, Schema: json.RawMessage(slices.Clone([]byte(r.Schema.Schema)))}
|
||||||
|
r.Schema = &s
|
||||||
|
}
|
||||||
|
return r
|
||||||
|
}
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFakeGenerateReturnsScriptInOrder(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
var f Fake
|
||||||
|
f.QueueGenerate(
|
||||||
|
GenerateResult{Response: Response{Text: "first", Model: "m1"}},
|
||||||
|
GenerateResult{Err: &RateLimitedError{RetryAfter: time.Second}},
|
||||||
|
)
|
||||||
|
f.QueueGenerate(GenerateResult{Response: Response{JSON: json.RawMessage(`{"ok":true}`), FinishReason: "stop"}})
|
||||||
|
|
||||||
|
req := Request{Messages: []Message{{Role: RoleUser, Content: "one"}}}
|
||||||
|
resp, err := f.Generate(ctx, req)
|
||||||
|
if err != nil || resp.Text != "first" || resp.Model != "m1" {
|
||||||
|
t.Fatalf("call 1 = %+v, %v", resp, err)
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(ctx, req); !errors.Is(err, ErrRateLimited) {
|
||||||
|
t.Fatalf("call 2 err = %v, want ErrRateLimited", err)
|
||||||
|
}
|
||||||
|
resp, err = f.Generate(ctx, req)
|
||||||
|
if err != nil || string(resp.JSON) != `{"ok":true}` || resp.FinishReason != "stop" {
|
||||||
|
t.Fatalf("call 3 = %+v, %v", resp, err)
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(ctx, req); err == nil {
|
||||||
|
t.Fatal("unscripted call 4 succeeded")
|
||||||
|
}
|
||||||
|
if g, l, v := f.Remaining(); g != 0 || l != 0 || v != 0 {
|
||||||
|
t.Fatalf("Remaining() = %d, %d, %d", g, l, v)
|
||||||
|
}
|
||||||
|
if n := len(f.Requests()); n != 4 {
|
||||||
|
t.Fatalf("recorded %d requests, want 4 (unscripted calls are recorded too)", n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFakeRecordsRequests(t *testing.T) {
|
||||||
|
var f Fake
|
||||||
|
f.QueueGenerate(GenerateResult{}, GenerateResult{})
|
||||||
|
temp := 0.2
|
||||||
|
first := Request{
|
||||||
|
System: "be brief",
|
||||||
|
Messages: []Message{{Role: RoleUser, Content: "plan dinners"}},
|
||||||
|
Temperature: &temp,
|
||||||
|
MaxTokens: 100,
|
||||||
|
Schema: &JSONSchema{Name: "plan", Schema: json.RawMessage(`{"type":"object"}`)},
|
||||||
|
Timeout: 5 * time.Second,
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(context.Background(), first); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(context.Background(), Request{System: "second"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Later changes by the caller must not alter what was recorded.
|
||||||
|
first.Messages[0].Content = "changed"
|
||||||
|
temp = 1
|
||||||
|
first.Schema.Schema[2] = 'X'
|
||||||
|
|
||||||
|
got := f.Requests()
|
||||||
|
if len(got) != 2 || got[1].System != "second" {
|
||||||
|
t.Fatalf("Requests() = %+v", got)
|
||||||
|
}
|
||||||
|
r := got[0]
|
||||||
|
if r.System != "be brief" || r.Messages[0].Content != "plan dinners" || *r.Temperature != 0.2 ||
|
||||||
|
r.MaxTokens != 100 || r.Schema.Name != "plan" || string(r.Schema.Schema) != `{"type":"object"}` ||
|
||||||
|
r.Timeout != 5*time.Second {
|
||||||
|
t.Fatalf("recorded request = %+v", r)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Changing a returned copy must not alter the record either.
|
||||||
|
got[0].Messages[0].Content = "mutated"
|
||||||
|
if f.Requests()[0].Messages[0].Content != "plan dinners" {
|
||||||
|
t.Fatal("Requests() returned shared state")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFakeListModelsAndVerify(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
var f Fake
|
||||||
|
f.QueueListModels(
|
||||||
|
ListModelsResult{Models: []Model{{ID: "a", DisplayName: "A"}, {ID: "b", DisplayName: "b"}}},
|
||||||
|
ListModelsResult{Err: ErrAuth},
|
||||||
|
)
|
||||||
|
f.QueueVerify(nil, fmt.Errorf("wrapped: %w", ErrAuth))
|
||||||
|
|
||||||
|
models, err := f.ListModels(ctx)
|
||||||
|
if err != nil || len(models) != 2 || models[0].ID != "a" {
|
||||||
|
t.Fatalf("ListModels 1 = %+v, %v", models, err)
|
||||||
|
}
|
||||||
|
if _, err := f.ListModels(ctx); !errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatalf("ListModels 2 err = %v, want ErrAuth", err)
|
||||||
|
}
|
||||||
|
if _, err := f.ListModels(ctx); err == nil {
|
||||||
|
t.Fatal("unscripted ListModels succeeded")
|
||||||
|
}
|
||||||
|
if err := f.Verify(ctx); err != nil {
|
||||||
|
t.Fatalf("Verify 1 = %v", err)
|
||||||
|
}
|
||||||
|
if err := f.Verify(ctx); !errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatalf("Verify 2 = %v, want ErrAuth", err)
|
||||||
|
}
|
||||||
|
if err := f.Verify(ctx); err == nil {
|
||||||
|
t.Fatal("unscripted Verify succeeded")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFakeConcurrentUse(t *testing.T) {
|
||||||
|
const n = 50
|
||||||
|
var f Fake
|
||||||
|
for i := range n {
|
||||||
|
f.QueueGenerate(GenerateResult{Response: Response{Text: fmt.Sprint(i)}})
|
||||||
|
}
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for range n {
|
||||||
|
wg.Go(func() {
|
||||||
|
if _, err := f.Generate(context.Background(), Request{}); err != nil {
|
||||||
|
t.Error(err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
if len(f.Requests()) != n {
|
||||||
|
t.Fatalf("recorded %d requests, want %d", len(f.Requests()), n)
|
||||||
|
}
|
||||||
|
if g, _, _ := f.Remaining(); g != 0 {
|
||||||
|
t.Fatalf("%d Generate results unused", g)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,132 @@
|
|||||||
|
// Package llm is the single interface MealQuest uses to talk to a language
|
||||||
|
// model. Every vendor adapter implements Provider over plain net/http (no
|
||||||
|
// vendor SDKs) and reports failures with this package's typed errors, so
|
||||||
|
// callers decide what to do from errors.Is and errors.As alone.
|
||||||
|
//
|
||||||
|
// Errors from this package and its adapters never contain an API key.
|
||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Provider is a language-model vendor.
|
||||||
|
type Provider interface {
|
||||||
|
// Generate sends one request and returns the model's reply.
|
||||||
|
Generate(ctx context.Context, req Request) (Response, error)
|
||||||
|
// ListModels returns the models the configured credentials can use.
|
||||||
|
ListModels(ctx context.Context) ([]Model, error)
|
||||||
|
// Verify checks that the configured credentials are accepted.
|
||||||
|
Verify(ctx context.Context) error
|
||||||
|
}
|
||||||
|
|
||||||
|
// Role is the author of a Message.
|
||||||
|
type Role string
|
||||||
|
|
||||||
|
// The roles a Message may have. The system prompt is Request.System, not a
|
||||||
|
// message.
|
||||||
|
const (
|
||||||
|
RoleUser Role = "user"
|
||||||
|
RoleAssistant Role = "assistant"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Message is one turn of a conversation.
|
||||||
|
type Message struct {
|
||||||
|
Role Role
|
||||||
|
Content string
|
||||||
|
}
|
||||||
|
|
||||||
|
// JSONSchema is a named structured-output contract: when a Request carries
|
||||||
|
// one, the model must reply with JSON matching Schema.
|
||||||
|
type JSONSchema struct {
|
||||||
|
// Name identifies the contract to the vendor, e.g. "side_quests".
|
||||||
|
Name string
|
||||||
|
// Schema is a JSON Schema object.
|
||||||
|
Schema json.RawMessage
|
||||||
|
}
|
||||||
|
|
||||||
|
// Request is one generation request.
|
||||||
|
type Request struct {
|
||||||
|
// System is the system prompt; empty means none.
|
||||||
|
System string
|
||||||
|
// Messages is the conversation, oldest first. It must not be empty.
|
||||||
|
Messages []Message
|
||||||
|
// Temperature is the sampling temperature; nil means the vendor default.
|
||||||
|
Temperature *float64
|
||||||
|
// MaxTokens caps the reply length; 0 means the vendor default.
|
||||||
|
MaxTokens int
|
||||||
|
// Schema, when set, requests structured output; the reply is then in
|
||||||
|
// Response.JSON.
|
||||||
|
Schema *JSONSchema
|
||||||
|
// Timeout bounds the request; 0 means only ctx bounds it.
|
||||||
|
Timeout time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate reports whether r is well formed. Adapters call it before
|
||||||
|
// contacting the vendor.
|
||||||
|
func (r Request) Validate() error {
|
||||||
|
if len(r.Messages) == 0 {
|
||||||
|
return errors.New("llm: request has no messages")
|
||||||
|
}
|
||||||
|
for i, m := range r.Messages {
|
||||||
|
if m.Role != RoleUser && m.Role != RoleAssistant {
|
||||||
|
return fmt.Errorf("llm: message %d has unknown role %q", i, m.Role)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if r.Temperature != nil && *r.Temperature < 0 {
|
||||||
|
return fmt.Errorf("llm: temperature %v is negative", *r.Temperature)
|
||||||
|
}
|
||||||
|
if r.MaxTokens < 0 {
|
||||||
|
return fmt.Errorf("llm: max tokens %d is negative", r.MaxTokens)
|
||||||
|
}
|
||||||
|
if r.Timeout < 0 {
|
||||||
|
return fmt.Errorf("llm: timeout %v is negative", r.Timeout)
|
||||||
|
}
|
||||||
|
if r.Schema != nil {
|
||||||
|
if r.Schema.Name == "" {
|
||||||
|
return errors.New("llm: schema has no name")
|
||||||
|
}
|
||||||
|
var obj map[string]json.RawMessage
|
||||||
|
if err := json.Unmarshal(r.Schema.Schema, &obj); err != nil || obj == nil {
|
||||||
|
return fmt.Errorf("llm: schema %q is not a JSON object", r.Schema.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Usage is the token count of one Generate call.
|
||||||
|
type Usage struct {
|
||||||
|
InputTokens int
|
||||||
|
OutputTokens int
|
||||||
|
}
|
||||||
|
|
||||||
|
// Total returns InputTokens + OutputTokens.
|
||||||
|
func (u Usage) Total() int { return u.InputTokens + u.OutputTokens }
|
||||||
|
|
||||||
|
// Response is the model's reply to a Request.
|
||||||
|
type Response struct {
|
||||||
|
// Text is the reply text.
|
||||||
|
Text string
|
||||||
|
// JSON is the structured reply when the Request carried a Schema, and
|
||||||
|
// nil otherwise.
|
||||||
|
JSON json.RawMessage
|
||||||
|
// Usage is the tokens the call consumed.
|
||||||
|
Usage Usage
|
||||||
|
// Model is the model that answered, as the vendor reports it.
|
||||||
|
Model string
|
||||||
|
// FinishReason is why generation stopped, as the vendor reports it.
|
||||||
|
FinishReason string
|
||||||
|
}
|
||||||
|
|
||||||
|
// Model is a model a Provider offers.
|
||||||
|
type Model struct {
|
||||||
|
// ID is the vendor's model identifier, used in requests.
|
||||||
|
ID string
|
||||||
|
// DisplayName is a human-readable name; it equals ID when the vendor
|
||||||
|
// gives none.
|
||||||
|
DisplayName string
|
||||||
|
}
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRequestValidate(t *testing.T) {
|
||||||
|
neg := -0.1
|
||||||
|
valid := func() Request {
|
||||||
|
return Request{
|
||||||
|
Messages: []Message{{Role: RoleUser, Content: "hi"}, {Role: RoleAssistant, Content: "hello"}},
|
||||||
|
Schema: &JSONSchema{Name: "plan", Schema: json.RawMessage(`{"type":"object"}`)},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := valid().Validate(); err != nil {
|
||||||
|
t.Fatalf("valid request: %v", err)
|
||||||
|
}
|
||||||
|
cases := map[string]func(*Request){
|
||||||
|
"no messages": func(r *Request) { r.Messages = nil },
|
||||||
|
"unknown role": func(r *Request) { r.Messages[0].Role = "system" },
|
||||||
|
"negative temp": func(r *Request) { r.Temperature = &neg },
|
||||||
|
"negative max": func(r *Request) { r.MaxTokens = -1 },
|
||||||
|
"negative time": func(r *Request) { r.Timeout = -1 },
|
||||||
|
"unnamed schema": func(r *Request) { r.Schema.Name = "" },
|
||||||
|
"schema array": func(r *Request) { r.Schema.Schema = json.RawMessage(`[]`) },
|
||||||
|
"schema null": func(r *Request) { r.Schema.Schema = json.RawMessage(`null`) },
|
||||||
|
"schema invalid": func(r *Request) { r.Schema.Schema = json.RawMessage(`{`) },
|
||||||
|
"schema empty": func(r *Request) { r.Schema.Schema = nil },
|
||||||
|
}
|
||||||
|
for name, mutate := range cases {
|
||||||
|
r := valid()
|
||||||
|
mutate(&r)
|
||||||
|
if err := r.Validate(); err == nil {
|
||||||
|
t.Errorf("%s: Validate() = nil", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUsageTotal(t *testing.T) {
|
||||||
|
if got := (Usage{InputTokens: 12, OutputTokens: 30}).Total(); got != 42 {
|
||||||
|
t.Fatalf("Total() = %d, want 42", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
+64
-13
@@ -68,6 +68,14 @@ var codePattern = regexp.MustCompile(`^[a-z][a-z0-9_]{0,63}$`)
|
|||||||
// ErrNotFound is returned by Get for an unknown job id.
|
// ErrNotFound is returned by Get for an unknown job id.
|
||||||
var ErrNotFound = errors.New("jobs: job not found")
|
var ErrNotFound = errors.New("jobs: job not found")
|
||||||
|
|
||||||
|
// ErrUnknownKind is returned by Runner.Enqueue for a kind with no
|
||||||
|
// registered handler.
|
||||||
|
var ErrUnknownKind = errors.New("jobs: unknown job kind")
|
||||||
|
|
||||||
|
// ErrInvalidPayload is returned by Runner.Enqueue for a payload that does
|
||||||
|
// not decode into the kind's payload type.
|
||||||
|
var ErrInvalidPayload = errors.New("jobs: payload does not decode")
|
||||||
|
|
||||||
// Failure is a handler error reporting a failure code the household may
|
// Failure is a handler error reporting a failure code the household may
|
||||||
// see. Code must be one the job's kind declared at registration; the runner
|
// see. Code must be one the job's kind declared at registration; the runner
|
||||||
// stores it as the job's error. Any other handler error, including a Failure
|
// stores it as the job's error. Any other handler error, including a Failure
|
||||||
@@ -106,6 +114,13 @@ type Job struct {
|
|||||||
// never quotes an error.
|
// never quotes an error.
|
||||||
type runFunc func(ctx context.Context, payload string, publish func(json.RawMessage) error) (result json.RawMessage, code, reason string)
|
type runFunc func(ctx context.Context, payload string, publish func(json.RawMessage) error) (result json.RawMessage, code, reason string)
|
||||||
|
|
||||||
|
// handler is a registered kind: run runs a job, and check reports whether
|
||||||
|
// a payload decodes into the kind's payload type.
|
||||||
|
type handler struct {
|
||||||
|
run runFunc
|
||||||
|
check func(payload []byte) bool
|
||||||
|
}
|
||||||
|
|
||||||
// Runner owns the job handlers and the worker pool.
|
// Runner owns the job handlers and the worker pool.
|
||||||
type Runner struct {
|
type Runner struct {
|
||||||
st *store.Store
|
st *store.Store
|
||||||
@@ -118,7 +133,7 @@ type Runner struct {
|
|||||||
claim func(ctx context.Context, now sql.NullString) (db.Job, error)
|
claim func(ctx context.Context, now sql.NullString) (db.Job, error)
|
||||||
|
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
handlers map[string]runFunc
|
handlers map[string]handler
|
||||||
started bool
|
started bool
|
||||||
|
|
||||||
wg sync.WaitGroup
|
wg sync.WaitGroup
|
||||||
@@ -141,7 +156,7 @@ func NewRunner(st *store.Store, logger *slog.Logger, workers int) (*Runner, erro
|
|||||||
workers: workers,
|
workers: workers,
|
||||||
wake: make(chan struct{}, workers),
|
wake: make(chan struct{}, workers),
|
||||||
claim: db.New(st.Writer).ClaimNextJob,
|
claim: db.New(st.Writer).ClaimNextJob,
|
||||||
handlers: make(map[string]runFunc),
|
handlers: make(map[string]handler),
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -179,11 +194,9 @@ func Register[P any, R Public](r *Runner, kind string, h Handler[P, R], failureC
|
|||||||
if _, ok := r.handlers[kind]; ok {
|
if _, ok := r.handlers[kind]; ok {
|
||||||
return nil, fmt.Errorf("jobs: register %q: kind already registered", kind)
|
return nil, fmt.Errorf("jobs: register %q: kind already registered", kind)
|
||||||
}
|
}
|
||||||
r.handlers[kind] = func(ctx context.Context, payload string, publish func(json.RawMessage) error) (json.RawMessage, string, string) {
|
run := func(ctx context.Context, payload string, publish func(json.RawMessage) error) (json.RawMessage, string, string) {
|
||||||
var p P
|
p, err := decodePayload[P]([]byte(payload))
|
||||||
dec := json.NewDecoder(bytes.NewReader([]byte(payload)))
|
if err != nil {
|
||||||
dec.DisallowUnknownFields()
|
|
||||||
if err := dec.Decode(&p); err != nil {
|
|
||||||
return nil, CodeInternal, "payload does not decode"
|
return nil, CodeInternal, "payload does not decode"
|
||||||
}
|
}
|
||||||
result, err := h(ctx, p, func(partial R) error {
|
result, err := h(ctx, p, func(partial R) error {
|
||||||
@@ -209,9 +222,23 @@ func Register[P any, R Public](r *Runner, kind string, h Handler[P, R], failureC
|
|||||||
}
|
}
|
||||||
return data, "", ""
|
return data, "", ""
|
||||||
}
|
}
|
||||||
|
check := func(payload []byte) bool {
|
||||||
|
_, err := decodePayload[P](payload)
|
||||||
|
return err == nil
|
||||||
|
}
|
||||||
|
r.handlers[kind] = handler{run: run, check: check}
|
||||||
return &Kind[P]{r: r, name: kind}, nil
|
return &Kind[P]{r: r, name: kind}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// decodePayload decodes payload into P, refusing unknown fields.
|
||||||
|
func decodePayload[P any](payload []byte) (P, error) {
|
||||||
|
var p P
|
||||||
|
dec := json.NewDecoder(bytes.NewReader(payload))
|
||||||
|
dec.DisallowUnknownFields()
|
||||||
|
err := dec.Decode(&p)
|
||||||
|
return p, err
|
||||||
|
}
|
||||||
|
|
||||||
// publicJSON returns v's PublicJSON, refusing anything that is not one JSON
|
// publicJSON returns v's PublicJSON, refusing anything that is not one JSON
|
||||||
// value. Its error never quotes PublicJSON's error.
|
// value. Its error never quotes PublicJSON's error.
|
||||||
func publicJSON(v Public) (json.RawMessage, error) {
|
func publicJSON(v Public) (json.RawMessage, error) {
|
||||||
@@ -229,19 +256,43 @@ func (k *Kind[P]) Enqueue(ctx context.Context, payload P) (string, error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("jobs: enqueue %s: encode payload: %w", k.name, err)
|
return "", fmt.Errorf("jobs: enqueue %s: encode payload: %w", k.name, err)
|
||||||
}
|
}
|
||||||
|
return k.r.enqueue(ctx, k.name, data)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enqueue stores a queued job of kind with its JSON payload and returns its
|
||||||
|
// id, as Kind.Enqueue does for a kind known at compile time. It returns
|
||||||
|
// ErrUnknownKind for a kind with no registered handler and
|
||||||
|
// ErrInvalidPayload for a payload that does not decode into the kind's
|
||||||
|
// payload type; neither stores a job.
|
||||||
|
func (r *Runner) Enqueue(ctx context.Context, kind string, payload json.RawMessage) (string, error) {
|
||||||
|
r.mu.Lock()
|
||||||
|
h, ok := r.handlers[kind]
|
||||||
|
r.mu.Unlock()
|
||||||
|
if !ok {
|
||||||
|
return "", ErrUnknownKind
|
||||||
|
}
|
||||||
|
if !h.check(payload) {
|
||||||
|
return "", ErrInvalidPayload
|
||||||
|
}
|
||||||
|
return r.enqueue(ctx, kind, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// enqueue stores a queued job of kind with its encoded payload and wakes a
|
||||||
|
// worker.
|
||||||
|
func (r *Runner) enqueue(ctx context.Context, kind string, payload []byte) (string, error) {
|
||||||
id := store.NewID()
|
id := store.NewID()
|
||||||
now := store.Now().String()
|
now := store.Now().String()
|
||||||
if err := db.New(k.r.st.Writer).CreateJob(ctx, db.CreateJobParams{
|
if err := db.New(r.st.Writer).CreateJob(ctx, db.CreateJobParams{
|
||||||
ID: id,
|
ID: id,
|
||||||
Kind: k.name,
|
Kind: kind,
|
||||||
Payload: string(data),
|
Payload: string(payload),
|
||||||
CreatedAt: now,
|
CreatedAt: now,
|
||||||
UpdatedAt: now,
|
UpdatedAt: now,
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
return "", fmt.Errorf("jobs: enqueue %s: %w", k.name, err)
|
return "", fmt.Errorf("jobs: enqueue %s: %w", kind, err)
|
||||||
}
|
}
|
||||||
select {
|
select {
|
||||||
case k.r.wake <- struct{}{}:
|
case r.wake <- struct{}{}:
|
||||||
default: // every worker already has a token and will look again
|
default: // every worker already has a token and will look again
|
||||||
}
|
}
|
||||||
return id, nil
|
return id, nil
|
||||||
@@ -365,7 +416,7 @@ func (r *Runner) run(ctx context.Context, job db.Job) {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
result, code, reason := h(ctx, job.Payload, publish)
|
result, code, reason := h.run(ctx, job.Payload, publish)
|
||||||
if code != "" {
|
if code != "" {
|
||||||
if ctx.Err() != nil {
|
if ctx.Err() != nil {
|
||||||
logger.WarnContext(write, "job interrupted by shutdown; it fails at the next start")
|
logger.WarnContext(write, "job interrupted by shutdown; it fails at the next start")
|
||||||
|
|||||||
@@ -568,3 +568,39 @@ func TestClaimErrorIsRetried(t *testing.T) {
|
|||||||
t.Fatalf("claim error text logged: %s", out)
|
t.Fatalf("claim error text logged: %s", out)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestRunnerEnqueueByKindName(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
st := migratedStore(t)
|
||||||
|
r := newRunner(t, st, 1)
|
||||||
|
register(t, r, "count", func(_ context.Context, p countPayload, _ func(countResult) error) (countResult, error) {
|
||||||
|
return countResult{Done: p.Steps}, nil
|
||||||
|
})
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
if _, err := r.Enqueue(ctx, "no_such_kind", json.RawMessage(`{}`)); !errors.Is(err, ErrUnknownKind) {
|
||||||
|
t.Errorf("unknown kind error = %v, want ErrUnknownKind", err)
|
||||||
|
}
|
||||||
|
for _, payload := range []string{`{"steps":"three"}`, `{"other":1}`, `[]`, ``} {
|
||||||
|
if _, err := r.Enqueue(ctx, "count", json.RawMessage(payload)); !errors.Is(err, ErrInvalidPayload) {
|
||||||
|
t.Errorf("payload %q error = %v, want ErrInvalidPayload", payload, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
var n int
|
||||||
|
if err := st.Reader.QueryRow("SELECT COUNT(*) FROM jobs").Scan(&n); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n != 0 {
|
||||||
|
t.Errorf("refused enqueues stored %d jobs, want 0", n)
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := r.Enqueue(ctx, "count", json.RawMessage(`{"steps":3}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
start(t, r)
|
||||||
|
job := waitFor(t, r, id, finished)
|
||||||
|
if job.Kind != "count" || job.Status != StatusSucceeded || string(job.Result) != `{"done":3}` {
|
||||||
|
t.Errorf("job = %+v (result %s), want succeeded with {\"done\":3}", job, job.Result)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,7 +1,11 @@
|
|||||||
package server
|
package server
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
|
||||||
@@ -10,6 +14,80 @@ import (
|
|||||||
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/jobs"
|
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/jobs"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// maxJobBody bounds a POST /jobs request body.
|
||||||
|
const maxJobBody = 64 << 10
|
||||||
|
|
||||||
|
// jobRequest is the POST /jobs body. Payload is required; a job kind that
|
||||||
|
// takes no input is sent the empty object {}.
|
||||||
|
type jobRequest struct {
|
||||||
|
Kind string `json:"kind"`
|
||||||
|
Payload json.RawMessage `json:"payload"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// jobCreated is the POST /jobs response.
|
||||||
|
type jobCreated struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// postJob serves POST /api/v1/jobs: it enqueues a job of a registered kind
|
||||||
|
// and returns 202 with its id, which GET /api/v1/jobs/{id} reports on.
|
||||||
|
func postJob(logger *slog.Logger, runner *jobs.Runner) http.HandlerFunc {
|
||||||
|
return func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
req, err := decodeJobRequest(http.MaxBytesReader(w, r.Body, maxJobBody))
|
||||||
|
if err != nil {
|
||||||
|
writeError(w, r, logger, http.StatusBadRequest, "invalid_request", err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
id, err := runner.Enqueue(r.Context(), req.Kind, req.Payload)
|
||||||
|
switch {
|
||||||
|
case errors.Is(err, jobs.ErrUnknownKind):
|
||||||
|
writeError(w, r, logger, http.StatusBadRequest, "unknown_kind", "unknown job kind")
|
||||||
|
return
|
||||||
|
case errors.Is(err, jobs.ErrInvalidPayload):
|
||||||
|
writeError(w, r, logger, http.StatusBadRequest, "invalid_payload", "payload does not match the job kind")
|
||||||
|
return
|
||||||
|
case err != nil:
|
||||||
|
logger.ErrorContext(r.Context(), "enqueue job", "kind", req.Kind, "error", err)
|
||||||
|
writeError(w, r, logger, http.StatusInternalServerError, "internal", "internal error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, r, logger, http.StatusAccepted, jobCreated{ID: id})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// decodeJobRequest reads a single JSON object with a non-empty kind and a
|
||||||
|
// payload, refusing unknown fields, an omitted or null payload and trailing
|
||||||
|
// data.
|
||||||
|
func decodeJobRequest(body io.Reader) (jobRequest, error) {
|
||||||
|
data, err := io.ReadAll(body)
|
||||||
|
if err != nil {
|
||||||
|
var tooLarge *http.MaxBytesError
|
||||||
|
if errors.As(err, &tooLarge) {
|
||||||
|
return jobRequest{}, fmt.Errorf("request body exceeds %d bytes", tooLarge.Limit)
|
||||||
|
}
|
||||||
|
return jobRequest{}, errors.New("cannot read request body")
|
||||||
|
}
|
||||||
|
var req jobRequest
|
||||||
|
dec := json.NewDecoder(bytes.NewReader(data))
|
||||||
|
dec.DisallowUnknownFields()
|
||||||
|
if err := dec.Decode(&req); err != nil {
|
||||||
|
return jobRequest{}, errors.New("request body must be a JSON object with kind and payload")
|
||||||
|
}
|
||||||
|
if _, err := dec.Token(); err != io.EOF {
|
||||||
|
return jobRequest{}, errors.New("request body must be a single JSON object")
|
||||||
|
}
|
||||||
|
if req.Kind == "" {
|
||||||
|
return jobRequest{}, errors.New("kind: required")
|
||||||
|
}
|
||||||
|
if req.Payload == nil {
|
||||||
|
return jobRequest{}, errors.New("payload: required")
|
||||||
|
}
|
||||||
|
if bytes.Equal(req.Payload, []byte("null")) {
|
||||||
|
return jobRequest{}, errors.New("payload: must not be null")
|
||||||
|
}
|
||||||
|
return req, nil
|
||||||
|
}
|
||||||
|
|
||||||
// getJob serves GET /api/v1/jobs/{id}: the job's id, kind, status, latest
|
// getJob serves GET /api/v1/jobs/{id}: the job's id, kind, status, latest
|
||||||
// result and error.
|
// result and error.
|
||||||
func getJob(logger *slog.Logger, runner *jobs.Runner) http.HandlerFunc {
|
func getJob(logger *slog.Logger, runner *jobs.Runner) http.HandlerFunc {
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"slices"
|
"slices"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -151,3 +152,95 @@ func TestGetJobRequiresSession(t *testing.T) {
|
|||||||
t.Errorf("status = %d, want 401", rec.Code)
|
t.Errorf("status = %d, want 401", rec.Code)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestPostJobThenGetProgressAndResult drives the whole flow over the
|
||||||
|
// router: POST enqueues a progressive "draft" job, GET reads its partial
|
||||||
|
// result while the handler is held mid-run, then its final result.
|
||||||
|
func TestPostJobThenGetProgressAndResult(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
release := make(chan struct{})
|
||||||
|
runner, _ := startedRunner(t, release, nil)
|
||||||
|
router := newTestRouterWith(t, "dev", newTestSettings(t), runner)
|
||||||
|
|
||||||
|
rec, body := serve(t, router, signedIn(http.MethodPost, "/api/v1/jobs", strings.NewReader(`{"kind":"draft","payload":{}}`)))
|
||||||
|
if rec.Code != http.StatusAccepted {
|
||||||
|
t.Fatalf("status = %d, want 202: %s", rec.Code, rec.Body)
|
||||||
|
}
|
||||||
|
id, _ := body["id"].(string)
|
||||||
|
if keys := slices.Sorted(maps.Keys(body)); !slices.Equal(keys, []string{"id"}) || id == "" {
|
||||||
|
t.Fatalf("body = %v, want only a job id", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
mid := pollJob(t, router, id, func(b map[string]any) bool { return resultDays(b) == 1 })
|
||||||
|
if mid["id"] != id || mid["kind"] != "draft" || mid["status"] != "running" || mid["error"] != nil {
|
||||||
|
t.Errorf("mid-run body = %v, want running draft job %s with no error", mid, id)
|
||||||
|
}
|
||||||
|
|
||||||
|
close(release)
|
||||||
|
done := pollJob(t, router, id, func(b map[string]any) bool { return b["status"] == "succeeded" })
|
||||||
|
if done["id"] != id || resultDays(done) != 7 || done["error"] != nil {
|
||||||
|
t.Errorf("finished body = %v, want result {\"days\":7} and no error", done)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPostJobRefusesBadRequests(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
runner, _ := startedRunner(t, nil, nil)
|
||||||
|
router := newTestRouterWith(t, "dev", newTestSettings(t), runner)
|
||||||
|
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name, body, code, message string
|
||||||
|
}{
|
||||||
|
{"unknown kind", `{"kind":"no_such_kind","payload":{}}`, "unknown_kind", "unknown job kind"},
|
||||||
|
{"payload for another type", `{"kind":"draft","payload":{"days":3}}`, "invalid_payload", "payload does not match the job kind"},
|
||||||
|
{"payload not an object", `{"kind":"draft","payload":[1]}`, "invalid_payload", "payload does not match the job kind"},
|
||||||
|
{"missing payload", `{"kind":"draft"}`, "invalid_request", "payload: required"},
|
||||||
|
{"null payload", `{"kind":"draft","payload":null}`, "invalid_request", "payload: must not be null"},
|
||||||
|
{"missing kind", `{"payload":{}}`, "invalid_request", "kind: required"},
|
||||||
|
{"empty body", ``, "invalid_request", "request body must be a JSON object with kind and payload"},
|
||||||
|
{"not an object", `["draft"]`, "invalid_request", "request body must be a JSON object with kind and payload"},
|
||||||
|
{"unknown field", `{"kind":"draft","payload":{},"extra":1}`, "invalid_request", "request body must be a JSON object with kind and payload"},
|
||||||
|
{"trailing data", `{"kind":"draft","payload":{}} {}`, "invalid_request", "request body must be a single JSON object"},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
router.ServeHTTP(rec, signedIn(http.MethodPost, "/api/v1/jobs", strings.NewReader(tc.body)))
|
||||||
|
if rec.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("status = %d, want 400: %s", rec.Code, rec.Body)
|
||||||
|
}
|
||||||
|
var body errorBody
|
||||||
|
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if want := (errorBody{Error: apiError{Code: tc.code, Message: tc.message}}); body != want {
|
||||||
|
t.Errorf("body = %+v, want %+v", body, want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPostJobNeedsSessionAndCSRF(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
router := newTestRouter(t, "dev")
|
||||||
|
const body = `{"kind":"noop","payload":{}}`
|
||||||
|
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/v1/jobs", strings.NewReader(body)))
|
||||||
|
if rec.Code != http.StatusUnauthorized || rec.Body.String() != unauthenticatedJSON {
|
||||||
|
t.Errorf("no session: %d %s, want 401 %s", rec.Code, rec.Body, unauthenticatedJSON)
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, header := range []string{"", "wrong-token"} {
|
||||||
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", strings.NewReader(body))
|
||||||
|
req.AddCookie(validSessionCookie)
|
||||||
|
if header != "" {
|
||||||
|
req.Header.Set(CSRFHeader, header)
|
||||||
|
}
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
router.ServeHTTP(rec, req)
|
||||||
|
if rec.Code != http.StatusForbidden || rec.Body.String() != csrfJSON {
|
||||||
|
t.Errorf("CSRF header %q: %d %s, want 403 %s", header, rec.Code, rec.Body, csrfJSON)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ type RouteMounter interface {
|
|||||||
// only from a peer in trusted (see resolveOrigin). version is reported by
|
// only from a peer in trusted (see resolveOrigin). version is reported by
|
||||||
// GET /health; auth mounts the /auth routes; sessions guards /api/v1 and
|
// GET /health; auth mounts the /auth routes; sessions guards /api/v1 and
|
||||||
// /media; settings serves /api/v1/settings/household; runner serves
|
// /media; settings serves /api/v1/settings/household; runner serves
|
||||||
// /api/v1/jobs/{id}; assets is the web build
|
// POST /api/v1/jobs and GET /api/v1/jobs/{id}; assets is the web build
|
||||||
// served for every other GET. /health, /auth and the web build are public.
|
// served for every other GET. /health, /auth and the web build are public.
|
||||||
// Requests to /api/v1 other than GET also need the session's CSRF token in CSRFHeader.
|
// Requests to /api/v1 other than GET also need the session's CSRF token in CSRFHeader.
|
||||||
func NewRouter(logger *slog.Logger, version string, trusted forwarded.Trusted, auth RouteMounter, sessions SessionLookup, settings *household.SettingsService, runner *jobs.Runner, assets fs.FS) (http.Handler, error) {
|
func NewRouter(logger *slog.Logger, version string, trusted forwarded.Trusted, auth RouteMounter, sessions SessionLookup, settings *household.SettingsService, runner *jobs.Runner, assets fs.FS) (http.Handler, error) {
|
||||||
@@ -68,6 +68,7 @@ func NewRouter(logger *slog.Logger, version string, trusted forwarded.Trusted, a
|
|||||||
apiV1.Get("/me", me(logger))
|
apiV1.Get("/me", me(logger))
|
||||||
apiV1.Get("/settings/household", getHouseholdSettings(logger, settings))
|
apiV1.Get("/settings/household", getHouseholdSettings(logger, settings))
|
||||||
apiV1.Patch("/settings/household", patchHouseholdSettings(logger, settings))
|
apiV1.Patch("/settings/household", patchHouseholdSettings(logger, settings))
|
||||||
|
apiV1.Post("/jobs", postJob(logger, runner))
|
||||||
apiV1.Get("/jobs/{id}", getJob(logger, runner))
|
apiV1.Get("/jobs/{id}", getJob(logger, runner))
|
||||||
r.Route("/api", func(r chi.Router) {
|
r.Route("/api", func(r chi.Router) {
|
||||||
r.MethodNotAllowed(methodNotAllowed)
|
r.MethodNotAllowed(methodNotAllowed)
|
||||||
|
|||||||
Reference in New Issue
Block a user