feat(sdk/go): add context.Context to Store interface methods

sdk/go/README.md

@@ -112,7 +112,7 @@ PostgreSQL:
 ```go
 import "github.com/mozilla-ai/cq/sdk/go/stores/postgres"
 
-store, err := postgres.New("postgres://user:pass@localhost:5432/cq")
+store, err := postgres.New(context.Background(), "postgres://user:pass@localhost:5432/cq")
 c, err := cq.NewClient(cq.WithStore(store))
 ```
 

sdk/go/client.go

@@ -93,12 +93,14 @@ func (c *Client) Confirm(ctx context.Context, ku KnowledgeUnit) (KnowledgeUnit,
 	ctx, cancel := c.operationContext(ctx)
 	defer cancel()
 
-	if err := ctx.Err(); err != nil {
-		return KnowledgeUnit{}, err
+	select {
+	case <-ctx.Done():
+		return KnowledgeUnit{}, ctx.Err()
+	default:
 	}
 
 	if !ku.Tier.IsRemote() {
-		stored, err := c.store.Unit(ku.ID)
+		stored, err := c.store.Unit(ctx, ku.ID)
 		if err != nil {
 			return KnowledgeUnit{}, fmt.Errorf("reading knowledge unit: %w", err)
 		}
@@ -108,7 +110,7 @@ func (c *Client) Confirm(ctx context.Context, ku KnowledgeUnit) (KnowledgeUnit,
 		}
 
 		updated := applyConfirmation(*stored)
-		if err := c.store.Update(updated); err != nil {
+		if err := c.store.Update(ctx, updated); err != nil {
 			return KnowledgeUnit{}, fmt.Errorf("updating knowledge unit: %w", err)
 		}
 
@@ -137,15 +139,17 @@ func (c *Client) Drain(ctx context.Context) (DrainResult, error) {
 	ctx, cancel := c.operationContext(ctx)
 	defer cancel()
 
-	if err := ctx.Err(); err != nil {
-		return DrainResult{}, err
+	select {
+	case <-ctx.Done():
+		return DrainResult{}, ctx.Err()
+	default:
 	}
 
 	if c.remote == nil {
 		return DrainResult{}, fmt.Errorf("no remote API configured")
 	}
 
-	units, err := c.store.All()
+	units, err := c.store.All(ctx)
 	if err != nil {
 		return DrainResult{}, fmt.Errorf("reading local units: %w", err)
 	}
@@ -169,7 +173,7 @@ func (c *Client) Drain(ctx context.Context) (DrainResult, error) {
 			continue
 		}
 
-		if err := c.store.Delete(ku.ID); err != nil {
+		if err := c.store.Delete(ctx, ku.ID); err != nil {
 			result.Warnings = append(result.Warnings, fmt.Errorf("deleting local %s: %w", ku.ID, err))
 
 			continue
@@ -183,13 +187,16 @@ func (c *Client) Drain(ctx context.Context) (DrainResult, error) {
 
 // DrainableCount returns the number of local units that Drain would push.
 func (c *Client) DrainableCount(ctx context.Context) (int, error) {
+	ctx, cancel := c.operationContext(ctx)
+	defer cancel()
+
 	select {
 	case <-ctx.Done():
 		return 0, ctx.Err()
 	default:
 	}
 
-	units, err := c.store.All()
+	units, err := c.store.All(ctx)
 	if err != nil {
 		return 0, fmt.Errorf("reading local units: %w", err)
 	}
@@ -216,8 +223,10 @@ func (c *Client) Flag(
 	ctx, cancel := c.operationContext(ctx)
 	defer cancel()
 
-	if err := ctx.Err(); err != nil {
-		return KnowledgeUnit{}, err
+	select {
+	case <-ctx.Done():
+		return KnowledgeUnit{}, ctx.Err()
+	default:
 	}
 
 	cfg := resolveFlagConfig(opts)
@@ -233,7 +242,7 @@ func (c *Client) Flag(
 	}
 
 	if !ku.Tier.IsRemote() {
-		stored, err := c.store.Unit(ku.ID)
+		stored, err := c.store.Unit(ctx, ku.ID)
 		if err != nil {
 			return KnowledgeUnit{}, fmt.Errorf("reading knowledge unit: %w", err)
 		}
@@ -243,7 +252,7 @@ func (c *Client) Flag(
 		}
 
 		updated := applyFlag(*stored, reason, cfg)
-		if err := c.store.Update(updated); err != nil {
+		if err := c.store.Update(ctx, updated); err != nil {
 			return KnowledgeUnit{}, fmt.Errorf("updating knowledge unit: %w", err)
 		}
 
@@ -283,8 +292,10 @@ func (c *Client) Propose(ctx context.Context, params ProposeParams) (KnowledgeUn
 	ctx, cancel := c.operationContext(ctx)
 	defer cancel()
 
-	if err := ctx.Err(); err != nil {
-		return KnowledgeUnit{}, err
+	select {
+	case <-ctx.Done():
+		return KnowledgeUnit{}, ctx.Err()
+	default:
 	}
 
 	if err := ValidateExtensionKeys(params.Extensions); err != nil {
@@ -335,7 +346,7 @@ func (c *Client) Propose(ctx context.Context, params ProposeParams) (KnowledgeUn
 	ku.Evidence.FirstObserved = &now
 	ku.Evidence.LastConfirmed = &now
 
-	if err := c.store.Insert(ku); err != nil {
+	if err := c.store.Insert(ctx, ku); err != nil {
 		insertErr := fmt.Errorf("inserting knowledge unit: %w", err)
 		if remoteErr != nil {
 			return KnowledgeUnit{}, fmt.Errorf(
@@ -359,8 +370,10 @@ func (c *Client) Query(ctx context.Context, params QueryParams) (QueryResult, er
 	ctx, cancel := c.operationContext(ctx)
 	defer cancel()
 
-	if err := ctx.Err(); err != nil {
-		return QueryResult{}, err
+	select {
+	case <-ctx.Done():
+		return QueryResult{}, ctx.Err()
+	default:
 	}
 
 	limit := params.Limit
@@ -375,7 +388,7 @@ func (c *Client) Query(ctx context.Context, params QueryParams) (QueryResult, er
 	storeParams := params
 	storeParams.Limit = limit
 
-	storeResult, err := c.store.Query(storeParams)
+	storeResult, err := c.store.Query(ctx, storeParams)
 	if err != nil {
 		return QueryResult{}, fmt.Errorf("querying store: %w", err)
 	}
@@ -422,17 +435,21 @@ func (c *Client) Status(ctx context.Context) (StoreStats, error) {
 	ctx, cancel := c.operationContext(ctx)
 	defer cancel()
 
-	if err := ctx.Err(); err != nil {
-		return StoreStats{}, err
+	select {
+	case <-ctx.Done():
+		return StoreStats{}, ctx.Err()
+	default:
 	}
 
-	stats, err := c.store.Stats(defaultRecentLimit)
+	stats, err := c.store.Stats(ctx, defaultRecentLimit)
 	if err != nil {
 		return StoreStats{}, fmt.Errorf("reading store stats: %w", err)
 	}
 
-	if err := ctx.Err(); err != nil {
-		return StoreStats{}, err
+	select {
+	case <-ctx.Done():
+		return StoreStats{}, ctx.Err()
+	default:
 	}
 
 	stats.TierCounts = map[Tier]int{Local: stats.TotalCount}
@@ -491,16 +508,12 @@ func (c *Client) Status(ctx context.Context) (StoreStats, error) {
 	return stats, nil
 }
 
-// operationContext ensures client operations consistently respect request timeout.
+// operationContext bounds an operation by the client timeout. The returned
+// context expires at the sooner of the caller's deadline and c.timeout, so
+// c.timeout is an upper bound, never an extension. A non-positive c.timeout
+// imposes no client-side cap.
+// NOTE: ctx must be non-nil.
 func (c *Client) operationContext(ctx context.Context) (context.Context, context.CancelFunc) {
-	if ctx == nil {
-		ctx = context.Background()
-	}
-
-	if _, hasDeadline := ctx.Deadline(); hasDeadline {
-		return ctx, func() {}
-	}
-
 	if c.timeout <= 0 {
 		return ctx, func() {}
 	}

sdk/go/options.go

@@ -121,7 +121,8 @@ func WithStore(s Store) ClientOption {
 	}
 }
 
-// WithTimeout overrides the default HTTP request timeout.
+// WithTimeout sets the maximum duration of a client operation. It is an upper
+// bound: when the caller's context has an earlier deadline, that deadline wins.
 func WithTimeout(d time.Duration) ClientOption {
 	return func(c *clientConfig) error {
 		if d <= 0 {

sdk/go/store.go

@@ -2,6 +2,7 @@ package cq
 
 import (
 	"cmp"
+	"context"
 	"errors"
 	"fmt"
 	"math"
@@ -50,32 +51,32 @@ var ErrStoreClosed = errors.New("store is closed")
 // NOTE: implementations must be safe for concurrent use.
 type Store interface {
 	// Unit returns the knowledge unit with the given ID, or nil when absent.
-	Unit(id string) (*KnowledgeUnit, error)
+	Unit(ctx context.Context, id string) (*KnowledgeUnit, error)
 
 	// All returns every knowledge unit in the store.
-	All() ([]KnowledgeUnit, error)
+	All(ctx context.Context) ([]KnowledgeUnit, error)
 
 	// Insert stores a new knowledge unit.
 	// NOTE: implementations must reject a duplicate ID and a unit whose
 	// domains are empty after normalization.
-	Insert(ku KnowledgeUnit) error
+	Insert(ctx context.Context, ku KnowledgeUnit) error
 
 	// Update replaces an existing knowledge unit.
 	// NOTE: implementations must error when the ID is absent and reject a
 	// unit whose domains are empty after normalization.
-	Update(ku KnowledgeUnit) error
+	Update(ctx context.Context, ku KnowledgeUnit) error
 
 	// Delete removes the knowledge unit with the given ID.
 	// NOTE: implementations must error when the ID is absent.
-	Delete(id string) error
+	Delete(ctx context.Context, id string) error
 
 	// Query returns knowledge units matching the parameters, ranked by
 	// relevance and confidence and truncated to the limit.
-	Query(params QueryParams) (StoreQueryResult, error)
+	Query(ctx context.Context, params QueryParams) (StoreQueryResult, error)
 
 	// Stats returns aggregated store statistics, including up to recentLimit
 	// most-recently-inserted units.
-	Stats(recentLimit int) (StoreStats, error)
+	Stats(ctx context.Context, recentLimit int) (StoreStats, error)
 
 	// Close releases the resources held by the store.
 	// NOTE: implementations must be safe to call more than once.

sdk/go/store_memory.go

@@ -1,6 +1,7 @@
 package cq
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"slices"
@@ -27,7 +28,13 @@ func NewInMemoryStore() Store {
 }
 
 // All returns every knowledge unit in insertion order.
-func (s *inMemoryStore) All() ([]KnowledgeUnit, error) {
+func (s *inMemoryStore) All(ctx context.Context) ([]KnowledgeUnit, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -55,7 +62,13 @@ func (s *inMemoryStore) Close() error {
 }
 
 // Delete removes a knowledge unit by ID.
-func (s *inMemoryStore) Delete(id string) error {
+func (s *inMemoryStore) Delete(ctx context.Context, id string) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -76,7 +89,13 @@ func (s *inMemoryStore) Delete(id string) error {
 }
 
 // Insert stores a new knowledge unit. Error if ID exists or domains empty after normalization.
-func (s *inMemoryStore) Insert(ku KnowledgeUnit) error {
+func (s *inMemoryStore) Insert(ctx context.Context, ku KnowledgeUnit) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -104,7 +123,13 @@ func (s *inMemoryStore) Insert(ku KnowledgeUnit) error {
 
 // Query returns units matching any of the requested domains, ranked by the
 // shared ranker. It does not run full-text search.
-func (s *inMemoryStore) Query(params QueryParams) (StoreQueryResult, error) {
+func (s *inMemoryStore) Query(ctx context.Context, params QueryParams) (StoreQueryResult, error) {
+	select {
+	case <-ctx.Done():
+		return StoreQueryResult{}, ctx.Err()
+	default:
+	}
+
 	norm, err := normalizeQueryParams(params)
 	if err != nil {
 		return StoreQueryResult{}, err
@@ -142,7 +167,13 @@ func (s *inMemoryStore) Query(params QueryParams) (StoreQueryResult, error) {
 }
 
 // Stats returns aggregated statistics, including up to recentLimit most-recently-inserted units.
-func (s *inMemoryStore) Stats(recentLimit int) (StoreStats, error) {
+func (s *inMemoryStore) Stats(ctx context.Context, recentLimit int) (StoreStats, error) {
+	select {
+	case <-ctx.Done():
+		return StoreStats{}, ctx.Err()
+	default:
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -190,7 +221,13 @@ func (s *inMemoryStore) Stats(recentLimit int) (StoreStats, error) {
 }
 
 // Unit retrieves a knowledge unit by ID. Returns nil, nil if not found.
-func (s *inMemoryStore) Unit(id string) (*KnowledgeUnit, error) {
+func (s *inMemoryStore) Unit(ctx context.Context, id string) (*KnowledgeUnit, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -209,7 +246,13 @@ func (s *inMemoryStore) Unit(id string) (*KnowledgeUnit, error) {
 }
 
 // Update replaces an existing knowledge unit.
-func (s *inMemoryStore) Update(ku KnowledgeUnit) error {
+func (s *inMemoryStore) Update(ctx context.Context, ku KnowledgeUnit) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()