diff --git a/internal/store/copy.go b/internal/store/copy.go index a9d0a9b..589e936 100644 --- a/internal/store/copy.go +++ b/internal/store/copy.go @@ -70,6 +70,8 @@ var copyOrder = []string{ // Who changed what is not lost by moving the store. "audit_events", "key_ttl", + // The servers, and the secrets their headers name are in cluster_secrets. + "mcp_servers", "files", "batches", "batch_results", diff --git a/internal/store/mcp.go b/internal/store/mcp.go new file mode 100644 index 0000000..1fa3aa7 --- /dev/null +++ b/internal/store/mcp.go @@ -0,0 +1,104 @@ +package store + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "time" +) + +// MCPServer is one of the company's MCP servers, served through llmgw's +// /mcp under its name: its tools are name.tool there. +type MCPServer struct { + Name string + URL string + Description string + // Headers are sent as written; HeaderRefs name a secret to read for + // each header (a file, env:, vault: or store:), so no token is stored + // in this row. + Headers map[string]string + HeaderRefs map[string]string + Enabled bool + // Timeout bounds one call; zero is the gateway's default. + Timeout time.Duration + CreatedAt time.Time + CreatedBy string +} + +// PutMCPServer adds a server, or replaces the one of that name. +func (s *DB) PutMCPServer(ctx context.Context, m MCPServer) error { + h, err := json.Marshal(nonNil(m.Headers)) + if err != nil { + return err + } + r, err := json.Marshal(nonNil(m.HeaderRefs)) + if err != nil { + return err + } + if m.CreatedAt.IsZero() { + m.CreatedAt = time.Now() + } + _, err = s.db.ExecContext(ctx, + `INSERT INTO mcp_servers (name, url, description, headers, header_refs, enabled, timeout_ms, created_at, created_by) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (name) DO UPDATE SET + url = excluded.url, description = excluded.description, headers = excluded.headers, + header_refs = excluded.header_refs, enabled = excluded.enabled, timeout_ms = excluded.timeout_ms`, + m.Name, m.URL, m.Description, string(h), string(r), boolInt(m.Enabled), m.Timeout.Milliseconds(), + m.CreatedAt.UnixMilli(), m.CreatedBy) + return err +} + +// MCPServers are every server registered, by name. +func (s *DB) MCPServers(ctx context.Context) ([]MCPServer, error) { + return queryAll(ctx, s.db, `SELECT name, url, description, headers, header_refs, enabled, timeout_ms, created_at, created_by + FROM mcp_servers ORDER BY name`, nil, scanMCPServer) +} + +// MCPServer is the server of that name, ErrNotFound when there is none. +func (s *DB) MCPServer(ctx context.Context, name string) (MCPServer, error) { + list, err := queryAll(ctx, s.db, `SELECT name, url, description, headers, header_refs, enabled, timeout_ms, created_at, created_by + FROM mcp_servers WHERE name = ?`, []any{name}, scanMCPServer) + if err != nil { + return MCPServer{}, err + } + if len(list) == 0 { + return MCPServer{}, ErrNotFound + } + return list[0], nil +} + +// DeleteMCPServer removes a server; ErrNotFound when there was none. +func (s *DB) DeleteMCPServer(ctx context.Context, name string) error { + res, err := s.db.ExecContext(ctx, `DELETE FROM mcp_servers WHERE name = ?`, name) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrNotFound + } + return nil +} + +func scanMCPServer(rows *sql.Rows) (MCPServer, error) { + var m MCPServer + var headers, refs string + var enabled int + var timeoutMs, created int64 + if err := rows.Scan(&m.Name, &m.URL, &m.Description, &headers, &refs, &enabled, &timeoutMs, &created, &m.CreatedBy); err != nil { + return m, err + } + m.Enabled, m.Timeout, m.CreatedAt = enabled != 0, time.Duration(timeoutMs)*time.Millisecond, time.UnixMilli(created) + if err := errors.Join(json.Unmarshal([]byte(headers), &m.Headers), json.Unmarshal([]byte(refs), &m.HeaderRefs)); err != nil { + return m, err + } + return m, nil +} + +func nonNil(m map[string]string) map[string]string { + if m == nil { + return map[string]string{} + } + return m +} diff --git a/internal/store/mcp_test.go b/internal/store/mcp_test.go new file mode 100644 index 0000000..6599915 --- /dev/null +++ b/internal/store/mcp_test.go @@ -0,0 +1,73 @@ +package store + +import ( + "context" + "errors" + "testing" + "time" +) + +// A company tool is granted, never assumed: none anywhere is none, and a +// level that names any narrows what the levels above it gave. +func TestACompanyToolIsTheCallersOnlyWhereGrantedAllTheWayDown(t *testing.T) { + s := openDB(t) + ctx := context.Background() + project := tree(t, s) + _, k := mustKey(t, s, "key1", project) + + e, err := s.Resolve(ctx, k.ID) + if err != nil { + t.Fatal(err) + } + if e.AllowsTool("jira.search") { + t.Fatal("a tool nobody granted is allowed") + } + + must(t, s.PutLimits(ctx, Limits{Scope: ScopeTeam, ScopeID: "team1", AllowedTools: []string{"jira.*", "wiki.read"}})) + e, _ = s.Resolve(ctx, k.ID) + if !e.AllowsTool("jira.search") || !e.AllowsTool("wiki.read") || e.AllowsTool("db.query") || e.AllowsTool("wiki.write") { + t.Errorf("the team's grant: %v", e.ToolGrants) + } + // The key narrows it; it cannot add what the team withheld. + must(t, s.PutLimits(ctx, Limits{Scope: ScopeKey, ScopeID: k.ID, AllowedTools: []string{"jira.search", "db.query"}})) + e, _ = s.Resolve(ctx, k.ID) + if !e.AllowsTool("jira.search") || e.AllowsTool("jira.create") || e.AllowsTool("db.query") { + t.Errorf("narrowed by the key: %v", e.ToolGrants) + } + // Read back as written. + l, err := s.LimitsFor(ctx, ScopeKey, k.ID) + if err != nil || len(l.AllowedTools) != 2 { + t.Errorf("limits %+v %v", l, err) + } + all, _ := s.AllLimits(ctx) + if got := all["team:team1"].AllowedTools; len(got) != 2 { + t.Errorf("AllLimits: %v", got) + } +} + +func TestACompanyMCPServerIsKeptReplacedAndRemoved(t *testing.T) { + s := openDB(t) + ctx := context.Background() + m := MCPServer{Name: "jira", URL: "https://jira.internal/mcp", Description: "tickets", + Headers: map[string]string{"X-Org": "acme"}, HeaderRefs: map[string]string{"Authorization": "store:mcp/jira/Authorization"}, + Enabled: true, Timeout: 30 * time.Second, CreatedBy: "ilker"} + must(t, s.PutMCPServer(ctx, m)) + got, err := s.MCPServer(ctx, "jira") + if err != nil || got.URL != m.URL || got.Headers["X-Org"] != "acme" || got.HeaderRefs["Authorization"] == "" || + !got.Enabled || got.Timeout != 30*time.Second || got.CreatedBy != "ilker" { + t.Fatalf("%+v %v", got, err) + } + m.Enabled, m.URL = false, "https://jira2.internal/mcp" + must(t, s.PutMCPServer(ctx, m)) + list, _ := s.MCPServers(ctx) + if len(list) != 1 || list[0].Enabled || list[0].URL != m.URL { + t.Errorf("replaced: %+v", list) + } + must(t, s.DeleteMCPServer(ctx, "jira")) + if _, err := s.MCPServer(ctx, "jira"); !errors.Is(err, ErrNotFound) { + t.Errorf("after removal: %v", err) + } + if err := s.DeleteMCPServer(ctx, "jira"); !errors.Is(err, ErrNotFound) { + t.Errorf("removing twice: %v", err) + } +} diff --git a/internal/store/migrations/postgres/0049_mcp_servers.sql b/internal/store/migrations/postgres/0049_mcp_servers.sql new file mode 100644 index 0000000..7173596 --- /dev/null +++ b/internal/store/migrations/postgres/0049_mcp_servers.sql @@ -0,0 +1,21 @@ +-- Ported from the SQLite migration of the same number. +-- The company's own MCP servers, served through llmgw's /mcp: registered +-- once by an administrator, their credentials kept here rather than in every +-- client. headers are sent as written; header_refs name a secret to read for +-- each (a file, env:, vault: or store:), so a token is never in this table. +CREATE TABLE mcp_servers ( + name TEXT PRIMARY KEY, + url TEXT NOT NULL, + description TEXT NOT NULL DEFAULT '', + headers TEXT NOT NULL DEFAULT '{}', + header_refs TEXT NOT NULL DEFAULT '{}', + enabled INTEGER NOT NULL DEFAULT 1, + timeout_ms INTEGER NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + created_by TEXT NOT NULL DEFAULT '' +); + +-- Which of those servers' tools a scope may call, as patterns (jira.*, +-- db.query). Unlike allowed_models, none set anywhere means none: a tool +-- reaches into a company system, so it is granted, never assumed. +ALTER TABLE limits ADD COLUMN allowed_tools TEXT; diff --git a/internal/store/migrations/sqlite/0049_mcp_servers.sql b/internal/store/migrations/sqlite/0049_mcp_servers.sql new file mode 100644 index 0000000..e65e006 --- /dev/null +++ b/internal/store/migrations/sqlite/0049_mcp_servers.sql @@ -0,0 +1,20 @@ +-- The company's own MCP servers, served through llmgw's /mcp: registered +-- once by an administrator, their credentials kept here rather than in every +-- client. headers are sent as written; header_refs name a secret to read for +-- each (a file, env:, vault: or store:), so a token is never in this table. +CREATE TABLE mcp_servers ( + name TEXT PRIMARY KEY, + url TEXT NOT NULL, + description TEXT NOT NULL DEFAULT '', + headers TEXT NOT NULL DEFAULT '{}', + header_refs TEXT NOT NULL DEFAULT '{}', + enabled INTEGER NOT NULL DEFAULT 1, + timeout_ms INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + created_by TEXT NOT NULL DEFAULT '' +); + +-- Which of those servers' tools a scope may call, as patterns (jira.*, +-- db.query). Unlike allowed_models, none set anywhere means none: a tool +-- reaches into a company system, so it is granted, never assumed. +ALTER TABLE limits ADD COLUMN allowed_tools TEXT; diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index ecdda69..a77787e 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -267,29 +267,33 @@ func (s *DB) PutLimits(ctx context.Context, l Limits) error { if err != nil { return err } + tools, err := encodeModels(l.AllowedTools) + if err != nil { + return err + } _, err = s.db.ExecContext(ctx, - `INSERT INTO limits (scope_type, scope_id, rpm, tpm, budget_micros, budget_period, allowed_models) - VALUES (?, ?, ?, ?, ?, ?, ?) + `INSERT INTO limits (scope_type, scope_id, rpm, tpm, budget_micros, budget_period, allowed_models, allowed_tools) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT (scope_type, scope_id) DO UPDATE SET rpm = excluded.rpm, tpm = excluded.tpm, budget_micros = excluded.budget_micros, budget_period = excluded.budget_period, - allowed_models = excluded.allowed_models`, + allowed_models = excluded.allowed_models, allowed_tools = excluded.allowed_tools`, string(l.Scope), l.ScopeID, nullInt(l.RPM), nullInt(l.TPM), - nullInt64(l.Budget), nullString(string(l.BudgetPeriod)), models) + nullInt64(l.Budget), nullString(string(l.BudgetPeriod)), models, tools) return err } // LimitsFor reads the limits set at one scope, without inheritance. func (s *DB) LimitsFor(ctx context.Context, scope Scope, id string) (Limits, error) { row := s.db.QueryRowContext(ctx, - `SELECT rpm, tpm, budget_micros, budget_period, allowed_models + `SELECT rpm, tpm, budget_micros, budget_period, allowed_models, allowed_tools FROM limits WHERE scope_type = ? AND scope_id = ?`, string(scope), id) l := Limits{Scope: scope, ScopeID: id} var rpm, tpm sql.NullInt64 var budget sql.NullInt64 - var period, models sql.NullString - if err := row.Scan(&rpm, &tpm, &budget, &period, &models); err != nil { + var period, models, tools sql.NullString + if err := row.Scan(&rpm, &tpm, &budget, &period, &models, &tools); err != nil { if errors.Is(err, sql.ErrNoRows) { return Limits{}, ErrNotFound } @@ -312,6 +316,9 @@ func (s *DB) LimitsFor(ctx context.Context, scope Scope, id string) (Limits, err if l.AllowedModels, err = decodeModels(models); err != nil { return Limits{}, err } + if l.AllowedTools, err = decodeModels(tools); err != nil { + return Limits{}, err + } return l, nil } @@ -409,6 +416,9 @@ func (e *Effective) tighten(l Limits) { e.BudgetPeriod = l.BudgetPeriod e.BudgetOwner = owner } + if l.AllowedTools != nil { + e.ToolGrants = append(e.ToolGrants, slices.Clone(l.AllowedTools)) + } if l.AllowedModels != nil { if e.AllowedModels == nil { e.AllowedModels = slices.Clone(l.AllowedModels) diff --git a/internal/store/store.go b/internal/store/store.go index d7067dd..5fac888 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -18,6 +18,7 @@ import ( "context" "encoding/json" "errors" + "path" "strings" "time" ) @@ -130,6 +131,10 @@ type Limits struct { BudgetPeriod Period // AllowedModels nil means every model. AllowedModels []string + // AllowedTools are the company MCP servers' tools this scope may call, + // as patterns: jira.*, db.query, *. Unlike models, nil says nothing + // here, and a tool no level grants is not granted. + AllowedTools []string } // Period is how often a budget resets. @@ -165,6 +170,10 @@ type Effective struct { BudgetPeriod Period // AllowedModels is the intersection along the chain. nil means every model. AllowedModels []string + // ToolGrants are AllowedTools of every level that set them, widest + // first: a tool is the caller's when each of them grants it, and at + // least one level said something. + ToolGrants [][]string // ExcludedLicences are the licences the caller's person has said not to // use, as LicenceKey gives them; none for a key without a person. ExcludedLicences []string @@ -228,6 +237,30 @@ func (e Effective) Allows(model string) bool { return false } +// AllowsTool reports whether a company MCP server's tool, server.tool, may +// be called with this key: granted somewhere along the chain, and by every +// level that names any. +func (e Effective) AllowsTool(name string) bool { + if len(e.ToolGrants) == 0 { + return false + } + for _, grants := range e.ToolGrants { + if !matchesTool(grants, name) { + return false + } + } + return true +} + +func matchesTool(patterns []string, name string) bool { + for _, p := range patterns { + if ok, err := path.Match(p, name); err == nil && ok { + return true + } + } + return false +} + // Usage is one request's accounting record. type Usage struct { At time.Time diff --git a/internal/store/tenancy.go b/internal/store/tenancy.go index cbba867..ec9c131 100644 --- a/internal/store/tenancy.go +++ b/internal/store/tenancy.go @@ -128,15 +128,15 @@ func (s *DB) nameIsFree(ctx context.Context, q string, args ...any) error { // otherwise be a query each. func (s *DB) AllLimits(ctx context.Context) (map[string]Limits, error) { rows, err := queryAll(ctx, s.db, ` - SELECT scope_type, scope_id, rpm, tpm, budget_micros, budget_period, allowed_models + SELECT scope_type, scope_id, rpm, tpm, budget_micros, budget_period, allowed_models, allowed_tools FROM limits`, nil, func(rows *sql.Rows) (Limits, error) { var l Limits var scope string - var models sql.NullString + var models, tools sql.NullString var rpm, tpm sql.NullInt64 var budget sql.NullInt64 var period sql.NullString - if err := rows.Scan(&scope, &l.ScopeID, &rpm, &tpm, &budget, &period, &models); err != nil { + if err := rows.Scan(&scope, &l.ScopeID, &rpm, &tpm, &budget, &period, &models, &tools); err != nil { return l, err } l.Scope = Scope(scope) @@ -155,6 +155,10 @@ func (s *DB) AllLimits(ctx context.Context) (map[string]Limits, error) { l.BudgetPeriod = Period(period.String) am, err := decodeModels(models) l.AllowedModels = am + if err != nil { + return l, err + } + l.AllowedTools, err = decodeModels(tools) return l, err }) if err != nil {