feat: add per-user tenant scoping
CI / build-and-test (push) Failing after 1m53s
CI / build-and-test (pull_request) Failing after 1m43s

This commit is contained in:
2026-07-15 04:27:44 +02:00
parent 5718685c40
commit 4b3f0b1b55
9 changed files with 202 additions and 35 deletions
+10 -5
View File
@@ -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
}
}
+44
View File
@@ -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)
}
}
+6 -5
View File
@@ -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)))
+17 -9
View File
@@ -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)
}
+34
View File
@@ -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 ""
}
+26 -15
View File
@@ -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)))
+30
View File
@@ -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 != ""
}
+22
View File
@@ -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);
+13 -1
View File
@@ -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 {