Merge pull request #40 from petal-labs/feat/sdk-reliability-hardening

feat: SDK reliability hardening (errors, structured output, timeouts)

Author: erikhoward

core/batch.go

@@ -2,6 +2,7 @@ package core
 
 import (
 	"context"
+	"errors"
 	"time"
 )
 
@@ -79,6 +80,10 @@ func AsBatchProvider(p Provider) (BatchProvider, bool) {
 }
 
 // BatchWaiter provides utilities for waiting on batch completion.
+//
+// BatchWaiter polls BatchProvider directly and does not go through
+// core.Client, so core.WithTimeout does not apply; bound Wait via
+// WithMaxWait and/or a context deadline passed by the caller.
 type BatchWaiter struct {
 	provider     BatchProvider
 	pollInterval time.Duration
@@ -109,14 +114,36 @@ func (w *BatchWaiter) WithMaxWait(d time.Duration) *BatchWaiter {
 
 // Wait blocks until the batch completes or the context is cancelled.
 // Returns the final batch info, or an error if the wait times out or fails.
+//
+// Each poll is individually bounded by the remaining maxWait budget: a
+// GetBatchStatus call that hangs (e.g. a stuck network request) cannot
+// exceed the overall deadline, since the per-poll context is derived from
+// that deadline rather than from the caller's context alone.
 func (w *BatchWaiter) Wait(ctx context.Context, id BatchID) (*BatchInfo, error) {
 	deadline := time.Now().Add(w.maxWait)
 	ticker := time.NewTicker(w.pollInterval)
 	defer ticker.Stop()
 
 	for {
-		info, err := w.provider.GetBatchStatus(ctx, id)
+		// Bound this poll by the remaining maxWait budget (capped further
+		// by the caller's own ctx deadline, if any). context.WithDeadline
+		// automatically uses the earlier of the two deadlines, so a
+		// caller-supplied deadline shorter than maxWait still wins.
+		pollCtx, cancel := context.WithDeadline(ctx, deadline)
+		info, err := w.provider.GetBatchStatus(pollCtx, id)
+		cancel()
+
 		if err != nil {
+			// The caller's own context expiring/cancelling takes priority
+			// over our internally imposed deadline.
+			if ctx.Err() != nil {
+				return nil, ctx.Err()
+			}
+			// pollCtx timed out against our maxWait deadline (this is what
+			// bounds a stuck/hanging GetBatchStatus call).
+			if errors.Is(err, context.DeadlineExceeded) {
+				return nil, ErrBatchTimeout
+			}
 			return nil, err
 		}
 

core/batch_test.go

@@ -116,6 +116,33 @@ func (m *mockBatchProvider) ListBatches(_ context.Context, _ int) ([]BatchInfo,
 	return nil, nil
 }
 
+// stuckBatchProvider implements BatchProvider with a GetBatchStatus that
+// blocks until its context is cancelled, simulating a hung network call.
+// Used to verify that BatchWaiter.Wait bounds each poll by the remaining
+// maxWait budget instead of hanging indefinitely.
+type stuckBatchProvider struct{}
+
+func (s *stuckBatchProvider) CreateBatch(_ context.Context, _ []BatchRequest) (BatchID, error) {
+	return "batch_123", nil
+}
+
+func (s *stuckBatchProvider) GetBatchStatus(ctx context.Context, _ BatchID) (*BatchInfo, error) {
+	<-ctx.Done()
+	return nil, ctx.Err()
+}
+
+func (s *stuckBatchProvider) GetBatchResults(_ context.Context, _ BatchID) ([]BatchResult, error) {
+	return nil, nil
+}
+
+func (s *stuckBatchProvider) CancelBatch(_ context.Context, _ BatchID) error {
+	return nil
+}
+
+func (s *stuckBatchProvider) ListBatches(_ context.Context, _ int) ([]BatchInfo, error) {
+	return nil, nil
+}
+
 func TestAsBatchProvider(t *testing.T) {
 	t.Run("supports batch", func(t *testing.T) {
 		mock := &fullBatchProvider{}
@@ -246,6 +273,25 @@ func TestBatchWaiter(t *testing.T) {
 	})
 }
 
+func TestBatchWaiterDoesNotHangOnStuckPoll(t *testing.T) {
+	p := &stuckBatchProvider{}
+	w := NewBatchWaiter(p).
+		WithPollInterval(10 * time.Millisecond).
+		WithMaxWait(100 * time.Millisecond)
+
+	start := time.Now()
+	_, err := w.Wait(context.Background(), "batch-1")
+	if err == nil {
+		t.Fatal("want timeout error")
+	}
+	if elapsed := time.Since(start); elapsed > 2*time.Second {
+		t.Fatalf("Wait hung for %v; maxWait budget not enforced per-poll", elapsed)
+	}
+	if !errors.Is(err, ErrBatchTimeout) {
+		t.Errorf("Wait() error = %v, want ErrBatchTimeout", err)
+	}
+}
+
 func TestBatchWaiterWaitAndCollect(t *testing.T) {
 	t.Run("waits and collects results", func(t *testing.T) {
 		mock := &mockBatchProvider{

core/client.go

@@ -349,18 +349,41 @@ func (b *ChatBuilder) ResponseJSON() *ChatBuilder {
 // This enables structured output mode where the model produces JSON conforming to the schema.
 // The schema parameter defines the structure the output must conform to.
 //
+// This always forces schema.Strict = true. Strict mode requires the schema to
+// set "additionalProperties": false and list every property in "required" at
+// every object node; validate() rejects non-compliant schemas with
+// ErrInvalidSchema before the request is sent. Use ResponseJSONSchemaNonStrict
+// to opt out of strict mode for schemas that cannot meet those constraints.
+//
 // Example:
 //
 //	schema := &core.JSONSchemaDefinition{
 //	    Name:   "person",
-//	    Strict: true,
-//	    Schema: json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"},"age":{"type":"integer"}},"required":["name","age"]}`),
+//	    Schema: json.RawMessage(`{"type":"object","additionalProperties":false,"properties":{"name":{"type":"string"},"age":{"type":"integer"}},"required":["name","age"]}`),
 //	}
 //	resp, err := client.Chat(model).
 //	    User("Extract: John is 30 years old").
 //	    ResponseJSONSchema(schema).
 //	    GetResponse(ctx)
 func (b *ChatBuilder) ResponseJSONSchema(schema *JSONSchemaDefinition) *ChatBuilder {
+	if schema != nil {
+		schema.Strict = true
+	}
+	b.req.ResponseFormat = ResponseFormatJSONSchema
+	b.req.JSONSchema = schema
+	return b
+}
+
+// ResponseJSONSchemaNonStrict constrains the model output to match a specific
+// JSON Schema without enforcing strict mode. This forces schema.Strict = false,
+// skipping the strict-schema validation that ResponseJSONSchema applies. Use
+// this when a schema cannot satisfy strict mode's requirements (every object
+// node needs "additionalProperties": false and a "required" array covering
+// all declared properties) and the provider still accepts loose schemas.
+func (b *ChatBuilder) ResponseJSONSchemaNonStrict(schema *JSONSchemaDefinition) *ChatBuilder {
+	if schema != nil {
+		schema.Strict = false
+	}
 	b.req.ResponseFormat = ResponseFormatJSONSchema
 	b.req.JSONSchema = schema
 	return b
@@ -477,6 +500,51 @@ func (b *ChatBuilder) validate() error {
 		}
 	}
 
+	// Capability gate: reject schema-based structured output requests
+	// (ResponseFormatJSONSchema) against a provider/model that does not
+	// support core.FeatureStructuredOutput. This must run before the
+	// strict-schema check below so an unsupported provider fails with
+	// ErrStructuredOutputUnsupported rather than a schema validation error.
+	//
+	// Plain JSON mode (ResponseFormatJSON / json_object) is intentionally
+	// NOT hard-gated here: it has no schema/shape contract to silently
+	// violate, and many models support json_object without being tagged
+	// with FeatureStructuredOutput (e.g. OpenAI's gpt-3.5-turbo, gpt-4, and
+	// gpt-4-turbo). Gating it would reject requests that previously worked.
+	// If a provider genuinely can't do json_object, its descriptive API
+	// error surfaces instead.
+	if b.req.ResponseFormat == ResponseFormatJSONSchema {
+		if !b.client.provider.Supports(FeatureStructuredOutput) {
+			return fmt.Errorf("%w: provider %s model %s",
+				ErrStructuredOutputUnsupported, b.client.provider.ID(), b.req.Model)
+		}
+
+		// Model-level gate: the provider supports structured output overall,
+		// but the specific requested model may not (e.g. OpenAI supports it,
+		// but gpt-3.5-turbo / gpt-4 do not). If the model is present in the
+		// provider's catalog and lacks the capability, reject it. If the
+		// model is NOT present in the catalog (unknown, brand-new, or a
+		// custom deployment not yet reflected in the static catalog), fall
+		// back to allowing it since the provider-level check already passed.
+		for _, m := range b.client.provider.Models() {
+			if m.ID == b.req.Model {
+				if !m.HasCapability(FeatureStructuredOutput) {
+					return fmt.Errorf("%w: provider %s model %s",
+						ErrStructuredOutputUnsupported, b.client.provider.ID(), b.req.Model)
+				}
+				break
+			}
+		}
+	}
+
+	// Strict structured output requires a schema shape the provider can
+	// enforce exactly.
+	if b.req.ResponseFormat == ResponseFormatJSONSchema && b.req.JSONSchema != nil && b.req.JSONSchema.Strict {
+		if err := validateStrictSchema(b.req.JSONSchema.Schema); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }