From 39c0427eda67a429a663dd92c103b70de31ddfae Mon Sep 17 00:00:00 2001 From: Moshe Immerman Date: Tue, 11 Aug 2026 09:16:24 +0300 Subject: [PATCH 1/2] feat(aichat): resolve hierarchical runtime profiles Claude-Session-Id: 019fef59-46c9-7b83-bc02-6bd4246d90de --- .../aimock_lifecycle_integration_test.go | 4 +- pkg/aichat/approval_execution.go | 6 +- pkg/aichat/execution.go | 1 + pkg/aichat/execution_authority_ginkgo_test.go | 10 + pkg/aichat/execution_database_authority.go | 18 +- pkg/aichat/messages.go | 35 +- pkg/aichat/provider_config.go | 83 +++++ pkg/aichat/runtime_profile_ginkgo_test.go | 80 +++++ pkg/aichat/runtime_settings.go | 35 +- pkg/aichat/service.go | 75 +++-- pkg/aichat/service_ginkgo_test.go | 42 +-- pkg/api/spec_layers.go | 299 ++++++++++++++++++ pkg/api/spec_layers_ginkgo_test.go | 100 ++++++ pkg/cli/serve_chat.go | 15 +- 14 files changed, 712 insertions(+), 91 deletions(-) create mode 100644 pkg/aichat/runtime_profile_ginkgo_test.go create mode 100644 pkg/api/spec_layers.go create mode 100644 pkg/api/spec_layers_ginkgo_test.go diff --git a/pkg/aichat/aimock_lifecycle_integration_test.go b/pkg/aichat/aimock_lifecycle_integration_test.go index a95ad0d1..b0c7db5d 100644 --- a/pkg/aichat/aimock_lifecycle_integration_test.go +++ b/pkg/aichat/aimock_lifecycle_integration_test.go @@ -132,8 +132,8 @@ var _ = Describe("Mocked Captain chat lifecycle", func() { var approvedInput map[string]any service := aichat.NewService(aichat.ServiceOptions{ Resolver: realChatResolver{}, Threads: aichat.FixedThreadStore(store), Authority: authority, - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{ProviderConfig: api.Config{ + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return aichat.RuntimeProfile{ProviderConfig: api.Config{ APIURL: mock.apiURL, APIKey: aimock.DummyKey, }}, nil }), diff --git a/pkg/aichat/approval_execution.go b/pkg/aichat/approval_execution.go index 578b29e2..36a63db0 100644 --- a/pkg/aichat/approval_execution.go +++ b/pkg/aichat/approval_execution.go @@ -16,9 +16,9 @@ func (s *Service) resumeToolApproval(ctx context.Context, threadID string, conti } execution := continuation.Execution defer closeExecution(execution) - settings, err := s.runtimeSettings(ctx) + profile, err := s.runtimeProfile(ctx) if err != nil { - return fmt.Errorf("load chat runtime settings: %w", err) + return fmt.Errorf("load chat runtime profile: %w", err) } set, err := s.loadTools(ctx) if err != nil { @@ -28,7 +28,7 @@ func (s *Service) resumeToolApproval(ctx context.Context, threadID string, conti if err != nil { return err } - config := settings.ProviderConfig + config := profile.ProviderConfig config.Model = continuation.Spec.Model config.Budget = continuation.Spec.Budget config.SessionID = continuation.Spec.SessionID diff --git a/pkg/aichat/execution.go b/pkg/aichat/execution.go index 9a5a0304..a36c338c 100644 --- a/pkg/aichat/execution.go +++ b/pkg/aichat/execution.go @@ -15,6 +15,7 @@ type ExecutionRequest struct { RequestID string Title string Spec api.Spec + Profile api.ResolvedSpec Definitions []api.ToolDefinition } diff --git a/pkg/aichat/execution_authority_ginkgo_test.go b/pkg/aichat/execution_authority_ginkgo_test.go index b7ea15c9..62b36401 100644 --- a/pkg/aichat/execution_authority_ginkgo_test.go +++ b/pkg/aichat/execution_authority_ginkgo_test.go @@ -171,6 +171,12 @@ var _ = Describe("Authoritative aichat execution", func() { } service := aichat.NewService(aichat.ServiceOptions{ Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), Authority: authority, + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return mustRuntimeProfile(api.SpecLayer{ + Name: "accounts", Scope: api.SpecLayerContext, + Spec: api.Spec{Prompt: api.Prompt{System: "Use account policy."}}, + }), nil + }), Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ Name: "account_edit", DefaultPermission: api.ToolModeAsk, Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, @@ -191,6 +197,10 @@ var _ = Describe("Authoritative aichat execution", func() { Expect(authority.begins).To(HaveLen(1)) Expect(authority.begins[0].ThreadID).To(Equal(thread.ID)) Expect(authority.begins[0].RequestID).To(Equal("user-message-1")) + Expect(authority.begins[0].Profile.Trace).To(HaveLen(2)) + Expect(authority.begins[0].Profile.Trace[0].Name).To(Equal("accounts")) + Expect(authority.begins[0].Profile.Spec).To(Equal(authority.begins[0].Spec)) + Expect(authority.begins[0].Profile.Spec.Prompt.System).To(Equal("Use account policy.")) Expect(provider.specs).To(HaveLen(1)) Expect(execution.observed).To(ContainElement( MatchFields(IgnoreExtras, Fields{"Kind": Equal(api.EventResult)}), diff --git a/pkg/aichat/execution_database_authority.go b/pkg/aichat/execution_database_authority.go index 1a38f9e7..b3789e6d 100644 --- a/pkg/aichat/execution_database_authority.go +++ b/pkg/aichat/execution_database_authority.go @@ -46,7 +46,7 @@ func (a *DatabaseExecutionAuthority) Begin( if session.Source != "aichat" { return nil, fmt.Errorf("chat thread %s has incompatible source %q", request.ThreadID, session.Source) } - renderedSpec, err := renderedSpecMap(request.Spec) + renderedSpec, err := renderedSpecMap(request.Spec, request.Profile) if err != nil { return nil, err } @@ -278,7 +278,7 @@ func runtimeSelection(model api.Model) database.PromptRunRuntimeSelection { } } -func renderedSpecMap(spec api.Spec) (map[string]any, error) { +func renderedSpecMap(spec api.Spec, profile api.ResolvedSpec) (map[string]any, error) { raw, err := json.Marshal(spec) if err != nil { return nil, fmt.Errorf("encode authoritative chat spec: %w", err) @@ -287,6 +287,20 @@ func renderedSpecMap(spec api.Spec) (map[string]any, error) { if err := json.Unmarshal(raw, &rendered); err != nil { return nil, fmt.Errorf("decode authoritative chat spec: %w", err) } + if len(profile.Trace) > 0 { + resolution, err := json.Marshal(struct { + Constraints api.RuntimeConstraints `json:"constraints"` + Trace []api.SpecLayer `json:"trace"` + }{Constraints: profile.Constraints, Trace: profile.Trace}) + if err != nil { + return nil, fmt.Errorf("encode authoritative chat profile: %w", err) + } + var value map[string]any + if err := json.Unmarshal(resolution, &value); err != nil { + return nil, fmt.Errorf("decode authoritative chat profile: %w", err) + } + rendered["resolution"] = value + } return rendered, nil } diff --git a/pkg/aichat/messages.go b/pkg/aichat/messages.go index 0665549f..3529ecb3 100644 --- a/pkg/aichat/messages.go +++ b/pkg/aichat/messages.go @@ -68,19 +68,31 @@ func (s *Service) resolveAttachments(ctx context.Context, messages []UIMessage) return resolved, nil } -func requestSpec(request ChatRequest, settings RuntimeSettings, attachments map[partLocation]api.AttachmentRef) (api.Spec, error) { - override, err := chatModel(request, settings.Spec.Model) +func requestSpec(request ChatRequest, profile RuntimeProfile, attachments map[partLocation]api.AttachmentRef) (api.ResolvedSpec, error) { + if len(profile.Resolved.Trace) == 0 && (!api.IsEmpty(profile.Resolved.Spec) || !api.IsEmpty(profile.Resolved.Constraints)) { + return api.ResolvedSpec{}, fmt.Errorf("chat runtime profile must include its resolution trace") + } + override, err := chatModel(request, profile.Resolved.Spec.Model) if err != nil { - return api.Spec{}, err + return api.ResolvedSpec{}, err } - spec := settings.Spec.Merge(api.Spec{ + user := api.SpecLayer{Name: "chat request", Scope: api.SpecLayerUser, Spec: api.Spec{ Model: override, Budget: request.Budget, ToolPreferences: request.ToolPreferences, ToolApproval: request.ToolApproval, Permissions: api.Permissions{Mode: request.PermissionMode}, SessionID: request.ProviderSessionID, - }) + }} + layers := append([]api.SpecLayer(nil), profile.Resolved.Trace...) + resolved, err := api.ResolveSpecLayers(append(layers, user)...) + if err != nil { + return api.ResolvedSpec{}, fmt.Errorf("resolve chat runtime profile: %w", err) + } + spec := resolved.Spec + baseSystem := strings.TrimSpace(strings.Join([]string{ + profile.System, spec.Prompt.System, spec.Prompt.AppendSystem, + }, "\n\n")) spec.Prompt.User = "" spec.Prompt.System = "" spec.Prompt.AppendSystem = "" @@ -88,16 +100,16 @@ func requestSpec(request ChatRequest, settings RuntimeSettings, attachments map[ if request.ToolApproval == nil { messages, err := canonicalMessages(request.Messages, attachments) if err != nil { - return api.Spec{}, err + return api.ResolvedSpec{}, err } - system, err := requestSystem(settings.System, request) + system, err := requestSystem(baseSystem, request) if err != nil { - return api.Spec{}, err + return api.ResolvedSpec{}, err } if isAgentBackend(spec.Backend) { user, promptAttachments, err := agentPrompt(messages, request.ProviderSessionID != "") if err != nil { - return api.Spec{}, err + return api.ResolvedSpec{}, err } spec.Messages = nil spec.Prompt.System = system @@ -113,9 +125,10 @@ func requestSpec(request ChatRequest, settings RuntimeSettings, attachments map[ spec.Messages = nil } if err := spec.Validate(); err != nil { - return api.Spec{}, err + return api.ResolvedSpec{}, err } - return spec, nil + resolved.Spec = spec + return resolved, nil } func chatModel(request ChatRequest, fallback api.Model) (api.Model, error) { diff --git a/pkg/aichat/provider_config.go b/pkg/aichat/provider_config.go index 2d37b888..3547e2fb 100644 --- a/pkg/aichat/provider_config.go +++ b/pkg/aichat/provider_config.go @@ -4,11 +4,94 @@ import ( "context" "fmt" "reflect" + "strings" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" ) +func annotateProfileModels(resolved api.ResolvedSpec, models ModelCatalogResponse) { + if len(resolved.Constraints.Models) == 0 { + return + } + for index := range models { + if resolved.AllowsModel(models[index].Runtime) { + continue + } + layer := modelRestrictionLayer(resolved, models[index].Runtime) + models[index].Configured = false + models[index].Default = false + models[index].Availability = restrictedAvailability(layer, resolved.Constraints.Models) + } +} + +func annotateProfileRuntimes(resolved api.ResolvedSpec, runtimes []api.RuntimeFamily) { + if len(resolved.Constraints.Models) == 0 { + return + } + for familyIndex := range runtimes { + for modeIndex := range runtimes[familyIndex].Modes { + mode := &runtimes[familyIndex].Modes[modeIndex] + backend := api.Backend(mode.Backend) + if runtimeAllowed(resolved.Constraints.Models, backend) { + continue + } + layer := runtimeRestrictionLayer(resolved, backend) + mode.Disabled = true + if layer != nil { + mode.DisabledReason = fmt.Sprintf("%s layer %s", layer.Scope, layer.Name) + } + mode.Availability = restrictedAvailability(layer, resolved.Constraints.Models) + } + } +} + +func restrictedAvailability(layer *api.SpecLayer, allowed []string) api.Availability { + reason := "Unavailable because the resolved runtime profile restricts the model catalog." + if layer != nil { + reason = fmt.Sprintf("Unavailable because %s layer %q restricts the model catalog.", layer.Scope, layer.Name) + } + return api.Availability{ + State: api.AvailabilityDisabled, + Reason: reason, + Remediation: "Select one of the allowed models: " + strings.Join(allowed, ", ") + ".", + } +} + +func modelRestrictionLayer(resolved api.ResolvedSpec, model api.Model) *api.SpecLayer { + for index := len(resolved.Trace) - 1; index >= 0; index-- { + layer := &resolved.Trace[index] + if len(layer.Constraints.Models) > 0 && !(api.ResolvedSpec{Constraints: layer.Constraints}).AllowsModel(model) { + return layer + } + } + return nil +} + +func runtimeRestrictionLayer(resolved api.ResolvedSpec, backend api.Backend) *api.SpecLayer { + for index := len(resolved.Trace) - 1; index >= 0; index-- { + layer := &resolved.Trace[index] + if len(layer.Constraints.Models) > 0 && !runtimeAllowed(layer.Constraints.Models, backend) { + return layer + } + } + return nil +} + +func runtimeAllowed(models []string, backend api.Backend) bool { + providerPrefix := ai.BackendToProvider(backend) + "/" + for _, selector := range models { + if strings.HasPrefix(selector, providerPrefix) { + return true + } + model, err := (api.Model{Name: selector}).Expand() + if err == nil && model.Backend == backend { + return true + } + } + return false +} + // ProviderConfigRequest carries the canonically resolved model and the runtime // config assembled by the chat service. type ProviderConfigRequest struct { diff --git a/pkg/aichat/runtime_profile_ginkgo_test.go b/pkg/aichat/runtime_profile_ginkgo_test.go new file mode 100644 index 00000000..fecc247d --- /dev/null +++ b/pkg/aichat/runtime_profile_ginkgo_test.go @@ -0,0 +1,80 @@ +package aichat_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Resolved runtime profiles", func() { + It("annotates models outside the effective catalog with actionable provenance", func() { + resolver := &fakeResolver{models: aichat.ModelCatalogResponse{ + {ID: "openai/gpt-5.6-sol", Provider: "openai", Label: "GPT", Runtime: api.Model{Name: "gpt-5.6-sol", Backend: api.BackendOpenAI}, Configured: true, Availability: api.Available()}, + {ID: "anthropic/claude-sonnet-5", Provider: "anthropic", Label: "Claude", Runtime: api.Model{Name: "claude-sonnet-5", Backend: api.BackendAnthropic}, Configured: true, Availability: api.Available()}, + }} + profile := mustRuntimeProfile(api.SpecLayer{ + Name: "claims", Scope: api.SpecLayerContext, + Constraints: api.RuntimeConstraints{Models: []string{"gpt-5.6-sol"}}, + }) + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return profile, nil + }), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/api/chat/models", nil)) + + Expect(response.Code).To(Equal(http.StatusOK)) + var models aichat.ModelCatalogResponse + Expect(json.Unmarshal(response.Body.Bytes(), &models)).To(Succeed()) + Expect(models[0].Availability).To(Equal(api.Available())) + Expect(models[1].Configured).To(BeFalse()) + Expect(models[1].Availability.State).To(Equal(api.AvailabilityDisabled)) + Expect(models[1].Availability.Reason).To(ContainSubstring(`context layer "claims"`)) + Expect(models[1].Availability.Remediation).To(ContainSubstring("allowed models")) + }) + + It("checks every named quota and reports the independently exhausted allowance", func() { + resolver := &fakeResolver{provider: &fakeStreamingProvider{}} + profile := mustRuntimeProfile( + api.SpecLayer{ + Name: "platform", Scope: api.SpecLayerGlobal, + Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}, + Constraints: api.RuntimeConstraints{Quotas: []api.RuntimeQuota{{Name: "platform-monthly", TokenLimit: 100, TokensUsed: 10}}}, + }, + api.SpecLayer{ + Name: "claims", Scope: api.SpecLayerContext, + Constraints: api.RuntimeConstraints{Quotas: []api.RuntimeQuota{{Name: "claims-monthly", CostLimitUSD: 20, CostUsedUSD: 20}}}, + }, + ) + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return profile, nil + }), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "hello"}}}}, + })) + + Expect(response.Code).To(Equal(http.StatusPaymentRequired)) + Expect(response.Body.String()).To(ContainSubstring(`context quota "claims-monthly" from layer "claims"`)) + Expect(resolver.configs).To(BeEmpty()) + }) +}) + +func mustRuntimeProfile(layers ...api.SpecLayer) aichat.RuntimeProfile { + resolved, err := api.ResolveSpecLayers(layers...) + Expect(err).NotTo(HaveOccurred()) + return aichat.RuntimeProfile{Resolved: resolved} +} diff --git a/pkg/aichat/runtime_settings.go b/pkg/aichat/runtime_settings.go index 0e9fff1f..5229f2b5 100644 --- a/pkg/aichat/runtime_settings.go +++ b/pkg/aichat/runtime_settings.go @@ -4,6 +4,8 @@ import ( "encoding/json" "fmt" "net/http" + + "github.com/flanksource/captain/pkg/api" ) type requestError struct { @@ -20,20 +22,23 @@ func requestErrorStatus(err error) int { return http.StatusBadRequest } -func enforceRuntimeSettings(request ChatRequest, settings RuntimeSettings) error { - if settings.MonthlyBudgetUSD > 0 && settings.CurrentMonthCostUSD >= settings.MonthlyBudgetUSD { - return requestError{status: http.StatusPaymentRequired, text: fmt.Sprintf( - "chat monthly cost budget exhausted: $%.4f used of $%.4f", - settings.CurrentMonthCostUSD, settings.MonthlyBudgetUSD, - )} - } - if settings.MonthlyTokenBudget > 0 && settings.CurrentMonthTokens >= settings.MonthlyTokenBudget { - return requestError{status: http.StatusPaymentRequired, text: fmt.Sprintf( - "chat monthly token budget exhausted: %d used of %d", - settings.CurrentMonthTokens, settings.MonthlyTokenBudget, - )} +func enforceRuntimeProfile(request ChatRequest, resolved api.ResolvedSpec) error { + for _, quota := range resolved.Constraints.Quotas { + if quota.CostLimitUSD > 0 && quota.CostUsedUSD >= quota.CostLimitUSD { + return requestError{status: http.StatusPaymentRequired, text: fmt.Sprintf( + "chat %s quota %q from layer %q exhausted: $%.4f used of $%.4f", + quota.Scope, quota.Name, quota.Layer, quota.CostUsedUSD, quota.CostLimitUSD, + )} + } + if quota.TokenLimit > 0 && quota.TokensUsed >= quota.TokenLimit { + return requestError{status: http.StatusPaymentRequired, text: fmt.Sprintf( + "chat %s quota %q from layer %q exhausted: %d tokens used of %d", + quota.Scope, quota.Name, quota.Layer, quota.TokensUsed, quota.TokenLimit, + )} + } } - if settings.MaxInputTokens <= 0 { + maxInputTokens := resolved.Constraints.Limits.MaxInputTokens + if maxInputTokens <= 0 { return nil } raw, err := json.Marshal(struct { @@ -45,10 +50,10 @@ func enforceRuntimeSettings(request ChatRequest, settings RuntimeSettings) error return fmt.Errorf("estimate chat input tokens: %w", err) } estimated := (len(raw) + 3) / 4 - if estimated > settings.MaxInputTokens { + if estimated > maxInputTokens { return requestError{status: http.StatusRequestEntityTooLarge, text: fmt.Sprintf( "chat input is about %d tokens, exceeding the configured per-turn limit of %d", - estimated, settings.MaxInputTokens, + estimated, maxInputTokens, )} } return nil diff --git a/pkg/aichat/service.go b/pkg/aichat/service.go index d8ec5944..118342e8 100644 --- a/pkg/aichat/service.go +++ b/pkg/aichat/service.go @@ -17,34 +17,23 @@ import ( var serviceLog = logger.GetLogger("aichat") -// RuntimeSettings are application-owned defaults and provider construction -// settings evaluated for each request. -// RuntimeSettings is the request-scoped application configuration for a chat. -// -// The default model lives in Spec.Model — there is deliberately no DefaultModel -// string beside it. A bare name next to a structured Spec is the lossy pattern: -// it cannot carry a backend/mode/effort, so whatever it named got re-inferred, and -// when both were set they could silently disagree. Spec.Model can say -// {Name: "sol", Mode: ModeAgent} and mean it. -type RuntimeSettings struct { - System string - Spec api.Spec - ProviderConfig api.Config - MaxInputTokens int - MonthlyTokenBudget int - CurrentMonthTokens int - MonthlyBudgetUSD float64 - CurrentMonthCostUSD float64 +// RuntimeProfile is the request-scoped, hierarchically resolved application +// configuration for a chat. Resolved carries the effective Spec, constraints, +// and ordered provenance; provider credentials remain runtime-only. +type RuntimeProfile struct { + System string + Resolved api.ResolvedSpec + ProviderConfig api.Config } -// RuntimeSettingsProvider supplies request-scoped application settings. -type RuntimeSettingsProvider interface { - RuntimeSettings(context.Context) (RuntimeSettings, error) +// RuntimeProfileProvider supplies request-scoped application profiles. +type RuntimeProfileProvider interface { + RuntimeProfile(context.Context) (RuntimeProfile, error) } -type RuntimeSettingsProviderFunc func(context.Context) (RuntimeSettings, error) +type RuntimeProfileProviderFunc func(context.Context) (RuntimeProfile, error) -func (f RuntimeSettingsProviderFunc) RuntimeSettings(ctx context.Context) (RuntimeSettings, error) { +func (f RuntimeProfileProviderFunc) RuntimeProfile(ctx context.Context) (RuntimeProfile, error) { return f(ctx) } @@ -71,7 +60,7 @@ func FixedThreadStore(store ThreadStore) ThreadStoreProvider { type ServiceOptions struct { Resolver Resolver ProviderConfig ProviderConfigSource - Settings RuntimeSettingsProvider + Profile RuntimeProfileProvider Tools ToolProvider MCP ToolProvider Attachments AttachmentResolver @@ -106,6 +95,11 @@ func (s *Service) Handler() http.Handler { } func (s *Service) handleRuntimes(w http.ResponseWriter, request *http.Request) { + profile, err := s.runtimeProfile(request.Context()) + if err != nil { + http.Error(w, fmt.Sprintf("load chat runtime profile: %v", err), http.StatusInternalServerError) + return + } runtimes, err := s.resolver.Runtimes(request.Context()) if err != nil { http.Error(w, err.Error(), http.StatusServiceUnavailable) @@ -115,12 +109,18 @@ func (s *Service) handleRuntimes(w http.ResponseWriter, request *http.Request) { http.Error(w, err.Error(), http.StatusServiceUnavailable) return } + annotateProfileRuntimes(profile.Resolved, runtimes) if err := writeJSON(w, http.StatusOK, runtimes); err != nil { serviceLog.Errorf("write chat runtimes response: %v", err) } } func (s *Service) handleModels(w http.ResponseWriter, request *http.Request) { + profile, err := s.runtimeProfile(request.Context()) + if err != nil { + http.Error(w, fmt.Sprintf("load chat runtime profile: %v", err), http.StatusInternalServerError) + return + } models, err := s.resolver.Models(request.Context()) if err != nil { http.Error(w, err.Error(), http.StatusServiceUnavailable) @@ -130,12 +130,17 @@ func (s *Service) handleModels(w http.ResponseWriter, request *http.Request) { http.Error(w, err.Error(), http.StatusServiceUnavailable) return } + annotateProfileModels(profile.Resolved, models) if err := writeJSON(w, http.StatusOK, models); err != nil { serviceLog.Errorf("write chat models response: %v", err) } } func (s *Service) handleTools(w http.ResponseWriter, request *http.Request) { + if _, err := s.runtimeProfile(request.Context()); err != nil { + http.Error(w, fmt.Sprintf("load chat runtime profile: %v", err), http.StatusInternalServerError) + return + } set, err := s.loadTools(request.Context()) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) @@ -157,12 +162,12 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { http.Error(w, err.Error(), http.StatusBadRequest) return } - settings, err := s.runtimeSettings(request.Context()) + profile, err := s.runtimeProfile(request.Context()) if err != nil { - http.Error(w, fmt.Sprintf("load chat runtime settings: %v", err), http.StatusInternalServerError) + http.Error(w, fmt.Sprintf("load chat runtime profile: %v", err), http.StatusInternalServerError) return } - if err := enforceRuntimeSettings(chat, settings); err != nil { + if err := enforceRuntimeProfile(chat, profile.Resolved); err != nil { http.Error(w, err.Error(), requestErrorStatus(err)) return } @@ -183,11 +188,12 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { return } } - spec, err := requestSpec(chat, settings, attachments) + resolved, err := requestSpec(chat, profile, attachments) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } + spec := resolved.Spec set, err := s.loadTools(request.Context()) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) @@ -205,7 +211,7 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { definitions = append(definitions, s.sessionTitleTool(chat.ThreadID)) appendSessionTitleInstruction(&spec) } - config := settings.ProviderConfig + config := profile.ProviderConfig config.Model = spec.Model config.Budget = spec.Budget config.SessionID = spec.SessionID @@ -217,6 +223,7 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { return } spec.Model = config.Model + resolved.Spec = spec var execution Execution var callerToolEvents <-chan api.Event if s.options.Authority != nil && chat.ThreadID != "" { @@ -226,7 +233,7 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { } execution, err = s.options.Authority.Begin(request.Context(), ExecutionRequest{ ThreadID: chat.ThreadID, RequestID: turnID, Title: title, - Spec: spec, Definitions: definitions, + Spec: spec, Profile: resolved, Definitions: definitions, }) if err != nil { http.Error(w, fmt.Sprintf("admit chat execution: %v", err), http.StatusInternalServerError) @@ -365,11 +372,11 @@ func validateThreadTurn(request ChatRequest, thread *Thread) error { return nil } -func (s *Service) runtimeSettings(ctx context.Context) (RuntimeSettings, error) { - if s.options.Settings == nil { - return RuntimeSettings{}, nil +func (s *Service) runtimeProfile(ctx context.Context) (RuntimeProfile, error) { + if s.options.Profile == nil { + return RuntimeProfile{}, nil } - return s.options.Settings.RuntimeSettings(ctx) + return s.options.Profile.RuntimeProfile(ctx) } func (s *Service) resolveThreadSession(ctx context.Context, request *ChatRequest) (*Thread, error) { diff --git a/pkg/aichat/service_ginkgo_test.go b/pkg/aichat/service_ginkgo_test.go index 2af07692..a7c00ae2 100644 --- a/pkg/aichat/service_ginkgo_test.go +++ b/pkg/aichat/service_ginkgo_test.go @@ -153,8 +153,8 @@ var _ = Describe("Captain aichat service", func() { } service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, ProviderConfig: source, - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{Spec: api.Spec{Model: api.Model{Name: "api:gpt-5.4"}}}, nil + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return mustRuntimeProfile(api.SpecLayer{Name: "application", Scope: api.SpecLayerGlobal, Spec: api.Spec{Model: api.Model{Name: "api:gpt-5.4"}}}), nil }), }) @@ -180,8 +180,8 @@ var _ = Describe("Captain aichat service", func() { }} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, ProviderConfig: source, - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{Spec: api.Spec{Model: api.Model{Name: "api:gpt-5.4"}}}, nil + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return mustRuntimeProfile(api.SpecLayer{Name: "application", Scope: api.SpecLayerGlobal, Spec: api.Spec{Model: api.Model{Name: "api:gpt-5.4"}}}), nil }), }) @@ -202,11 +202,11 @@ var _ = Describe("Captain aichat service", func() { resolver := &fakeResolver{provider: provider} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{Spec: api.Spec{ + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return mustRuntimeProfile(api.SpecLayer{Name: "application", Scope: api.SpecLayerGlobal, Spec: api.Spec{ Model: api.Model{Name: "openai/test-model"}, Budget: api.Budget{Cost: 5, MaxTokens: 2_000, MaxTurns: 3}, - }}, nil + }}), nil }), }) @@ -224,11 +224,14 @@ var _ = Describe("Captain aichat service", func() { resolver := &fakeResolver{provider: &fakeStreamingProvider{}} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{ - Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}, - MonthlyBudgetUSD: 10, CurrentMonthCostUSD: 10, - }, nil + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return mustRuntimeProfile(api.SpecLayer{ + Name: "application", Scope: api.SpecLayerGlobal, + Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}, + Constraints: api.RuntimeConstraints{Quotas: []api.RuntimeQuota{{ + Name: "application-monthly", CostLimitUSD: 10, CostUsedUSD: 10, + }}}, + }), nil }), }) @@ -288,12 +291,11 @@ var _ = Describe("Captain aichat service", func() { resolver := &fakeResolver{provider: provider} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{ - Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}, - System: "Use application tools.", - ProviderConfig: api.Config{APIURL: "https://example.com/ai", ProjectName: "tenant-x"}, - }, nil + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + profile := mustRuntimeProfile(api.SpecLayer{Name: "application", Scope: api.SpecLayerGlobal, Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}}) + profile.System = "Use application tools." + profile.ProviderConfig = api.Config{APIURL: "https://example.com/ai", ProjectName: "tenant-x"} + return profile, nil }), Attachments: fakeAttachmentResolver{}, Tools: aichat.ToolProviderFunc(func(context.Context) (aichat.ToolSet, error) { @@ -346,8 +348,8 @@ var _ = Describe("Captain aichat service", func() { resolver := &fakeResolver{provider: provider} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, Threads: aichat.FixedThreadStore(store), - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{System: "Use accounting tools."}, nil + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + return aichat.RuntimeProfile{System: "Use accounting tools."}, nil }), }) diff --git a/pkg/api/spec_layers.go b/pkg/api/spec_layers.go new file mode 100644 index 00000000..3320ffce --- /dev/null +++ b/pkg/api/spec_layers.go @@ -0,0 +1,299 @@ +package api + +import ( + "fmt" + "slices" + "strings" + "time" +) + +// SpecLayerScope identifies one deterministic level in a resolved runtime profile. +type SpecLayerScope string + +const ( + SpecLayerGlobal SpecLayerScope = "global" + SpecLayerContext SpecLayerScope = "context" + SpecLayerSurface SpecLayerScope = "surface" + SpecLayerUser SpecLayerScope = "user" +) + +// RuntimeLimits are ceilings applied after structural Spec defaults are layered. +type RuntimeLimits struct { + MaxInputTokens int `json:"maxInputTokens,omitempty" yaml:"maxInputTokens,omitempty"` + Budget Budget `json:"budget,omitempty" yaml:"budget,omitempty"` +} + +// RuntimeQuota is one independently enforced usage allowance. +type RuntimeQuota struct { + Name string `json:"name" yaml:"name"` + Scope SpecLayerScope `json:"scope" yaml:"scope"` + Layer string `json:"layer" yaml:"layer"` + TokenLimit int `json:"tokenLimit,omitempty" yaml:"tokenLimit,omitempty"` + TokensUsed int `json:"tokensUsed,omitempty" yaml:"tokensUsed,omitempty"` + CostLimitUSD float64 `json:"costLimitUsd,omitempty" yaml:"costLimitUsd,omitempty"` + CostUsedUSD float64 `json:"costUsedUsd,omitempty" yaml:"costUsedUsd,omitempty"` +} + +// RuntimeConstraints restrict values a later Spec layer may select. +type RuntimeConstraints struct { + Models []string `json:"models,omitempty" yaml:"models,omitempty"` + Limits RuntimeLimits `json:"limits,omitempty" yaml:"limits,omitempty"` + Quotas []RuntimeQuota `json:"quotas,omitempty" yaml:"quotas,omitempty"` +} + +// SpecLayer is one named source of runtime defaults and constraints. +type SpecLayer struct { + Name string `json:"name" yaml:"name"` + Scope SpecLayerScope `json:"scope" yaml:"scope"` + Spec Spec `json:"spec,omitempty" yaml:"spec,omitempty"` + Constraints RuntimeConstraints `json:"constraints,omitempty" yaml:"constraints,omitempty"` +} + +// ResolvedSpec is Captain's effective runtime profile plus ordered provenance. +type ResolvedSpec struct { + Spec Spec `json:"spec" yaml:"spec"` + Constraints RuntimeConstraints `json:"constraints" yaml:"constraints"` + Trace []SpecLayer `json:"trace" yaml:"trace"` +} + +// PromptSpecLayer adapts parsed .prompt frontmatter into the normal surface layer. +func PromptSpecLayer(name string, spec Spec) SpecLayer { + return SpecLayer{Name: name, Scope: SpecLayerSurface, Spec: spec} +} + +// ResolveSpecLayers deterministically overlays defaults and intersects constraints. +func ResolveSpecLayers(input ...SpecLayer) (ResolvedSpec, error) { + layers := append([]SpecLayer(nil), input...) + slices.SortStableFunc(layers, func(left, right SpecLayer) int { + return scopeRank(left.Scope) - scopeRank(right.Scope) + }) + + resolved := ResolvedSpec{Trace: make([]SpecLayer, 0, len(layers))} + for _, layer := range layers { + if err := validateSpecLayer(layer); err != nil { + return ResolvedSpec{}, err + } + resolved.Spec = resolved.Spec.Merge(layer.Spec) + if len(layer.Constraints.Models) > 0 { + resolved.Constraints.Models = intersectModels(resolved.Constraints.Models, layer.Constraints.Models) + if len(resolved.Constraints.Models) == 0 { + return ResolvedSpec{}, fmt.Errorf("spec layer %q leaves the effective model catalog empty", layer.Name) + } + } + limits, err := strictRuntimeLimits(resolved.Constraints.Limits, layer.Constraints.Limits) + if err != nil { + return ResolvedSpec{}, fmt.Errorf("spec layer %q limits: %w", layer.Name, err) + } + resolved.Constraints.Limits = limits + for _, quota := range layer.Constraints.Quotas { + quota.Name = strings.TrimSpace(quota.Name) + quota.Scope = layer.Scope + quota.Layer = layer.Name + resolved.Constraints.Quotas = append(resolved.Constraints.Quotas, quota) + } + resolved.Trace = append(resolved.Trace, cloneSpecLayer(layer)) + } + + budget, err := strictBudget(resolved.Spec.Budget, resolved.Constraints.Limits.Budget) + if err != nil { + return ResolvedSpec{}, fmt.Errorf("effective run budget: %w", err) + } + resolved.Spec.Budget = budget + if err := validateResolvedModels(resolved); err != nil { + return ResolvedSpec{}, err + } + return resolved, nil +} + +// AllowsModel reports whether a model belongs to the effective restrictive catalog. +func (resolved ResolvedSpec) AllowsModel(model Model) bool { + if len(resolved.Constraints.Models) == 0 { + return true + } + for _, allowed := range resolved.Constraints.Models { + if modelSelectorMatches(allowed, model) { + return true + } + } + return false +} + +func validateSpecLayer(layer SpecLayer) error { + if strings.TrimSpace(layer.Name) == "" { + return fmt.Errorf("spec layer name is required") + } + if scopeRank(layer.Scope) < 0 { + return fmt.Errorf("spec layer %q has invalid scope %q", layer.Name, layer.Scope) + } + if _, err := strictRuntimeLimits(RuntimeLimits{}, layer.Constraints.Limits); err != nil { + return fmt.Errorf("spec layer %q limits: %w", layer.Name, err) + } + seenModels := map[string]bool{} + for _, model := range layer.Constraints.Models { + model = strings.TrimSpace(model) + if model == "" { + return fmt.Errorf("spec layer %q model catalog contains an empty selector", layer.Name) + } + if seenModels[model] { + return fmt.Errorf("spec layer %q model catalog repeats %q", layer.Name, model) + } + seenModels[model] = true + } + seenQuotas := map[string]bool{} + for _, quota := range layer.Constraints.Quotas { + name := strings.TrimSpace(quota.Name) + if name == "" { + return fmt.Errorf("spec layer %q quota name is required", layer.Name) + } + if layer.Scope != SpecLayerGlobal && layer.Scope != SpecLayerContext { + return fmt.Errorf("spec layer %q quota %q requires global or context scope", layer.Name, name) + } + if seenQuotas[name] { + return fmt.Errorf("spec layer %q repeats quota %q", layer.Name, name) + } + seenQuotas[name] = true + if quota.TokenLimit < 0 || quota.TokensUsed < 0 || quota.CostLimitUSD < 0 || quota.CostUsedUSD < 0 { + return fmt.Errorf("spec layer %q quota %q cannot contain negative usage or limits", layer.Name, name) + } + } + return nil +} + +func scopeRank(scope SpecLayerScope) int { + switch scope { + case SpecLayerGlobal: + return 0 + case SpecLayerContext: + return 1 + case SpecLayerSurface: + return 2 + case SpecLayerUser: + return 3 + default: + return -1 + } +} + +func intersectModels(current, restrictive []string) []string { + if len(current) == 0 { + return append([]string(nil), restrictive...) + } + allowed := make(map[string]bool, len(restrictive)) + for _, model := range restrictive { + allowed[strings.TrimSpace(model)] = true + } + out := make([]string, 0, len(current)) + for _, model := range current { + if allowed[model] { + out = append(out, model) + } + } + return out +} + +func strictRuntimeLimits(current, next RuntimeLimits) (RuntimeLimits, error) { + if current.MaxInputTokens < 0 || next.MaxInputTokens < 0 { + return RuntimeLimits{}, fmt.Errorf("maxInputTokens must be non-negative") + } + budget, err := strictBudget(current.Budget, next.Budget) + if err != nil { + return RuntimeLimits{}, err + } + return RuntimeLimits{ + MaxInputTokens: strictPositiveInt(current.MaxInputTokens, next.MaxInputTokens), + Budget: budget, + }, nil +} + +func strictBudget(current, next Budget) (Budget, error) { + if err := current.Validate(); err != nil { + return Budget{}, err + } + if err := next.Validate(); err != nil { + return Budget{}, err + } + timeout, err := strictTimeout(current.Timeout, next.Timeout) + if err != nil { + return Budget{}, err + } + return Budget{ + Cost: strictPositiveFloat(current.Cost, next.Cost), + MaxTokens: strictPositiveInt(current.MaxTokens, next.MaxTokens), + MaxTurns: strictPositiveInt(current.MaxTurns, next.MaxTurns), + Timeout: timeout, + }, nil +} + +func strictPositiveInt(left, right int) int { + if left == 0 || right > 0 && right < left { + return right + } + return left +} + +func strictPositiveFloat(left, right float64) float64 { + if left == 0 || right > 0 && right < left { + return right + } + return left +} + +func strictTimeout(left, right string) (string, error) { + leftDuration, err := parseOptionalDuration(left) + if err != nil { + return "", err + } + rightDuration, err := parseOptionalDuration(right) + if err != nil { + return "", err + } + if leftDuration == 0 || rightDuration > 0 && rightDuration < leftDuration { + return strings.TrimSpace(right), nil + } + return strings.TrimSpace(left), nil +} + +func parseOptionalDuration(value string) (time.Duration, error) { + value = strings.TrimSpace(value) + if value == "" { + return 0, nil + } + duration, err := time.ParseDuration(value) + if err != nil || duration <= 0 { + return 0, fmt.Errorf("invalid positive timeout %q", value) + } + return duration, nil +} + +func validateResolvedModels(resolved ResolvedSpec) error { + model := resolved.Spec.Model + if strings.TrimSpace(model.Name) == "" { + return nil + } + if !resolved.AllowsModel(model) { + return fmt.Errorf("selected model %q is outside the effective model catalog", model.Name) + } + for _, fallback := range model.Fallbacks { + if !resolved.AllowsModel(fallback) { + return fmt.Errorf("fallback model %q is outside the effective model catalog", fallback.Name) + } + } + return nil +} + +func modelSelectorMatches(selector string, model Model) bool { + selector = strings.TrimSpace(selector) + if selector == model.Name || selector == model.ID { + return true + } + allowed, allowedErr := (Model{Name: selector}).Expand() + actual, actualErr := model.Expand() + return allowedErr == nil && actualErr == nil && allowed.Name == actual.Name && allowed.Backend == actual.Backend +} + +func cloneSpecLayer(layer SpecLayer) SpecLayer { + layer.Spec = Spec{}.Merge(layer.Spec) + layer.Constraints.Models = append([]string(nil), layer.Constraints.Models...) + layer.Constraints.Quotas = append([]RuntimeQuota(nil), layer.Constraints.Quotas...) + return layer +} diff --git a/pkg/api/spec_layers_ginkgo_test.go b/pkg/api/spec_layers_ginkgo_test.go new file mode 100644 index 00000000..93f537ad --- /dev/null +++ b/pkg/api/spec_layers_ginkgo_test.go @@ -0,0 +1,100 @@ +package api + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Hierarchical spec profiles", func() { + It("orders layers by scope and structurally overlays defaults without mutating inputs", func() { + global := SpecLayer{ + Name: "platform", Scope: SpecLayerGlobal, + Spec: Spec{Model: Model{Name: "claude-sonnet-5"}, Budget: Budget{MaxTokens: 8000}}, + } + context := SpecLayer{ + Name: "claims", Scope: SpecLayerContext, + Spec: Spec{Model: Model{Effort: EffortHigh}, Budget: Budget{MaxTurns: 6}}, + } + surface := PromptSpecLayer("triage.prompt", Spec{Prompt: Prompt{System: "Triage claims."}}) + user := SpecLayer{ + Name: "request", Scope: SpecLayerUser, + Spec: Spec{Model: Model{Effort: EffortLow}}, + } + + resolved, err := ResolveSpecLayers(user, surface, context, global) + + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Spec.Model.Name).To(Equal("claude-sonnet-5")) + Expect(resolved.Spec.Model.Effort).To(Equal(EffortLow)) + Expect(resolved.Spec.Budget).To(Equal(Budget{MaxTokens: 8000, MaxTurns: 6})) + Expect(resolved.Spec.Prompt.System).To(Equal("Triage claims.")) + Expect(resolved.Trace).To(HaveLen(4)) + Expect([]SpecLayerScope{ + resolved.Trace[0].Scope, resolved.Trace[1].Scope, resolved.Trace[2].Scope, resolved.Trace[3].Scope, + }).To(Equal([]SpecLayerScope{SpecLayerGlobal, SpecLayerContext, SpecLayerSurface, SpecLayerUser})) + Expect(context.Spec.Model.Name).To(BeEmpty()) + Expect(global.Spec.Model.Effort).To(BeEmpty()) + }) + + It("intersects restrictive model catalogs and validates every fallback", func() { + layers := []SpecLayer{ + { + Name: "platform", Scope: SpecLayerGlobal, + Constraints: RuntimeConstraints{Models: []string{"claude-sonnet-5", "gpt-5.6-sol", "gpt-5.4"}}, + }, + { + Name: "claims", Scope: SpecLayerContext, + Constraints: RuntimeConstraints{Models: []string{"gpt-5.6-sol", "claude-sonnet-5"}}, + }, + { + Name: "request", Scope: SpecLayerUser, + Spec: Spec{Model: Model{Name: "gpt-5.6-sol", Fallbacks: []Model{{Name: "claude-sonnet-5"}}}}, + }, + } + + resolved, err := ResolveSpecLayers(layers...) + + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Constraints.Models).To(Equal([]string{"claude-sonnet-5", "gpt-5.6-sol"})) + Expect(resolved.AllowsModel(Model{Name: "gpt-5.4"})).To(BeFalse()) + Expect(resolved.AllowsModel(Model{Name: "gpt-5.6-sol"})).To(BeTrue()) + + layers[2].Spec.Model.Fallbacks = []Model{{Name: "gpt-5.4"}} + _, err = ResolveSpecLayers(layers...) + Expect(err).To(MatchError(ContainSubstring(`fallback model "gpt-5.4" is outside the effective model catalog`))) + }) + + It("uses strict non-zero run ceilings and retains each named quota independently", func() { + resolved, err := ResolveSpecLayers( + SpecLayer{ + Name: "platform", Scope: SpecLayerGlobal, + Spec: Spec{Budget: Budget{Cost: 12, MaxTokens: 9000, MaxTurns: 10, Timeout: "10m"}}, + Constraints: RuntimeConstraints{ + Limits: RuntimeLimits{MaxInputTokens: 12000, Budget: Budget{Cost: 8, MaxTokens: 7000, MaxTurns: 8, Timeout: "8m"}}, + Quotas: []RuntimeQuota{{Name: "platform-monthly", TokenLimit: 1_000_000, TokensUsed: 10}}, + }, + }, + SpecLayer{ + Name: "claims", Scope: SpecLayerContext, + Spec: Spec{Budget: Budget{Cost: 10, MaxTokens: 8000, MaxTurns: 6, Timeout: "9m"}}, + Constraints: RuntimeConstraints{ + Limits: RuntimeLimits{MaxInputTokens: 4000, Budget: Budget{Cost: 5, MaxTokens: 6000, MaxTurns: 7, Timeout: "5m"}}, + Quotas: []RuntimeQuota{{Name: "claims-monthly", CostLimitUSD: 50, CostUsedUSD: 2}}, + }, + }, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Spec.Budget).To(Equal(Budget{Cost: 5, MaxTokens: 6000, MaxTurns: 6, Timeout: "5m"})) + Expect(resolved.Constraints.Limits.MaxInputTokens).To(Equal(4000)) + Expect(resolved.Constraints.Quotas).To(Equal([]RuntimeQuota{ + {Name: "platform-monthly", Scope: SpecLayerGlobal, Layer: "platform", TokenLimit: 1_000_000, TokensUsed: 10}, + {Name: "claims-monthly", Scope: SpecLayerContext, Layer: "claims", CostLimitUSD: 50, CostUsedUSD: 2}, + })) + duration, err := time.ParseDuration(resolved.Spec.Budget.Timeout) + Expect(err).NotTo(HaveOccurred()) + Expect(duration).To(Equal(5 * time.Minute)) + }) +}) diff --git a/pkg/cli/serve_chat.go b/pkg/cli/serve_chat.go index 84512c9b..72b6252a 100644 --- a/pkg/cli/serve_chat.go +++ b/pkg/cli/serve_chat.go @@ -37,14 +37,21 @@ func newCaptainChatService( return nil, nil, err } chat := aichat.NewService(aichat.ServiceOptions{ - Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { - return aichat.RuntimeSettings{ - System: "You are Captain's coding-agent launcher assistant. Use Captain and Clicky tools when useful, " + - "prefer read-only inspection unless the user explicitly asks for edits, and keep follow-up guidance concise.", + Profile: aichat.RuntimeProfileProviderFunc(func(context.Context) (aichat.RuntimeProfile, error) { + resolved, err := api.ResolveSpecLayers(api.SpecLayer{ + Name: "captain serve", Scope: api.SpecLayerGlobal, Spec: api.Spec{ Model: api.Model{Name: "sol", Mode: registry.ModeAgent}, Setup: &shell.Setup{Cwd: cwd}, }, + }) + if err != nil { + return aichat.RuntimeProfile{}, err + } + return aichat.RuntimeProfile{ + System: "You are Captain's coding-agent launcher assistant. Use Captain and Clicky tools when useful, " + + "prefer read-only inspection unless the user explicitly asks for edits, and keep follow-up guidance concise.", + Resolved: resolved, }, nil }), // Thread reads follow the request's database context; writes never reach From 1998d0970fe43c22fb7422335873e2ff756562ef Mon Sep 17 00:00:00 2001 From: Moshe Immerman Date: Tue, 11 Aug 2026 09:25:33 +0300 Subject: [PATCH 2/2] feat(database): add schema-scoped persistence Claude-Session-Id: 019fef59-46c9-7b83-bc02-6bd4246d90de --- go.mod | 2 +- go.sum | 4 +- migrations/concurrency_integration_test.go | 4 +- migrations/migrations.go | 87 +++++++++++++----- migrations/migrations_test.go | 32 +++---- migrations/schema.go | 42 +++++++++ migrations/schema_ginkgo_test.go | 46 ++++++++++ pkg/database/database.go | 66 +++++++++++--- pkg/database/database_test.go | 16 ++-- pkg/database/model_usage.go | 89 +++++++++++++++++++ pkg/database/options_ginkgo_test.go | 37 +++++++- .../schema_scoped_integration_test.go | 81 +++++++++++++++++ 12 files changed, 439 insertions(+), 67 deletions(-) create mode 100644 migrations/schema.go create mode 100644 migrations/schema_ginkgo_test.go create mode 100644 pkg/database/model_usage.go create mode 100644 pkg/database/schema_scoped_integration_test.go diff --git a/go.mod b/go.mod index 82bff7e2..d083cbeb 100644 --- a/go.mod +++ b/go.mod @@ -40,7 +40,7 @@ require ( ) require ( - github.com/flanksource/commons-db v0.1.26 + github.com/flanksource/commons-db v0.1.27-0.20260806180738-8eb0bca01132 github.com/gliderlabs/ssh v0.3.8 github.com/pelletier/go-toml/v2 v2.4.3 ) diff --git a/go.sum b/go.sum index 74e16fbe..53dd6220 100644 --- a/go.sum +++ b/go.sum @@ -280,8 +280,8 @@ github.com/flanksource/clicky/aichat v1.21.48 h1:f8Kvl96Lfp1qcqPuve1zsjaYN8ZcK1/ github.com/flanksource/clicky/aichat v1.21.48/go.mod h1:PGN/lVAgxpchRctciUCpR4YIuqWoDwRxDh339A6wi3w= github.com/flanksource/commons v1.55.0 h1:gj9zBY3V1qgAAnEiLaeGbkqCmNK0p1tJVQCDurdTZ2k= github.com/flanksource/commons v1.55.0/go.mod h1:gupTCRqGpgD8dd2ooE7bMDJxkfcVKvkPVuHg3cWkW+Q= -github.com/flanksource/commons-db v0.1.26 h1:NXAP0WvMs4ufyDfl1L2ryRxBv5qxV67GiI1nINd4YIw= -github.com/flanksource/commons-db v0.1.26/go.mod h1:i378WIxy8g9xOeLBvhh1y3FO99oCumUqNmfhuDF79kM= +github.com/flanksource/commons-db v0.1.27-0.20260806180738-8eb0bca01132 h1:dbBt8TmT2tEE9y3zEUJh6bGbB7gH5JisP4BbhXoYvmU= +github.com/flanksource/commons-db v0.1.27-0.20260806180738-8eb0bca01132/go.mod h1:i378WIxy8g9xOeLBvhh1y3FO99oCumUqNmfhuDF79kM= github.com/flanksource/gomplate/v3 v3.24.84 h1:UOE0yCJsczTIKRaHUvhD6tjCYrbNvOugAizuy0FVlhE= github.com/flanksource/gomplate/v3 v3.24.84/go.mod h1:NMMZkFsjbLy/8iY8Fip5N86Y0PP6lZeq+kmPwpVVIL0= github.com/flanksource/is-healthy v1.0.88 h1:ATQuKoNdp8Qfzf41/eMFajmT0qzOmZlZNG5eLK41RFo= diff --git a/migrations/concurrency_integration_test.go b/migrations/concurrency_integration_test.go index 354c5fa4..2b3881b8 100644 --- a/migrations/concurrency_integration_test.go +++ b/migrations/concurrency_integration_test.go @@ -17,7 +17,7 @@ func TestConcurrentApplySerializesCaptainMigrations(t *testing.T) { // Hold the same session lock before releasing a group of Apply calls. This // proves every caller enters through the advisory-lock boundary rather than // racing the Atlas inspect/diff/apply window. - blocker, err := acquireMigrationLock(t.Context(), dsn) + blocker, err := acquireMigrationLock(t.Context(), applyRequest{Connection: dsn, Schema: DefaultSchema}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, blocker.Close()) }) @@ -63,7 +63,7 @@ func TestConcurrentApplySerializesCaptainMigrations(t *testing.T) { // Bound the reacquisition to catch a leaked dedicated connection cleanly. reacquireCtx, cancel := context.WithTimeout(t.Context(), 2*time.Second) defer cancel() - reacquired, err := acquireMigrationLock(reacquireCtx, dsn) + reacquired, err := acquireMigrationLock(reacquireCtx, applyRequest{Connection: dsn, Schema: DefaultSchema}) require.NoError(t, err) require.NoError(t, reacquired.Close()) } diff --git a/migrations/migrations.go b/migrations/migrations.go index cf01b0f9..884cd220 100644 --- a/migrations/migrations.go +++ b/migrations/migrations.go @@ -5,8 +5,10 @@ package migrations import ( "context" + "crypto/sha256" "database/sql" "embed" + "encoding/binary" "errors" "fmt" "strings" @@ -18,6 +20,7 @@ import ( ) const Scope = "captain" +const DefaultSchema = "public" const ( // captainMigrationLockNamespace and captainMigrationLockKey are stable, @@ -42,16 +45,38 @@ type migrationLockHandle interface { } type applyDependencies struct { - acquireLock func(context.Context, string) (migrationLockHandle, error) - migrate func(context.Context, string) error - verify func(context.Context, string) error + acquireLock func(context.Context, applyRequest) (migrationLockHandle, error) + migrate func(context.Context, applyRequest) error + verify func(context.Context, applyRequest) error +} + +type applyRequest struct { + Connection string + Schema string +} + +type options struct { + schema string +} + +// Option configures Captain's migration bundle. +type Option func(*options) + +// WithSchema selects the schema that owns Captain's migration bundle. +func WithSchema(name string) Option { + return func(options *options) { options.schema = name } } var defaultApplyDependencies = applyDependencies{ acquireLock: acquireMigrationLock, - migrate: func(ctx context.Context, connection string) error { - return commonsmigrate.Apply(ctx, connection, schemaFS, + migrate: func(ctx context.Context, request applyRequest) error { + filesystem, err := schemaFilesystem(request.Schema) + if err != nil { + return err + } + return commonsmigrate.Apply(ctx, request.Connection, filesystem, commonsmigrate.WithName(Scope), + commonsmigrate.WithSchema(request.Schema), commonsmigrate.WithExclude("todo_*"), ) }, @@ -63,16 +88,25 @@ var defaultApplyDependencies = applyDependencies{ // migration bundle across processes. It is safe to call repeatedly and uses a // stable scope so Captain can share a database with other independently // migrated applications. -func Apply(ctx context.Context, connection string) error { - return apply(ctx, connection, defaultApplyDependencies) +func Apply(ctx context.Context, connection string, optionFns ...Option) error { + config := options{schema: DefaultSchema} + for _, option := range optionFns { + if option != nil { + option(&config) + } + } + return apply(ctx, applyRequest{Connection: connection, Schema: config.schema}, defaultApplyDependencies) } -func apply(ctx context.Context, connection string, deps applyDependencies) (resultErr error) { - if strings.TrimSpace(connection) == "" { +func apply(ctx context.Context, request applyRequest, deps applyDependencies) (resultErr error) { + if strings.TrimSpace(request.Connection) == "" { return errors.New("captain migration connection string is empty") } + if err := commonsmigrate.ValidateSchemaName(request.Schema); err != nil { + return fmt.Errorf("captain migration schema: %w", err) + } - lock, err := deps.acquireLock(ctx, connection) + lock, err := deps.acquireLock(ctx, request) if err != nil { return fmt.Errorf("acquire Captain migration lock: %w", err) } @@ -82,16 +116,20 @@ func apply(ctx context.Context, connection string, deps applyDependencies) (resu } }() - if err := deps.migrate(ctx, connection); err != nil { + if err := deps.migrate(ctx, request); err != nil { return fmt.Errorf("migrate Captain database: %w", err) } - if err := deps.verify(ctx, connection); err != nil { + if err := deps.verify(ctx, request); err != nil { return fmt.Errorf("verify Captain database: %w", err) } return nil } -func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultErr error) { +func verifyToolApprovalIdentity(ctx context.Context, request applyRequest) (resultErr error) { + connection, err := commonsmigrate.ConnectionForSchema(request.Connection, request.Schema) + if err != nil { + return fmt.Errorf("scope schema verification database: %w", err) + } db, err := commonsdb.NewDB(connection) if err != nil { return fmt.Errorf("open schema verification database: %w", err) @@ -109,11 +147,11 @@ func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultE FROM pg_constraint c JOIN pg_class relation ON relation.oid = c.conrelid JOIN pg_namespace namespace ON namespace.oid = relation.relnamespace - WHERE namespace.nspname = 'public' + WHERE namespace.nspname = $1 AND relation.relname = 'captain_turn_requests' AND c.conname = 'captain_turn_requests_tool_approval_identity' AND c.contype = 'c' - `).Scan(&validated, &definition) + `, request.Schema).Scan(&validated, &definition) if errors.Is(err, sql.ErrNoRows) { return errors.New("captain_turn_requests_tool_approval_identity constraint is missing") } @@ -138,13 +176,14 @@ func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultE type migrationLock struct { db *sql.DB conn *sql.Conn + key int32 once sync.Once err error } -func acquireMigrationLock(ctx context.Context, connection string) (migrationLockHandle, error) { - db, err := commonsdb.NewDB(connection) +func acquireMigrationLock(ctx context.Context, request applyRequest) (migrationLockHandle, error) { + db, err := commonsdb.NewDB(request.Connection) if err != nil { return nil, fmt.Errorf("open advisory-lock database: %w", err) } @@ -154,12 +193,12 @@ func acquireMigrationLock(ctx context.Context, connection string) (migrationLock return nil, fmt.Errorf("reserve advisory-lock connection: %w", err) } if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1, $2)`, - captainMigrationLockNamespace, captainMigrationLockKey); err != nil { + captainMigrationLockNamespace, migrationLockKey(request.Schema)); err != nil { _ = conn.Close() _ = db.Close() return nil, fmt.Errorf("lock Captain migration scope: %w", err) } - return &migrationLock{db: db, conn: conn}, nil + return &migrationLock{db: db, conn: conn, key: migrationLockKey(request.Schema)}, nil } func (lock *migrationLock) Close() error { @@ -172,7 +211,7 @@ func (lock *migrationLock) Close() error { ctx, cancel := context.WithTimeout(context.Background(), migrationUnlockTimeout) var unlocked bool if err := lock.conn.QueryRowContext(ctx, `SELECT pg_advisory_unlock($1, $2)`, - captainMigrationLockNamespace, captainMigrationLockKey).Scan(&unlocked); err != nil { + captainMigrationLockNamespace, lock.key).Scan(&unlocked); err != nil { cleanupErrors = append(cleanupErrors, fmt.Errorf("unlock Captain migration scope: %w", err)) } else if !unlocked { cleanupErrors = append(cleanupErrors, errors.New("captain migration advisory lock was not held")) @@ -191,3 +230,11 @@ func (lock *migrationLock) Close() error { }) return lock.err } + +func migrationLockKey(schemaName string) int32 { + if schemaName == DefaultSchema { + return captainMigrationLockKey + } + digest := sha256.Sum256([]byte(schemaName)) + return int32(binary.BigEndian.Uint32(digest[:4])) +} diff --git a/migrations/migrations_test.go b/migrations/migrations_test.go index e241d7e7..350be4a2 100644 --- a/migrations/migrations_test.go +++ b/migrations/migrations_test.go @@ -231,16 +231,16 @@ func TestApplyHoldsMigrationLockAcrossMigration(t *testing.T) { t.Parallel() var events []string - err := apply(t.Context(), "postgres://captain", applyDependencies{ - acquireLock: func(context.Context, string) (migrationLockHandle, error) { + err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{ + acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) { events = append(events, "lock") return &recordingMigrationLock{events: &events}, nil }, - migrate: func(context.Context, string) error { + migrate: func(context.Context, applyRequest) error { events = append(events, "migrate") return nil }, - verify: func(context.Context, string) error { + verify: func(context.Context, applyRequest) error { events = append(events, "verify") return nil }, @@ -256,16 +256,16 @@ func TestApplyReleasesMigrationLockOnVerificationFailure(t *testing.T) { var events []string verificationErr := errors.New("constraint drifted") - err := apply(t.Context(), "postgres://captain", applyDependencies{ - acquireLock: func(context.Context, string) (migrationLockHandle, error) { + err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{ + acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) { events = append(events, "lock") return &recordingMigrationLock{events: &events}, nil }, - migrate: func(context.Context, string) error { + migrate: func(context.Context, applyRequest) error { events = append(events, "migrate") return nil }, - verify: func(context.Context, string) error { + verify: func(context.Context, applyRequest) error { events = append(events, "verify") return verificationErr }, @@ -281,12 +281,12 @@ func TestApplyReleasesMigrationLockOnMigrationFailure(t *testing.T) { var events []string migrationErr := errors.New("atlas failed") - err := apply(t.Context(), "postgres://captain", applyDependencies{ - acquireLock: func(context.Context, string) (migrationLockHandle, error) { + err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{ + acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) { events = append(events, "lock") return &recordingMigrationLock{events: &events}, nil }, - migrate: func(context.Context, string) error { + migrate: func(context.Context, applyRequest) error { events = append(events, "migrate") return migrationErr }, @@ -303,8 +303,8 @@ func TestApplyReportsLockAcquisitionAndReleaseErrors(t *testing.T) { t.Run("acquire", func(t *testing.T) { t.Parallel() wantErr := errors.New("lock unavailable") - err := apply(t.Context(), "postgres://captain", applyDependencies{ - acquireLock: func(context.Context, string) (migrationLockHandle, error) { + err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{ + acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) { return nil, wantErr }, }) @@ -317,11 +317,11 @@ func TestApplyReportsLockAcquisitionAndReleaseErrors(t *testing.T) { t.Parallel() migrationErr := errors.New("migration failed") releaseErr := errors.New("unlock failed") - err := apply(t.Context(), "postgres://captain", applyDependencies{ - acquireLock: func(context.Context, string) (migrationLockHandle, error) { + err := apply(t.Context(), applyRequest{Connection: "postgres://captain", Schema: DefaultSchema}, applyDependencies{ + acquireLock: func(context.Context, applyRequest) (migrationLockHandle, error) { return &recordingMigrationLock{err: releaseErr}, nil }, - migrate: func(context.Context, string) error { return migrationErr }, + migrate: func(context.Context, applyRequest) error { return migrationErr }, }) if !errors.Is(err, migrationErr) || !errors.Is(err, releaseErr) { t.Fatalf("apply error = %v, want joined migration and release errors", err) diff --git a/migrations/schema.go b/migrations/schema.go new file mode 100644 index 00000000..dd864a38 --- /dev/null +++ b/migrations/schema.go @@ -0,0 +1,42 @@ +package migrations + +import ( + "fmt" + "io/fs" + "path" + "strings" + "testing/fstest" + + commonsmigrate "github.com/flanksource/commons-db/migrate" +) + +func schemaFilesystem(schemaName string) (fs.FS, error) { + if err := commonsmigrate.ValidateSchemaName(schemaName); err != nil { + return nil, fmt.Errorf("captain migration schema: %w", err) + } + if schemaName == DefaultSchema { + return schemaFS, nil + } + files := fstest.MapFS{} + err := fs.WalkDir(schemaFS, ".", func(name string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if entry.IsDir() { + return nil + } + content, err := fs.ReadFile(schemaFS, name) + if err != nil { + return fmt.Errorf("read Captain migration %s: %w", name, err) + } + if strings.EqualFold(path.Ext(name), ".sql") { + content = []byte(strings.ReplaceAll(string(content), DefaultSchema+".", schemaName+".")) + } + files[name] = &fstest.MapFile{Data: content} + return nil + }) + if err != nil { + return nil, fmt.Errorf("render Captain migration schema: %w", err) + } + return files, nil +} diff --git a/migrations/schema_ginkgo_test.go b/migrations/schema_ginkgo_test.go new file mode 100644 index 00000000..aeab0f94 --- /dev/null +++ b/migrations/schema_ginkgo_test.go @@ -0,0 +1,46 @@ +package migrations + +import ( + "io/fs" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("schema-scoped Captain migrations", func() { + It("leaves the public bundle unchanged", func() { + filesystem, err := schemaFilesystem(DefaultSchema) + Expect(err).NotTo(HaveOccurred()) + content, err := fs.ReadFile(filesystem, "51_state_triggers.sql") + Expect(err).NotTo(HaveOccurred()) + Expect(string(content)).To(ContainSubstring("public.captain_sessions")) + }) + + It("qualifies SQL objects with the selected schema while retaining portable HCL", func() { + const schemaName = "agent_namespace_context" + filesystem, err := schemaFilesystem(schemaName) + Expect(err).NotTo(HaveOccurred()) + + sqlContent, err := fs.ReadFile(filesystem, "51_state_triggers.sql") + Expect(err).NotTo(HaveOccurred()) + Expect(string(sqlContent)).To(ContainSubstring(schemaName + ".captain_sessions")) + Expect(string(sqlContent)).NotTo(ContainSubstring(DefaultSchema + ".captain_")) + + hclContent, err := fs.ReadFile(filesystem, "10_sessions.pg.hcl") + Expect(err).NotTo(HaveOccurred()) + Expect(string(hclContent)).To(ContainSubstring("schema.public")) + Expect(string(hclContent)).NotTo(ContainSubstring(schemaName)) + }) + + It("rejects invalid schemas", func() { + _, err := schemaFilesystem(strings.Repeat("x", 64)) + Expect(err).To(HaveOccurred()) + }) + + It("uses a stable schema-specific advisory lock", func() { + Expect(migrationLockKey(DefaultSchema)).To(Equal(captainMigrationLockKey)) + Expect(migrationLockKey("agent_namespace_one")).To(Equal(migrationLockKey("agent_namespace_one"))) + Expect(migrationLockKey("agent_namespace_one")).NotTo(Equal(migrationLockKey("agent_namespace_two"))) + }) +}) diff --git a/pkg/database/database.go b/pkg/database/database.go index 213afa72..fccb98d8 100644 --- a/pkg/database/database.go +++ b/pkg/database/database.go @@ -12,6 +12,7 @@ import ( "github.com/flanksource/captain/migrations" commonsdb "github.com/flanksource/commons-db/db" + commonsmigrate "github.com/flanksource/commons-db/migrate" "gorm.io/gorm" ) @@ -23,6 +24,7 @@ type openOptions struct { gorm *gorm.DB gormSet bool migrate bool + schema string maxOpenConns int } @@ -44,6 +46,11 @@ func WithMigrations() Option { return func(options *openOptions) { options.migrate = true } } +// WithSchema selects the validated PostgreSQL schema owned by this Captain handle. +func WithSchema(name string) Option { + return func(options *openOptions) { options.schema = name } +} + // WithMaxOpenConns caps a Captain-owned pool. Processes that hold several // database handles at once use it so the pools do not add up to an unreasonable // number of backends. Ignored for injected pools, which the host sizes. @@ -54,18 +61,21 @@ func WithMaxOpenConns(conns int) Option { // DB is a Captain database handle. It records pool ownership so a host can // safely share its application pool without Captain closing it. type DB struct { - gorm *gorm.DB - owned bool + gorm *gorm.DB + owned bool + schema string } type dependencies struct { - migrate func(context.Context, string) error + migrate func(context.Context, string, string) error open func(string, *gorm.Config) (*gorm.DB, error) } var defaultDependencies = dependencies{ - migrate: migrations.Apply, - open: commonsdb.NewGorm, + migrate: func(ctx context.Context, dsn, schemaName string) error { + return migrations.Apply(ctx, dsn, migrations.WithSchema(schemaName)) + }, + open: commonsdb.NewGorm, } // Open reuses an injected pool or opens a Captain-owned pool. It does not @@ -75,13 +85,16 @@ func Open(ctx context.Context, options ...Option) (*DB, error) { } func open(ctx context.Context, deps dependencies, optionFns ...Option) (*DB, error) { - var options openOptions + options := openOptions{schema: migrations.DefaultSchema} for _, option := range optionFns { if option != nil { option(&options) } } dsn := strings.TrimSpace(options.dsn) + if err := commonsmigrate.ValidateSchemaName(options.schema); err != nil { + return nil, fmt.Errorf("captain database schema: %w", err) + } if options.gormSet && options.gorm == nil { return nil, errors.New("captain database GORM pool is nil") } @@ -91,16 +104,27 @@ func open(ctx context.Context, deps dependencies, optionFns ...Option) (*DB, err if options.migrate && dsn == "" { return nil, errors.New("captain database migrations require a DSN") } + if options.gorm != nil && options.schema != migrations.DefaultSchema { + return nil, fmt.Errorf("captain database cannot select schema %q on a host-owned GORM pool", options.schema) + } if options.migrate { - if err := deps.migrate(ctx, dsn); err != nil { + if err := deps.migrate(ctx, dsn, options.schema); err != nil { return nil, err } } if options.gorm != nil { - return &DB{gorm: options.gorm}, nil + return &DB{gorm: options.gorm, schema: options.schema}, nil } - gormDB, err := deps.open(dsn, commonsdb.DefaultGormConfig()) + scopedDSN := dsn + if options.schema != migrations.DefaultSchema { + var err error + scopedDSN, err = commonsmigrate.ConnectionForSchema(dsn, options.schema) + if err != nil { + return nil, fmt.Errorf("scope Captain database: %w", err) + } + } + gormDB, err := deps.open(scopedDSN, commonsdb.DefaultGormConfig()) if err != nil { return nil, fmt.Errorf("open Captain database: %w", err) } @@ -112,7 +136,7 @@ func open(ctx context.Context, deps dependencies, optionFns ...Option) (*DB, err sqlDB.SetMaxOpenConns(options.maxOpenConns) sqlDB.SetMaxIdleConns(options.maxOpenConns) } - return &DB{gorm: gormDB, owned: true}, nil + return &DB{gorm: gormDB, owned: true, schema: options.schema}, nil } // Use wraps an already migrated shared GORM pool. The returned handle does not @@ -121,13 +145,19 @@ func Use(gormDB *gorm.DB) (*DB, error) { if gormDB == nil { return nil, errors.New("captain database GORM pool is nil") } - return &DB{gorm: gormDB}, nil + return &DB{gorm: gormDB, schema: migrations.DefaultSchema}, nil } // Migrate applies Captain's authoritative HCL and SQL migration bundle without // opening another application pool. -func Migrate(ctx context.Context, dsn string) error { - return migrations.Apply(ctx, dsn) +func Migrate(ctx context.Context, dsn string, options ...Option) error { + config := openOptions{schema: migrations.DefaultSchema} + for _, option := range options { + if option != nil { + option(&config) + } + } + return migrations.Apply(ctx, dsn, migrations.WithSchema(config.schema)) } // Gorm returns the application pool supplied to or opened by Captain. @@ -138,6 +168,14 @@ func (db *DB) Gorm() *gorm.DB { return db.gorm } +// Schema returns the PostgreSQL schema owned by the handle. +func (db *DB) Schema() string { + if db == nil { + return "" + } + return db.schema +} + // Transaction runs fn with a Captain handle backed by the same GORM // transaction. The scoped handle never owns or closes the underlying pool, so // hosts can atomically update Captain rows and their own rows in one database. @@ -149,7 +187,7 @@ func (db *DB) Transaction(ctx context.Context, fn func(*DB) error) error { return errors.New("captain database transaction callback is nil") } return db.gorm.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - return fn(&DB{gorm: tx}) + return fn(&DB{gorm: tx, schema: db.schema}) }) } diff --git a/pkg/database/database_test.go b/pkg/database/database_test.go index 6cc5e5dc..242480be 100644 --- a/pkg/database/database_test.go +++ b/pkg/database/database_test.go @@ -15,8 +15,8 @@ func TestOpenMigratesThenOpensStandalonePool(t *testing.T) { var calls []string opened := &gorm.DB{} db, err := open(t.Context(), dependencies{ - migrate: func(_ context.Context, dsn string) error { - calls = append(calls, "migrate:"+dsn) + migrate: func(_ context.Context, dsn, schemaName string) error { + calls = append(calls, "migrate:"+dsn+":"+schemaName) return nil }, open: func(dsn string, _ *gorm.Config) (*gorm.DB, error) { @@ -30,7 +30,7 @@ func TestOpenMigratesThenOpensStandalonePool(t *testing.T) { if db.Gorm() != opened || !db.owned { t.Fatalf("database = %+v, want owned standalone pool", db) } - want := []string{"migrate:postgres://captain", "open:postgres://captain"} + want := []string{"migrate:postgres://captain:public", "open:postgres://captain"} if !reflect.DeepEqual(calls, want) { t.Fatalf("calls = %v, want %v", calls, want) } @@ -42,8 +42,8 @@ func TestOpenMigratesThenReusesInjectedPool(t *testing.T) { shared := &gorm.DB{} var calls []string db, err := open(t.Context(), dependencies{ - migrate: func(_ context.Context, dsn string) error { - calls = append(calls, "migrate:"+dsn) + migrate: func(_ context.Context, dsn, schemaName string) error { + calls = append(calls, "migrate:"+dsn+":"+schemaName) return nil }, open: func(string, *gorm.Config) (*gorm.DB, error) { @@ -57,7 +57,7 @@ func TestOpenMigratesThenReusesInjectedPool(t *testing.T) { if db.Gorm() != shared || db.owned { t.Fatalf("database = %+v, want non-owning injected pool", db) } - if want := []string{"migrate:postgres://shared"}; !reflect.DeepEqual(calls, want) { + if want := []string{"migrate:postgres://shared:public"}; !reflect.DeepEqual(calls, want) { t.Fatalf("calls = %v, want %v", calls, want) } if err := db.Close(); err != nil { @@ -70,7 +70,7 @@ func TestOpenUsesPreMigratedInjectedPoolWithoutDSN(t *testing.T) { shared := &gorm.DB{} db, err := open(t.Context(), dependencies{ - migrate: func(context.Context, string) error { + migrate: func(context.Context, string, string) error { t.Fatal("pre-migrated injected pool must not run migrations") return nil }, @@ -92,7 +92,7 @@ func TestOpenStopsWhenMigrationFails(t *testing.T) { wantErr := errors.New("migration failed") _, err := open(t.Context(), dependencies{ - migrate: func(context.Context, string) error { return wantErr }, + migrate: func(context.Context, string, string) error { return wantErr }, open: func(string, *gorm.Config) (*gorm.DB, error) { t.Fatal("pool must not open after migration failure") return nil, nil diff --git a/pkg/database/model_usage.go b/pkg/database/model_usage.go new file mode 100644 index 00000000..6f1c7d26 --- /dev/null +++ b/pkg/database/model_usage.go @@ -0,0 +1,89 @@ +package database + +import ( + "context" + "fmt" + "math" + "strings" + "time" + + commonsmigrate "github.com/flanksource/commons-db/migrate" +) + +type ModelUsageTotals struct { + TotalTokens int + TotalCostUSD float64 +} + +// ModelUsageSince aggregates Captain model calls in the selected schemas. +// Configured schemas that have never been migrated have no usage and are +// skipped; a migrated schema missing Captain's ledger fails as corrupted. +func (db *DB) ModelUsageSince(ctx context.Context, since time.Time, schemas ...string) (ModelUsageTotals, error) { + if db == nil || db.gorm == nil { + return ModelUsageTotals{}, fmt.Errorf("captain database is not initialized") + } + if since.IsZero() { + return ModelUsageTotals{}, fmt.Errorf("model usage start time is required") + } + if len(schemas) == 0 { + schemas = []string{db.schema} + } + seen := map[string]struct{}{} + var total ModelUsageTotals + for _, schema := range schemas { + schema = strings.TrimSpace(schema) + if err := commonsmigrate.ValidateSchemaName(schema); err != nil { + return ModelUsageTotals{}, fmt.Errorf("model usage schema: %w", err) + } + if _, duplicate := seen[schema]; duplicate { + continue + } + seen[schema] = struct{}{} + usage, exists, err := db.modelUsageSince(ctx, since, schema) + if err != nil { + return ModelUsageTotals{}, err + } + if !exists { + continue + } + if usage.TotalTokens > math.MaxInt-total.TotalTokens { + return ModelUsageTotals{}, fmt.Errorf("model usage token total overflows int") + } + total.TotalTokens += usage.TotalTokens + total.TotalCostUSD += usage.TotalCostUSD + } + return total, nil +} + +func (db *DB) modelUsageSince(ctx context.Context, since time.Time, schema string) (ModelUsageTotals, bool, error) { + var state struct { + SchemaExists bool `gorm:"column:schema_exists"` + LedgerExists bool `gorm:"column:ledger_exists"` + } + err := db.gorm.WithContext(ctx).Raw(` + SELECT + EXISTS (SELECT 1 FROM information_schema.schemata WHERE schema_name = ?) AS schema_exists, + EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = ? AND table_name = 'captain_model_calls') AS ledger_exists`, + schema, schema, + ).Scan(&state).Error + if err != nil { + return ModelUsageTotals{}, false, fmt.Errorf("inspect Captain usage schema %q: %w", schema, err) + } + if !state.SchemaExists { + return ModelUsageTotals{}, false, nil + } + if !state.LedgerExists { + return ModelUsageTotals{}, false, fmt.Errorf("captain usage schema %q is missing captain_model_calls", schema) + } + var usage ModelUsageTotals + err = db.gorm.WithContext(ctx). + Table(schema+".captain_model_calls"). + Where("created_at >= ?", since). + Select(`COALESCE(SUM(input_tokens + output_tokens), 0) AS total_tokens, + COALESCE(SUM(input_cost + output_cost + reasoning_cost + cache_read_cost + cache_write_cost), 0) AS total_cost_usd`). + Scan(&usage).Error + if err != nil { + return ModelUsageTotals{}, false, fmt.Errorf("aggregate Captain model usage for schema %q: %w", schema, err) + } + return usage, true, nil +} diff --git a/pkg/database/options_ginkgo_test.go b/pkg/database/options_ginkgo_test.go index c487c8f9..2f9d3221 100644 --- a/pkg/database/options_ginkgo_test.go +++ b/pkg/database/options_ginkgo_test.go @@ -20,8 +20,8 @@ var _ = Describe("Open options", func() { calls = nil opened = &gorm.DB{} deps = dependencies{ - migrate: func(_ context.Context, dsn string) error { - calls = append(calls, "migrate:"+dsn) + migrate: func(_ context.Context, dsn, schemaName string) error { + calls = append(calls, "migrate:"+dsn+":"+schemaName) return nil }, open: func(dsn string, _ *gorm.Config) (*gorm.DB, error) { @@ -44,7 +44,36 @@ var _ = Describe("Open options", func() { Expect(err).NotTo(HaveOccurred()) Expect(db.Gorm()).To(BeIdenticalTo(opened)) - Expect(calls).To(Equal([]string{"migrate:postgres://captain", "open:postgres://captain"})) + Expect(calls).To(Equal([]string{"migrate:postgres://captain:public", "open:postgres://captain"})) + }) + + It("scopes migrations and Captain-owned pools to the selected schema", func(ctx SpecContext) { + db, err := open(ctx, deps, + WithDSN("postgres://captain/database?sslmode=disable"), + WithSchema("agent_namespace_context"), + WithMigrations(), + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(db.Schema()).To(Equal("agent_namespace_context")) + Expect(calls).To(Equal([]string{ + "migrate:postgres://captain/database?sslmode=disable:agent_namespace_context", + "open:postgres://captain/database?search_path=agent_namespace_context&sslmode=disable", + })) + }) + + It("rejects an invalid explicit schema before migrations or pool creation", func(ctx SpecContext) { + _, err := open(ctx, deps, WithDSN("postgres://captain"), WithSchema(""), WithMigrations()) + + Expect(err).To(MatchError(ContainSubstring("captain database schema"))) + Expect(calls).To(BeEmpty()) + }) + + It("rejects a non-public schema on a host-owned pool", func(ctx SpecContext) { + _, err := open(ctx, deps, WithGorm(&gorm.DB{}), WithSchema("agent_namespace_context")) + + Expect(err).To(MatchError(ContainSubstring("host-owned GORM pool"))) + Expect(calls).To(BeEmpty()) }) It("reuses an injected pool without a DSN by default", func(ctx SpecContext) { @@ -65,7 +94,7 @@ var _ = Describe("Open options", func() { It("does not open after an explicit migration fails", func(ctx SpecContext) { migrationErr := errors.New("migration failed") - deps.migrate = func(context.Context, string) error { return migrationErr } + deps.migrate = func(context.Context, string, string) error { return migrationErr } _, err := open(ctx, deps, WithDSN("postgres://captain"), WithMigrations()) diff --git a/pkg/database/schema_scoped_integration_test.go b/pkg/database/schema_scoped_integration_test.go new file mode 100644 index 00000000..39bf315b --- /dev/null +++ b/pkg/database/schema_scoped_integration_test.go @@ -0,0 +1,81 @@ +package database + +import ( + "time" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/commons-db/dbtest" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Schema-scoped Captain storage", func() { + It("isolates identical session identities and usage aggregates between schemas", func(ctx SpecContext) { + handle := dbtest.ForGinkgo(dbtest.Options{Name: "captain_schema_scoped"}) + schemaNames := []string{"agent_tenant_a_context", "agent_tenant_b_context"} + titles := []string{"Tenant A context", "Tenant B context"} + inputTokens := []int{17, 29} + sessionID := uuid.MustParse("00000000-0000-0000-0000-000000000501") + promptRunID := uuid.MustParse("00000000-0000-0000-0000-000000001501") + databases := make([]*DB, len(schemaNames)) + + for index, schemaName := range schemaNames { + db, err := Open(ctx, WithDSN(handle.DSN()), WithSchema(schemaName), WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + databases[index] = db + DeferCleanup(func() { Expect(db.Close()).To(Succeed()) }) + + session, err := db.CreateOrGetSession(ctx, CreateSessionInput{ + ID: sessionID, ProviderSessionID: "shared-provider-session", Source: "codex", + Provider: "openai", HostID: "schema-test", Title: titles[index], + }) + Expect(err).NotTo(HaveOccurred()) + + turn, created, err := db.CreateChatTurn(ctx, CreateChatTurnInput{ + SessionID: session.ID, ProviderTurnID: "shared-provider-turn", + }) + Expect(err).NotTo(HaveOccurred()) + Expect(created).To(BeTrue()) + + run, err := db.CreatePromptRun(ctx, CreatePromptRunInput{ + ID: promptRunID, SessionID: session.ID, TurnID: &turn.ID, + }) + Expect(err).NotTo(HaveOccurred()) + + callID, err := db.CreateChatModelCall(ctx, CreateChatModelCallInput{ + TurnID: turn.ID, PromptRunID: run.ID, Model: "model-schema-test", Backend: "codex", + }) + Expect(err).NotTo(HaveOccurred()) + Expect(db.FinishChatModelCall(ctx, FinishChatModelCallInput{ + ID: callID, Status: ModelCallStatusSucceeded, StopReason: "end_turn", + Event: api.Event{Usage: &api.Usage{InputTokens: inputTokens[index], OutputTokens: 3}}, + })).To(Succeed()) + + costs, err := db.ListThreadCosts(ctx, session.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(costs).To(ConsistOf(SatisfyAll( + HaveField("SessionID", sessionID), + HaveField("InputTokens", int64(inputTokens[index])), + HaveField("TotalTokens", int64(inputTokens[index]+3)), + ))) + } + + for index, db := range databases { + session, err := db.GetSession(ctx, sessionID) + Expect(err).NotTo(HaveOccurred()) + Expect(session.Title).To(Equal(titles[index])) + } + + usageSchemas := append([]string{}, schemaNames...) + usageSchemas = append(usageSchemas, "agent_context_not_opened", schemaNames[0]) + usage, err := databases[0].ModelUsageSince(ctx, time.Now().Add(-time.Hour), usageSchemas...) + Expect(err).NotTo(HaveOccurred()) + Expect(usage.TotalTokens).To(Equal(52)) + Expect(usage.TotalCostUSD).To(Equal(0.0)) + + Expect(databases[0].Gorm().WithContext(ctx).Exec("CREATE SCHEMA agent_context_incomplete").Error).NotTo(HaveOccurred()) + _, err = databases[0].ModelUsageSince(ctx, time.Now().Add(-time.Hour), "agent_context_incomplete") + Expect(err).To(MatchError(ContainSubstring("missing captain_model_calls"))) + }) +})