feat(batch): async wait/reconcile primitives for the Batches API

client/batch_async.go

@@ -0,0 +1,217 @@
+package client
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"math/rand"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+)
+
+// Async batch reconciliation, completing the BatchClient surface: a
+// wait-until-terminal loop with exponential backoff + jitter + Retry-After
+// honoring, per-request result retrieval from the JSONL results endpoint,
+// and a reconcile primitive that tolerates eventual consistency — an
+// individual request's result row can lag the batch reaching a terminal
+// state, so callers must keep polling for the row itself instead of
+// trusting the batch status alone.
+
+// PollOptions configures wait loops.
+type PollOptions struct {
+	// InitialInterval is the first sleep between polls (default 2s).
+	InitialInterval time.Duration
+	// MaxInterval caps the exponential growth of the sleep (default 30s).
+	MaxInterval time.Duration
+	// Timeout bounds total wall-clock waiting (default 10m).
+	Timeout time.Duration
+	// JitterFraction randomizes each sleep by ± this fraction (default 0.2).
+	JitterFraction float64
+}
+
+func (o PollOptions) withDefaults() PollOptions {
+	if o.InitialInterval <= 0 {
+		o.InitialInterval = 2 * time.Second
+	}
+	if o.MaxInterval <= 0 {
+		o.MaxInterval = 30 * time.Second
+	}
+	if o.Timeout <= 0 {
+		o.Timeout = 10 * time.Minute
+	}
+	if o.JitterFraction <= 0 || o.JitterFraction >= 1 {
+		o.JitterFraction = 0.2
+	}
+	return o
+}
+
+var terminalBatchStates = map[string]bool{
+	"ended":     true,
+	"completed": true,
+	"failed":    true,
+	"expired":   true,
+	"canceled":  true,
+	"cancelled": true,
+}
+
+func isTerminalBatchState(s string) bool { return terminalBatchStates[strings.ToLower(s)] }
+
+// backoffDelay computes attempt-th exponential delay with jitter, capped by
+// MaxInterval; a Retry-After seconds header on resp overrides it when larger.
+func backoffDelay(attempt int, retryAfter string, o PollOptions) time.Duration {
+	d := o.InitialInterval << uint(min(attempt, 16))
+	if d > o.MaxInterval || d <= 0 {
+		d = o.MaxInterval
+	}
+	j := 1 - o.JitterFraction + rand.Float64()*2*o.JitterFraction
+	d = time.Duration(float64(d) * j)
+	if ra, err := strconv.Atoi(strings.TrimSpace(retryAfter)); err == nil && ra > 0 {
+		rd := time.Duration(ra) * time.Second
+		if rd > d {
+			return rd
+		}
+	}
+	return d
+}
+
+// WaitUntilDone polls the batch until it reaches a terminal state or the
+// timeout elapses. Non-terminal responses keep polling; 429/5xx responses
+// are retried with Retry-After-aware backoff; other non-200s fail fast.
+func (bc *BatchClient) WaitUntilDone(ctx context.Context, batchID string, opts PollOptions) (*BatchResult, error) {
+	opts = opts.withDefaults()
+	deadline := time.Now().Add(opts.Timeout)
+	attempt := 0
+	for {
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, bc.baseURL+"/v1/messages/batches/"+batchID, nil)
+		if err != nil {
+			return nil, err
+		}
+		req.Header.Set("X-Api-Key", bc.apiKey)
+		req.Header.Set("Anthropic-Version", "2023-06-01")
+		req.Header.Set("Anthropic-Beta", "message-batches-2024-09-24")
+		resp, err := bc.httpClient.Do(req)
+		if err != nil {
+			return nil, fmt.Errorf("eyrie: batch wait: %w", err)
+		}
+		switch {
+		case resp.StatusCode == http.StatusOK:
+			var res BatchResult
+			derr := json.NewDecoder(resp.Body).Decode(&res)
+			_ = resp.Body.Close()
+			if derr != nil {
+				return nil, derr
+			}
+			if isTerminalBatchState(res.Status) {
+				return &res, nil
+			}
+		case resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500:
+			// transient: retried below with Retry-After-aware backoff
+		default:
+			errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
+			_ = resp.Body.Close()
+			return nil, fmt.Errorf("eyrie: batch wait error %d: %s", resp.StatusCode, strings.TrimSpace(string(errBody)))
+		}
+
+		if time.Now().After(deadline) {
+			return nil, fmt.Errorf("eyrie: batch %s not terminal within %s", batchID, opts.Timeout)
+		}
+		delay := backoffDelay(attempt, resp.Header.Get("Retry-After"), opts)
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(delay):
+		}
+		attempt++
+	}
+}
+
+// BatchRequestResult is one row of the batch results JSONL: the raw provider
+// payload is preserved byte-exact so hosts decode per-provider shapes.
+type BatchRequestResult struct {
+	CustomID string          `json:"custom_id"`
+	Result   json.RawMessage `json:"result,omitempty"`
+	Error    json.RawMessage `json:"error,omitempty"`
+}
+
+// ErrResultNotVisible reports that a request's result row was absent even
+// though polling continued past batch completion — callers may retry.
+var ErrResultNotVisible = errors.New("eyrie: batch result row not yet visible")
+
+// RequestResults fetches and parses the newline-delimited results document.
+func (bc *BatchClient) RequestResults(ctx context.Context, batchID string) ([]BatchRequestResult, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, bc.baseURL+"/v1/messages/batches/"+batchID+"/results", nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("X-Api-Key", bc.apiKey)
+	req.Header.Set("Anthropic-Version", "2023-06-01")
+	req.Header.Set("Anthropic-Beta", "message-batches-2024-09-24")
+	resp, err := bc.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("eyrie: batch results: %w", err)
+	}
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			slog.Warn("batch: close results body", "error", err)
+		}
+	}()
+	if resp.StatusCode != http.StatusOK {
+		errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
+		return nil, fmt.Errorf("eyrie: batch results error %d: %s", resp.StatusCode, strings.TrimSpace(string(errBody)))
+	}
+	var out []BatchRequestResult
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Buffer(make([]byte, 64*1024), 4*1024*1024)
+	for scanner.Scan() {
+		line := bytes.TrimSpace(scanner.Bytes())
+		if len(line) == 0 {
+			continue
+		}
+		var r BatchRequestResult
+		if err := json.Unmarshal(line, &r); err != nil {
+			return nil, fmt.Errorf("eyrie: parse results line: %w", err)
+		}
+		out = append(out, r)
+	}
+	return out, scanner.Err()
+}
+
+// PollRequestResult reconciles one request: waits until the batch is terminal
+// AND the custom_id's row appears, tolerating eventual consistency where the
+// results endpoint lags batch completion. After the batch completes, up to
+// opts.InitialInterval-scaled extra polls are made before surfacing
+// ErrResultNotVisible.
+func (bc *BatchClient) PollRequestResult(ctx context.Context, batchID, customID string, opts PollOptions) (*BatchRequestResult, error) {
+	opts = opts.withDefaults()
+	if _, err := bc.WaitUntilDone(ctx, batchID, opts); err != nil {
+		return nil, err
+	}
+	const maxExtra = 5
+	for extra := 0; ; extra++ {
+		rows, err := bc.RequestResults(ctx, batchID)
+		if err != nil {
+			return nil, err
+		}
+		for i := range rows {
+			if rows[i].CustomID == customID {
+				return &rows[i], nil
+			}
+		}
+		if extra >= maxExtra {
+			return nil, fmt.Errorf("%w after batch completion: %s", ErrResultNotVisible, customID)
+		}
+		delay := backoffDelay(extra, "", opts)
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(delay):
+		}
+	}
+}

client/batch_async_test.go

@@ -0,0 +1,173 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func fastOpts() PollOptions {
+	return PollOptions{InitialInterval: 5 * time.Millisecond, MaxInterval: 10 * time.Millisecond, Timeout: 2 * time.Second}
+}
+
+func TestWaitUntilDoneTerminal(t *testing.T) {
+	var polls int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n := atomic.AddInt32(&polls, 1)
+		status := "in_progress"
+		if n >= 3 {
+			status = "ended"
+		}
+		fmt.Fprintf(w, `{"id":"b1","status":%q}`, status)
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	res, err := bc.WaitUntilDone(context.Background(), "b1", fastOpts())
+	if err != nil {
+		t.Fatalf("WaitUntilDone: %v", err)
+	}
+	if res.Status != "ended" {
+		t.Fatalf("status = %q", res.Status)
+	}
+	if atomic.LoadInt32(&polls) < 3 {
+		t.Fatalf("expected >=3 polls, got %d", polls)
+	}
+}
+
+func TestWaitUntilDoneTimeout(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		fmt.Fprint(w, `{"id":"b1","status":"in_progress"}`)
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	opts := PollOptions{InitialInterval: 5 * time.Millisecond, MaxInterval: 10 * time.Millisecond, Timeout: 80 * time.Millisecond}
+	if _, err := bc.WaitUntilDone(context.Background(), "b1", opts); err == nil || !strings.Contains(err.Error(), "not terminal") {
+		t.Fatalf("err = %v", err)
+	}
+}
+
+func TestWaitUntilDoneHonorsRetryAfterOn429(t *testing.T) {
+	var saw429 int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if atomic.CompareAndSwapInt32(&saw429, 0, 1) {
+			w.Header().Set("Retry-After", "1")
+			w.WriteHeader(http.StatusTooManyRequests)
+			return
+		}
+		fmt.Fprint(w, `{"id":"b1","status":"ended"}`)
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	start := time.Now()
+	res, err := bc.WaitUntilDone(context.Background(), "b1", fastOpts())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res.Status != "ended" {
+		t.Fatal("wrong status")
+	}
+	// Retry-After: 1 second must have been honored (>= ~900ms elapsed).
+	if elapsed := time.Since(start); elapsed < 900*time.Millisecond {
+		t.Fatalf("Retry-After not honored; elapsed=%v", elapsed)
+	}
+}
+
+func TestWaitUntilDoneFailsFastOn401(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusUnauthorized)
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	_, err := bc.WaitUntilDone(context.Background(), "b1", fastOpts())
+	if err == nil || !strings.Contains(err.Error(), "401") {
+		t.Fatalf("err = %v", err)
+	}
+}
+
+const resultsJSONL = "{\"custom_id\":\"r1\",\"result\":{\"type\":\"succeeded\"}}\n" +
+	"{\"custom_id\":\"r2\",\"error\":{\"type\":\"invalid_request\"}}\n"
+
+func TestRequestResultsParsesJSONL(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		fmt.Fprint(w, resultsJSONL)
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	rows, err := bc.RequestResults(context.Background(), "b1")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(rows) != 2 {
+		t.Fatalf("rows = %d", len(rows))
+	}
+	if rows[0].CustomID != "r1" || len(rows[0].Result) == 0 {
+		t.Fatalf("row0 = %+v", rows[0])
+	}
+	if rows[1].CustomID != "r2" || len(rows[1].Error) == 0 {
+		t.Fatalf("row1 = %+v", rows[1])
+	}
+}
+
+func TestPollRequestResultToleratesLag(t *testing.T) {
+	// First results fetch returns empty (eventual consistency), second has the row.
+	var fetches int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case strings.HasSuffix(r.URL.Path, "/results"):
+			if atomic.AddInt32(&fetches, 1) == 1 {
+				fmt.Fprint(w, "")
+				return
+			}
+			fmt.Fprint(w, resultsJSONL)
+		default:
+			fmt.Fprint(w, `{"id":"b1","status":"ended"}`)
+		}
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	row, err := bc.PollRequestResult(context.Background(), "b1", "r2", fastOpts())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if row.CustomID != "r2" || len(row.Error) == 0 {
+		t.Fatalf("row = %+v", row)
+	}
+	if atomic.LoadInt32(&fetches) < 2 {
+		t.Fatal("expected a retry after empty results")
+	}
+}
+
+func TestPollRequestResultNotVisibleAfterExtraAttempts(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.HasSuffix(r.URL.Path, "/results") {
+			fmt.Fprint(w, `{"custom_id":"other","result":{}}`)
+			return
+		}
+		fmt.Fprint(w, `{"id":"b1","status":"ended"}`)
+	}))
+	defer srv.Close()
+	bc := NewBatchClient("k", srv.URL)
+	_, err := bc.PollRequestResult(context.Background(), "b1", "missing", fastOpts())
+	if !errors.Is(err, ErrResultNotVisible) {
+		t.Fatalf("err = %v, want ErrResultNotVisible", err)
+	}
+}
+
+func TestBackoffDelayRetryAfterOverride(t *testing.T) {
+	o := fastOpts()
+	d := backoffDelay(0, strconv.Itoa(60), o)
+	if d < 59*time.Second {
+		t.Fatalf("Retry-After override ignored: %v", d)
+	}
+	// Without header: exponential with cap.
+	if got := backoffDelay(20, "", o); got > o.MaxInterval*2 {
+		t.Fatalf("cap exceeded: %v", got)
+	}
+}