From f1418e1ecca1a3b797fc0fb5fd80b14da5a62d24 Mon Sep 17 00:00:00 2001 From: Joshua Gilman Date: Tue, 25 Aug 2026 14:55:47 -0700 Subject: [PATCH] test(mcp): prove composite capability programs --- builder_test.go | 138 ++++++++++++++++++ internal/catalog/catalog_test.go | 75 +++++++++- mcpserver/e2e_test.go | 231 +++++++++++++++++++++++++++++++ mcpserver/server_test.go | 4 + moon.yml | 2 +- server_test.go | 210 ++++++++++++++++++++++++++++ 6 files changed, 656 insertions(+), 4 deletions(-) diff --git a/builder_test.go b/builder_test.go index 75446d3..564339c 100644 --- a/builder_test.go +++ b/builder_test.go @@ -2,6 +2,7 @@ package codemode_test import ( "context" + "encoding/json" "math" "os/exec" "strings" @@ -32,6 +33,44 @@ type invalidBuilderInput struct { Count int32 `json:"count"` } +// interfaceBuilderOutput is an unsupported interface/any nested output. +type interfaceBuilderOutput struct { + // Value is intentionally an unconstrained interface. + Value any `json:"value"` +} + +// rawMessageBuilderOutput is an unsupported [json.RawMessage] nested output. +type rawMessageBuilderOutput struct { + // Value is intentionally opaque JSON. + Value json.RawMessage `json:"value"` +} + +// cyclicBuilderOutput is an unsupported cyclic nested output. +type cyclicBuilderOutput struct { + // Next is a self-referential pointer that must fail registration. + Next *cyclicBuilderOutput `json:"next"` +} + +// marshalerBuilderValue is a custom JSON marshaler used only as a rejected fixture. +type marshalerBuilderValue struct{} + +// MarshalJSON exists so registration rejects custom marshalers. +func (marshalerBuilderValue) MarshalJSON() ([]byte, error) { + return []byte(`""`), nil +} + +// marshalerBuilderOutput is an unsupported custom-marshaler nested output. +type marshalerBuilderOutput struct { + // Value is intentionally a custom JSON marshaler. + Value marshalerBuilderValue `json:"value"` +} + +// mapKeyBuilderOutput is an unsupported non-string map key nested output. +type mapKeyBuilderOutput struct { + // Value is intentionally keyed by integers. + Value map[int]string `json:"value"` +} + // nilPolicy is a typed-nil authorization implementation used to test Build validation. type nilPolicy struct{} @@ -133,6 +172,105 @@ func TestRegisterRejectsInvalidContractsBeforeRetention(t *testing.T) { )), codemode.ErrInvalidRegistration) } +// TestRegisterRejectsUnsupportedNestedOutputsBeforeRetention proves representative +// unsupportable output graphs never enter the builder. +func TestRegisterRejectsUnsupportedNestedOutputsBeforeRetention(t *testing.T) { + tests := []struct { + // name identifies the unsupported nested output. + name string + + // register attempts one invalid public registration. + register func(*codemode.Builder) error + }{ + { + name: "interface any", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[builderInput, interfaceBuilderOutput]{ + ID: "cap.interface", + Name: "records.interface", + Summary: "Unsupported interface output.", + Description: "Rejected before retention.", + Handler: func(context.Context, authz.Subject, builderInput) (interfaceBuilderOutput, error) { + return interfaceBuilderOutput{}, nil + }, + }) + }, + }, + { + name: "json.RawMessage", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[builderInput, rawMessageBuilderOutput]{ + ID: "cap.raw", + Name: "records.raw", + Summary: "Unsupported raw message output.", + Description: "Rejected before retention.", + Handler: func(context.Context, authz.Subject, builderInput) (rawMessageBuilderOutput, error) { + return rawMessageBuilderOutput{}, nil + }, + }) + }, + }, + { + name: "cyclic type", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[builderInput, cyclicBuilderOutput]{ + ID: "cap.cycle", + Name: "records.cycle", + Summary: "Unsupported cyclic output.", + Description: "Rejected before retention.", + Handler: func(context.Context, authz.Subject, builderInput) (cyclicBuilderOutput, error) { + return cyclicBuilderOutput{}, nil + }, + }) + }, + }, + { + name: "custom marshaler", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[builderInput, marshalerBuilderOutput]{ + ID: "cap.marshaler", + Name: "records.marshaler", + Summary: "Unsupported marshaler output.", + Description: "Rejected before retention.", + Handler: func(context.Context, authz.Subject, builderInput) (marshalerBuilderOutput, error) { + return marshalerBuilderOutput{}, nil + }, + }) + }, + }, + { + name: "non-string map key", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[builderInput, mapKeyBuilderOutput]{ + ID: "cap.mapkey", + Name: "records.mapkey", + Summary: "Unsupported map key output.", + Description: "Rejected before retention.", + Handler: func(context.Context, authz.Subject, builderInput) (mapKeyBuilderOutput, error) { + return mapKeyBuilderOutput{}, nil + }, + }) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: codemode.DefaultLimits()}) + + err := tt.register(builder) + + require.ErrorIs(t, err, codemode.ErrInvalidRegistration) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.retained", "records.retained"))) + server, buildErr := builder.Build() + require.NoError(t, buildErr) + results, searchErr := server.Search("unsupported") + require.NoError(t, searchErr) + assert.Empty(t, results) + }) + } +} + // TestRegisterRejectsObviousDuplicates proves duplicate stable and model-facing identities fail immediately. func TestRegisterRejectsObviousDuplicates(t *testing.T) { builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: codemode.DefaultLimits()}) diff --git a/internal/catalog/catalog_test.go b/internal/catalog/catalog_test.go index 487b6fe..a67266c 100644 --- a/internal/catalog/catalog_test.go +++ b/internal/catalog/catalog_test.go @@ -33,6 +33,24 @@ type ExportedOutput struct { Name string `json:"name"` } +// NamedItem is a nested catalog result whose Go identifier must stay host-local. +type NamedItem struct { + // ID is the nested row identifier. + ID string `json:"id"` + + // Active reports whether the row survives filtering. + Active bool `json:"active"` + + // Score is the nested finite floating-point field. + Score float64 `json:"score"` +} + +// compositeTestOutput is a list-valued catalog result used to lock derived notation. +type compositeTestOutput struct { + // Items is the compiled list of nested objects. + Items []NamedItem `json:"items"` +} + // TestBuildValidatesEveryRegistrationBeforeFiltering proves disabled entries cannot hide invalid contracts. func TestBuildValidatesEveryRegistrationBeforeFiltering(t *testing.T) { registration := validRegistration("cap.bad", "records.bad", "Bad record") @@ -163,12 +181,26 @@ func TestBuildFiltersOnceAndDerivesEverySurface(t *testing.T) { validRegistration("cap.zeta", "teams.zeta", "Zeta team"), validRegistration("cap.disabled", "records.disabled", "Disabled record"), validRegistration("cap.alpha", "records.alpha", "Alpha record"), + { + ID: "cap.composite", + Name: "records.composite", + Summary: "Composite record", + Description: "Composite record full description.", + Plan: mustCompileCompositePlan(), + Invoke: func(context.Context, authz.Subject, any) (any, error) { + return compositeTestOutput{}, nil + }, + }, }, testOptions("cap.disabled")) require.NoError(t, err) entries := catalog.Entries() - require.Len(t, entries, 2) - assert.Equal(t, []string{"records.alpha", "teams.zeta"}, []string{entries[0].Name, entries[1].Name}) + require.Len(t, entries, 3) + assert.Equal(t, []string{"records.alpha", "records.composite", "teams.zeta"}, []string{ + entries[0].Name, + entries[1].Name, + entries[2].Name, + }) _, foundByName := catalog.Lookup("records.disabled") _, foundByID := catalog.LookupID("cap.disabled") _, foundDisabledDescription := catalog.Describe("records.disabled") @@ -179,8 +211,16 @@ func TestBuildFiltersOnceAndDerivesEverySurface(t *testing.T) { assert.False(t, foundDisabledDescription) assert.Empty(t, disabledSearch) + recordSearch, err := catalog.Search("record") + require.NoError(t, err) + require.Len(t, recordSearch, 2) + assert.Equal(t, []string{"records.alpha", "records.composite"}, []string{ + recordSearch[0].Name, + recordSearch[1].Name, + }) + bindings := catalog.NamespaceBindings() - require.Len(t, bindings, 2) + require.Len(t, bindings, 3) assert.Equal(t, []string{"records"}, bindings[0].Segments) assert.Equal(t, "alpha", bindings[0].Function) assert.Same(t, entries[0].Plan, bindings[0].Capability.Plan) @@ -201,6 +241,26 @@ func TestBuildFiltersOnceAndDerivesEverySurface(t *testing.T) { assert.Equal(t, "name", description.Output[0].Name) assert.Equal(t, "str", description.Output[0].Type) assert.True(t, description.Output[0].Required) + + compositeDescription, foundComposite := catalog.Describe("records.composite") + require.True(t, foundComposite) + assert.Equal(t, "records.composite(*, org: str, limit: int | None)", compositeDescription.Signature) + require.Len(t, compositeDescription.Output, 1) + assert.Equal(t, "items", compositeDescription.Output[0].Name) + assert.Equal(t, "list[{id: str, active: bool, score: float}]", compositeDescription.Output[0].Type) + assert.True(t, compositeDescription.Output[0].Required) + assertDescriptionOmitsOutputTypeNames(t, compositeDescription, "NamedItem", "compositeTestOutput") + compositeDescription.Input[0].Name = "mutated" + compositeDescription.Input[0].Type = "mutated" + compositeDescription.Output[0].Name = "mutated" + compositeDescription.Output[0].Type = "mutated" + freshComposite, foundComposite := catalog.Describe("records.composite") + require.True(t, foundComposite) + assert.Equal(t, "org", freshComposite.Input[0].Name) + assert.Equal(t, "str", freshComposite.Input[0].Type) + assert.Equal(t, "items", freshComposite.Output[0].Name) + assert.Equal(t, "list[{id: str, active: bool, score: float}]", freshComposite.Output[0].Type) + description.Input[0].Name = "mutated" freshDescription, foundDescription := catalog.Describe("records.alpha") require.True(t, foundDescription) @@ -395,6 +455,15 @@ func mustCompileExportedOutputPlan() *binding.Plan { return plan } +// mustCompileCompositePlan compiles the catalog test input with a list-valued output. +func mustCompileCompositePlan() *binding.Plan { + plan, err := binding.CompileFor[testInput, compositeTestOutput]() + if err != nil { + panic(err) + } + return plan +} + // assertSearchOmitsOutputTypeNames requires a Search result to omit host Go output identifiers. func assertSearchOmitsOutputTypeNames(t *testing.T, result SearchResult, forbidden ...string) { t.Helper() diff --git a/mcpserver/e2e_test.go b/mcpserver/e2e_test.go index d063b36..240a047 100644 --- a/mcpserver/e2e_test.go +++ b/mcpserver/e2e_test.go @@ -82,6 +82,57 @@ type StatusResult struct { State string `json:"state"` } +// searchInput is the widened records.search argument contract. +type searchInput struct { + // Count is the required integer argument. + Count int64 `json:"count"` + + // Active is the required Boolean argument. + Active bool `json:"active"` + + // Score is the required floating-point argument. + Score float64 `json:"score"` + + // Label is the optional string omitted from the canonical map when absent. + Label *string `json:"label,omitempty"` +} + +// NamedItem is one nested search row whose Go identifier must not cross MCP. +type NamedItem struct { + // ID is the nested row identifier. + ID string `json:"id"` + + // Active reports whether the row survives the program filter. + Active bool `json:"active"` + + // Score is one finite floating-point row value. + Score float64 `json:"score"` +} + +// searchOutput is the composite records.search handler output. +type searchOutput struct { + // Items is the compiled list of nested objects. + Items []NamedItem `json:"items"` +} + +// compositeDigest is the only value the composite program may return. +type compositeDigest struct { + // Count is the number of active rows. + Count int64 `json:"count"` + + // Score is the sum of active row scores. + Score float64 `json:"score"` + + // IDs are the active row identifiers in encounter order. + IDs []string `json:"ids"` +} + +// compositeExecuteEnvelope is the exact structured execute payload for the digest. +type compositeExecuteEnvelope struct { + // Result is main's final converted digest. + Result compositeDigest `json:"result"` +} + // executeEnvelope is the exact structured execute payload. type executeEnvelope struct { // Result is main's final converted value. @@ -179,6 +230,42 @@ func (recorder *lookupRecorder) snapshot() (int, []authz.Subject, []lookupInput, recorder.sawCanary } +// searchRecorder is the enabled composite handler that records trusted dispatch. +type searchRecorder struct { + // mu protects recorded handler observations. + mu sync.Mutex + + // calls is the number of handler dispatches. + calls int + + // inputs are cloned typed arguments observed at dispatch. + inputs []searchInput +} + +// invoke records typed arguments and returns multiple nested rows in one call. +func (recorder *searchRecorder) invoke( + _ context.Context, + _ authz.Subject, + input searchInput, +) (searchOutput, error) { + recorder.mu.Lock() + defer recorder.mu.Unlock() + recorder.calls++ + recorder.inputs = append(recorder.inputs, cloneSearchInput(input)) + return searchOutput{Items: []NamedItem{ + {ID: "keep-a", Active: true, Score: 1.5}, + {ID: "drop-b", Active: false, Score: 10}, + {ID: "keep-c", Active: true, Score: 2.25}, + }}, nil +} + +// snapshot returns a copy of recorded composite handler observations. +func (recorder *searchRecorder) snapshot() (int, []searchInput) { + recorder.mu.Lock() + defer recorder.mu.Unlock() + return recorder.calls, append([]searchInput(nil), recorder.inputs...) +} + // TestActualMCPSecureLoop proves the official in-memory MCP boundary preserves the secure loop. func TestActualMCPSecureLoop(t *testing.T) { authorizer := &recordingAuthorizer{} @@ -407,6 +494,140 @@ def main(): assert.NotEqual(t, string(attackerSubjectID), string(authorizations[1].Subject.ID)) } +// TestActualMCPCompositeProgram proves one native call can return nested rows +// that a Starlark program loops, filters, and reduces to a digest. +func TestActualMCPCompositeProgram(t *testing.T) { + authorizer := &recordingAuthorizer{} + search := &searchRecorder{} + builder := codemode.New(codemode.Options{ + Authorizer: authorizer, + Limits: codemode.DefaultLimits(), + }) + require.NoError(t, codemode.Register(builder, codemode.Capability[searchInput, searchOutput]{ + ID: "records.entry.search", + Name: "records.search", + Summary: "Search records and return nested items.", + Description: "Returns multiple nested rows for one widened search call.", + Handler: search.invoke, + })) + root, err := builder.Build() + require.NoError(t, err) + + mcpServer, err := mcpserver.New(root, contextResolver{}) + require.NoError(t, err) + + trustedCtx := withInvocationIdentity(t.Context(), invocationIdentity{ + Subject: authz.Subject{ID: trustedSubjectID}, + Canary: credentialCanary, + }) + serverTransport, clientTransport := mcp.NewInMemoryTransports() + serverSession, err := mcpServer.Connect(trustedCtx, serverTransport, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = serverSession.Close() }) + + client := mcp.NewClient(&mcp.Implementation{Name: "codemode-e2e", Version: "test"}, nil) + session, err := client.Connect(t.Context(), clientTransport, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = session.Close() }) + + const ( + compositeSignature = "records.search(*, count: int, active: bool, score: float, label: str | None)" + compositeItemsType = "list[{id: str, active: bool, score: float}]" + ) + + searched, err := session.CallTool(t.Context(), &mcp.CallToolParams{ + Name: "search_api", + Arguments: map[string]any{"query": "record"}, + }) + require.NoError(t, err) + assertSuccessfulTool(t, searched) + searchResults := decodeStructured[[]codemode.SearchResult](t, searched) + require.Len(t, searchResults, 1) + assert.Equal(t, "records.search", searchResults[0].Name) + assert.Equal(t, compositeSignature, searchResults[0].Signature) + assert.Equal(t, "Search records and return nested items.", searchResults[0].Summary) + assertDiscoveryOmitsGoTypeNames(t, searched, "NamedItem", "searchOutput") + + described, err := session.CallTool(t.Context(), &mcp.CallToolParams{ + Name: "describe_api", + Arguments: map[string]any{"name": "records.search"}, + }) + require.NoError(t, err) + assertSuccessfulTool(t, described) + description := decodeStructured[codemode.Description](t, described) + assert.Equal(t, "records.search", description.Name) + assert.Equal(t, compositeSignature, description.Signature) + require.Len(t, description.Input, 4) + assert.Equal(t, "count", description.Input[0].Name) + assert.Equal(t, "int", description.Input[0].Type) + assert.True(t, description.Input[0].Required) + assert.Equal(t, "active", description.Input[1].Name) + assert.Equal(t, "bool", description.Input[1].Type) + assert.True(t, description.Input[1].Required) + assert.Equal(t, "score", description.Input[2].Name) + assert.Equal(t, "float", description.Input[2].Type) + assert.True(t, description.Input[2].Required) + assert.Equal(t, "label", description.Input[3].Name) + assert.Equal(t, "str | None", description.Input[3].Type) + assert.False(t, description.Input[3].Required) + require.Len(t, description.Output, 1) + assert.Equal(t, "items", description.Output[0].Name) + assert.Equal(t, compositeItemsType, description.Output[0].Type) + assert.True(t, description.Output[0].Required) + assertDiscoveryOmitsGoTypeNames(t, described, "NamedItem", "searchOutput") + requireNonNullDescribeFieldArrays(t, described) + + executed, err := session.CallTool(t.Context(), &mcp.CallToolParams{ + Name: "execute", + Arguments: map[string]any{ + "source": ` +def main(): + response = records.search(count=3, active=True, score=1.5) + ids = [] + total = 0.0 + count = 0 + for item in response["items"]: + if item["active"]: + ids.append(item["id"]) + total += item["score"] + count += 1 + return {"count": count, "score": total, "ids": ids} +`, + }, + }) + require.NoError(t, err) + assertSuccessfulTool(t, executed) + wantDigest := compositeExecuteEnvelope{Result: compositeDigest{ + Count: 2, + Score: 3.75, + IDs: []string{"keep-a", "keep-c"}, + }} + assert.Equal(t, wantDigest, decodeStructured[compositeExecuteEnvelope](t, executed)) + assertExactExecuteEnvelope(t, executed.StructuredContent) + resultObject := requireJSONObject(t, requireJSONObject(t, executed.StructuredContent)["result"]) + require.Len(t, resultObject, 3) + requireJSONTextMirror(t, executed, wantDigest) + + handlerCalls, inputs := search.snapshot() + require.Equal(t, 1, handlerCalls) + require.Len(t, inputs, 1) + assert.Equal(t, searchInput{Count: 3, Active: true, Score: 1.5}, inputs[0]) + assert.Nil(t, inputs[0].Label) + + authorizations := authorizer.snapshot() + require.Len(t, authorizations, 1) + assert.Equal(t, authz.Subject{ID: trustedSubjectID}, authorizations[0].Subject) + assert.Equal(t, "records.entry.search", authorizations[0].CapabilityID) + assert.Equal(t, "records.search", authorizations[0].CapabilityName) + assert.Equal(t, map[string]any{ + "count": int64(3), + "active": true, + "score": 1.5, + }, authorizations[0].Arguments) + _, hasLabel := authorizations[0].Arguments["label"] + assert.False(t, hasLabel, "omitted optional string must not appear in the canonical map") +} + // withInvocationIdentity stores trusted identity on the server-owned context. func withInvocationIdentity(ctx context.Context, identity invocationIdentity) context.Context { return context.WithValue(ctx, invocationContextKey{}, identity) @@ -438,6 +659,16 @@ func cloneLookupInput(input lookupInput) lookupInput { return cloned } +// cloneSearchInput copies one typed search input including the optional label. +func cloneSearchInput(input searchInput) searchInput { + cloned := input + if input.Label != nil { + label := *input.Label + cloned.Label = &label + } + return cloned +} + // decodeStructured decodes MCP structured content into a typed value. func decodeStructured[T any](t *testing.T, result *mcp.CallToolResult) T { t.Helper() diff --git a/mcpserver/server_test.go b/mcpserver/server_test.go index 8c3403a..a432809 100644 --- a/mcpserver/server_test.go +++ b/mcpserver/server_test.go @@ -745,7 +745,11 @@ func requireFieldShapeItemSchema(t *testing.T, items any) { item := requireJSONObject(t, items) requireNonNullJSONType(t, item, "object") requireRequiredNames(t, item, "name", "type", "required") + required, ok := item["required"].([]any) + require.True(t, ok, "schema required must be an array") + require.Len(t, required, 3, "FieldShape must require exactly name, type, and required") properties := requireJSONObject(t, item["properties"]) + require.Len(t, properties, 3, "FieldShape must stay a flat name/type/required object") requireNonNullJSONType(t, requireJSONObject(t, properties["name"]), "string") requireNonNullJSONType(t, requireJSONObject(t, properties["type"]), "string") requireNonNullJSONType(t, requireJSONObject(t, properties["required"]), "boolean") diff --git a/moon.yml b/moon.yml index 86feed5..2143940 100644 --- a/moon.yml +++ b/moon.yml @@ -72,7 +72,7 @@ tasks: # trailing single-dash `-flag=value` into two arguments; `script` runs the line # through the shell verbatim. mcp-smoke: - script: "go test ./mcpserver -run '^TestActualMCPSecureLoop$' -count=1" + script: "go test ./mcpserver -run '^(TestActualMCPSecureLoop|TestActualMCPCompositeProgram)$' -count=1" inputs: - '@group(goSources)' options: diff --git a/server_test.go b/server_test.go index 9fd4fb5..e0990b3 100644 --- a/server_test.go +++ b/server_test.go @@ -3,6 +3,7 @@ package codemode_test import ( "context" "errors" + "math" "sync" "sync/atomic" "testing" @@ -394,6 +395,215 @@ def main(): assert.Equal(t, int64(2), handlerCalls.Load()) } +// compositeProjectionInput is the empty input used by public composite projections. +type compositeProjectionInput struct{} + +// nanProjectionOutput covers nested non-finite floating-point results. +type nanProjectionOutput struct { + // Score is a finite floating-point field. + Score float64 `json:"score"` +} + +// overflowProjectionOutput covers unsigned values above MaxInt64. +type overflowProjectionOutput struct { + // Count is an unsigned 64-bit integer. + Count uint64 `json:"count"` +} + +// nestedProjectionItem is one nested object used for depth and budget cases. +type nestedProjectionItem struct { + // ID is the nested row identifier. + ID string `json:"id"` +} + +// nestedProjectionOutput is a composite result whose nesting exceeds MaxValueDepth 2. +type nestedProjectionOutput struct { + // Items is the compiled list of nested objects. + Items []nestedProjectionItem `json:"items"` +} + +// TestServerExecuteProjectsCompositeOutputFailures proves nested NaN, Inf, and +// unsigned overflow become the bare public capability-failure sentinel. +func TestServerExecuteProjectsCompositeOutputFailures(t *testing.T) { + tests := []struct { + // name identifies the invalid handler output. + name string + + // capabilityName is the Starlark capability invoked by main. + capabilityName string + + // register retains one capability that returns an invalid composite value. + register func(*codemode.Builder) error + }{ + { + name: "nested NaN", + capabilityName: "records.nan", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[compositeProjectionInput, nanProjectionOutput]{ + ID: "cap.nan", + Name: "records.nan", + Summary: "Return a NaN score.", + Description: "Projects non-finite floats as capability failure.", + Handler: func(context.Context, authz.Subject, compositeProjectionInput) (nanProjectionOutput, error) { + return nanProjectionOutput{Score: math.NaN()}, nil + }, + }) + }, + }, + { + name: "nested Inf", + capabilityName: "records.inf", + register: func(builder *codemode.Builder) error { + return codemode.Register(builder, codemode.Capability[compositeProjectionInput, nanProjectionOutput]{ + ID: "cap.inf", + Name: "records.inf", + Summary: "Return an infinite score.", + Description: "Projects non-finite floats as capability failure.", + Handler: func(context.Context, authz.Subject, compositeProjectionInput) (nanProjectionOutput, error) { + return nanProjectionOutput{Score: math.Inf(1)}, nil + }, + }) + }, + }, + { + name: "uint overflow", + capabilityName: "records.overflow", + register: func(builder *codemode.Builder) error { + return codemode.Register( + builder, + codemode.Capability[compositeProjectionInput, overflowProjectionOutput]{ + ID: "cap.overflow", + Name: "records.overflow", + Summary: "Return an overflowing unsigned count.", + Description: "Projects unsigned overflow as capability failure.", + Handler: func(context.Context, authz.Subject, compositeProjectionInput) (overflowProjectionOutput, error) { + return overflowProjectionOutput{Count: uint64(1) << 63}, nil + }, + }, + ) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: codemode.DefaultLimits()}) + require.NoError(t, tt.register(builder)) + server, err := builder.Build() + require.NoError(t, err) + + _, err = server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, codemode.Program(` +def main(): + return `+tt.capabilityName+`() +`)) + + require.ErrorIs(t, err, codemode.ErrCapabilityFailure) + assert.Equal(t, codemode.ErrCapabilityFailure, err) + }) + } +} + +// TestServerExecuteProjectsCompositeValueLimits proves composite depth and +// independent per-value/aggregate budgets project only the bare resource sentinel. +func TestServerExecuteProjectsCompositeValueLimits(t *testing.T) { + t.Run("max value depth", func(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxValueDepth = 2 + limits.MaxValueBytes = 1024 + limits.MaxIntermediateValueBytes = 1024 + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: limits}) + require.NoError( + t, + codemode.Register(builder, codemode.Capability[compositeProjectionInput, nestedProjectionOutput]{ + ID: "cap.depth", + Name: "records.depth", + Summary: "Return nested items.", + Description: "Projects composite depth exhaustion as a resource limit.", + Handler: func(context.Context, authz.Subject, compositeProjectionInput) (nestedProjectionOutput, error) { + return nestedProjectionOutput{Items: []nestedProjectionItem{{ID: "a"}}}, nil + }, + }), + ) + server, err := builder.Build() + require.NoError(t, err) + + _, err = server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.depth() +`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit, err) + }) + + t.Run("per-value bytes", func(t *testing.T) { + const body = `{"items":[{"id":"xx"}]}` + limits := codemode.DefaultLimits() + limits.MaxValueBytes = len(body) - 1 + limits.MaxIntermediateValueBytes = 1024 + var handlerCalls atomic.Int64 + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: limits}) + require.NoError( + t, + codemode.Register(builder, codemode.Capability[compositeProjectionInput, nestedProjectionOutput]{ + ID: "cap.pervalue", + Name: "records.pervalue", + Summary: "Return one nested item.", + Description: "Projects one oversized composite result independently of the aggregate budget.", + Handler: func(context.Context, authz.Subject, compositeProjectionInput) (nestedProjectionOutput, error) { + handlerCalls.Add(1) + return nestedProjectionOutput{Items: []nestedProjectionItem{{ID: "xx"}}}, nil + }, + }), + ) + server, err := builder.Build() + require.NoError(t, err) + + _, err = server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.pervalue() +`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit, err) + assert.Equal(t, int64(1), handlerCalls.Load()) + }) + + t.Run("aggregate intermediate bytes", func(t *testing.T) { + const body = `{"items":[{"id":"xx"}]}` + limits := codemode.DefaultLimits() + limits.MaxValueBytes = 1024 + limits.MaxIntermediateValueBytes = len(body)*2 - 1 + var handlerCalls atomic.Int64 + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: limits}) + require.NoError( + t, + codemode.Register(builder, codemode.Capability[compositeProjectionInput, nestedProjectionOutput]{ + ID: "cap.aggregate", + Name: "records.aggregate", + Summary: "Return one nested item.", + Description: "Projects composite aggregate exhaustion independently of MaxValueBytes.", + Handler: func(context.Context, authz.Subject, compositeProjectionInput) (nestedProjectionOutput, error) { + handlerCalls.Add(1) + return nestedProjectionOutput{Items: []nestedProjectionItem{{ID: "xx"}}}, nil + }, + }), + ) + server, err := builder.Build() + require.NoError(t, err) + + _, err = server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + records.aggregate() + return records.aggregate() +`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit, err) + assert.Equal(t, int64(2), handlerCalls.Load()) + }) +} + // TestServerExecuteReturnsOnlyMainResult proves top-level values and printed text do not escape. func TestServerExecuteReturnsOnlyMainResult(t *testing.T) { server := buildTestServer(