refactor: address kb connector review (timeout, composite conflict, tests)

CodeRabbit review follow-ups on the kb/* connectors:

- Bound pgx.Connect with a 15s connect timeout (Major) so a stalled
  Postgres handshake can't hang a step that has no timeout of its own; it
  only shortens an existing deadline.
- conflict_target accepts a comma-separated composite key (each part
  validated as an identifier), not just a single column.
- kbDistanceOperator error message lists the l2/euclidean and
  inner_product/ip aliases it accepts.
- Tests: both-metadata-forms, both-vector-forms, and top_k: 0 error cases.
- Schema example: note that only cosine is indexed; l2/inner_product need
  a matching index or they seq-scan.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79

Author: claude

packages/engine/internal/connector/kb.go

@@ -7,6 +7,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/jackc/pgx/v5"
 )
@@ -29,6 +30,10 @@ const (
 	kbDefaultMetadataColumn  = "metadata"
 	kbDefaultTopK            = 5
 	kbMaxTopK                = 1000
+	// kbConnectTimeout bounds the Postgres connection handshake so a stalled
+	// server can't hang a step that has no timeout of its own. It only shortens
+	// an existing deadline, never extends one.
+	kbConnectTimeout = 15 * time.Second
 )
 
 // kbIdentParam reads an identifier param, applies a default when absent, and
@@ -162,10 +167,17 @@ func prepareUpsert(params map[string]any) (string, []any, error) {
 
 	var conflict string
 	if ct, ok := params["conflict_target"].(string); ok && ct != "" {
-		if !kbIdentRe.MatchString(ct) {
-			return "", nil, fmt.Errorf("conflict_target %q is not a valid SQL identifier", ct)
+		// Accept a single column or a comma-separated list (composite key),
+		// validating each part as an identifier.
+		parts := strings.Split(ct, ",")
+		for i, p := range parts {
+			p = strings.TrimSpace(p)
+			if !kbIdentRe.MatchString(p) {
+				return "", nil, fmt.Errorf("conflict_target %q is not a valid SQL identifier", ct)
+			}
+			parts[i] = p
 		}
-		conflict = fmt.Sprintf(" ON CONFLICT (%s) DO NOTHING", ct)
+		conflict = fmt.Sprintf(" ON CONFLICT (%s) DO NOTHING", strings.Join(parts, ", "))
 	}
 
 	var sb strings.Builder
@@ -193,7 +205,7 @@ func kbDistanceOperator(metric string) (string, error) {
 	case "inner_product", "ip":
 		return "<#>", nil
 	default:
-		return "", fmt.Errorf("unknown metric %q (use cosine, l2, or inner_product)", metric)
+		return "", fmt.Errorf("unknown metric %q (use cosine, l2/euclidean, or inner_product/ip)", metric)
 	}
 }
 
@@ -293,7 +305,9 @@ func (c *KBUpsertConnector) Execute(ctx context.Context, params map[string]any)
 		return nil, fmt.Errorf("kb/upsert: %w", err)
 	}
 
-	conn, err := pgx.Connect(ctx, connURL)
+	connectCtx, cancel := context.WithTimeout(ctx, kbConnectTimeout)
+	conn, err := pgx.Connect(connectCtx, connURL)
+	cancel()
 	if err != nil {
 		return nil, fmt.Errorf("kb/upsert: connecting: %w", err)
 	}
@@ -321,7 +335,9 @@ func (c *KBQueryConnector) Execute(ctx context.Context, params map[string]any) (
 		return nil, fmt.Errorf("kb/query: %w", err)
 	}
 
-	conn, err := pgx.Connect(ctx, connURL)
+	connectCtx, cancel := context.WithTimeout(ctx, kbConnectTimeout)
+	conn, err := pgx.Connect(connectCtx, connURL)
+	cancel()
 	if err != nil {
 		return nil, fmt.Errorf("kb/query: connecting: %w", err)
 	}

packages/engine/internal/connector/kb_test.go

@@ -108,6 +108,8 @@ func TestPrepareUpsert_Errors(t *testing.T) {
 		{"both content forms", map[string]any{"table": "t", "content": "x", "contents": []any{"y"}, "vector": "[1]"}},
 		{"count mismatch", map[string]any{"table": "t", "contents": []any{"a", "b"}, "vectors": []any{"[1]"}}},
 		{"metadata mismatch", map[string]any{"table": "t", "content": "a", "vector": "[1]", "metadatas": []any{map[string]any{}, map[string]any{}}}},
+		{"both metadata forms", map[string]any{"table": "t", "content": "a", "vector": "[1]", "metadata": map[string]any{}, "metadatas": []any{map[string]any{}}}},
+		{"both vector forms", map[string]any{"table": "t", "content": "a", "vector": "[1]", "vectors": []any{"[1]"}}},
 		{"sql injection in table", map[string]any{"table": "kb; DROP TABLE users;--", "content": "x", "vector": "[1]"}},
 		{"sql injection in column", map[string]any{"table": "t", "content_column": "c)); DROP", "content": "x", "vector": "[1]"}},
 		{"sql injection in conflict", map[string]any{"table": "t", "content": "x", "vector": "[1]", "conflict_target": "a) DO UPDATE"}},
@@ -180,6 +182,7 @@ func TestPrepareQuery_Errors(t *testing.T) {
 		{"missing vector", map[string]any{"table": "t"}},
 		{"bad metric", map[string]any{"table": "t", "vector": "[1]", "metric": "manhattan"}},
 		{"negative top_k", map[string]any{"table": "t", "vector": "[1]", "top_k": -1}},
+		{"zero top_k", map[string]any{"table": "t", "vector": "[1]", "top_k": 0}},
 		{"sql injection in table", map[string]any{"table": "t; DROP TABLE x", "vector": "[1]"}},
 		{"sql injection in column", map[string]any{"table": "t", "vector": "[1]", "columns": []any{"content", "x); DROP"}}},
 	}

packages/site/src/content/examples/rag-kb-schema.sql

@@ -27,7 +27,9 @@ CREATE TABLE IF NOT EXISTS kb_documents (
 CREATE UNIQUE INDEX IF NOT EXISTS idx_kb_documents_dedupe
     ON kb_documents (dedupe_key);
 
--- Approximate-nearest-neighbour index for cosine distance (pgvector >= 0.5).
--- For small corpora you can omit this and rely on an exact scan.
+-- Approximate-nearest-neighbour index for cosine distance (pgvector >= 0.5),
+-- which is kb/query's default metric. For small corpora you can omit this and
+-- rely on an exact scan. If you query with metric: l2 or inner_product, add a
+-- matching index (vector_l2_ops / vector_ip_ops) or those queries seq-scan.
 CREATE INDEX IF NOT EXISTS idx_kb_documents_embedding
     ON kb_documents USING hnsw (embedding vector_cosine_ops);