Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions cli/search.go
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,7 @@ func runSearch(cmd *cobra.Command, args []string) error {
st = gobStore
case "postgres":
var err error
st, err = store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions())
st, err = store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions(), cfg.Embedder.CacheNamespace())
if err != nil {
return fmt.Errorf("failed to connect to postgres: %w", err)
}
Expand Down Expand Up @@ -520,7 +520,7 @@ func SearchJSON(projectRoot string, query string, limit int) ([]store.SearchResu
st = gobStore
case "postgres":
var err error
st, err = store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions())
st, err = store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions(), cfg.Embedder.CacheNamespace())
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -577,7 +577,7 @@ func runWorkspaceSearch(ctx context.Context, query string, projects []string, pa

switch ws.Store.Backend {
case "postgres":
st, err = store.NewPostgresStore(ctx, ws.Store.Postgres.DSN, projectID, ws.Embedder.GetDimensions())
st, err = store.NewPostgresStore(ctx, ws.Store.Postgres.DSN, projectID, ws.Embedder.GetDimensions(), ws.Embedder.CacheNamespace())
if err != nil {
return fmt.Errorf("failed to connect to postgres: %w", err)
}
Expand Down
2 changes: 1 addition & 1 deletion cli/status.go
Original file line number Diff line number Diff line change
Expand Up @@ -460,7 +460,7 @@ func runStatus(cmd *cobra.Command, args []string) error {
st = gobStore
case "postgres":
var err error
st, err = store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions())
st, err = store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions(), cfg.Embedder.CacheNamespace())
if err != nil {
return fmt.Errorf("failed to connect to postgres: %w", err)
}
Expand Down
4 changes: 2 additions & 2 deletions cli/watch.go
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,7 @@ func initializeStore(ctx context.Context, cfg *config.Config, projectRoot string
}
return gobStore, nil
case "postgres":
return store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions())
return store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, projectRoot, cfg.Embedder.GetDimensions(), cfg.Embedder.CacheNamespace())
case "qdrant":
collectionName := cfg.Store.Qdrant.Collection
if collectionName == "" {
Expand Down Expand Up @@ -2870,7 +2870,7 @@ func initializeWorkspaceStore(ctx context.Context, ws *config.Workspace) (store.

switch ws.Store.Backend {
case "postgres":
return store.NewPostgresStore(ctx, ws.Store.Postgres.DSN, projectID, ws.Embedder.GetDimensions())
return store.NewPostgresStore(ctx, ws.Store.Postgres.DSN, projectID, ws.Embedder.GetDimensions(), ws.Embedder.CacheNamespace())
case "qdrant":
collectionName := ws.Store.Qdrant.Collection
if collectionName == "" {
Expand Down
12 changes: 12 additions & 0 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,18 @@ func (e *EmbedderConfig) GetDimensions() int {
}
}

// CacheNamespace returns the semantic namespace used for embedding-cache reuse.
// It intentionally excludes secrets such as API keys.
func (e *EmbedderConfig) CacheNamespace() string {
return fmt.Sprintf(
"embedding-cache-v2:provider=%s:model=%s:dimensions=%d:endpoint=%s",
strings.TrimSpace(e.Provider),
strings.TrimSpace(e.Model),
e.GetDimensions(),
strings.TrimSpace(e.Endpoint),
)
}

func DefaultEmbedderForProvider(provider string) EmbedderConfig {
switch provider {
case "synthetic":
Expand Down
36 changes: 36 additions & 0 deletions config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,42 @@ func TestDefaultEmbedderForProvider(t *testing.T) {
}
}

func TestEmbedderConfigCacheNamespace(t *testing.T) {
dim := 2048
cfg := EmbedderConfig{
Provider: " openai ",
Model: " text-embedding-3-small ",
Endpoint: " https://api.openai.com/v1 ",
Dimensions: &dim,
}

got := cfg.CacheNamespace()
want := "embedding-cache-v2:provider=openai:model=text-embedding-3-small:dimensions=2048:endpoint=https://api.openai.com/v1"
if got != want {
t.Fatalf("CacheNamespace() = %q, want %q", got, want)
}
}

func TestEmbedderConfigCacheNamespaceUsesResolvedDimensions(t *testing.T) {
cfg := EmbedderConfig{
Provider: "openai",
Model: OpenAIEmbeddingModelLarge,
}

if got, want := cfg.CacheNamespace(), "embedding-cache-v2:provider=openai:model=text-embedding-3-large:dimensions=3072:endpoint="; got != want {
t.Fatalf("CacheNamespace() = %q, want %q", got, want)
}
}

func TestEmbedderConfigCacheNamespaceChangesForModel(t *testing.T) {
small := EmbedderConfig{Provider: "openai", Model: DefaultOpenAIEmbeddingModel}
large := EmbedderConfig{Provider: "openai", Model: OpenAIEmbeddingModelLarge}

if small.CacheNamespace() == large.CacheNamespace() {
t.Fatalf("cache namespace must differ across models: %q", small.CacheNamespace())
}
}

func TestDefaultStoreForBackend(t *testing.T) {
postgres := DefaultStoreForBackend("postgres")
if postgres.Backend != "postgres" || postgres.Postgres.DSN != DefaultPostgresDSN {
Expand Down
4 changes: 2 additions & 2 deletions mcp/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -1004,7 +1004,7 @@ func (s *Server) createWorkspaceStore(ctx context.Context, ws *config.Workspace)

switch ws.Store.Backend {
case "postgres":
return store.NewPostgresStore(ctx, ws.Store.Postgres.DSN, projectID, ws.Embedder.GetDimensions())
return store.NewPostgresStore(ctx, ws.Store.Postgres.DSN, projectID, ws.Embedder.GetDimensions(), ws.Embedder.CacheNamespace())
case "qdrant":
collectionName := ws.Store.Qdrant.Collection
if collectionName == "" {
Expand Down Expand Up @@ -2046,7 +2046,7 @@ func (s *Server) createStore(ctx context.Context, cfg *config.Config) (store.Vec
}
return gobStore, nil
case "postgres":
return store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, s.projectRoot, cfg.Embedder.GetDimensions())
return store.NewPostgresStore(ctx, cfg.Store.Postgres.DSN, s.projectRoot, cfg.Embedder.GetDimensions(), cfg.Embedder.CacheNamespace())
case "qdrant":
collectionName := cfg.Store.Qdrant.Collection
if collectionName == "" {
Expand Down
36 changes: 23 additions & 13 deletions store/postgres.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,21 +11,29 @@ import (
)

type PostgresStore struct {
pool *pgxpool.Pool
projectID string
dimensions int
pool *pgxpool.Pool
projectID string
dimensions int
embeddingCacheNamespace string
}

func NewPostgresStore(ctx context.Context, dsn string, projectID string, vectorDimensions int) (*PostgresStore, error) {
const lookupByContentHashSQL = `SELECT vector FROM chunks WHERE project_id = $1 AND content_hash = $2 AND embedding_cache_namespace = $3 AND vector IS NOT NULL LIMIT 1`

func NewPostgresStore(ctx context.Context, dsn string, projectID string, vectorDimensions int, embeddingCacheNamespace ...string) (*PostgresStore, error) {
pool, err := pgxpool.New(ctx, dsn)
if err != nil {
return nil, fmt.Errorf("failed to connect to postgres: %w", err)
}
cacheNamespace := fmt.Sprintf("embedding-cache-v2:dimensions=%d", vectorDimensions)
if len(embeddingCacheNamespace) > 0 {
cacheNamespace = embeddingCacheNamespace[0]
}

store := &PostgresStore{
pool: pool,
projectID: projectID,
dimensions: vectorDimensions,
pool: pool,
projectID: projectID,
dimensions: vectorDimensions,
embeddingCacheNamespace: cacheNamespace,
}

if err := store.ensureSchema(ctx); err != nil {
Expand Down Expand Up @@ -61,7 +69,8 @@ func (s *PostgresStore) ensureSchema(ctx context.Context) error {
PRIMARY KEY (project_id, path)
)`,
`ALTER TABLE chunks ADD COLUMN IF NOT EXISTS content_hash TEXT DEFAULT ''`,
`CREATE INDEX IF NOT EXISTS idx_chunks_content_hash ON chunks(content_hash) WHERE content_hash != ''`,
`ALTER TABLE chunks ADD COLUMN IF NOT EXISTS embedding_cache_namespace TEXT NOT NULL DEFAULT ''`,
`CREATE INDEX IF NOT EXISTS idx_chunks_project_content_cache ON chunks(project_id, content_hash, embedding_cache_namespace) WHERE content_hash != ''`,
buildEnsureVectorSQL(s.dimensions),
// Migrate chunks primary key from (id) to (project_id, id) so that
// worktrees sharing the same database get their own chunk rows instead
Expand Down Expand Up @@ -99,8 +108,8 @@ func (s *PostgresStore) SaveChunks(ctx context.Context, chunks []Chunk) error {
for _, chunk := range chunks {
vec := pgvector.NewVector(chunk.Vector)
batch.Queue(
`INSERT INTO chunks (id, project_id, file_path, start_line, end_line, content, vector, hash, content_hash, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
`INSERT INTO chunks (id, project_id, file_path, start_line, end_line, content, vector, hash, content_hash, embedding_cache_namespace, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
ON CONFLICT (project_id, id) DO UPDATE SET
file_path = EXCLUDED.file_path,
start_line = EXCLUDED.start_line,
Expand All @@ -109,9 +118,10 @@ func (s *PostgresStore) SaveChunks(ctx context.Context, chunks []Chunk) error {
vector = EXCLUDED.vector,
hash = EXCLUDED.hash,
content_hash = EXCLUDED.content_hash,
embedding_cache_namespace = EXCLUDED.embedding_cache_namespace,
updated_at = EXCLUDED.updated_at`,
chunk.ID, s.projectID, chunk.FilePath, chunk.StartLine, chunk.EndLine,
chunk.Content, vec, chunk.Hash, chunk.ContentHash, chunk.UpdatedAt,
chunk.Content, vec, chunk.Hash, chunk.ContentHash, s.embeddingCacheNamespace, chunk.UpdatedAt,
)
}

Expand Down Expand Up @@ -381,8 +391,8 @@ func (s *PostgresStore) LookupByContentHash(ctx context.Context, contentHash str

var vec pgvector.Vector
err := s.pool.QueryRow(ctx,
`SELECT vector FROM chunks WHERE content_hash = $1 AND vector IS NOT NULL LIMIT 1`,
contentHash,
lookupByContentHashSQL,
s.projectID, contentHash, s.embeddingCacheNamespace,
).Scan(&vec)

if err == pgx.ErrNoRows {
Expand Down
82 changes: 82 additions & 0 deletions store/postgres_integration_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
package store

import (
"context"
"os"
"testing"
"time"
)

func TestPostgresStoreLookupByContentHashScopesCacheIdentity(t *testing.T) {
dsn := os.Getenv("GREPAI_POSTGRES_TEST_DSN")
if dsn == "" {
t.Skip("GREPAI_POSTGRES_TEST_DSN is not set")
}

ctx := context.Background()
victim, err := NewPostgresStore(ctx, dsn, "proj_victim", 3, "model=text-embedding-3-large")
if err != nil {
t.Fatalf("failed to create victim store: %v", err)
}
defer victim.Close()

attackerProject, err := NewPostgresStore(ctx, dsn, "proj_attacker", 3, "model=text-embedding-3-large")
if err != nil {
t.Fatalf("failed to create attacker project store: %v", err)
}
defer attackerProject.Close()

attackerModel, err := NewPostgresStore(ctx, dsn, "proj_victim", 3, "model=text-embedding-3-small")
if err != nil {
t.Fatalf("failed to create attacker model store: %v", err)
}
defer attackerModel.Close()

if _, err := victim.pool.Exec(ctx, `TRUNCATE TABLE chunks, documents`); err != nil {
t.Fatalf("failed to clean tables: %v", err)
}

contentHash := "sha256:shared-content"
victimVector := []float32{9.9, 8.8, 7.7}
if err := victim.SaveChunks(ctx, []Chunk{{
ID: "README.md_0",
FilePath: "README.md",
StartLine: 1,
EndLine: 1,
Content: "File: README.md\n\nShared content",
Vector: victimVector,
Hash: "victim-row",
ContentHash: contentHash,
UpdatedAt: time.Now().UTC(),
}}); err != nil {
t.Fatalf("failed to save victim chunk: %v", err)
}

if _, found, err := attackerProject.LookupByContentHash(ctx, contentHash); err != nil {
t.Fatalf("cross-project lookup failed: %v", err)
} else if found {
t.Fatal("different project should not reuse victim project's cached vector")
}

if _, found, err := attackerModel.LookupByContentHash(ctx, contentHash); err != nil {
t.Fatalf("cross-model lookup failed: %v", err)
} else if found {
t.Fatal("different embedding cache namespace should not reuse victim model's cached vector")
}

got, found, err := victim.LookupByContentHash(ctx, contentHash)
if err != nil {
t.Fatalf("victim lookup failed: %v", err)
}
if !found {
t.Fatal("victim project should find its own cached vector")
}
if len(got) != len(victimVector) {
t.Fatalf("victim vector length mismatch: got %d, want %d", len(got), len(victimVector))
}
for i := range got {
if got[i] != victimVector[i] {
t.Fatalf("victim vector mismatch at %d: got %v, want %v", i, got, victimVector)
}
}
}
20 changes: 20 additions & 0 deletions store/postgres_sql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,26 @@ func TestBuildEnsureVectorSQL_ContainsExpectedFragments(t *testing.T) {
}
}

func TestLookupByContentHashSQLScopesCacheIdentity(t *testing.T) {
expected := []string{
"SELECT vector FROM chunks",
"project_id = $1",
"content_hash = $2",
"embedding_cache_namespace = $3",
"vector IS NOT NULL",
}

for _, frag := range expected {
if !strings.Contains(lookupByContentHashSQL, frag) {
t.Fatalf("expected lookup SQL to contain %q, got: %q", frag, lookupByContentHashSQL)
}
}

if strings.Contains(lookupByContentHashSQL, "WHERE content_hash = $1") {
t.Fatalf("lookup SQL must not use content_hash as the first and only scope: %q", lookupByContentHashSQL)
}
}

// strconvItoa converts an int to string without importing strconv.
// This keeps the test dependency surface minimal.
func strconvItoa(n int) string {
Expand Down
2 changes: 1 addition & 1 deletion store/search_path_filter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ func TestPostgresStoreSearchWithPathPrefix(t *testing.T) {
}

ctx := context.Background()
store, err := NewPostgresStore(ctx, dsn, "test-project-"+randomString(8), 3)
store, err := NewPostgresStore(ctx, dsn, "test-project-"+randomString(8), 3, "test-cache-namespace")
if err != nil {
t.Skip("could not connect to PostgreSQL:", err)
}
Expand Down