diff --git a/internal/auth/middleware.go b/internal/auth/middleware.go index bccd62a..836d7d4 100644 --- a/internal/auth/middleware.go +++ b/internal/auth/middleware.go @@ -11,6 +11,7 @@ import ( "git.warky.dev/wdevs/amcs/internal/config" "git.warky.dev/wdevs/amcs/internal/observability" "git.warky.dev/wdevs/amcs/internal/requestip" + "git.warky.dev/wdevs/amcs/internal/tenancy" ) type contextKey string @@ -50,6 +51,10 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg ) } } + withTenant := func(ctx context.Context, keyID string) context.Context { + ctx = context.WithValue(ctx, keyIDContextKey, keyID) + return tenancy.WithTenantKey(ctx, keyID) + } return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { remoteAddr := requestip.FromRequest(r) @@ -63,7 +68,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg return } recordAccess(r, keyID) - next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID))) + next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID))) return } } @@ -73,14 +78,14 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg if tokenStore != nil { if keyID, ok := tokenStore.Lookup(bearer); ok { recordAccess(r, keyID) - next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID))) + next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID))) return } } if keyring != nil { if keyID, ok := keyring.Lookup(bearer); ok { recordAccess(r, keyID) - next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID))) + next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID))) return } } @@ -103,7 +108,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg return } recordAccess(r, keyID) - next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID))) + next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID))) return } @@ -117,7 +122,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg return } recordAccess(r, keyID) - next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID))) + next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID))) return } } diff --git a/internal/auth/tenancy_test.go b/internal/auth/tenancy_test.go new file mode 100644 index 0000000..83a2e72 --- /dev/null +++ b/internal/auth/tenancy_test.go @@ -0,0 +1,44 @@ +package auth + +import ( + "log/slog" + "net/http" + "net/http/httptest" + "testing" + + "git.warky.dev/wdevs/amcs/internal/config" + "git.warky.dev/wdevs/amcs/internal/tenancy" +) + +func TestMiddlewareAddsTenantKeyToContext(t *testing.T) { + keyring, err := NewKeyring([]config.APIKey{{ID: "user-a", Value: "secret-a"}}) + if err != nil { + t.Fatalf("NewKeyring error = %v", err) + } + + var gotKeyID, gotTenant string + handler := Middleware(config.AuthConfig{}, keyring, nil, nil, nil, slog.Default())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var ok bool + gotKeyID, ok = KeyIDFromContext(r.Context()) + if !ok { + t.Fatal("KeyIDFromContext ok = false") + } + gotTenant, ok = tenancy.KeyFromContext(r.Context()) + if !ok { + t.Fatal("tenancy.KeyFromContext ok = false") + } + w.WriteHeader(http.StatusNoContent) + })) + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + req.Header.Set("x-brain-key", "secret-a") + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if rr.Code != http.StatusNoContent { + t.Fatalf("status = %d, want %d", rr.Code, http.StatusNoContent) + } + if gotKeyID != "user-a" || gotTenant != "user-a" { + t.Fatalf("keyID=%q tenant=%q, want user-a/user-a", gotKeyID, gotTenant) + } +} diff --git a/internal/store/files.go b/internal/store/files.go index 4841ef6..279e512 100644 --- a/internal/store/files.go +++ b/internal/store/files.go @@ -15,10 +15,10 @@ import ( func (db *DB) InsertStoredFile(ctx context.Context, file thoughttypes.StoredFile) (thoughttypes.StoredFile, error) { row := db.pool.QueryRow(ctx, ` - insert into stored_files (thought_id, project_id, name, media_type, kind, encoding, size_bytes, sha256, content) - values ($1, $2, $3, $4, $5, $6, $7, $8, $9) + insert into stored_files (thought_id, project_id, tenant_key, name, media_type, kind, encoding, size_bytes, sha256, content) + values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) returning id, guid, thought_id, project_id, name, media_type, kind, encoding, size_bytes, sha256, created_at, updated_at - `, file.ThoughtID, file.ProjectID, file.Name, file.MediaType, file.Kind, file.Encoding, file.SizeBytes, file.SHA256, file.Content) + `, file.ThoughtID, file.ProjectID, tenantKeyPtr(ctx), file.Name, file.MediaType, file.Kind, file.Encoding, file.SizeBytes, file.SHA256, file.Content) var model generatedmodels.ModelPublicStoredFiles if err := row.Scan( @@ -42,11 +42,11 @@ func (db *DB) InsertStoredFile(ctx context.Context, file thoughttypes.StoredFile } func (db *DB) GetStoredFile(ctx context.Context, id uuid.UUID) (thoughttypes.StoredFile, error) { + args := []any{id} row := db.pool.QueryRow(ctx, ` select id, guid, thought_id, project_id, name, media_type, kind, encoding, size_bytes, sha256, content, created_at, updated_at from stored_files - where guid = $1 - `, id) + where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) var model generatedmodels.ModelPublicStoredFiles if err := row.Scan( @@ -77,6 +77,7 @@ func (db *DB) ListStoredFiles(ctx context.Context, filter thoughttypes.StoredFil args := make([]any, 0, 4) conditions := make([]string, 0, 3) + addTenantCondition(ctx, &args, &conditions, "tenant_key") if filter.ThoughtID != nil { args = append(args, *filter.ThoughtID) conditions = append(conditions, fmt.Sprintf("thought_id = $%d", len(args))) diff --git a/internal/store/projects.go b/internal/store/projects.go index 0c2190a..24d5b1f 100644 --- a/internal/store/projects.go +++ b/internal/store/projects.go @@ -14,10 +14,10 @@ import ( func (db *DB) CreateProject(ctx context.Context, name, description string) (thoughttypes.Project, error) { row := db.pool.QueryRow(ctx, ` - insert into projects (name, description) - values ($1, $2) + insert into projects (name, description, tenant_key) + values ($1, $2, $3) returning id, guid, name, description, created_at, last_active_at - `, name, description) + `, name, description, tenantKeyPtr(ctx)) var model generatedmodels.ModelPublicProjects if err := row.Scan(&model.ID, &model.GUID, &model.Name, &model.Description, &model.CreatedAt, &model.LastActiveAt); err != nil { @@ -45,20 +45,20 @@ func (db *DB) GetProject(ctx context.Context, nameOrID string) (thoughttypes.Pro } func (db *DB) getProjectByGUID(ctx context.Context, id uuid.UUID) (thoughttypes.Project, error) { + args := []any{id} row := db.pool.QueryRow(ctx, ` select id, guid, name, description, created_at, last_active_at from projects - where guid = $1 - `, id) + where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) return scanProject(row) } func (db *DB) getProjectByName(ctx context.Context, name string) (thoughttypes.Project, error) { + args := []any{name} row := db.pool.QueryRow(ctx, ` select id, guid, name, description, created_at, last_active_at from projects - where name = $1 - `, name) + where name = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) return scanProject(row) } @@ -74,13 +74,20 @@ func scanProject(row pgx.Row) (thoughttypes.Project, error) { } func (db *DB) ListProjects(ctx context.Context) ([]thoughttypes.ProjectSummary, error) { + args := []any{} + where := "" + if key, ok := tenantKey(ctx); ok { + args = append(args, key) + where = "where p.tenant_key = $1" + } rows, err := db.pool.Query(ctx, ` select p.id, p.guid, p.name, p.description, p.created_at, p.last_active_at, count(t.id) as thought_count from projects p left join thoughts t on t.project_id = p.id and t.archived_at is null + `+where+` group by p.id, p.guid, p.name, p.description, p.created_at, p.last_active_at order by p.last_active_at desc, p.created_at desc - `) + `, args...) if err != nil { return nil, fmt.Errorf("list projects: %w", err) } @@ -105,7 +112,8 @@ func (db *DB) ListProjects(ctx context.Context) ([]thoughttypes.ProjectSummary, } func (db *DB) TouchProject(ctx context.Context, id int64) error { - tag, err := db.pool.Exec(ctx, `update projects set last_active_at = now() where id = $1`, id) + args := []any{id} + tag, err := db.pool.Exec(ctx, `update projects set last_active_at = now() where id = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) if err != nil { return fmt.Errorf("touch project: %w", err) } diff --git a/internal/store/tenancy.go b/internal/store/tenancy.go new file mode 100644 index 0000000..2b6ec9a --- /dev/null +++ b/internal/store/tenancy.go @@ -0,0 +1,34 @@ +package store + +import ( + "context" + "fmt" + + "git.warky.dev/wdevs/amcs/internal/tenancy" +) + +func tenantKeyPtr(ctx context.Context) *string { + if key, ok := tenancy.KeyFromContext(ctx); ok { + return &key + } + return nil +} + +func tenantKey(ctx context.Context) (string, bool) { + return tenancy.KeyFromContext(ctx) +} + +func addTenantCondition(ctx context.Context, args *[]any, conditions *[]string, column string) { + if key, ok := tenancy.KeyFromContext(ctx); ok { + *args = append(*args, key) + *conditions = append(*conditions, fmt.Sprintf("%s = $%d", column, len(*args))) + } +} + +func tenantSQL(ctx context.Context, args *[]any, column string) string { + if key, ok := tenancy.KeyFromContext(ctx); ok { + *args = append(*args, key) + return fmt.Sprintf(" and %s = $%d", column, len(*args)) + } + return "" +} diff --git a/internal/store/thoughts.go b/internal/store/thoughts.go index 942a8dc..fdd2f9f 100644 --- a/internal/store/thoughts.go +++ b/internal/store/thoughts.go @@ -31,10 +31,10 @@ func (db *DB) InsertThought(ctx context.Context, thought thoughttypes.Thought, e }() row := tx.QueryRow(ctx, ` - insert into thoughts (content, metadata, project_id) - values ($1, $2::jsonb, $3) + insert into thoughts (content, metadata, project_id, tenant_key) + values ($1, $2::jsonb, $3, $4) returning id, guid, created_at, updated_at - `, thought.Content, metadata, thought.ProjectID) + `, thought.Content, metadata, thought.ProjectID, tenantKeyPtr(ctx)) created := thought created.Embedding = nil @@ -107,6 +107,7 @@ func (db *DB) ListThoughts(ctx context.Context, filter thoughttypes.ListFilter) args := make([]any, 0, 6) conditions := []string{} + addTenantCondition(ctx, &args, &conditions, "tenant_key") if !filter.IncludeArchived { conditions = append(conditions, "archived_at is null") } @@ -170,11 +171,14 @@ func (db *DB) ListThoughts(ctx context.Context, filter thoughttypes.ListFilter) func (db *DB) Stats(ctx context.Context) (thoughttypes.ThoughtStats, error) { var total int - if err := db.pool.QueryRow(ctx, `select count(*) from thoughts where archived_at is null`).Scan(&total); err != nil { + statsArgs := []any{} + statsConditions := []string{"archived_at is null"} + addTenantCondition(ctx, &statsArgs, &statsConditions, "tenant_key") + if err := db.pool.QueryRow(ctx, `select count(*) from thoughts where `+strings.Join(statsConditions, " and "), statsArgs...).Scan(&total); err != nil { return thoughttypes.ThoughtStats{}, fmt.Errorf("count thoughts: %w", err) } - rows, err := db.pool.Query(ctx, `select metadata from thoughts where archived_at is null`) + rows, err := db.pool.Query(ctx, `select metadata from thoughts where `+strings.Join(statsConditions, " and "), statsArgs...) if err != nil { return thoughttypes.ThoughtStats{}, fmt.Errorf("query stats metadata: %w", err) } @@ -217,11 +221,11 @@ func (db *DB) Stats(ctx context.Context) (thoughttypes.ThoughtStats, error) { } func (db *DB) GetThought(ctx context.Context, id uuid.UUID) (thoughttypes.Thought, error) { + args := []any{id} row := db.pool.QueryRow(ctx, ` select id, guid, content, metadata, project_id, archived_at, created_at, updated_at from thoughts - where guid = $1 - `, id) + where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) var model generatedmodels.ModelPublicThoughts if err := row.Scan(&model.ID, &model.GUID, &model.Content, &model.Metadata, &model.ProjectID, &model.ArchivedAt, &model.CreatedAt, &model.UpdatedAt); err != nil { @@ -240,11 +244,11 @@ func (db *DB) GetThought(ctx context.Context, id uuid.UUID) (thoughttypes.Though } func (db *DB) GetThoughtByID(ctx context.Context, id int64) (thoughttypes.Thought, error) { + args := []any{id} row := db.pool.QueryRow(ctx, ` select id, guid, content, metadata, project_id, archived_at, created_at, updated_at from thoughts - where id = $1 - `, id) + where id = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) var model generatedmodels.ModelPublicThoughts if err := row.Scan(&model.ID, &model.GUID, &model.Content, &model.Metadata, &model.ProjectID, &model.ArchivedAt, &model.CreatedAt, &model.UpdatedAt); err != nil { @@ -276,14 +280,14 @@ func (db *DB) UpdateThought(ctx context.Context, id uuid.UUID, content string, e _ = tx.Rollback(ctx) }() + args := []any{id, content, metadataBytes, projectID} tag, err := tx.Exec(ctx, ` update thoughts set content = $2, metadata = $3::jsonb, project_id = $4, updated_at = now() - where guid = $1 - `, id, content, metadataBytes, projectID) + where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) if err != nil { return thoughttypes.Thought{}, fmt.Errorf("update thought: %w", err) } @@ -317,12 +321,12 @@ func (db *DB) UpdateThoughtMetadata(ctx context.Context, id int64, metadata thou return thoughttypes.Thought{}, fmt.Errorf("marshal updated metadata: %w", err) } + args := []any{id, metadataBytes} tag, err := db.pool.Exec(ctx, ` update thoughts set metadata = $2::jsonb, updated_at = now() - where id = $1 - `, id, metadataBytes) + where id = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) if err != nil { return thoughttypes.Thought{}, fmt.Errorf("update thought metadata: %w", err) } @@ -334,7 +338,8 @@ func (db *DB) UpdateThoughtMetadata(ctx context.Context, id int64, metadata thou } func (db *DB) DeleteThought(ctx context.Context, id uuid.UUID) error { - tag, err := db.pool.Exec(ctx, `delete from thoughts where guid = $1`, id) + args := []any{id} + tag, err := db.pool.Exec(ctx, `delete from thoughts where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) if err != nil { return fmt.Errorf("delete thought: %w", err) } @@ -345,7 +350,8 @@ func (db *DB) DeleteThought(ctx context.Context, id uuid.UUID) error { } func (db *DB) ArchiveThought(ctx context.Context, id uuid.UUID) error { - tag, err := db.pool.Exec(ctx, `update thoughts set archived_at = now(), updated_at = now() where guid = $1`, id) + args := []any{id} + tag, err := db.pool.Exec(ctx, `update thoughts set archived_at = now(), updated_at = now() where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...) if err != nil { return fmt.Errorf("archive thought: %w", err) } @@ -424,6 +430,7 @@ func (db *DB) SearchSimilarThoughts(ctx context.Context, embedding []float32, em "1 - (e.embedding <=> $1) > $2", "e.model = $3", } + addTenantCondition(ctx, &args, &conditions, "t.tenant_key") if projectID != nil { args = append(args, *projectID) conditions = append(conditions, fmt.Sprintf("t.project_id = $%d", len(args))) @@ -472,6 +479,7 @@ func (db *DB) HasEmbeddingsForModel(ctx context.Context, model string, projectID "e.model = $1", "t.archived_at is null", } + addTenantCondition(ctx, &args, &conditions, "t.tenant_key") if projectID != nil { args = append(args, *projectID) conditions = append(conditions, fmt.Sprintf("t.project_id = $%d", len(args))) @@ -490,6 +498,7 @@ func (db *DB) HasEmbeddingsForModel(ctx context.Context, model string, projectID func (db *DB) ListThoughtsMissingEmbedding(ctx context.Context, model string, limit int, projectID *int64, includeArchived bool, olderThanDays int) ([]thoughttypes.Thought, error) { args := []any{model} conditions := []string{"e.id is null"} + addTenantCondition(ctx, &args, &conditions, "t.tenant_key") if !includeArchived { conditions = append(conditions, "t.archived_at is null") @@ -539,6 +548,7 @@ func (db *DB) ListThoughtsMissingEmbedding(ctx context.Context, model string, li func (db *DB) ListThoughtsForMetadataReparse(ctx context.Context, limit int, projectID *int64, includeArchived bool, olderThanDays int) ([]thoughttypes.Thought, error) { args := make([]any, 0, 3) conditions := make([]string, 0, 4) + addTenantCondition(ctx, &args, &conditions, "tenant_key") if !includeArchived { conditions = append(conditions, "archived_at is null") @@ -608,6 +618,7 @@ func (db *DB) SearchThoughtsText(ctx context.Context, query string, limit int, p "t.archived_at is null", "(to_tsvector('simple', t.content) || to_tsvector('simple', coalesce(p.name, ''))) @@ websearch_to_tsquery('simple', $1)", } + addTenantCondition(ctx, &args, &conditions, "t.tenant_key") if projectID != nil { args = append(args, *projectID) conditions = append(conditions, fmt.Sprintf("t.project_id = $%d", len(args))) diff --git a/internal/tenancy/tenancy.go b/internal/tenancy/tenancy.go new file mode 100644 index 0000000..27b1604 --- /dev/null +++ b/internal/tenancy/tenancy.go @@ -0,0 +1,30 @@ +package tenancy + +import ( + "context" + "strings" +) + +type contextKey string + +const tenantKeyContextKey contextKey = "tenancy.tenant_key" + +// WithTenantKey returns a context scoped to the authenticated tenant boundary. +// The tenant key is intentionally opaque; today it is the authenticated API key +// or OAuth client id, and callers should not interpret it as a human username. +func WithTenantKey(ctx context.Context, tenantKey string) context.Context { + tenantKey = strings.TrimSpace(tenantKey) + if tenantKey == "" { + return ctx + } + return context.WithValue(ctx, tenantKeyContextKey, tenantKey) +} + +func KeyFromContext(ctx context.Context) (string, bool) { + if ctx == nil { + return "", false + } + value, ok := ctx.Value(tenantKeyContextKey).(string) + value = strings.TrimSpace(value) + return value, ok && value != "" +} diff --git a/migrations/110_per_user_tenancy.sql b/migrations/110_per_user_tenancy.sql new file mode 100644 index 0000000..16455fd --- /dev/null +++ b/migrations/110_per_user_tenancy.sql @@ -0,0 +1,22 @@ +-- Per-user tenancy: authenticate key/client id becomes an opaque tenant boundary. +-- Existing rows remain in the legacy unscoped tenant (NULL) for single-tenant installs. + +alter table projects add column if not exists tenant_key text; +alter table thoughts add column if not exists tenant_key text; +alter table stored_files add column if not exists tenant_key text; +alter table learnings add column if not exists tenant_key text; +alter table plans add column if not exists tenant_key text; +alter table chat_histories add column if not exists tenant_key text; + +-- Project names are now unique inside a tenant rather than globally. +drop index if exists projects_name_key; +create unique index if not exists projects_tenant_key_name_idx + on projects (coalesce(tenant_key, ''), name); + +create index if not exists projects_tenant_key_idx on projects (tenant_key); +create index if not exists thoughts_tenant_key_idx on thoughts (tenant_key); +create index if not exists thoughts_tenant_key_project_id_idx on thoughts (tenant_key, project_id); +create index if not exists stored_files_tenant_key_idx on stored_files (tenant_key); +create index if not exists learnings_tenant_key_idx on learnings (tenant_key); +create index if not exists plans_tenant_key_idx on plans (tenant_key); +create index if not exists chat_histories_tenant_key_idx on chat_histories (tenant_key); diff --git a/schema/core.dbml b/schema/core.dbml index 8790a87..3006be7 100644 --- a/schema/core.dbml +++ b/schema/core.dbml @@ -6,16 +6,28 @@ Table thoughts { created_at timestamptz [default: `now()`] updated_at timestamptz [default: `now()`] project_id bigint [ref: > projects.id] + tenant_key text archived_at timestamptz + + indexes { + tenant_key + (tenant_key, project_id) + } } Table projects { id bigserial [pk] guid uuid [unique, not null, default: `gen_random_uuid()`] - name text [unique, not null] + name text [not null] description text + tenant_key text created_at timestamptz [default: `now()`] last_active_at timestamptz [default: `now()`] + + indexes { + (tenant_key, name) [unique] + tenant_key + } } Table thought_links {