From 9a5db87ee6bd8f9853f246c9c2bd938afa65f9f9 Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:06:58 +0100
Subject: [PATCH 1/8] feat(sdk/go): add context.Context to Store interface
 methods
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Store interface methods accepted no context, forcing the
PostgreSQL adapter to use context.Background() for every pgx
call. A stalled connection could block indefinitely while
holding the store mutex with no way for callers to cancel.

Add context.Context as the first parameter to Unit, All, Insert,
Update, Delete, Query, and Stats. Close is unchanged — it does
no cancellable I/O. All three implementations (memory, SQLite,
PostgreSQL), the Client, the conformance suite, and tests are
updated in lockstep.

The SQLite store now uses context-aware database/sql methods
(QueryContext, ExecContext, BeginTx, PrepareContext).

Fixes #468
---
 sdk/go/client.go                   |  20 +--
 sdk/go/store.go                    |  15 +-
 sdk/go/store_memory.go             |  15 +-
 sdk/go/store_memory_test.go        |  11 +-
 sdk/go/store_sqlite.go             |  77 +++++-----
 sdk/go/store_sqlite_test.go        | 223 +++++++++++++++--------------
 sdk/go/stores/postgres/postgres.go |  25 ++--
 sdk/go/storetest/storetest.go      |  77 +++++-----
 8 files changed, 244 insertions(+), 219 deletions(-)

diff --git a/sdk/go/client.go b/sdk/go/client.go
index 4efb9d29..b0a5d67b 100644
--- a/sdk/go/client.go
+++ b/sdk/go/client.go
@@ -98,7 +98,7 @@ func (c *Client) Confirm(ctx context.Context, ku KnowledgeUnit) (KnowledgeUnit,
 	}
 
 	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 +108,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)
 		}
 
@@ -145,7 +145,7 @@ func (c *Client) Drain(ctx context.Context) (DrainResult, error) {
 		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 +169,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
@@ -189,7 +189,7 @@ func (c *Client) DrainableCount(ctx context.Context) (int, error) {
 	default:
 	}
 
-	units, err := c.store.All()
+	units, err := c.store.All(ctx)
 	if err != nil {
 		return 0, fmt.Errorf("reading local units: %w", err)
 	}
@@ -233,7 +233,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 +243,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)
 		}
 
@@ -335,7 +335,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(
@@ -375,7 +375,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)
 	}
@@ -426,7 +426,7 @@ func (c *Client) Status(ctx context.Context) (StoreStats, error) {
 		return StoreStats{}, err
 	}
 
-	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)
 	}
diff --git a/sdk/go/store.go b/sdk/go/store.go
index 422cf32b..9e391fa1 100644
--- a/sdk/go/store.go
+++ b/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.
diff --git a/sdk/go/store_memory.go b/sdk/go/store_memory.go
index d78516fa..3761d71e 100644
--- a/sdk/go/store_memory.go
+++ b/sdk/go/store_memory.go
@@ -1,6 +1,7 @@
 package cq
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"slices"
@@ -27,7 +28,7 @@ func NewInMemoryStore() Store {
 }
 
 // All returns every knowledge unit in insertion order.
-func (s *inMemoryStore) All() ([]KnowledgeUnit, error) {
+func (s *inMemoryStore) All(_ context.Context) ([]KnowledgeUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -55,7 +56,7 @@ func (s *inMemoryStore) Close() error {
 }
 
 // Delete removes a knowledge unit by ID.
-func (s *inMemoryStore) Delete(id string) error {
+func (s *inMemoryStore) Delete(_ context.Context, id string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -76,7 +77,7 @@ 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(_ context.Context, ku KnowledgeUnit) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -104,7 +105,7 @@ 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(_ context.Context, params QueryParams) (StoreQueryResult, error) {
 	norm, err := normalizeQueryParams(params)
 	if err != nil {
 		return StoreQueryResult{}, err
@@ -142,7 +143,7 @@ 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(_ context.Context, recentLimit int) (StoreStats, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -190,7 +191,7 @@ 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(_ context.Context, id string) (*KnowledgeUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -209,7 +210,7 @@ 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(_ context.Context, ku KnowledgeUnit) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
diff --git a/sdk/go/store_memory_test.go b/sdk/go/store_memory_test.go
index bf15f695..69d41f46 100644
--- a/sdk/go/store_memory_test.go
+++ b/sdk/go/store_memory_test.go
@@ -1,6 +1,9 @@
 package cq
 
-import "testing"
+import (
+	"context"
+	"testing"
+)
 
 // TestInMemoryStoreMatchesDomainsOnly verifies the in-memory store selects
 // candidates by domain tag and does not consult full-text: a term that appears
@@ -21,11 +24,11 @@ func TestInMemoryStoreMatchesDomainsOnly(t *testing.T) {
 		Evidence: Evidence{Confidence: 0.7, Confirmations: 1},
 		Tier:     Local,
 	}
-	if err := s.Insert(ku); err != nil {
+	if err := s.Insert(context.Background(), ku); err != nil {
 		t.Fatalf("Insert: %s", err)
 	}
 
-	got, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: 5})
+	got, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: 5})
 	if err != nil {
 		t.Fatalf("Query by domain: %s", err)
 	}
@@ -35,7 +38,7 @@ func TestInMemoryStoreMatchesDomainsOnly(t *testing.T) {
 
 	// "payments" appears only in the summary, never as a domain; the
 	// full-text-free store must not surface it.
-	none, err := s.Query(QueryParams{Domains: []string{"payments"}, Limit: 5})
+	none, err := s.Query(context.Background(), QueryParams{Domains: []string{"payments"}, Limit: 5})
 	if err != nil {
 		t.Fatalf("Query by summary term: %s", err)
 	}
diff --git a/sdk/go/store_sqlite.go b/sdk/go/store_sqlite.go
index dbd40da0..c533ebf2 100644
--- a/sdk/go/store_sqlite.go
+++ b/sdk/go/store_sqlite.go
@@ -1,6 +1,7 @@
 package cq
 
 import (
+	"context"
 	"database/sql"
 	"encoding/json"
 	"errors"
@@ -89,7 +90,7 @@ func newSQLiteStore(dbPath string) (*sqliteStore, error) {
 }
 
 // All returns every knowledge unit in the store.
-func (s *sqliteStore) All() ([]KnowledgeUnit, error) {
+func (s *sqliteStore) All(ctx context.Context) ([]KnowledgeUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -97,7 +98,7 @@ func (s *sqliteStore) All() ([]KnowledgeUnit, error) {
 		return nil, ErrStoreClosed
 	}
 
-	rows, err := s.db.Query("SELECT data FROM knowledge_units")
+	rows, err := s.db.QueryContext(ctx, "SELECT data FROM knowledge_units")
 	if err != nil {
 		return nil, fmt.Errorf("querying all units: %w", err)
 	}
@@ -140,7 +141,7 @@ func (s *sqliteStore) Close() error {
 }
 
 // Delete removes a knowledge unit by ID.
-func (s *sqliteStore) Delete(unitID string) error {
+func (s *sqliteStore) Delete(ctx context.Context, unitID string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -148,18 +149,17 @@ func (s *sqliteStore) Delete(unitID string) error {
 		return ErrStoreClosed
 	}
 
-	tx, err := s.db.Begin()
+	tx, err := s.db.BeginTx(ctx, nil)
 	if err != nil {
 		return fmt.Errorf("beginning transaction: %w", err)
 	}
 	defer func() { _ = tx.Rollback() }()
 
-	// Delete FTS entry first; virtual tables don't support CASCADE.
-	if _, err := tx.Exec("DELETE FROM knowledge_units_fts WHERE id = ?", unitID); err != nil {
+	if _, err := tx.ExecContext(ctx, "DELETE FROM knowledge_units_fts WHERE id = ?", unitID); err != nil {
 		return fmt.Errorf("deleting FTS entry: %w", err)
 	}
 
-	res, err := tx.Exec("DELETE FROM knowledge_units WHERE id = ?", unitID)
+	res, err := tx.ExecContext(ctx, "DELETE FROM knowledge_units WHERE id = ?", unitID)
 	if err != nil {
 		return fmt.Errorf("deleting unit: %w", err)
 	}
@@ -176,7 +176,7 @@ func (s *sqliteStore) Delete(unitID string) error {
 }
 
 // Insert stores a knowledge unit. Error if ID exists or domains empty after normalization.
-func (s *sqliteStore) Insert(ku KnowledgeUnit) error {
+func (s *sqliteStore) Insert(ctx context.Context, ku KnowledgeUnit) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -198,21 +198,26 @@ func (s *sqliteStore) Insert(ku KnowledgeUnit) error {
 		return fmt.Errorf("marshalling unit: %w", err)
 	}
 
-	tx, err := s.db.Begin()
+	tx, err := s.db.BeginTx(ctx, nil)
 	if err != nil {
 		return fmt.Errorf("beginning transaction: %w", err)
 	}
 	defer func() { _ = tx.Rollback() }()
 
-	if _, err := tx.Exec("INSERT INTO knowledge_units (id, data) VALUES (?, ?)", ku.ID, string(data)); err != nil {
+	if _, err := tx.ExecContext(
+		ctx,
+		"INSERT INTO knowledge_units (id, data) VALUES (?, ?)",
+		ku.ID,
+		string(data),
+	); err != nil {
 		return fmt.Errorf("inserting unit: %w", err)
 	}
 
-	if err := insertDomains(tx, ku.ID, domains); err != nil {
+	if err := insertDomains(ctx, tx, ku.ID, domains); err != nil {
 		return err
 	}
 
-	if err := insertFTS(tx, stored); err != nil {
+	if err := insertFTS(ctx, tx, stored); err != nil {
 		return err
 	}
 
@@ -220,7 +225,7 @@ func (s *sqliteStore) Insert(ku KnowledgeUnit) error {
 }
 
 // Query searches by domain with FTS and relevance ranking.
-func (s *sqliteStore) Query(params QueryParams) (StoreQueryResult, error) {
+func (s *sqliteStore) Query(ctx context.Context, params QueryParams) (StoreQueryResult, error) {
 	norm, err := normalizeQueryParams(params)
 	if err != nil {
 		return StoreQueryResult{}, err
@@ -253,7 +258,7 @@ func (s *sqliteStore) Query(params QueryParams) (StoreQueryResult, error) {
 		"JOIN knowledge_unit_domains d ON k.id = d.unit_id " +
 		"WHERE d.domain IN (" + strings.Join(placeholders, ",") + ")"
 
-	rows, err := s.db.Query(domainQuery, args...)
+	rows, err := s.db.QueryContext(ctx, domainQuery, args...)
 	if err != nil {
 		return StoreQueryResult{}, fmt.Errorf("querying by domain: %w", err)
 	}
@@ -279,7 +284,7 @@ func (s *sqliteStore) Query(params QueryParams) (StoreQueryResult, error) {
 
 	matchExpr := buildFTSMatchExpr(domains)
 	if matchExpr != "" {
-		ftsRows, ftsErr := s.db.Query(
+		ftsRows, ftsErr := s.db.QueryContext(ctx,
 			"SELECT k.data FROM knowledge_units k "+
 				"JOIN knowledge_units_fts f ON k.id = f.id "+
 				"WHERE knowledge_units_fts MATCH ?",
@@ -324,7 +329,7 @@ func (s *sqliteStore) Query(params QueryParams) (StoreQueryResult, error) {
 }
 
 // Stats returns aggregated store statistics.
-func (s *sqliteStore) Stats(recentLimit int) (StoreStats, error) {
+func (s *sqliteStore) Stats(ctx context.Context, recentLimit int) (StoreStats, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -337,12 +342,12 @@ func (s *sqliteStore) Stats(recentLimit int) (StoreStats, error) {
 	}
 
 	var totalCount int
-	if err := s.db.QueryRow("SELECT COUNT(*) FROM knowledge_units").Scan(&totalCount); err != nil {
+	if err := s.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM knowledge_units").Scan(&totalCount); err != nil {
 		return StoreStats{}, fmt.Errorf("counting units: %w", err)
 	}
 
 	domainCounts := make(map[string]int)
-	domainRows, err := s.db.Query("SELECT domain, COUNT(*) FROM knowledge_unit_domains GROUP BY domain")
+	domainRows, err := s.db.QueryContext(ctx, "SELECT domain, COUNT(*) FROM knowledge_unit_domains GROUP BY domain")
 	if err != nil {
 		return StoreStats{}, fmt.Errorf("querying domain counts: %w", err)
 	}
@@ -361,7 +366,11 @@ func (s *sqliteStore) Stats(recentLimit int) (StoreStats, error) {
 	}
 
 	// Fetch recent units ordered by rowid descending (insertion order).
-	recentRows, err := s.db.Query("SELECT data FROM knowledge_units ORDER BY rowid DESC LIMIT ?", recentLimit)
+	recentRows, err := s.db.QueryContext(
+		ctx,
+		"SELECT data FROM knowledge_units ORDER BY rowid DESC LIMIT ?",
+		recentLimit,
+	)
 	if err != nil {
 		return StoreStats{}, fmt.Errorf("querying recent units: %w", err)
 	}
@@ -391,7 +400,7 @@ func (s *sqliteStore) Stats(recentLimit int) (StoreStats, error) {
 		buckets[label] = 0
 	}
 
-	allRows, err := s.db.Query("SELECT data FROM knowledge_units")
+	allRows, err := s.db.QueryContext(ctx, "SELECT data FROM knowledge_units")
 	if err != nil {
 		return StoreStats{}, fmt.Errorf("querying for confidence distribution: %w", err)
 	}
@@ -427,7 +436,7 @@ func (s *sqliteStore) Stats(recentLimit int) (StoreStats, error) {
 }
 
 // Unit retrieves a knowledge unit by ID. Returns nil, nil if not found.
-func (s *sqliteStore) Unit(id string) (*KnowledgeUnit, error) {
+func (s *sqliteStore) Unit(ctx context.Context, id string) (*KnowledgeUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -436,7 +445,7 @@ func (s *sqliteStore) Unit(id string) (*KnowledgeUnit, error) {
 	}
 
 	var data string
-	err := s.db.QueryRow("SELECT data FROM knowledge_units WHERE id = ?", id).Scan(&data)
+	err := s.db.QueryRowContext(ctx, "SELECT data FROM knowledge_units WHERE id = ?", id).Scan(&data)
 	if errors.Is(err, sql.ErrNoRows) {
 		return nil, nil
 	}
@@ -453,7 +462,7 @@ func (s *sqliteStore) Unit(id string) (*KnowledgeUnit, error) {
 }
 
 // Update replaces an existing knowledge unit.
-func (s *sqliteStore) Update(ku KnowledgeUnit) error {
+func (s *sqliteStore) Update(ctx context.Context, ku KnowledgeUnit) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -474,13 +483,13 @@ func (s *sqliteStore) Update(ku KnowledgeUnit) error {
 		return fmt.Errorf("marshalling unit: %w", err)
 	}
 
-	tx, err := s.db.Begin()
+	tx, err := s.db.BeginTx(ctx, nil)
 	if err != nil {
 		return fmt.Errorf("beginning transaction: %w", err)
 	}
 	defer func() { _ = tx.Rollback() }()
 
-	res, err := tx.Exec("UPDATE knowledge_units SET data = ? WHERE id = ?", string(data), ku.ID)
+	res, err := tx.ExecContext(ctx, "UPDATE knowledge_units SET data = ? WHERE id = ?", string(data), ku.ID)
 	if err != nil {
 		return fmt.Errorf("updating unit: %w", err)
 	}
@@ -494,18 +503,18 @@ func (s *sqliteStore) Update(ku KnowledgeUnit) error {
 	}
 
 	// Replace domain rows.
-	if _, err := tx.Exec("DELETE FROM knowledge_unit_domains WHERE unit_id = ?", ku.ID); err != nil {
+	if _, err := tx.ExecContext(ctx, "DELETE FROM knowledge_unit_domains WHERE unit_id = ?", ku.ID); err != nil {
 		return fmt.Errorf("clearing domains: %w", err)
 	}
-	if err := insertDomains(tx, ku.ID, domains); err != nil {
+	if err := insertDomains(ctx, tx, ku.ID, domains); err != nil {
 		return err
 	}
 
 	// Replace FTS entry.
-	if _, err := tx.Exec("DELETE FROM knowledge_units_fts WHERE id = ?", ku.ID); err != nil {
+	if _, err := tx.ExecContext(ctx, "DELETE FROM knowledge_units_fts WHERE id = ?", ku.ID); err != nil {
 		return fmt.Errorf("clearing FTS entry: %w", err)
 	}
-	if err := insertFTS(tx, stored); err != nil {
+	if err := insertFTS(ctx, tx, stored); err != nil {
 		return err
 	}
 
@@ -547,15 +556,15 @@ func ensureMetadata(db *sql.DB) error {
 }
 
 // insertDomains writes domain rows for a unit within an existing transaction.
-func insertDomains(tx *sql.Tx, unitID string, domains []string) error {
-	stmt, err := tx.Prepare("INSERT INTO knowledge_unit_domains (unit_id, domain) VALUES (?, ?)")
+func insertDomains(ctx context.Context, tx *sql.Tx, unitID string, domains []string) error {
+	stmt, err := tx.PrepareContext(ctx, "INSERT INTO knowledge_unit_domains (unit_id, domain) VALUES (?, ?)")
 	if err != nil {
 		return fmt.Errorf("preparing domain insert: %w", err)
 	}
 	defer func() { _ = stmt.Close() }()
 
 	for _, d := range domains {
-		if _, err := stmt.Exec(unitID, d); err != nil {
+		if _, err := stmt.ExecContext(ctx, unitID, d); err != nil {
 			return fmt.Errorf("inserting domain %q: %w", d, err)
 		}
 	}
@@ -564,8 +573,8 @@ func insertDomains(tx *sql.Tx, unitID string, domains []string) error {
 }
 
 // insertFTS writes an FTS entry for a unit within an existing transaction.
-func insertFTS(tx *sql.Tx, ku KnowledgeUnit) error {
-	_, err := tx.Exec(
+func insertFTS(ctx context.Context, tx *sql.Tx, ku KnowledgeUnit) error {
+	_, err := tx.ExecContext(ctx,
 		"INSERT INTO knowledge_units_fts (id, summary, detail, action) VALUES (?, ?, ?, ?)",
 		ku.ID,
 		ku.Insight.Summary,
diff --git a/sdk/go/store_sqlite_test.go b/sdk/go/store_sqlite_test.go
index c627ad41..dbcf9125 100644
--- a/sdk/go/store_sqlite_test.go
+++ b/sdk/go/store_sqlite_test.go
@@ -1,6 +1,7 @@
 package cq
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"path/filepath"
@@ -94,21 +95,21 @@ func TestClose(t *testing.T) {
 		_ = s.Close()
 
 		ku := newFakeKU(t, []string{"testing"})
-		require.ErrorIs(t, s.Insert(ku), ErrStoreClosed)
+		require.ErrorIs(t, s.Insert(context.Background(), ku), ErrStoreClosed)
 
-		_, err = s.Unit("any")
+		_, err = s.Unit(context.Background(), "any")
 		require.ErrorIs(t, err, ErrStoreClosed)
 
-		_, err = s.All()
+		_, err = s.All(context.Background())
 		require.ErrorIs(t, err, ErrStoreClosed)
 
-		require.ErrorIs(t, s.Delete("any"), ErrStoreClosed)
-		require.ErrorIs(t, s.Update(ku), ErrStoreClosed)
+		require.ErrorIs(t, s.Delete(context.Background(), "any"), ErrStoreClosed)
+		require.ErrorIs(t, s.Update(context.Background(), ku), ErrStoreClosed)
 
-		_, err = s.Query(QueryParams{Domains: []string{"testing"}, Limit: 10})
+		_, err = s.Query(context.Background(), QueryParams{Domains: []string{"testing"}, Limit: 10})
 		require.ErrorIs(t, err, ErrStoreClosed)
 
-		_, err = s.Stats(5)
+		_, err = s.Stats(context.Background(), 5)
 		require.ErrorIs(t, err, ErrStoreClosed)
 	})
 }
@@ -121,9 +122,9 @@ func TestInsert(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"databases"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.NotNil(t, got)
 		require.Equal(t, ku.ID, got.ID)
@@ -134,8 +135,8 @@ func TestInsert(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"databases"})
-		require.NoError(t, s.Insert(ku))
-		require.Error(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
+		require.Error(t, s.Insert(context.Background(), ku))
 	})
 
 	t.Run("rejects empty domains", func(t *testing.T) {
@@ -143,7 +144,7 @@ func TestInsert(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, nil)
-		require.Error(t, s.Insert(ku))
+		require.Error(t, s.Insert(context.Background(), ku))
 	})
 
 	t.Run("rejects whitespace-only domains", func(t *testing.T) {
@@ -151,7 +152,7 @@ func TestInsert(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"  ", "\t", ""})
-		require.Error(t, s.Insert(ku))
+		require.Error(t, s.Insert(context.Background(), ku))
 	})
 }
 
@@ -162,7 +163,7 @@ func TestUnit(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		got, err := s.Unit("ku_nonexistent")
+		got, err := s.Unit(context.Background(), "ku_nonexistent")
 		require.NoError(t, err)
 		require.Nil(t, got)
 	})
@@ -198,9 +199,9 @@ func TestUnit(t *testing.T) {
 			},
 		}
 
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.Equal(t, ku.ID, got.ID)
 		require.Equal(t, ku.Domains, got.Domains)
@@ -226,7 +227,7 @@ func TestAll(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		all, err := s.All()
+		all, err := s.All(context.Background())
 		require.NoError(t, err)
 		require.Empty(t, all)
 	})
@@ -237,10 +238,10 @@ func TestAll(t *testing.T) {
 
 		ku1 := newFakeKU(t, []string{"api"})
 		ku2 := newFakeKU(t, []string{"databases"})
-		require.NoError(t, s.Insert(ku1))
-		require.NoError(t, s.Insert(ku2))
+		require.NoError(t, s.Insert(context.Background(), ku1))
+		require.NoError(t, s.Insert(context.Background(), ku2))
 
-		all, err := s.All()
+		all, err := s.All(context.Background())
 		require.NoError(t, err)
 		require.Len(t, all, 2)
 
@@ -258,11 +259,11 @@ func TestDelete(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		require.NoError(t, s.Delete(ku.ID))
+		require.NoError(t, s.Delete(context.Background(), ku.ID))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.Nil(t, got)
 	})
@@ -271,7 +272,7 @@ func TestDelete(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		require.Error(t, s.Delete("ku_nonexistent"))
+		require.Error(t, s.Delete(context.Background(), "ku_nonexistent"))
 	})
 
 	t.Run("domain rows cascaded", func(t *testing.T) {
@@ -279,11 +280,14 @@ func TestDelete(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"unique-domain-for-cascade"})
-		require.NoError(t, s.Insert(ku))
-		require.NoError(t, s.Delete(ku.ID))
+		require.NoError(t, s.Insert(context.Background(), ku))
+		require.NoError(t, s.Delete(context.Background(), ku.ID))
 
 		// Verify domain is no longer queryable.
-		results, err := s.Query(QueryParams{Domains: []string{"unique-domain-for-cascade"}, Limit: 10})
+		results, err := s.Query(
+			context.Background(),
+			QueryParams{Domains: []string{"unique-domain-for-cascade"}, Limit: 10},
+		)
 		require.NoError(t, err)
 		require.Empty(t, results.KUs)
 	})
@@ -297,16 +301,16 @@ func TestUpdate(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		ku.Insight = Insight{
 			Summary: "updated summary",
 			Detail:  "updated detail",
 			Action:  "updated action",
 		}
-		require.NoError(t, s.Update(ku))
+		require.NoError(t, s.Update(context.Background(), ku))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.Equal(t, "updated summary", got.Insight.Summary)
 	})
@@ -316,7 +320,7 @@ func TestUpdate(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.Error(t, s.Update(ku))
+		require.Error(t, s.Update(context.Background(), ku))
 	})
 
 	t.Run("rejects empty domains", func(t *testing.T) {
@@ -324,10 +328,10 @@ func TestUpdate(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		ku.Domains = nil
-		require.Error(t, s.Update(ku))
+		require.Error(t, s.Update(context.Background(), ku))
 	})
 
 	t.Run("refreshes domain tags", func(t *testing.T) {
@@ -335,18 +339,18 @@ func TestUpdate(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"old-domain"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		ku.Domains = []string{"new-domain"}
-		require.NoError(t, s.Update(ku))
+		require.NoError(t, s.Update(context.Background(), ku))
 
 		// Old domain should no longer match.
-		results, err := s.Query(QueryParams{Domains: []string{"old-domain"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"old-domain"}, Limit: 10})
 		require.NoError(t, err)
 		require.Empty(t, results.KUs)
 
 		// New domain should match.
-		results, err = s.Query(QueryParams{Domains: []string{"new-domain"}, Limit: 10})
+		results, err = s.Query(context.Background(), QueryParams{Domains: []string{"new-domain"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 		require.Equal(t, ku.ID, results.KUs[0].ID)
@@ -361,9 +365,9 @@ func TestDomainNormalization(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"API", "Databases"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 	})
@@ -373,9 +377,9 @@ func TestDomainNormalization(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{"API"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"API"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 	})
@@ -385,9 +389,9 @@ func TestDomainNormalization(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api", "API", "Api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		stats, err := s.Stats(0)
+		stats, err := s.Stats(context.Background(), 0)
 		require.NoError(t, err)
 		// Only one domain "api" should exist.
 		require.Equal(t, 1, stats.DomainCounts["api"])
@@ -402,9 +406,9 @@ func TestQuery(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 		require.Equal(t, ku.ID, results.KUs[0].ID)
@@ -415,9 +419,9 @@ func TestQuery(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{"databases"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"databases"}, Limit: 10})
 		require.NoError(t, err)
 		require.Empty(t, results.KUs)
 	})
@@ -426,7 +430,7 @@ func TestQuery(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		results, err := s.Query(QueryParams{Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Limit: 10})
 		require.NoError(t, err)
 		require.Empty(t, results.KUs)
 	})
@@ -436,9 +440,9 @@ func TestQuery(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: 0})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: 0})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 	})
@@ -447,7 +451,7 @@ func TestQuery(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		_, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: -1})
+		_, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: -1})
 		require.Error(t, err)
 	})
 
@@ -462,7 +466,7 @@ func TestQuery(t *testing.T) {
 			domains = append(domains, fmt.Sprintf("domain%d", i))
 		}
 
-		_, err := s.Query(QueryParams{Domains: domains, Limit: 10})
+		_, err := s.Query(context.Background(), QueryParams{Domains: domains, Limit: 10})
 		require.Error(t, err)
 		require.Contains(t, err.Error(), "maximum number of domains")
 	})
@@ -473,10 +477,10 @@ func TestQuery(t *testing.T) {
 
 		for i := 0; i < 5; i++ {
 			ku := newFakeKU(t, []string{"api"})
-			require.NoError(t, s.Insert(ku))
+			require.NoError(t, s.Insert(context.Background(), ku))
 		}
 
-		results, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: 3})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: 3})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 3)
 	})
@@ -488,13 +492,13 @@ func TestQuery(t *testing.T) {
 		// Unit with two matching domains should rank higher.
 		kuBetter := newFakeKU(t, []string{"api", "payments"})
 		kuBetter.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuBetter))
+		require.NoError(t, s.Insert(context.Background(), kuBetter))
 
 		kuWorse := newFakeKU(t, []string{"api"})
 		kuWorse.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuWorse))
+		require.NoError(t, s.Insert(context.Background(), kuWorse))
 
-		results, err := s.Query(QueryParams{Domains: []string{"api", "payments"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api", "payments"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 2)
 		require.Equal(t, kuBetter.ID, results.KUs[0].ID)
@@ -507,14 +511,17 @@ func TestQuery(t *testing.T) {
 		kuGo := newFakeKU(t, []string{"api"})
 		kuGo.Context = Context{Languages: []string{"go"}}
 		kuGo.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuGo))
+		require.NoError(t, s.Insert(context.Background(), kuGo))
 
 		kuPython := newFakeKU(t, []string{"api"})
 		kuPython.Context = Context{Languages: []string{"python"}}
 		kuPython.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuPython))
+		require.NoError(t, s.Insert(context.Background(), kuPython))
 
-		results, err := s.Query(QueryParams{Domains: []string{"api"}, Languages: []string{"go"}, Limit: 10})
+		results, err := s.Query(
+			context.Background(),
+			QueryParams{Domains: []string{"api"}, Languages: []string{"go"}, Limit: 10},
+		)
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 2)
 		require.Equal(t, kuGo.ID, results.KUs[0].ID)
@@ -527,15 +534,15 @@ func TestQuery(t *testing.T) {
 		kuMatch := newFakeKU(t, []string{"api"})
 		kuMatch.Context = Context{Languages: []string{"python"}}
 		kuMatch.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuMatch))
+		require.NoError(t, s.Insert(context.Background(), kuMatch))
 
 		kuNoMatch := newFakeKU(t, []string{"api"})
 		kuNoMatch.Context = Context{Languages: []string{"rust"}}
 		kuNoMatch.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuNoMatch))
+		require.NoError(t, s.Insert(context.Background(), kuNoMatch))
 
 		// Querying with multiple languages; "python" overlaps with kuMatch.
-		results, err := s.Query(QueryParams{
+		results, err := s.Query(context.Background(), QueryParams{
 			Domains:   []string{"api"},
 			Languages: []string{"python", "go"},
 			Limit:     10,
@@ -552,15 +559,15 @@ func TestQuery(t *testing.T) {
 		kuMatch := newFakeKU(t, []string{"api"})
 		kuMatch.Context = Context{Frameworks: []string{"grpc"}}
 		kuMatch.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuMatch))
+		require.NoError(t, s.Insert(context.Background(), kuMatch))
 
 		kuNoMatch := newFakeKU(t, []string{"api"})
 		kuNoMatch.Context = Context{Frameworks: []string{"django"}}
 		kuNoMatch.Evidence.Confidence = 0.8
-		require.NoError(t, s.Insert(kuNoMatch))
+		require.NoError(t, s.Insert(context.Background(), kuNoMatch))
 
 		// Querying with multiple frameworks; "grpc" overlaps with kuMatch.
-		results, err := s.Query(QueryParams{
+		results, err := s.Query(context.Background(), QueryParams{
 			Domains:    []string{"api"},
 			Frameworks: []string{"grpc", "http"},
 			Limit:      10,
@@ -576,13 +583,13 @@ func TestQuery(t *testing.T) {
 
 		kuHigh := newFakeKU(t, []string{"api"})
 		kuHigh.Evidence.Confidence = 0.9
-		require.NoError(t, s.Insert(kuHigh))
+		require.NoError(t, s.Insert(context.Background(), kuHigh))
 
 		kuLow := newFakeKU(t, []string{"api"})
 		kuLow.Evidence.Confidence = 0.3
-		require.NoError(t, s.Insert(kuLow))
+		require.NoError(t, s.Insert(context.Background(), kuLow))
 
-		results, err := s.Query(QueryParams{Domains: []string{"api"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 2)
 		require.Equal(t, kuHigh.ID, results.KUs[0].ID)
@@ -602,10 +609,10 @@ func TestFTS(t *testing.T) {
 			Detail:  "some detail",
 			Action:  "some action",
 		}
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		// Query with a domain that matches the summary text via FTS.
-		results, err := s.Query(QueryParams{Domains: []string{"retries"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"retries"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 		require.Equal(t, ku.ID, results.KUs[0].ID)
@@ -621,9 +628,9 @@ func TestFTS(t *testing.T) {
 			Detail:  "Exponential backoff prevents thundering herd problems",
 			Action:  "some action",
 		}
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{"exponential"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"exponential"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 		require.Equal(t, ku.ID, results.KUs[0].ID)
@@ -639,10 +646,10 @@ func TestFTS(t *testing.T) {
 			Detail:  "detail",
 			Action:  "action",
 		}
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		// "retries" matches both domain and FTS but should appear only once.
-		results, err := s.Query(QueryParams{Domains: []string{"retries"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"retries"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 	})
@@ -657,23 +664,23 @@ func TestFTS(t *testing.T) {
 			Detail:  "original detail",
 			Action:  "original action",
 		}
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		ku.Insight = Insight{
 			Summary: "completely new unique-fts-keyword summary",
 			Detail:  "updated detail",
 			Action:  "updated action",
 		}
-		require.NoError(t, s.Update(ku))
+		require.NoError(t, s.Update(context.Background(), ku))
 
 		// FTS should find the updated text.
-		results, err := s.Query(QueryParams{Domains: []string{"unique-fts-keyword"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"unique-fts-keyword"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 		require.Equal(t, ku.ID, results.KUs[0].ID)
 
 		// FTS should not find the old text.
-		results, err = s.Query(QueryParams{Domains: []string{"original"}, Limit: 10})
+		results, err = s.Query(context.Background(), QueryParams{Domains: []string{"original"}, Limit: 10})
 		require.NoError(t, err)
 		require.Empty(t, results.KUs)
 	})
@@ -683,9 +690,12 @@ func TestFTS(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"safe-domain"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{`bad"domain`, "safe-domain"}, Limit: 10})
+		results, err := s.Query(
+			context.Background(),
+			QueryParams{Domains: []string{`bad"domain`, "safe-domain"}, Limit: 10},
+		)
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 	})
@@ -695,9 +705,12 @@ func TestFTS(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"safe-domain"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{`") OR (id:`, "safe-domain"}, Limit: 10})
+		results, err := s.Query(
+			context.Background(),
+			QueryParams{Domains: []string{`") OR (id:`, "safe-domain"}, Limit: 10},
+		)
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 	})
@@ -707,9 +720,9 @@ func TestFTS(t *testing.T) {
 		s := newTestStore(t)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
-		results, err := s.Query(QueryParams{Domains: []string{""}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{""}, Limit: 10})
 		require.NoError(t, err)
 		require.Empty(t, results.KUs)
 	})
@@ -722,7 +735,7 @@ func TestStats(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		stats, err := s.Stats(5)
+		stats, err := s.Stats(context.Background(), 5)
 		require.NoError(t, err)
 		require.Equal(t, 0, stats.TotalCount)
 		require.Empty(t, stats.DomainCounts)
@@ -733,10 +746,10 @@ func TestStats(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		require.NoError(t, s.Insert(newFakeKU(t, []string{"api"})))
-		require.NoError(t, s.Insert(newFakeKU(t, []string{"databases"})))
+		require.NoError(t, s.Insert(context.Background(), newFakeKU(t, []string{"api"})))
+		require.NoError(t, s.Insert(context.Background(), newFakeKU(t, []string{"databases"})))
 
-		stats, err := s.Stats(10)
+		stats, err := s.Stats(context.Background(), 10)
 		require.NoError(t, err)
 		require.Equal(t, 2, stats.TotalCount)
 	})
@@ -745,10 +758,10 @@ func TestStats(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		require.NoError(t, s.Insert(newFakeKU(t, []string{"api", "payments"})))
-		require.NoError(t, s.Insert(newFakeKU(t, []string{"api", "databases"})))
+		require.NoError(t, s.Insert(context.Background(), newFakeKU(t, []string{"api", "payments"})))
+		require.NoError(t, s.Insert(context.Background(), newFakeKU(t, []string{"api", "databases"})))
 
-		stats, err := s.Stats(10)
+		stats, err := s.Stats(context.Background(), 10)
 		require.NoError(t, err)
 		require.Equal(t, 2, stats.DomainCounts["api"])
 		require.Equal(t, 1, stats.DomainCounts["payments"])
@@ -760,10 +773,10 @@ func TestStats(t *testing.T) {
 		s := newTestStore(t)
 
 		for i := 0; i < 5; i++ {
-			require.NoError(t, s.Insert(newFakeKU(t, []string{"api"})))
+			require.NoError(t, s.Insert(context.Background(), newFakeKU(t, []string{"api"})))
 		}
 
-		stats, err := s.Stats(3)
+		stats, err := s.Stats(context.Background(), 3)
 		require.NoError(t, err)
 		require.Len(t, stats.Recent, 3)
 	})
@@ -772,7 +785,7 @@ func TestStats(t *testing.T) {
 		t.Parallel()
 		s := newTestStore(t)
 
-		_, err := s.Stats(-1)
+		_, err := s.Stats(context.Background(), -1)
 		require.Error(t, err)
 	})
 
@@ -785,10 +798,10 @@ func TestStats(t *testing.T) {
 		for _, conf := range []float64{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.95} {
 			ku := newFakeKU(t, []string{"api"})
 			ku.Evidence.Confidence = conf
-			require.NoError(t, s.Insert(ku))
+			require.NoError(t, s.Insert(context.Background(), ku))
 		}
 
-		stats, err := s.Stats(10)
+		stats, err := s.Stats(context.Background(), 10)
 		require.NoError(t, err)
 		require.Equal(t, 2, stats.ConfidenceDistribution["0.0-0.3"]) // 0.1, 0.2
 		require.Equal(t, 2, stats.ConfidenceDistribution["0.3-0.5"]) // 0.3, 0.4
@@ -806,28 +819,28 @@ func TestEndToEnd(t *testing.T) {
 
 		ku := newFakeKU(t, []string{"api", "payments"})
 		ku.Evidence.Confidence = 0.6
-		require.NoError(t, s.Insert(ku))
+		require.NoError(t, s.Insert(context.Background(), ku))
 
 		// Confirm the unit.
 		confirmed := applyConfirmation(ku)
-		require.NoError(t, s.Update(confirmed))
+		require.NoError(t, s.Update(context.Background(), confirmed))
 
 		// Verify confidence increased.
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.Greater(t, got.Evidence.Confidence, 0.6)
 
 		// Query should find it.
-		results, err := s.Query(QueryParams{Domains: []string{"api", "payments"}, Limit: 10})
+		results, err := s.Query(context.Background(), QueryParams{Domains: []string{"api", "payments"}, Limit: 10})
 		require.NoError(t, err)
 		require.Len(t, results.KUs, 1)
 		require.Equal(t, ku.ID, results.KUs[0].ID)
 
 		// Flag the unit.
 		flagged := applyFlag(*got, Stale, flagConfig{})
-		require.NoError(t, s.Update(flagged))
+		require.NoError(t, s.Update(context.Background(), flagged))
 
-		got, err = s.Unit(ku.ID)
+		got, err = s.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.Len(t, got.Flags, 1)
 		require.Equal(t, Stale, got.Flags[0].Reason)
@@ -841,14 +854,14 @@ func TestEndToEnd(t *testing.T) {
 		require.NoError(t, err)
 
 		ku := newFakeKU(t, []string{"api"})
-		require.NoError(t, s1.Insert(ku))
+		require.NoError(t, s1.Insert(context.Background(), ku))
 		_ = s1.Close()
 
 		s2, err := newSQLiteStore(dbPath)
 		require.NoError(t, err)
 		t.Cleanup(func() { _ = s2.Close() })
 
-		got, err := s2.Unit(ku.ID)
+		got, err := s2.Unit(context.Background(), ku.ID)
 		require.NoError(t, err)
 		require.NotNil(t, got)
 		require.Equal(t, ku.ID, got.ID)
diff --git a/sdk/go/stores/postgres/postgres.go b/sdk/go/stores/postgres/postgres.go
index 5e316ac1..360ec2d8 100644
--- a/sdk/go/stores/postgres/postgres.go
+++ b/sdk/go/stores/postgres/postgres.go
@@ -109,13 +109,13 @@ func New(connString string) (*Store, error) {
 }
 
 // All returns every knowledge unit in the store.
-func (s *Store) All() ([]cq.KnowledgeUnit, error) {
+func (s *Store) All(ctx context.Context) ([]cq.KnowledgeUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
 		return nil, cq.ErrStoreClosed
 	}
-	return s.scanUnits(context.Background(), sqlSelectAll)
+	return s.scanUnits(ctx, sqlSelectAll)
 }
 
 // Close releases the connection pool. Safe to call more than once.
@@ -132,13 +132,13 @@ func (s *Store) Close() error {
 
 // Delete removes the knowledge unit with the given ID.
 // Returns an error if no unit with that ID exists.
-func (s *Store) Delete(id string) error {
+func (s *Store) Delete(ctx context.Context, id string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
 		return cq.ErrStoreClosed
 	}
-	ct, err := s.pool.Exec(context.Background(), sqlDeleteUnit, id)
+	ct, err := s.pool.Exec(ctx, sqlDeleteUnit, id)
 	if err != nil {
 		return err
 	}
@@ -151,7 +151,7 @@ func (s *Store) Delete(id string) error {
 // Insert stores a new knowledge unit.
 // Domains are normalized before storage.
 // Returns an error on duplicate ID or empty domains after normalization.
-func (s *Store) Insert(ku cq.KnowledgeUnit) error {
+func (s *Store) Insert(ctx context.Context, ku cq.KnowledgeUnit) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
@@ -166,7 +166,6 @@ func (s *Store) Insert(ku cq.KnowledgeUnit) error {
 	if err != nil {
 		return err
 	}
-	ctx := context.Background()
 	tx, err := s.pool.Begin(ctx)
 	if err != nil {
 		return err
@@ -184,7 +183,7 @@ func (s *Store) Insert(ku cq.KnowledgeUnit) error {
 
 // Query returns knowledge units whose domain tags overlap with the query,
 // ranked by relevance and confidence, truncated to the limit.
-func (s *Store) Query(params cq.QueryParams) (cq.StoreQueryResult, error) {
+func (s *Store) Query(ctx context.Context, params cq.QueryParams) (cq.StoreQueryResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
@@ -197,7 +196,7 @@ func (s *Store) Query(params cq.QueryParams) (cq.StoreQueryResult, error) {
 	if len(nq.Domains) == 0 {
 		return cq.StoreQueryResult{}, nil
 	}
-	candidates, err := s.scanUnits(context.Background(), sqlQueryByDomains, nq.Domains)
+	candidates, err := s.scanUnits(ctx, sqlQueryByDomains, nq.Domains)
 	if err != nil {
 		return cq.StoreQueryResult{}, err
 	}
@@ -208,7 +207,7 @@ func (s *Store) Query(params cq.QueryParams) (cq.StoreQueryResult, error) {
 // Stats returns aggregated store statistics including unit counts per domain,
 // the most recently inserted units, and confidence distribution across the
 // canonical buckets.
-func (s *Store) Stats(recentLimit int) (cq.StoreStats, error) {
+func (s *Store) Stats(ctx context.Context, recentLimit int) (cq.StoreStats, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
@@ -217,7 +216,6 @@ func (s *Store) Stats(recentLimit int) (cq.StoreStats, error) {
 	if recentLimit < 0 {
 		return cq.StoreStats{}, fmt.Errorf("recent limit must be non-negative: %d", recentLimit)
 	}
-	ctx := context.Background()
 	totalCount, err := s.countUnits(ctx)
 	if err != nil {
 		return cq.StoreStats{}, err
@@ -244,14 +242,14 @@ func (s *Store) Stats(recentLimit int) (cq.StoreStats, error) {
 }
 
 // Unit returns the knowledge unit with the given ID, or nil when absent.
-func (s *Store) Unit(id string) (*cq.KnowledgeUnit, error) {
+func (s *Store) Unit(ctx context.Context, id string) (*cq.KnowledgeUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
 		return nil, cq.ErrStoreClosed
 	}
 	var data []byte
-	err := s.pool.QueryRow(context.Background(), sqlSelectByID, id).Scan(&data)
+	err := s.pool.QueryRow(ctx, sqlSelectByID, id).Scan(&data)
 	if errors.Is(err, pgx.ErrNoRows) {
 		return nil, nil
 	}
@@ -268,7 +266,7 @@ func (s *Store) Unit(id string) (*cq.KnowledgeUnit, error) {
 // Update replaces an existing knowledge unit.
 // Domains are re-normalized and the domain index is rebuilt.
 // Returns an error if no unit with that ID exists.
-func (s *Store) Update(ku cq.KnowledgeUnit) error {
+func (s *Store) Update(ctx context.Context, ku cq.KnowledgeUnit) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.closed {
@@ -283,7 +281,6 @@ func (s *Store) Update(ku cq.KnowledgeUnit) error {
 	if err != nil {
 		return err
 	}
-	ctx := context.Background()
 	tx, err := s.pool.Begin(ctx)
 	if err != nil {
 		return err
diff --git a/sdk/go/storetest/storetest.go b/sdk/go/storetest/storetest.go
index b1951f02..b90c0777 100644
--- a/sdk/go/storetest/storetest.go
+++ b/sdk/go/storetest/storetest.go
@@ -10,6 +10,7 @@
 package storetest
 
 import (
+	"context"
 	"testing"
 
 	cq "github.com/mozilla-ai/cq/sdk/go"
@@ -27,9 +28,9 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_00000000000000000000000000000001", []string{"api"}, 0.8)
-		noErr(t, s.Insert(ku))
+		noErr(t, s.Insert(context.Background(), ku))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		noErr(t, err)
 		if got == nil {
 			t.Fatal("Get returned nil for an inserted unit")
@@ -46,7 +47,7 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		s := newStore()
 		t.Cleanup(func() { _ = s.Close() })
 
-		got, err := s.Unit("ku_0000000000000000000000000000dead")
+		got, err := s.Unit(context.Background(), "ku_0000000000000000000000000000dead")
 		noErr(t, err)
 		if got != nil {
 			t.Fatalf("Get for a missing ID = %+v, want nil", got)
@@ -58,8 +59,8 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_00000000000000000000000000000002", []string{"api"}, 0.5)
-		noErr(t, s.Insert(ku))
-		wantErr(t, s.Insert(ku), "duplicate Insert")
+		noErr(t, s.Insert(context.Background(), ku))
+		wantErr(t, s.Insert(context.Background(), ku), "duplicate Insert")
 	})
 
 	t.Run("insert rejects empty domains", func(t *testing.T) {
@@ -67,7 +68,7 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_00000000000000000000000000000003", []string{"  ", ""}, 0.5)
-		wantErr(t, s.Insert(ku), "Insert with empty domains")
+		wantErr(t, s.Insert(context.Background(), ku), "Insert with empty domains")
 	})
 
 	t.Run("update existing persists", func(t *testing.T) {
@@ -75,12 +76,12 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_00000000000000000000000000000004", []string{"api"}, 0.5)
-		noErr(t, s.Insert(ku))
+		noErr(t, s.Insert(context.Background(), ku))
 
 		ku.Insight.Summary = "updated summary"
-		noErr(t, s.Update(ku))
+		noErr(t, s.Update(context.Background(), ku))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		noErr(t, err)
 		if got.Insight.Summary != "updated summary" {
 			t.Fatalf("summary after Update = %s, want updated summary", got.Insight.Summary)
@@ -92,7 +93,7 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_00000000000000000000000000000005", []string{"api"}, 0.5)
-		wantErr(t, s.Update(ku), "Update of a missing unit")
+		wantErr(t, s.Update(context.Background(), ku), "Update of a missing unit")
 	})
 
 	t.Run("delete existing then missing", func(t *testing.T) {
@@ -100,17 +101,17 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_00000000000000000000000000000006", []string{"api"}, 0.5)
-		noErr(t, s.Insert(ku))
+		noErr(t, s.Insert(context.Background(), ku))
 
-		noErr(t, s.Delete(ku.ID))
+		noErr(t, s.Delete(context.Background(), ku.ID))
 
-		got, err := s.Unit(ku.ID)
+		got, err := s.Unit(context.Background(), ku.ID)
 		noErr(t, err)
 		if got != nil {
 			t.Fatalf("Get after Delete = %+v, want nil", got)
 		}
 
-		wantErr(t, s.Delete(ku.ID), "Delete of a missing unit")
+		wantErr(t, s.Delete(context.Background(), ku.ID), "Delete of a missing unit")
 	})
 
 	t.Run("all returns inserted units", func(t *testing.T) {
@@ -119,10 +120,10 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 
 		ku1 := newKU("ku_00000000000000000000000000000007", []string{"api"}, 0.5)
 		ku2 := newKU("ku_00000000000000000000000000000008", []string{"db"}, 0.5)
-		noErr(t, s.Insert(ku1))
-		noErr(t, s.Insert(ku2))
+		noErr(t, s.Insert(context.Background(), ku1))
+		noErr(t, s.Insert(context.Background(), ku2))
 
-		all, err := s.All()
+		all, err := s.All(context.Background())
 		noErr(t, err)
 		if len(all) != 2 {
 			t.Fatalf("All returned %d units, want 2", len(all))
@@ -143,12 +144,12 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 
 		// Two matching domains and high confidence ranks first.
 		better := newKU("ku_00000000000000000000000000000009", []string{"api", "payments"}, 0.9)
-		noErr(t, s.Insert(better))
+		noErr(t, s.Insert(context.Background(), better))
 
 		worse := newKU("ku_0000000000000000000000000000000a", []string{"api"}, 0.3)
-		noErr(t, s.Insert(worse))
+		noErr(t, s.Insert(context.Background(), worse))
 
-		res, err := s.Query(cq.QueryParams{Domains: []string{"api", "payments"}, Limit: 10})
+		res, err := s.Query(context.Background(), cq.QueryParams{Domains: []string{"api", "payments"}, Limit: 10})
 		noErr(t, err)
 		if len(res.KUs) != 2 {
 			t.Fatalf("Query returned %d units, want 2", len(res.KUs))
@@ -168,10 +169,10 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 			"ku_0000000000000000000000000000000d",
 		} {
 			ku := newKU(id, []string{"api"}, 0.5+float64(i)*0.1)
-			noErr(t, s.Insert(ku))
+			noErr(t, s.Insert(context.Background(), ku))
 		}
 
-		res, err := s.Query(cq.QueryParams{Domains: []string{"api"}, Limit: 2})
+		res, err := s.Query(context.Background(), cq.QueryParams{Domains: []string{"api"}, Limit: 2})
 		noErr(t, err)
 		if len(res.KUs) != 2 {
 			t.Fatalf("Query with Limit 2 returned %d units, want 2", len(res.KUs))
@@ -183,9 +184,9 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		t.Cleanup(func() { _ = s.Close() })
 
 		ku := newKU("ku_0000000000000000000000000000000e", []string{"unique-domain"}, 0.5)
-		noErr(t, s.Insert(ku))
+		noErr(t, s.Insert(context.Background(), ku))
 
-		res, err := s.Query(cq.QueryParams{Domains: []string{"unique-domain"}, Limit: 10})
+		res, err := s.Query(context.Background(), cq.QueryParams{Domains: []string{"unique-domain"}, Limit: 10})
 		noErr(t, err)
 		if len(res.KUs) != 1 {
 			t.Fatalf("Query returned %d units, want 1", len(res.KUs))
@@ -194,7 +195,7 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 			t.Fatalf("matched ID = %s, want %s", res.KUs[0].ID, ku.ID)
 		}
 
-		none, err := s.Query(cq.QueryParams{Domains: []string{"no-such-domain"}, Limit: 10})
+		none, err := s.Query(context.Background(), cq.QueryParams{Domains: []string{"no-such-domain"}, Limit: 10})
 		noErr(t, err)
 		if len(none.KUs) != 0 {
 			t.Fatalf("Query for an absent domain returned %d units, want 0", len(none.KUs))
@@ -214,17 +215,17 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 			"ku_00000000000000000000000000000035",
 			"ku_00000000000000000000000000000036",
 		} {
-			noErr(t, s.Insert(newKU(id, []string{"api"}, 0.5+float64(i)*0.01)))
+			noErr(t, s.Insert(context.Background(), newKU(id, []string{"api"}, 0.5+float64(i)*0.01)))
 		}
 
 		// A zero limit means "unset": the store falls back to its default (5).
-		res, err := s.Query(cq.QueryParams{Domains: []string{"api"}, Limit: 0})
+		res, err := s.Query(context.Background(), cq.QueryParams{Domains: []string{"api"}, Limit: 0})
 		noErr(t, err)
 		if len(res.KUs) != 5 {
 			t.Fatalf("Query with Limit 0 returned %d units, want default 5", len(res.KUs))
 		}
 
-		_, err = s.Query(cq.QueryParams{Domains: []string{"api"}, Limit: -1})
+		_, err = s.Query(context.Background(), cq.QueryParams{Domains: []string{"api"}, Limit: -1})
 		wantErr(t, err, "Query with a negative limit")
 	})
 
@@ -246,10 +247,10 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		}
 		for i, conf := range confidences {
 			ku := newKU(ids[i], []string{"api"}, conf)
-			noErr(t, s.Insert(ku))
+			noErr(t, s.Insert(context.Background(), ku))
 		}
 
-		stats, err := s.Stats(3)
+		stats, err := s.Stats(context.Background(), 3)
 		noErr(t, err)
 		if stats.TotalCount != len(confidences) {
 			t.Fatalf("TotalCount = %d, want %d", stats.TotalCount, len(confidences))
@@ -272,7 +273,7 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		s := newStore()
 		t.Cleanup(func() { _ = s.Close() })
 
-		_, err := s.Stats(-1)
+		_, err := s.Stats(context.Background(), -1)
 		wantErr(t, err, "Stats with a negative recent limit")
 	})
 
@@ -283,20 +284,20 @@ func RunConformance(t *testing.T, newStore func() cq.Store) {
 		noErr(t, s.Close())
 
 		ku := newKU("ku_00000000000000000000000000000020", []string{"api"}, 0.5)
-		wantErr(t, s.Insert(ku), "Insert after Close")
-		wantErr(t, s.Update(ku), "Update after Close")
-		wantErr(t, s.Delete(ku.ID), "Delete after Close")
+		wantErr(t, s.Insert(context.Background(), ku), "Insert after Close")
+		wantErr(t, s.Update(context.Background(), ku), "Update after Close")
+		wantErr(t, s.Delete(context.Background(), ku.ID), "Delete after Close")
 
-		_, err := s.Unit(ku.ID)
+		_, err := s.Unit(context.Background(), ku.ID)
 		wantErr(t, err, "Get after Close")
 
-		_, err = s.All()
+		_, err = s.All(context.Background())
 		wantErr(t, err, "All after Close")
 
-		_, err = s.Query(cq.QueryParams{Domains: []string{"api"}, Limit: 10})
+		_, err = s.Query(context.Background(), cq.QueryParams{Domains: []string{"api"}, Limit: 10})
 		wantErr(t, err, "Query after Close")
 
-		_, err = s.Stats(5)
+		_, err = s.Stats(context.Background(), 5)
 		wantErr(t, err, "Stats after Close")
 	})
 }

From 62c6e0e74d435d89e0d3065277f78f49549e2471 Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:07:28 +0100
Subject: [PATCH 2/8] feat(sdk/go): add context.Context to postgres.New
 constructor
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

postgres.New used context.Background() for pool creation, server
ping, and schema migration — the same indefinite-blocking problem
the Store interface change addressed. Thread caller-provided
context through New, ensureSchema, and stampWriter so callers
can set timeouts on connection setup.
---
 sdk/go/stores/postgres/postgres.go      | 18 ++++++++----------
 sdk/go/stores/postgres/postgres_test.go |  3 ++-
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/sdk/go/stores/postgres/postgres.go b/sdk/go/stores/postgres/postgres.go
index 360ec2d8..278abb34 100644
--- a/sdk/go/stores/postgres/postgres.go
+++ b/sdk/go/stores/postgres/postgres.go
@@ -2,7 +2,7 @@
 //
 // Construct with New and pass to the client via cq.WithStore:
 //
-//	store, err := postgres.New("postgres://localhost/cq")
+//	store, err := postgres.New(ctx, "postgres://localhost/cq")
 //	client, err := cq.NewClient(cq.WithStore(store))
 package postgres
 
@@ -84,7 +84,7 @@ type Store struct {
 // The connection string must be a valid PostgreSQL URL or DSN.
 // New validates the string, connects, pings the server, and ensures the
 // schema exists before returning.
-func New(connString string) (*Store, error) {
+func New(ctx context.Context, connString string) (*Store, error) {
 	if connString == "" {
 		return nil, fmt.Errorf("connection string must not be empty")
 	}
@@ -92,16 +92,16 @@ func New(connString string) (*Store, error) {
 	if err != nil {
 		return nil, fmt.Errorf("invalid connection string: %w", err)
 	}
-	pool, err := pgxpool.NewWithConfig(context.Background(), cfg)
+	pool, err := pgxpool.NewWithConfig(ctx, cfg)
 	if err != nil {
 		return nil, fmt.Errorf("creating connection pool: %w", err)
 	}
-	if err := pool.Ping(context.Background()); err != nil {
+	if err := pool.Ping(ctx); err != nil {
 		pool.Close()
 		return nil, fmt.Errorf("connecting to server: %w", err)
 	}
 	s := &Store{pool: pool}
-	if err := s.ensureSchema(); err != nil {
+	if err := s.ensureSchema(ctx); err != nil {
 		pool.Close()
 		return nil, fmt.Errorf("ensuring schema: %w", err)
 	}
@@ -342,12 +342,11 @@ func (s *Store) countUnits(ctx context.Context) (int, error) {
 // ensureSchema creates the tables and indexes if they do not exist, then
 // stamps the writer metadata so cross-SDK diagnostics can identify the
 // last SDK that wrote to the database.
-func (s *Store) ensureSchema() error {
-	ctx := context.Background()
+func (s *Store) ensureSchema(ctx context.Context) error {
 	if _, err := s.pool.Exec(ctx, schemaDDL); err != nil {
 		return err
 	}
-	return s.stampWriter()
+	return s.stampWriter(ctx)
 }
 
 // queryDomainCounts returns the number of knowledge units per domain tag.
@@ -394,8 +393,7 @@ func (s *Store) scanUnits(ctx context.Context, sql string, args ...any) ([]cq.Kn
 
 // stampWriter records the SDK version and timestamp in the metadata table
 // so operators can identify which SDK last modified the database.
-func (s *Store) stampWriter() error {
-	ctx := context.Background()
+func (s *Store) stampWriter(ctx context.Context) error {
 	tag := fmt.Sprintf(writerTagFmt, runtime.Version())
 	now := time.Now().UTC().Format(time.RFC3339)
 	batch := &pgx.Batch{}
diff --git a/sdk/go/stores/postgres/postgres_test.go b/sdk/go/stores/postgres/postgres_test.go
index 41a025af..88808ca8 100644
--- a/sdk/go/stores/postgres/postgres_test.go
+++ b/sdk/go/stores/postgres/postgres_test.go
@@ -1,6 +1,7 @@
 package postgres_test
 
 import (
+	"context"
 	"encoding/json"
 	"testing"
 	"time"
@@ -37,7 +38,7 @@ func TestNew(t *testing.T) {
 	for _, tc := range tcs {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			_, err := postgres.New(tc.connString)
+			_, err := postgres.New(context.Background(), tc.connString)
 			require.Error(t, err)
 			require.Contains(t, err.Error(), tc.wantErr)
 		})

From aa025e585fe00f78db8cd1b3f8aa539cbe24420f Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:42:09 +0100
Subject: [PATCH 3/8] fix(sdk/go): use bounded cleanup context for postgres tx
 rollback

The deferred rollback in Insert and Update used the caller's ctx.
When that ctx is cancelled or times out mid-operation, the rollback
cannot run, and pgx churns the pooled connection instead of
releasing it cleanly.

Extract a rollback helper that discards the transaction with a
short background-derived context, independent of caller
cancellation but still bounded so it cannot block indefinitely.
---
 sdk/go/stores/postgres/postgres.go | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/sdk/go/stores/postgres/postgres.go b/sdk/go/stores/postgres/postgres.go
index 278abb34..730a2e47 100644
--- a/sdk/go/stores/postgres/postgres.go
+++ b/sdk/go/stores/postgres/postgres.go
@@ -17,6 +17,7 @@ import (
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+
 	cq "github.com/mozilla-ai/cq/sdk/go"
 )
 
@@ -26,6 +27,10 @@ const (
 	writerTagFmt   = "cq-go-sdk/postgres go/%s"
 )
 
+// rollbackTimeout bounds the deferred rollback cleanup so it cannot block
+// indefinitely on an unresponsive server.
+const rollbackTimeout = 5 * time.Second
+
 const schemaDDL = `
 CREATE TABLE IF NOT EXISTS knowledge_units (
     rowid BIGINT GENERATED ALWAYS AS IDENTITY,
@@ -170,8 +175,7 @@ func (s *Store) Insert(ctx context.Context, ku cq.KnowledgeUnit) error {
 	if err != nil {
 		return err
 	}
-	// Rollback is a no-op after Commit; the returned error is not actionable.
-	defer func() { _ = tx.Rollback(ctx) }()
+	defer rollback(tx)
 	if _, err := tx.Exec(ctx, sqlInsertUnit, ku.ID, data); err != nil {
 		return fmt.Errorf("inserting unit %s: %w", ku.ID, err)
 	}
@@ -285,8 +289,7 @@ func (s *Store) Update(ctx context.Context, ku cq.KnowledgeUnit) error {
 	if err != nil {
 		return err
 	}
-	// Rollback is a no-op after Commit; the returned error is not actionable.
-	defer func() { _ = tx.Rollback(ctx) }()
+	defer rollback(tx)
 	ct, err := tx.Exec(ctx, sqlUpdateUnit, data, ku.ID)
 	if err != nil {
 		return err
@@ -412,6 +415,16 @@ func insertDomains(ctx context.Context, tx pgx.Tx, unitID string, domains []stri
 	return nil
 }
 
+// rollback discards tx using a context independent of the caller's, so cleanup
+// runs even after the caller cancels — otherwise a cancelled rollback churns
+// the pooled connection instead of releasing it. No-op after a successful
+// Commit; the returned error is not actionable.
+func rollback(tx pgx.Tx) {
+	ctx, cancel := context.WithTimeout(context.Background(), rollbackTimeout)
+	defer cancel()
+	_ = tx.Rollback(ctx)
+}
+
 // marshal serializes a KnowledgeUnit to JSON for JSONB storage.
 func marshal(ku cq.KnowledgeUnit) ([]byte, error) {
 	return json.Marshal(ku)

From 9748745068209da8c598fdbfb0ea2790bae4d944 Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:44:26 +0100
Subject: [PATCH 4/8] test(sdk/go): cover cancelled context in postgres.New

Guard against future drift back to context.Background() inside New
by asserting that an already-cancelled context aborts connection
setup. Asserts errors.Is(err, context.Canceled) so it holds
regardless of which setup stage observes the cancellation.
---
 sdk/go/stores/postgres/postgres_test.go | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/sdk/go/stores/postgres/postgres_test.go b/sdk/go/stores/postgres/postgres_test.go
index 88808ca8..c7b1cc06 100644
--- a/sdk/go/stores/postgres/postgres_test.go
+++ b/sdk/go/stores/postgres/postgres_test.go
@@ -6,9 +6,11 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	cq "github.com/mozilla-ai/cq/sdk/go"
+
 	"github.com/mozilla-ai/cq/sdk/go/stores/postgres"
-	"github.com/stretchr/testify/require"
 )
 
 func TestNew(t *testing.T) {
@@ -45,6 +47,17 @@ func TestNew(t *testing.T) {
 	}
 }
 
+func TestNewRespectsCancelledContext(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err := postgres.New(ctx, "postgres://localhost:1/nonexistent")
+	require.Error(t, err)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 func TestMarshalUnmarshalRoundTrip(t *testing.T) {
 	t.Parallel()
 
@@ -174,4 +187,3 @@ func TestUnmarshalInvalidJSON(t *testing.T) {
 		})
 	}
 }
-

From 2c89aff41658b2b72641c47a170ea39e70376b1a Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:46:19 +0100
Subject: [PATCH 5/8] fix(sdk/go): honor context cancellation in in-memory
 store

The in-memory store discarded ctx, so a cancelled call could still
acquire the mutex and walk the store while the SQLite and Postgres
backends would return early. Add a pre-flight select/default check
to each context-taking method so cancellation behaves consistently
across backends. The store does no blocking I/O, so per-iteration
checks are unnecessary.
---
 sdk/go/store_memory.go | 56 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 49 insertions(+), 7 deletions(-)

diff --git a/sdk/go/store_memory.go b/sdk/go/store_memory.go
index 3761d71e..4ef4d300 100644
--- a/sdk/go/store_memory.go
+++ b/sdk/go/store_memory.go
@@ -28,7 +28,13 @@ func NewInMemoryStore() Store {
 }
 
 // All returns every knowledge unit in insertion order.
-func (s *inMemoryStore) All(_ context.Context) ([]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()
 
@@ -56,7 +62,13 @@ func (s *inMemoryStore) Close() error {
 }
 
 // Delete removes a knowledge unit by ID.
-func (s *inMemoryStore) Delete(_ context.Context, 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()
 
@@ -77,7 +89,13 @@ func (s *inMemoryStore) Delete(_ context.Context, id string) error {
 }
 
 // Insert stores a new knowledge unit. Error if ID exists or domains empty after normalization.
-func (s *inMemoryStore) Insert(_ context.Context, 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()
 
@@ -105,7 +123,13 @@ func (s *inMemoryStore) Insert(_ context.Context, 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(_ context.Context, 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
@@ -143,7 +167,13 @@ func (s *inMemoryStore) Query(_ context.Context, params QueryParams) (StoreQuery
 }
 
 // Stats returns aggregated statistics, including up to recentLimit most-recently-inserted units.
-func (s *inMemoryStore) Stats(_ context.Context, 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()
 
@@ -191,7 +221,13 @@ func (s *inMemoryStore) Stats(_ context.Context, recentLimit int) (StoreStats, e
 }
 
 // Unit retrieves a knowledge unit by ID. Returns nil, nil if not found.
-func (s *inMemoryStore) Unit(_ context.Context, 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()
 
@@ -210,7 +246,13 @@ func (s *inMemoryStore) Unit(_ context.Context, id string) (*KnowledgeUnit, erro
 }
 
 // Update replaces an existing knowledge unit.
-func (s *inMemoryStore) Update(_ context.Context, 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()
 

From 0098f58e3509242587973d44c1906722afd11f38 Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:49:42 +0100
Subject: [PATCH 6/8] refactor(sdk/go): treat client timeout as an upper bound
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

operationContext skipped the client timeout entirely when the
caller's context already carried a deadline, making c.timeout a
default-when-unset rather than a ceiling. Always derive from the
caller's context with WithTimeout, which expires at the sooner of
the two — so c.timeout caps every operation without ever extending
a shorter caller deadline.

Also drop the nil-context guard: passing a nil context is a caller
bug and panicking is the idiomatic Go outcome.
---
 sdk/go/client.go  | 14 +++++---------
 sdk/go/options.go |  3 ++-
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/sdk/go/client.go b/sdk/go/client.go
index b0a5d67b..89fe7648 100644
--- a/sdk/go/client.go
+++ b/sdk/go/client.go
@@ -491,16 +491,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() {}
 	}
diff --git a/sdk/go/options.go b/sdk/go/options.go
index 3eabdfbb..72e7c8e7 100644
--- a/sdk/go/options.go
+++ b/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 {

From 8a2c03d4b1e8c026980c6ffbf92915d72f870d15 Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 20:52:41 +0100
Subject: [PATCH 7/8] refactor(sdk/go): standardize client ctx guards on
 select/default

Convert the pre-flight cancellation guards across the client entry
points from ctx.Err() to a uniform select/default form, and route
DrainableCount through operationContext like every other entry
point so it honors the client timeout instead of using the raw
caller context. The mid-loop check in Drain stays on ctx.Err()
because a select there would need a labeled break.
---
 sdk/go/client.go | 45 +++++++++++++++++++++++++++++++--------------
 1 file changed, 31 insertions(+), 14 deletions(-)

diff --git a/sdk/go/client.go b/sdk/go/client.go
index 89fe7648..d0d8407c 100644
--- a/sdk/go/client.go
+++ b/sdk/go/client.go
@@ -93,8 +93,10 @@ 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() {
@@ -137,8 +139,10 @@ 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 {
@@ -183,6 +187,9 @@ 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()
@@ -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)
@@ -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 {
@@ -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
@@ -422,8 +435,10 @@ 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(ctx, defaultRecentLimit)
@@ -431,8 +446,10 @@ func (c *Client) Status(ctx context.Context) (StoreStats, error) {
 		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}

From b48c13fa85c9dafeb32170f938c61446b9681c7e Mon Sep 17 00:00:00 2001
From: Peter Wilson <peter@mozilla.ai>
Date: Thu, 25 Jun 2026 21:10:05 +0100
Subject: [PATCH 8/8] docs(sdk/go): update postgres.New examples for context
 parameter

The README usage snippets still showed the context-free
postgres.New(connString) form. Update both to the new
postgres.New(ctx, connString) signature so the docs match the API.
---
 sdk/go/README.md                 | 2 +-
 sdk/go/stores/postgres/README.md | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/sdk/go/README.md b/sdk/go/README.md
index 0ca79229..60716cc3 100644
--- a/sdk/go/README.md
+++ b/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))
 ```
 
diff --git a/sdk/go/stores/postgres/README.md b/sdk/go/stores/postgres/README.md
index 0948ac9c..cd230cc6 100644
--- a/sdk/go/stores/postgres/README.md
+++ b/sdk/go/stores/postgres/README.md
@@ -17,11 +17,13 @@ core SDK.
 
 ```go
 import (
+    "context"
+
     cq "github.com/mozilla-ai/cq/sdk/go"
     "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")
 if err != nil {
     log.Fatal(err)
 }