From 24e7670f08a3bb917c52a1410e40eaf869273c67 Mon Sep 17 00:00:00 2001 From: Joshua Gilman Date: Tue, 25 Aug 2026 19:49:10 -0700 Subject: [PATCH] feat: add bounded relevance-ranked search --- README.md | 12 +- SECURITY.md | 6 +- builder.go | 1 + builder_test.go | 32 +- capability.go | 7 + docs/docs/explanation/security-model.md | 17 + docs/docs/how-to/disable-capabilities.md | 12 +- docs/docs/reference/mcp-tools.md | 132 +++++- docs/docs/reference/public-api.md | 144 +++++- docs/docs/tutorials/first-server.md | 16 +- internal/catalog/build.go | 134 ++++-- internal/catalog/catalog.go | 14 +- internal/catalog/catalog_test.go | 47 +- internal/catalog/doc.go | 2 +- internal/catalog/search.go | 532 ++++++++++++++++++++++- internal/catalog/search_test.go | 385 ++++++++++++++++ limits_test.go | 6 +- mcpserver/e2e_test.go | 59 ++- mcpserver/mocks/service.go | 18 +- mcpserver/server.go | 21 +- mcpserver/server_test.go | 99 ++++- mcpserver/service.go | 4 +- server.go | 17 +- server_test.go | 43 +- 24 files changed, 1532 insertions(+), 228 deletions(-) create mode 100644 internal/catalog/search_test.go diff --git a/README.md b/README.md index fed1afc..ec7d0b0 100644 --- a/README.md +++ b/README.md @@ -33,9 +33,10 @@ func main() { builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll()}) codemode.Register(builder, codemode.Capability[lookupInput, lookupOutput]{ - Name: "records.lookup", - Summary: "Look up one record by key.", - Handler: lookup, + Name: "records.lookup", + Summary: "Look up one record by key.", + SearchTerms: []string{"fetch entry"}, + Handler: lookup, }) server, err := builder.Build() @@ -50,6 +51,11 @@ func main() { } ``` +`SearchTerms` adds task and resource vocabulary for discovery only. Search +terms are not returned or callable, but callers can infer them by probing. +Never put secrets, credentials, policy facts, tenant identifiers, or sensitive +examples in search terms. + The repository also contains shorter, compile-checked examples: - [`example_test.go`](example_test.go) — typed registration with default limits and an explicit subject, plus direct execution diff --git a/SECURITY.md b/SECURITY.md index 7616b38..3c65f8e 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -32,8 +32,10 @@ attempted native calls, concurrent workers, and the depth and encoded size of every value that crosses the worker boundary. `MaxIntermediateValueBytes` is the cumulative encoded size of successful parent-to-child native-result value bodies per execution, independent of -the per-value `MaxValueBytes` bound. Search query bytes and -search result counts are bounded separately in the parent. An elapsed deadline +the per-value `MaxValueBytes` bound. Search query bytes, result counts, and the +structured search response size are bounded separately in the parent. The +structured-response bound excludes the surrounding JSON-RPC envelope and the +MCP SDK's JSON text mirror. An elapsed deadline or request cancellation kills and reaps the worker. Each native call whose arguments bind successfully is rebound in the parent and passes through the host-supplied `authz.Authorizer` before its handler runs. The MCP adapter diff --git a/builder.go b/builder.go index 294af48..cf13f81 100644 --- a/builder.go +++ b/builder.go @@ -103,6 +103,7 @@ func Register[Input, Output any](builder *Builder, capability Capability[Input, Name: string(capability.Name), Summary: capability.Summary, Description: description, + SearchTerms: slices.Clone(capability.SearchTerms), Plan: plan, Invoke: func(ctx context.Context, subject authz.Subject, input any) (any, error) { typed, ok := input.(Input) diff --git a/builder_test.go b/builder_test.go index 3db1612..c63bca3 100644 --- a/builder_test.go +++ b/builder_test.go @@ -109,9 +109,37 @@ func TestBuilderCopiesStaticFilteringOptions(t *testing.T) { server, err := builder.Build() require.NoError(t, err) - results, err := server.Search("disabled") + response, err := server.Search("disabled") require.NoError(t, err) - assert.Empty(t, results) + require.NotNil(t, response.Results) + assert.Empty(t, response.Results) + assert.False(t, response.Truncated) +} + +// TestBuilderCopiesSearchTerms proves caller search-term mutation cannot alter discovery. +func TestBuilderCopiesSearchTerms(t *testing.T) { + terms := []string{"open ticket"} + capability := validBuilderCapability("cap.alpha", "records.alpha") + capability.SearchTerms = terms + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll()}) + codemode.Register(builder, capability) + terms[0] = "mutated after register" + + server, err := builder.Build() + require.NoError(t, err) + + response, err := server.Search("ticket") + require.NoError(t, err) + require.NotNil(t, response.Results) + require.Len(t, response.Results, 1) + assert.Equal(t, "records.alpha", response.Results[0].Name) + assert.False(t, response.Truncated) + + mutated, err := server.Search("mutated") + require.NoError(t, err) + require.NotNil(t, mutated.Results) + assert.Empty(t, mutated.Results) + assert.False(t, mutated.Truncated) } // TestBuilderRejectsNilAuthorizersAndNegativeLimits proves required construction policy fails closed. diff --git a/capability.go b/capability.go index 053b81c..aa2028d 100644 --- a/capability.go +++ b/capability.go @@ -32,6 +32,13 @@ type Capability[Input, Output any] struct { // requests. An empty Description defaults to Summary. Description string + // SearchTerms contains alternative task vocabulary used only for discovery. + // Terms are not callable aliases and are not accepted by Describe or Execute. + // They are not returned in search results, but callers can infer indexed + // vocabulary by probing. Do not put secrets, policy facts, credentials, + // tenant identifiers, or sensitive examples in search terms. + SearchTerms []string + // Handler executes the capability after binding and authorization succeed. Handler Handler[Input, Output] } diff --git a/docs/docs/explanation/security-model.md b/docs/docs/explanation/security-model.md index 1b81a62..929bc36 100644 --- a/docs/docs/explanation/security-model.md +++ b/docs/docs/explanation/security-model.md @@ -84,6 +84,23 @@ See [Use Rego for authorization](../how-to/use-rego-authorization.md) for config Static filtering is useful for deployment-wide availability, but it is not dynamic authorization. It cannot express subject-specific or argument-specific decisions. Conversely, authorization alone does not hide a capability's metadata from discovery. Use static filtering to remove a capability from the deployment surface and authorization to decide whether an enabled native call may dispatch. +## Discovery metadata is observable + +Search indexes each enabled capability's name, `SearchTerms`, summary, and +description. A search response returns only the exact name, signature, and +summary, but omission is not secrecy. A caller can submit different queries +and infer whether particular vocabulary changes the ranked results. + +Treat `SearchTerms` as model-visible discovery metadata even though the terms +are not returned directly. Do not put secrets, credentials, policy facts, +tenant identifiers, or sensitive examples in them. Search terms do not create +callable aliases and are not accepted by exact `Describe` or by execution. + +Static filtering removes a disabled capability before the search index is +built, so its metadata does not contribute tokens or ranking. Authorization is +different: search does not run the per-capability authorizer and does not hide +enabled discovery metadata based on the resolved subject. + ## Client errors are intentionally coarse Detailed causes exist only on the trusted side of the public boundary. Internal packages and host authorizers, resolvers, and handlers can hold or log those causes. A direct call to `authz/rego.Authorize` can return an ordinary error that identifies an undefined or non-Boolean decision, or carries an OPA evaluation or builtin failure. diff --git a/docs/docs/how-to/disable-capabilities.md b/docs/docs/how-to/disable-capabilities.md index 4e48cc0..dcccd58 100644 --- a/docs/docs/how-to/disable-capabilities.md +++ b/docs/docs/how-to/disable-capabilities.md @@ -14,10 +14,11 @@ explicit ID before creating a deployment filter: ```go codemode.Register(builder, codemode.Capability[lookupInput, lookupOutput]{ - ID: "records.entry.lookup", - Name: "records.lookup", - Summary: "Look up one record by key.", - Handler: lookup, + ID: "records.entry.lookup", + Name: "records.lookup", + Summary: "Look up one record by key.", + SearchTerms: []string{"fetch entry"}, + Handler: lookup, }) ``` @@ -66,7 +67,8 @@ go build -o codemode-first-server . ## Verify the filter -1. Call `search_api` with `{"query":"records.lookup"}`. The result is `[]` and no longer lists `records.lookup`. +1. Call `search_api` with `{"query":"records.lookup"}`. The result is + `{"results":[],"truncated":false}` and no longer lists `records.lookup`. 2. Call `describe_api` with `{"name":"records.lookup"}`. The call returns the tool error `capability not found`. diff --git a/docs/docs/reference/mcp-tools.md b/docs/docs/reference/mcp-tools.md index e1ab1f5..1bfe78e 100644 --- a/docs/docs/reference/mcp-tools.md +++ b/docs/docs/reference/mcp-tools.md @@ -17,7 +17,10 @@ On success, `CallToolResult.StructuredContent` contains the value described by t ## `search_api` -Search enabled names and summaries with a short literal substring. Retry an empty result with a shorter term. +Search enabled capabilities using task, resource, or exact-name vocabulary. +Results are relevance-ranked. Pass the exact returned name to `describe_api`. +If `truncated` is `true` and no result fits, submit a more specific +task/resource query. ### Input @@ -34,37 +37,122 @@ Search enabled names and summaries with a short literal substring. Retry an empt } ``` -The raw query is limited by `MaxSearchQueryBytes` before trimming or case normalization. Whitespace padding counts. CodeMode then trims surrounding whitespace and normalizes case. Matching is a short literal substring over capability names and summaries. A blank normalized query or any other empty result is `[]`, not `null`. Retry an empty result with a shorter term. Search does not add fuzzy matching, aliases, or extra query rewriting. - -Results are sorted by exact dotted name and limited by `MaxSearchResults`. Static filtering happens before search, so disabled capabilities never appear. +The raw query is limited by `MaxSearchQueryBytes` before trimming or +tokenization. Whitespace padding counts. CodeMode trims surrounding Unicode +whitespace, treats every rune that is not a Unicode letter or digit as a +separator, splits camel-case, acronym-to-word, and letter/number transitions, +and lowercases each token. Dots, underscores, and hyphens are therefore +separators. + +For example, `GitHub.Pulls.createReview` becomes `github`, `pulls`, `create`, +`review`; `pull_request` becomes `pull`, `request`; and `sql` and `mysql` +remain different tokens. + +Search removes the connector tokens `a`, `an`, `and`, `by`, `for`, `from`, +`in`, `of`, `on`, `or`, `the`, `to`, and `with`, then deduplicates the +remaining query tokens. A query with more than 16 distinct normalized tokens +returns `resource limit exceeded`. + +Search compares the distinct query tokens with tokens from each enabled +capability's name, registered `SearchTerms`, summary, and description. Exact +token matches are supported. A query token of at least three Unicode characters +can also match the prefix of a capability token. Arbitrary infix matching, +fuzzy matching, stemming, and built-in synonym expansion are not supported. + +For one query token within one field, an exact match ranks above a prefix +match. Search retains only the strongest contribution for that query token. +The field precedence used for weighting is the capability name, then +`SearchTerms`, then the summary, then the description. Terms found in fewer +enabled capabilities contribute more than catalog-wide terms. Numeric scoring +weights are internal. + +Eligibility depends on the number `q` of distinct normalized query tokens: + +| Query tokens | Required matched tokens | +| ---: | ---: | +| 1 | 1 | +| 2 | 2 | +| 3 or more | `ceil(2q / 3)` | + +Search ranks every eligible capability before applying output bounds. A +case-insensitive exact dotted-name query, after trimming surrounding +whitespace, ranks first. Remaining ordering is relevance score descending, +then exact dotted name ascending. Static filtering happens before indexing, so +disabled capabilities cannot match or affect ranking. ### Successful structured output ```json { - "type": "array", - "items": { - "type": "object", - "required": ["name", "signature", "summary"], - "additionalProperties": false, - "properties": { - "name": { "type": "string" }, - "signature": { "type": "string" }, - "summary": { "type": "string" } + "type": "object", + "required": ["results", "truncated"], + "additionalProperties": false, + "properties": { + "results": { + "type": "array", + "items": { + "type": "object", + "required": ["name", "signature", "summary"], + "additionalProperties": false, + "properties": { + "name": { "type": "string" }, + "signature": { "type": "string" }, + "summary": { "type": "string" } + } + } + }, + "truncated": { + "type": "boolean" } } } ``` -The structured content itself is the array described above, not an object that wraps the array. When there are no matches, it is `[]`, not `null`. +A populated successful value is: + +```json +{ + "results": [ + { + "name": "records.lookup", + "signature": "records.lookup(*, key: str, limit: int | None)", + "summary": "Look up one record by key." + } + ], + "truncated": false +} +``` + +A blank query, a separator- or connector-only query, and a query with no +eligible matches all succeed with this exact object: + +```json +{ + "results": [], + "truncated": false +} +``` + +`results` is always a non-null array on success. CodeMode packs the +highest-ranked prefix under `MaxSearchResults` and an internal structured +response-byte bound. `truncated` is `true` when either bound omits at least one +eligible capability. It does not expose a total count or provide pagination. +The byte bound covers the compact JSON representation of the structured +`{results, truncated}` response. The surrounding JSON-RPC envelope and the MCP +SDK's JSON `TextContent` mirror are outside that cap. | Field | Meaning | | --- | --- | -| `name` | Exact enabled dotted capability name, such as `records.lookup`. | -| `signature` | Invocation-only keyword signature. It ends after the parameter list and never contains a Go output type. | -| `summary` | Registered compact summary. | +| `results[].name` | Exact enabled dotted capability name, such as `records.lookup`. | +| `results[].signature` | Invocation-only keyword signature. It ends after the parameter list and never contains a Go output type. | +| `results[].summary` | Registered compact summary. | +| `truncated` | Whether at least one eligible result was omitted by the result-count or structured-response byte bound. | -`signature` contains the dotted name, a `*` keyword-only marker when there are parameters, and the ordered input fields with their type notations. It ends at `)`. The exact forms are `records.lookup(*, key: str, limit: int | None)` and `records.status()`. The result contract is `describe_api.output`, not `signature`. +`signature` contains the dotted name, a `*` keyword-only marker when there are +parameters, and the ordered input fields with their type notations. It ends at +`)`. The exact forms are +`records.lookup(*, key: str, limit: int | None)` and `records.status()`. The +result contract is `describe_api.output`, not `signature`. ## `describe_api` @@ -326,7 +414,7 @@ Only the final converted value from the worker process is exposed in the success The listed descriptions above are the model-facing contract. Recovery uses the same fixed coarse errors on this page. When recording or reporting a failed call, keep the coarse text and the recovery action; do not echo the failed source, arguments, credentials, or unknown requested name. -- Search with a short literal substring over enabled names and summaries. If the result is empty, retry with a shorter term. +- Search with task, resource, or exact-name vocabulary. If `truncated` is `true`, use a more specific task/resource query. Pass an exact returned `name` to `describe_api`. - After `capability not found`, search again and pass `describe_api` an exact returned `name`, without whitespace or case changes. - After `invalid capability arguments`, compare the call with the published `signature` and `input` field shapes. - After `invalid program`, check the program against these requirements: @@ -335,9 +423,9 @@ The listed descriptions above are the model-facing contract. Recovery uses the s - Return the final value from `main`. - Use `describe_api.output` as the result contract. Do not parse a type name from `signature`. - After `resource limit exceeded`, reduce the applicable bounded quantity: - query or source bytes, execution steps or time, native calls, crossing-value - depth or per-value encoded size, or the cumulative encoded size of successful - native results. Then retry the call. + query or source bytes, distinct normalized query tokens, execution steps or + time, native calls, crossing-value depth or per-value encoded size, or the + cumulative encoded size of successful native results. Then retry the call. - After `permission denied` or `authorization policy failure`, contact the host if the access was expected. One allowed or denied input cannot establish whether the policy is default-open, default-deny, complete, or incomplete. ## Errors diff --git a/docs/docs/reference/public-api.md b/docs/docs/reference/public-api.md index a40efc8..edabf00 100644 --- a/docs/docs/reference/public-api.md +++ b/docs/docs/reference/public-api.md @@ -57,10 +57,28 @@ to perform that setup; `IsWorker` does not serve worker mode. | --- | --- | | `ID CapabilityID` | Stable deployment and policy identity. An empty ID defaults to `Name`; explicit IDs must have no surrounding whitespace and must be unique. An explicit `ID` preserves policy and filter identity across `Name` changes. | | `Name CapabilityName` | A unique dotted Starlark name. A complete capability name cannot also be another capability's namespace. | -| `Summary string` | Non-empty compact text searched by `Search`, with no surrounding whitespace. | -| `Description string` | Detail returned by `Describe`. An empty value defaults to `Summary`; an explicit value must have no surrounding whitespace. | +| `Summary string` | Non-empty compact text used by `Search` and returned by discovery, with no surrounding whitespace. | +| `Description string` | Detail returned by `Describe` and used by `Search`. An empty value defaults to `Summary`; an explicit value must have no surrounding whitespace. | +| `SearchTerms []string` | Optional alternative task and resource phrases used only by `Search`. | | `Handler Handler[Input, Output]` | A non-nil function called after argument binding and authorization. | +`SearchTerms` is discovery-only metadata. `Register` clones the slice. Search +terms are not returned by `Search` or `Describe`, are not callable aliases, and +are not accepted as names by `Describe` or `Execute`. Callers can still infer +indexed vocabulary by probing search results. Do not include secrets, +credentials, policy facts, tenant identifiers, or sensitive examples. + +A capability can register at most 16 non-empty search-term phrases with no +surrounding whitespace. Their combined raw size must not exceed 1,024 bytes. +For example: + +```go +SearchTerms: []string{ + "fetch entry", + "find stored item", +}, +``` + `Handler[Input, Output]` has this signature: ```text @@ -169,9 +187,12 @@ Output-depth, per-value byte, and aggregate intermediate-byte exhaustion map to `Build` returns all capability-specific registration failures as one joined error. It also rejects a missing authorizer, negative signed limits, namespace -collisions, duplicate disabled IDs, and disabled IDs that do not match a -registered capability. It re-executes the current binary and performs a fixed -five-second private worker probe only after construction validation succeeds. +collisions, duplicate disabled IDs, disabled IDs that do not match a registered +capability, more than 4,096 registrations, aggregate searchable metadata that +exceeds the internal build budget, and any enabled compact `SearchResult` that +cannot fit the internal structured-response cap. It re-executes the current +binary and performs a fixed five-second private worker probe only after +construction validation succeeds. The probe detects an absent or nonfunctional worker entry and returns `ErrInvalidRegistration` with a host-wiring diagnostic that identifies the missing or misplaced `ServeWorkerAndExit` call. It cannot detect ordinary host @@ -205,8 +226,8 @@ Static filtering is build-scoped deployment configuration. Subject- or argument- | `MaxValueDepth int` | 32 | Inclusive nesting depth of any value crossing the worker boundary. | | `MaxValueBytes int` | 1,048,576 bytes (1 MiB) | Type-preserving encoding of any value crossing the worker boundary. | | `MaxIntermediateValueBytes int` | 8,388,608 bytes (8 MiB) | Cumulative encoded successful parent-to-child native-result value bodies in one `Execute` call. | -| `MaxSearchQueryBytes int` | 256 bytes | Raw search query before trimming or case normalization. Whitespace padding counts. | -| `MaxSearchResults int` | 20 | Search results returned. | +| `MaxSearchQueryBytes int` | 256 bytes | Raw search query before trimming or tokenization. Whitespace padding counts. | +| `MaxSearchResults int` | 20 | Maximum number of entries in `SearchResponse.Results`. | | `MaxConcurrentExecutions int` | 8 | Concurrent spawn attempts and live worker processes. | `Build` replaces every zero-valued field with the corresponding value from @@ -216,6 +237,27 @@ unrelated budgets. Negative signed fields return `ErrInvalidRegistration`. Calling `Limits.Validate()` directly still rejects zero and otherwise non-positive fields because it validates an already resolved limit set. +The structured search response also has an internal byte bound. This bound +applies to the compact JSON representation of `SearchResponse`; its exact value +is not part of the public configuration contract. The surrounding JSON-RPC +envelope and the MCP SDK's JSON text mirror are outside this bound. + +The registration ceiling, aggregate searchable-metadata budget, and +single-result fit check are hard build constraints rather than configurable +`Limits` fields. The aggregate covers the raw bytes of every registered +capability's name, summary, description, and `SearchTerms`; it is checked +before static filtering, so disabling a capability does not reduce that +accounting. The single-result check applies after filtering to each enabled +compact object containing its name, generated signature, and summary. +Exceeding any of these constraints returns `ErrInvalidRegistration`. + +To reduce registration count or aggregate metadata, remove capabilities from +the server build, split them across servers, or shorten their discovery +metadata. To make one compact result fit, shorten its capability name, summary, +or input field names that form the generated signature. The internal +structured-response cap is not public configuration and cannot be raised +through `Limits`. + `MaxValueDepth` is inclusive. A scalar or `None` is depth 1. Each tuple, list, or dictionary wrapper adds one. A scalar with limit 1 succeeds, a one-level container with limit 2 succeeds, and one more wrapper with limit 2 fails. @@ -260,22 +302,96 @@ effects. See [Understanding CodeMode's security model](../explanation/security-m #### `Search` ```text -Search(query string) ([]SearchResult, error) +Search(query string) (SearchResponse, error) ``` -`Search` first enforces `MaxSearchQueryBytes` on the raw query, then trims surrounding whitespace and normalizes case. Whitespace padding counts toward the byte budget. It performs substring matching against enabled capability names and summaries. Results are sorted by exact capability name and capped by `MaxSearchResults`. A blank normalized query returns an empty, non-nil result. +`SearchResponse` and `SearchResult` have these JSON shapes: -`SearchResult` contains these JSON fields: +```go +type SearchResponse struct { + Results []SearchResult `json:"results"` + Truncated bool `json:"truncated"` +} -| Field | Type | Meaning | +type SearchResult struct { + Name string `json:"name"` + Signature string `json:"signature"` + Summary string `json:"summary"` +} +``` + +`Search` first enforces `MaxSearchQueryBytes` on the raw query. Whitespace +padding counts. It then: + +1. trims surrounding Unicode whitespace; +2. treats every rune that is not a Unicode letter or digit as a separator; +3. splits camel-case, acronym-to-word, and letter/number transitions; +4. lowercases each token; +5. removes `a`, `an`, `and`, `by`, `for`, `from`, `in`, `of`, `on`, `or`, + `the`, `to`, and `with`; and +6. deduplicates the remaining query tokens. + +A query with more than 16 distinct normalized tokens returns +`ErrResourceLimit`. A blank query, or one containing only separators and +removed connector tokens, succeeds with: + +```json +{"results":[],"truncated":false} +``` + +Search compares the distinct query tokens with tokens from each enabled +capability's name, `SearchTerms`, summary, and description. Exact token matches +are supported. A query token of at least three Unicode characters can also +match the prefix of a capability token. Arbitrary infix matching, fuzzy +matching, stemming, and built-in synonym expansion are not supported. Hosts +must put alternative task and resource vocabulary in `SearchTerms` or other +descriptive metadata. + +For one query token within one field, an exact match ranks above a prefix +match. Search retains only the strongest contribution for that query token. +The field precedence used for weighting is the capability name, then +`SearchTerms`, then the summary, then the description. Terms that occur in +fewer enabled capabilities contribute more than catalog-wide terms. The +numeric scoring weights are internal. + +Eligibility depends on the number `q` of distinct normalized query tokens: + +| Query tokens | Required matched tokens | +| ---: | ---: | +| 1 | 1 | +| 2 | 2 | +| 3 or more | `ceil(2q / 3)` | + +Search ranks every eligible capability before applying output bounds. A +case-insensitive exact dotted-name query, after trimming surrounding +whitespace, ranks first. Remaining ordering is relevance score descending, +then exact dotted name ascending. Static filtering happens before indexing, so +disabled capabilities cannot match or affect ranking. + +`Results` is always a non-nil array on success. CodeMode packs the +highest-ranked prefix under `MaxSearchResults` and the internal structured +response-byte bound. `Truncated` is `true` if either bound omits at least one +eligible capability; it does not expose a total count or provide pagination. +The response-byte bound covers the compact JSON representation of +`SearchResponse`, not a surrounding JSON-RPC envelope or the MCP SDK's JSON +text mirror. + +`SearchResult` contains: + +| JSON field | Type | Meaning | | --- | --- | --- | | `name` | string | Exact enabled dotted name. | | `signature` | string | Invocation-only keyword signature. It ends after the parameter list and never contains a Go output type. | | `summary` | string | Registered summary. | -`signature` contains the dotted name, a `*` keyword-only marker when there are parameters, and the ordered input fields with their type notations. It ends at `)`. The exact forms are `records.lookup(*, key: str, limit: int | None)` and `records.status()`. The result contract is `Description.Output`, not `signature`. +`signature` contains the dotted name, a `*` keyword-only marker when there are +parameters, and the ordered input fields with their type notations. It ends at +`)`. The exact forms are +`records.lookup(*, key: str, limit: int | None)` and `records.status()`. The +result contract is `Description.Output`, not `signature`. -An oversized query returns `ErrResourceLimit`. An unexpected server-state failure returns `ErrInternal`. +An oversized or over-tokenized query returns `ErrResourceLimit`. An unexpected +server-state failure returns `ErrInternal`. #### `Describe` @@ -489,7 +605,7 @@ See [Use Rego for authorization](../how-to/use-rego-authorization.md) for server `Service` is the adapter's application port: ```text -Search(query string) ([]codemode.SearchResult, error) +Search(query string) (codemode.SearchResponse, error) Describe(name codemode.CapabilityName) (codemode.Description, error) Execute(context.Context, authz.Subject, codemode.Program) (any, error) ``` diff --git a/docs/docs/tutorials/first-server.md b/docs/docs/tutorials/first-server.md index 4277f0a..731a0ce 100644 --- a/docs/docs/tutorials/first-server.md +++ b/docs/docs/tutorials/first-server.md @@ -71,9 +71,10 @@ func main() { builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll()}) codemode.Register(builder, codemode.Capability[lookupInput, lookupOutput]{ - Name: "records.lookup", - Summary: "Look up one record by key.", - Handler: lookup, + Name: "records.lookup", + Summary: "Look up one record by key.", + SearchTerms: []string{"fetch entry"}, + Handler: lookup, }) server, err := builder.Build() @@ -100,8 +101,9 @@ process is the authentication boundary. A multi-user host must not use `Build` supplies bounded defaults for each zero-valued `Limits` field. It also reports all invalid registrations together. An omitted capability `ID` defaults to `Name`; set `ID` explicitly before writing authorization policy or deployment -filters that must survive a capability rename. An omitted `Description` -defaults to `Summary`. +filters that must survive a capability rename. `SearchTerms` supplies +discovery-only task and resource vocabulary; it does not create callable names. +An omitted `Description` defaults to `Summary`. ## Build and configure the server @@ -129,7 +131,9 @@ Agents that use the `mcpServers` configuration shape accept: Restart or reload the agent's MCP servers. -Ask the agent to search for a record capability. +Ask the agent to search for a capability that can `fetch entry`. This phrase +comes from `SearchTerms`, so the agent can discover `records.lookup` even though +the phrase does not appear in its name or summary. Ask the agent to call `records.lookup` with `key="alpha"` and `limit=2`. The agent can use `search_api`, `describe_api`, and `execute`. The final structured diff --git a/internal/catalog/build.go b/internal/catalog/build.go index c7f0cd4..509bd24 100644 --- a/internal/catalog/build.go +++ b/internal/catalog/build.go @@ -3,6 +3,7 @@ package catalog import ( "errors" "fmt" + "slices" "sort" "strings" @@ -26,6 +27,9 @@ var ( type candidate struct { // entry contains copied metadata, the supplied plan, and the handler adapter. entry Entry + + // searchTerms is the cloned discovery vocabulary excluded from disabled documents. + searchTerms []string } // Build validates every registration, applies static filtering once, and returns an immutable catalog. @@ -36,64 +40,114 @@ func Build(registrations []Registration, options Options) (*Catalog, error) { if options.MaxSearchResults <= 0 { return nil, fmt.Errorf("%w: MaxSearchResults must be positive", ErrInvalidRegistration) } + if len(registrations) > maxSupportedRegistrations { + return nil, fmt.Errorf( + "%w: %d registrations exceed the %d capability limit", + ErrInvalidRegistration, + len(registrations), + maxSupportedRegistrations, + ) + } + + candidates, ids, names, err := collectCandidates(registrations) + if err != nil { + return nil, err + } + if collisionErr := validateNamespaceCollisions(names); collisionErr != nil { + return nil, collisionErr + } + disabled, err := validateDisabled(options.DisabledCapabilities, ids) + if err != nil { + return nil, err + } + enabled, searchTerms := filterEnabled(candidates, disabled) + search, err := compileSearchIndex(enabled, searchTerms) + if err != nil { + return nil, err + } + catalog := &Catalog{ + enabled: enabled, + byName: make(map[string]int, len(enabled)), + byID: make(map[string]int, len(enabled)), + namespaces: make([]NamespaceBinding, 0, len(enabled)), + search: search, + maxSearchQueryBytes: options.MaxSearchQueryBytes, + maxSearchResults: options.MaxSearchResults, + } + indexCatalog(catalog) + return catalog, nil +} +// collectCandidates validates registrations, enforces searchable-metadata and identity uniqueness, and copies owned entries. +func collectCandidates( + registrations []Registration, +) ([]candidate, map[string]struct{}, map[string]struct{}, error) { candidates := make([]candidate, 0, len(registrations)) ids := make(map[string]struct{}, len(registrations)) names := make(map[string]struct{}, len(registrations)) + searchableBytes := 0 for index, registration := range registrations { if err := ValidateRegistration(registration); err != nil { - return nil, fmt.Errorf("%w: registration %d: %w", ErrInvalidRegistration, index, err) + return nil, nil, nil, fmt.Errorf("%w: registration %d: %w", ErrInvalidRegistration, index, err) + } + searchableBytes += searchableMetadataBytes(registration) + if searchableBytes > maxSearchableMetadataBytes { + return nil, nil, nil, fmt.Errorf( + "%w: searchable metadata exceeds %d bytes", + ErrInvalidRegistration, + maxSearchableMetadataBytes, + ) } if _, duplicate := ids[registration.ID]; duplicate { - return nil, fmt.Errorf("%w: duplicate ID %q", ErrInvalidRegistration, registration.ID) + return nil, nil, nil, fmt.Errorf("%w: duplicate ID %q", ErrInvalidRegistration, registration.ID) } if _, duplicate := names[registration.Name]; duplicate { - return nil, fmt.Errorf("%w: duplicate name %q", ErrInvalidRegistration, registration.Name) + return nil, nil, nil, fmt.Errorf("%w: duplicate name %q", ErrInvalidRegistration, registration.Name) } ids[registration.ID] = struct{}{} names[registration.Name] = struct{}{} - candidates = append(candidates, candidate{entry: Entry{ - ID: registration.ID, - Name: registration.Name, - Summary: registration.Summary, - Description: registration.Description, - Plan: registration.Plan, - Invoke: registration.Invoke, - signature: registration.Plan.Signature(registration.Name), - searchName: strings.ToLower(registration.Name), - searchSummary: strings.ToLower(registration.Summary), - inputShape: registration.Plan.InputShape(), - outputShape: registration.Plan.OutputShape(), - }}) - } - if err := validateNamespaceCollisions(names); err != nil { - return nil, err + candidates = append(candidates, candidate{ + entry: Entry{ + ID: registration.ID, + Name: registration.Name, + Summary: registration.Summary, + Description: registration.Description, + Plan: registration.Plan, + Invoke: registration.Invoke, + signature: registration.Plan.Signature(registration.Name), + inputShape: registration.Plan.InputShape(), + outputShape: registration.Plan.OutputShape(), + }, + searchTerms: slices.Clone(registration.SearchTerms), + }) } + return candidates, ids, names, nil +} - disabled, err := validateDisabled(options.DisabledCapabilities, ids) - if err != nil { - return nil, err - } - enabled := make([]Entry, 0, len(candidates)-len(disabled)) - for _, candidate := range candidates { - if _, isDisabled := disabled[candidate.entry.ID]; isDisabled { +// filterEnabled drops disabled candidates and returns name-sorted enabled entries aligned with search terms. +func filterEnabled(candidates []candidate, disabled map[string]struct{}) ([]Entry, [][]string) { + enabledCandidates := make([]candidate, 0, len(candidates)-len(disabled)) + for _, current := range candidates { + if _, isDisabled := disabled[current.entry.ID]; isDisabled { continue } - enabled = append(enabled, candidate.entry) + enabledCandidates = append(enabledCandidates, current) } - sort.Slice(enabled, func(left int, right int) bool { - return enabled[left].Name < enabled[right].Name + sort.Slice(enabledCandidates, func(left int, right int) bool { + return enabledCandidates[left].entry.Name < enabledCandidates[right].entry.Name }) - - catalog := &Catalog{ - enabled: enabled, - byName: make(map[string]int, len(enabled)), - byID: make(map[string]int, len(enabled)), - namespaces: make([]NamespaceBinding, 0, len(enabled)), - maxSearchQueryBytes: options.MaxSearchQueryBytes, - maxSearchResults: options.MaxSearchResults, + enabled := make([]Entry, len(enabledCandidates)) + searchTerms := make([][]string, len(enabledCandidates)) + for index, current := range enabledCandidates { + enabled[index] = current.entry + searchTerms[index] = current.searchTerms } - for index, entry := range enabled { + return enabled, searchTerms +} + +// indexCatalog fills exact-name, exact-ID, and namespace indexes from enabled entries. +func indexCatalog(catalog *Catalog) { + for index, entry := range catalog.enabled { catalog.byName[entry.Name] = index catalog.byID[entry.ID] = index segments := strings.Split(entry.Name, ".") @@ -103,7 +157,6 @@ func Build(registrations []Registration, options Options) (*Catalog, error) { Capability: entry, }) } - return catalog, nil } // ValidateRegistration reports whether one copied registration has valid metadata, a compiled plan, and a handler. @@ -126,6 +179,9 @@ func ValidateRegistration(registration Registration) error { if registration.Invoke == nil { return errors.New("handler must not be nil") } + if err := validateSearchTerms(registration.SearchTerms); err != nil { + return err + } return nil } diff --git a/internal/catalog/catalog.go b/internal/catalog/catalog.go index 973125a..17c2cba 100644 --- a/internal/catalog/catalog.go +++ b/internal/catalog/catalog.go @@ -19,12 +19,15 @@ type Registration struct { // Name is the dotted Starlark and model-facing capability name. Name string - // Summary is the compact text searched with the capability name. + // Summary is the compact description used by capability search. Summary string // Description is the full exact-description text. Description string + // SearchTerms contains alternative task vocabulary used only for discovery. + SearchTerms []string + // Plan is the caller-compiled input, output, canonical-argument, and signature plan. Plan *binding.Plan @@ -67,12 +70,6 @@ type Entry struct { // signature is generated once from the immutable binding plan. signature string - // searchName is the registration-time normalized capability name. - searchName string - - // searchSummary is the registration-time normalized capability summary. - searchSummary string - // inputShape is the registration-time compiled model-facing input shape. inputShape []binding.FieldShape @@ -106,6 +103,9 @@ type Catalog struct { // namespaces is derived only from enabled and remains name-sorted. namespaces []NamespaceBinding + // search is the immutable enabled-document search index. + search searchIndex + // maxSearchQueryBytes is the positive search input budget. maxSearchQueryBytes int diff --git a/internal/catalog/catalog_test.go b/internal/catalog/catalog_test.go index a67266c..59693f2 100644 --- a/internal/catalog/catalog_test.go +++ b/internal/catalog/catalog_test.go @@ -209,15 +209,17 @@ func TestBuildFiltersOnceAndDerivesEverySurface(t *testing.T) { assert.False(t, foundByName) assert.False(t, foundByID) assert.False(t, foundDisabledDescription) - assert.Empty(t, disabledSearch) + assert.Empty(t, disabledSearch.Results) + assert.False(t, disabledSearch.Truncated) recordSearch, err := catalog.Search("record") require.NoError(t, err) - require.Len(t, recordSearch, 2) + require.Len(t, recordSearch.Results, 2) assert.Equal(t, []string{"records.alpha", "records.composite"}, []string{ - recordSearch[0].Name, - recordSearch[1].Name, + recordSearch.Results[0].Name, + recordSearch.Results[1].Name, }) + assert.False(t, recordSearch.Truncated) bindings := catalog.NamespaceBindings() require.Len(t, bindings, 3) @@ -283,24 +285,27 @@ func TestSearchIsSortedNormalizedAndBounded(t *testing.T) { }, options) require.NoError(t, err) - results, err := catalog.Search(" ALPHA ") + response, err := catalog.Search(" ALPHA ") require.NoError(t, err) - require.Len(t, results, 2) - assert.Equal(t, []string{"records.alpha", "records.beta"}, []string{results[0].Name, results[1].Name}) - assert.Equal(t, "records.alpha(*, org: str, limit: int | None)", results[0].Signature) - - nonContiguous, err := catalog.Search("alpha record") - require.NoError(t, err) - assert.Empty(t, nonContiguous) + require.Len(t, response.Results, 2) + assert.Equal(t, []string{"records.alpha", "records.beta"}, []string{ + response.Results[0].Name, + response.Results[1].Name, + }) + assert.True(t, response.Truncated) + assert.Equal(t, "records.alpha(*, org: str, limit: int | None)", response.Results[0].Signature) - crossBoundary, err := catalog.Search("alpha alpha") + compound, err := catalog.Search("alpha record") require.NoError(t, err) - assert.Empty(t, crossBoundary) + require.NotEmpty(t, compound.Results) + assert.Equal(t, "records.alpha", compound.Results[0].Name) empty, err := catalog.Search(" ") require.NoError(t, err) - assert.Empty(t, empty) + require.NotNil(t, empty.Results) + assert.Empty(t, empty.Results) + assert.False(t, empty.Truncated) _, err = catalog.Search("query exceeding the configured byte budget") require.Error(t, err) @@ -330,9 +335,9 @@ func TestSearchAndDescribeOmitHostOutputTypeNames(t *testing.T) { results, err := catalog.Search("alpha") require.NoError(t, err) - require.Len(t, results, 1) - assert.Equal(t, "records.alpha(*, org: str, limit: int | None)", results[0].Signature) - assertSearchOmitsOutputTypeNames(t, results[0], unexportedName, exportedName) + require.Len(t, results.Results, 1) + assert.Equal(t, "records.alpha(*, org: str, limit: int | None)", results.Results[0].Signature) + assertSearchOmitsOutputTypeNames(t, results.Results[0], unexportedName, exportedName) description, found := catalog.Describe("records.alpha") require.True(t, found) @@ -344,9 +349,9 @@ func TestSearchAndDescribeOmitHostOutputTypeNames(t *testing.T) { statusResults, err := catalog.Search("health") require.NoError(t, err) - require.Len(t, statusResults, 1) - assert.Equal(t, "health.status()", statusResults[0].Signature) - assertSearchOmitsOutputTypeNames(t, statusResults[0], unexportedName, exportedName) + require.Len(t, statusResults.Results, 1) + assert.Equal(t, "health.status()", statusResults.Results[0].Signature) + assertSearchOmitsOutputTypeNames(t, statusResults.Results[0], unexportedName, exportedName) statusDescription, found := catalog.Describe("health.status") require.True(t, found) diff --git a/internal/catalog/doc.go b/internal/catalog/doc.go index b163c3d..ad8d934 100644 --- a/internal/catalog/doc.go +++ b/internal/catalog/doc.go @@ -1,2 +1,2 @@ -// Package catalog validates, filters, and indexes immutable native capability registrations. +// Package catalog validates, filters, and compiles immutable native capability registrations. package catalog diff --git a/internal/catalog/search.go b/internal/catalog/search.go index f12476a..e99879b 100644 --- a/internal/catalog/search.go +++ b/internal/catalog/search.go @@ -1,16 +1,47 @@ package catalog import ( + "encoding/json" "errors" "fmt" "slices" "strings" + "unicode" + "unicode/utf8" "github.com/meigma/codemode/internal/binding" ) +const ( + maxSupportedRegistrations = 4096 + maxSearchableMetadataBytes = 32 * 1024 * 1024 + maxSearchTermPhrases = 16 + maxSearchTermBytes = 1024 + maxDistinctQueryTokens = 16 + maxSearchResponseBytes = 64 * 1024 + minPrefixQueryTokenLength = 3 + fieldWeightName = 12 + fieldWeightSearchTerms = 10 + fieldWeightSummary = 5 + fieldWeightDescription = 2 + matchQualityExact = 100 + matchQualityPrefix = 72 + idfBucketCount = 8 + idfMinFactor = 1 + tokenizeScratchCap = 8 + tokensPerSearchPhrase = 2 + twoTokenQuery = 2 + coverageMatchNumerator = 2 + coverageMatchDenominator = 3 + coverageMatchRoundUp = 2 + queryTokenOverflowSlot = 1 + searchResponsePrefixJSONBytes = len(`{"results":[`) + searchResponseTruncatedFalse = len(`],"truncated":false}`) + searchResponseCommaJSONBytes = 1 +) + var ( - // ErrSearchQueryLimit classifies a search query that exceeds its byte budget. + // ErrSearchQueryLimit classifies a search query that exceeds its byte or token budget. ErrSearchQueryLimit = errors.New("search query limit exceeded") ) @@ -26,6 +57,15 @@ type SearchResult struct { Summary string `json:"summary"` } +// SearchResponse is one bounded ranked discovery result set. +type SearchResponse struct { + // Results contains the packed ranked prefix of eligible capabilities. + Results []SearchResult `json:"results"` + + // Truncated reports whether at least one eligible result was omitted. + Truncated bool `json:"truncated"` +} + // Description is one exact model-facing capability description. type Description struct { // Name is the enabled capability's exact dotted name. @@ -47,37 +87,91 @@ type Description struct { Output []binding.FieldShape `json:"output"` } -// Search performs a deterministic case-normalized linear scan over enabled names and summaries. -func (catalog *Catalog) Search(query string) ([]SearchResult, error) { +// searchIndex is the immutable owned document slice compiled after static filtering. +type searchIndex struct { + // documents is aligned by index with Catalog.enabled. + documents []searchDocument +} + +// searchDocument is one enabled capability's owned searchable tokens. +type searchDocument struct { + // normalizedName is the lowercase dotted capability name used for exact-name priority. + normalizedName string + + // nameTokens contains distinct tokens compiled from the capability name. + nameTokens []searchToken + + // searchTermTokens contains distinct tokens compiled from explicit search terms. + searchTermTokens []searchToken + + // summaryTokens contains distinct tokens compiled from the compact summary. + summaryTokens []searchToken + + // descriptionTokens contains distinct tokens compiled from the full description. + descriptionTokens []searchToken + + // resultJSONBytes is the compact encoding/json size of the projected SearchResult. + resultJSONBytes int +} + +// searchToken is one retained field token with a monotone document-frequency factor. +type searchToken struct { + // text is the normalized token text. + text string + + // idf is the integer rarity factor assigned at catalog build. + idf uint64 +} + +// searchCandidate is one eligible document retained until packing. +type searchCandidate struct { + // document is the index into Catalog.enabled and searchIndex.documents. + document int + + // score is the integer ranking score after coverage adjustment. + score uint64 + + // exact reports whether the normalized query equals the normalized dotted name. + exact bool +} + +// Search ranks enabled capabilities for one bounded query. +func (catalog *Catalog) Search(query string) (SearchResponse, error) { if len(query) > catalog.maxSearchQueryBytes { - return nil, fmt.Errorf( + return SearchResponse{}, fmt.Errorf( "%w: query is %d bytes; maximum is %d", ErrSearchQueryLimit, len(query), catalog.maxSearchQueryBytes, ) } - normalized := strings.ToLower(strings.TrimSpace(query)) - if normalized == "" { - return []SearchResult{}, nil + tokens, err := tokenizeQuery(query) + if err != nil { + return SearchResponse{}, err } - - results := make([]SearchResult, 0, min(catalog.maxSearchResults, len(catalog.enabled))) - for _, entry := range catalog.enabled { - if !strings.Contains(entry.searchName, normalized) && - !strings.Contains(entry.searchSummary, normalized) { + if len(tokens) == 0 { + return emptySearchResponse(), nil + } + normalizedName := strings.ToLower(strings.TrimSpace(query)) + candidates := make([]searchCandidate, 0, len(catalog.search.documents)) + for index, document := range catalog.search.documents { + score, matched := scoreDocument(document, tokens) + if matched < requiredMatches(len(tokens)) { continue } - results = append(results, SearchResult{ - Name: entry.Name, - Signature: entry.signature, - Summary: entry.Summary, + candidates = append(candidates, searchCandidate{ + document: index, + score: score * matched / uintCount(len(tokens)), + exact: document.normalizedName == normalizedName, }) - if len(results) == catalog.maxSearchResults { - break - } } - return results, nil + slices.SortFunc(candidates, func(left searchCandidate, right searchCandidate) int { + if order := compareSearchCandidates(left, right); order != 0 { + return order + } + return strings.Compare(catalog.enabled[left.document].Name, catalog.enabled[right.document].Name) + }) + return catalog.packSearchResponse(candidates), nil } // Describe returns the exact description of one enabled capability without fuzzy expansion. @@ -95,3 +189,401 @@ func (catalog *Catalog) Describe(name string) (Description, bool) { Output: slices.Clone(entry.outputShape), }, true } + +// emptySearchResponse returns a successful empty discovery payload. +func emptySearchResponse() SearchResponse { + return SearchResponse{Results: []SearchResult{}} +} + +// tokenizeQuery normalizes and deduplicates query tokens after the raw-byte check. +func tokenizeQuery(query string) ([]string, error) { + tokens := tokenize(query) + unique := make([]string, 0, min(len(tokens), maxDistinctQueryTokens+queryTokenOverflowSlot)) + for _, token := range tokens { + if slices.Contains(unique, token) { + continue + } + unique = append(unique, token) + if len(unique) > maxDistinctQueryTokens { + return nil, fmt.Errorf( + "%w: query has %d distinct tokens; maximum is %d", + ErrSearchQueryLimit, + len(unique), + maxDistinctQueryTokens, + ) + } + } + return unique, nil +} + +// uniqueTokens retains first-seen tokens while preserving order. +func uniqueTokens(tokens []string) []string { + if len(tokens) == 0 { + return nil + } + seen := make(map[string]struct{}, len(tokens)) + unique := make([]string, 0, len(tokens)) + for _, token := range tokens { + if _, exists := seen[token]; exists { + continue + } + seen[token] = struct{}{} + unique = append(unique, token) + } + return unique +} + +// tokenize splits value on punctuation, case, and letter/number boundaries and drops connectors. +func tokenize(value string) []string { + value = strings.TrimSpace(value) + if value == "" { + return nil + } + tokens := make([]string, 0, tokenizeScratchCap) + var current strings.Builder + flush := func() { + if current.Len() == 0 { + return + } + token := strings.ToLower(current.String()) + current.Reset() + if isConnectorToken(token) { + return + } + tokens = append(tokens, token) + } + runes := []rune(value) + startedUpper := false + for index, currentRune := range runes { + if !unicode.IsLetter(currentRune) && !unicode.IsDigit(currentRune) { + flush() + startedUpper = false + continue + } + if current.Len() > 0 && shouldSplitBefore(runes, index, startedUpper) { + flush() + startedUpper = unicode.IsUpper(currentRune) + } else if current.Len() == 0 { + startedUpper = unicode.IsUpper(currentRune) + } + current.WriteRune(currentRune) + } + flush() + return tokens +} + +// shouldSplitBefore reports a camel, acronym, or letter/number boundary before runes[index]. +// startedUpper is true when the current token began with an uppercase letter, which keeps MySQL and GitHub whole. +func shouldSplitBefore(runes []rune, index int, startedUpper bool) bool { + previous := runes[index-1] + current := runes[index] + if unicode.IsLetter(previous) && unicode.IsDigit(current) { + return true + } + if unicode.IsDigit(previous) && unicode.IsLetter(current) { + return true + } + if unicode.IsLower(previous) && unicode.IsUpper(current) { + return !startedUpper + } + if unicode.IsUpper(previous) && unicode.IsUpper(current) && + index+1 < len(runes) && unicode.IsLower(runes[index+1]) { + return true + } + return false +} + +// isConnectorToken reports whether token is a dropped grammatical connector. +func isConnectorToken(token string) bool { + switch token { + case "a", "an", "and", "by", "for", "from", "in", "of", "on", "or", "the", "to", "with": + return true + default: + return false + } +} + +// uintCount converts a non-negative int count to uint64. +func uintCount(value int) uint64 { + if value < 0 { + return 0 + } + return uint64(value) +} + +// scoreDocument returns the best per-token contributions and the distinct match count. +func scoreDocument(document searchDocument, queryTokens []string) (uint64, uint64) { + var score uint64 + var matched uint64 + for _, queryToken := range queryTokens { + contribution := bestTokenContribution(document, queryToken) + if contribution == 0 { + continue + } + matched++ + score += contribution + } + return score, matched +} + +// bestTokenContribution returns the highest field/quality contribution for one query token. +func bestTokenContribution(document searchDocument, queryToken string) uint64 { + best := tokenFieldContribution(document.nameTokens, queryToken, fieldWeightName) + if contribution := tokenFieldContribution( + document.searchTermTokens, + queryToken, + fieldWeightSearchTerms, + ); contribution > best { + best = contribution + } + if contribution := tokenFieldContribution( + document.summaryTokens, + queryToken, + fieldWeightSummary, + ); contribution > best { + best = contribution + } + if contribution := tokenFieldContribution( + document.descriptionTokens, + queryToken, + fieldWeightDescription, + ); contribution > best { + best = contribution + } + return best +} + +// tokenFieldContribution returns the best exact-then-prefix contribution in one field. +// Exact quality dominates prefix quality; rarity is compared only among equal-quality matches. +func tokenFieldContribution(tokens []searchToken, queryToken string, weight uint64) uint64 { + var bestQuality uint64 + var bestIDF uint64 + for _, token := range tokens { + quality := tokenMatchQuality(queryToken, token.text) + if quality == 0 { + continue + } + if quality > bestQuality || (quality == bestQuality && token.idf > bestIDF) { + bestQuality = quality + bestIDF = token.idf + } + } + if bestQuality == 0 { + return 0 + } + return weight * bestIDF * bestQuality +} + +// tokenMatchQuality scores exact token equality above a bounded prefix match. +func tokenMatchQuality(queryToken string, documentToken string) uint64 { + if queryToken == documentToken { + return matchQualityExact + } + if utf8.RuneCountInString(queryToken) < minPrefixQueryTokenLength { + return 0 + } + if strings.HasPrefix(documentToken, queryToken) { + return matchQualityPrefix + } + return 0 +} + +// requiredMatches returns the strict eligibility threshold for q distinct query tokens. +func requiredMatches(queryTokens int) uint64 { + switch { + case queryTokens <= 1: + return uintCount(queryTokens) + case queryTokens == twoTokenQuery: + return twoTokenQuery + default: + return uintCount((coverageMatchNumerator*queryTokens + coverageMatchRoundUp) / coverageMatchDenominator) + } +} + +// compareSearchCandidates orders exact-name hits first, then score descending, then name ascending. +func compareSearchCandidates(left searchCandidate, right searchCandidate) int { + switch { + case left.exact != right.exact: + if left.exact { + return -1 + } + return 1 + case left.score != right.score: + if left.score > right.score { + return -1 + } + return 1 + default: + return 0 + } +} + +// packSearchResponse walks ranked candidates under count and compact-response byte bounds. +func (catalog *Catalog) packSearchResponse(candidates []searchCandidate) SearchResponse { + if len(candidates) == 0 { + return emptySearchResponse() + } + results := make([]SearchResult, 0, min(catalog.maxSearchResults, len(candidates))) + used := searchResponsePrefixJSONBytes + for _, candidate := range candidates { + if len(results) == catalog.maxSearchResults { + return SearchResponse{Results: results, Truncated: true} + } + resultBytes := catalog.search.documents[candidate.document].resultJSONBytes + separator := 0 + if len(results) > 0 { + separator = searchResponseCommaJSONBytes + } + if used+separator+resultBytes+searchResponseTruncatedFalse > maxSearchResponseBytes { + return SearchResponse{Results: results, Truncated: true} + } + entry := catalog.enabled[candidate.document] + results = append(results, SearchResult{ + Name: entry.Name, + Signature: entry.signature, + Summary: entry.Summary, + }) + used += separator + resultBytes + } + return SearchResponse{Results: results, Truncated: false} +} + +// compileSearchIndex builds owned documents and document-frequency factors after filtering. +func compileSearchIndex(enabled []Entry, searchTerms [][]string) (searchIndex, error) { + documents := make([]searchDocument, len(enabled)) + documentFrequency := make(map[string]int) + for index, entry := range enabled { + nameTokens := uniqueTokens(tokenize(entry.Name)) + termTokens := uniqueTokens(tokenizeJoined(searchTerms[index])) + summaryTokens := uniqueTokens(tokenize(entry.Summary)) + descriptionTokens := uniqueTokens(tokenize(entry.Description)) + seen := make(map[string]struct{}) + markDocumentTokens(seen, nameTokens) + markDocumentTokens(seen, termTokens) + markDocumentTokens(seen, summaryTokens) + markDocumentTokens(seen, descriptionTokens) + for token := range seen { + documentFrequency[token]++ + } + result := SearchResult{ + Name: entry.Name, + Signature: entry.signature, + Summary: entry.Summary, + } + encoded, err := json.Marshal(result) + if err != nil { + return searchIndex{}, fmt.Errorf("%w: search result %q: %w", ErrInvalidRegistration, entry.Name, err) + } + if !resultFitsResponseCap(len(encoded)) { + return searchIndex{}, fmt.Errorf( + "%w: search result %q is %d bytes and cannot fit the %d-byte response cap", + ErrInvalidRegistration, + entry.Name, + len(encoded), + maxSearchResponseBytes, + ) + } + documents[index] = searchDocument{ + normalizedName: strings.ToLower(entry.Name), + nameTokens: tokensWithoutIDF(nameTokens), + searchTermTokens: tokensWithoutIDF(termTokens), + summaryTokens: tokensWithoutIDF(summaryTokens), + descriptionTokens: tokensWithoutIDF(descriptionTokens), + resultJSONBytes: len(encoded), + } + } + documentCount := len(enabled) + assignIDF(documents, documentFrequency, documentCount) + return searchIndex{documents: documents}, nil +} + +// tokenizeJoined tokenizes each phrase and concatenates the tokens. +func tokenizeJoined(phrases []string) []string { + tokens := make([]string, 0, len(phrases)*tokensPerSearchPhrase) + for _, phrase := range phrases { + tokens = append(tokens, tokenize(phrase)...) + } + return tokens +} + +// markDocumentTokens records each token as present in the current document. +func markDocumentTokens(seen map[string]struct{}, tokens []string) { + for _, token := range tokens { + seen[token] = struct{}{} + } +} + +// tokensWithoutIDF constructs search tokens before document-frequency assignment. +func tokensWithoutIDF(tokens []string) []searchToken { + if len(tokens) == 0 { + return nil + } + converted := make([]searchToken, len(tokens)) + for index, token := range tokens { + converted[index] = searchToken{text: token} + } + return converted +} + +// assignIDF writes a small monotone integer rarity factor onto every retained token. +func assignIDF(documents []searchDocument, documentFrequency map[string]int, documentCount int) { + factors := make(map[string]uint64, len(documentFrequency)) + for token, frequency := range documentFrequency { + factors[token] = idfFactor(frequency, documentCount) + } + applyIDF := func(tokens []searchToken) { + for index, token := range tokens { + tokens[index].idf = factors[token.text] + } + } + for index := range documents { + applyIDF(documents[index].nameTokens) + applyIDF(documents[index].searchTermTokens) + applyIDF(documents[index].summaryTokens) + applyIDF(documents[index].descriptionTokens) + } +} + +// idfFactor maps document frequency onto a small increasing rarity scale. +func idfFactor(frequency int, documentCount int) uint64 { + if frequency <= 0 || documentCount <= 0 { + return idfMinFactor + } + numerator := (documentCount - frequency) * (idfBucketCount - idfMinFactor) + return idfMinFactor + uintCount(numerator/documentCount) +} + +// resultFitsResponseCap reports whether one compact result can occupy a successful response. +func resultFitsResponseCap(resultJSONBytes int) bool { + return searchResponsePrefixJSONBytes+resultJSONBytes+searchResponseTruncatedFalse <= maxSearchResponseBytes +} + +// searchableMetadataBytes returns the aggregate raw searchable metadata size of one registration. +func searchableMetadataBytes(registration Registration) int { + total := len(registration.Name) + len(registration.Summary) + len(registration.Description) + for _, term := range registration.SearchTerms { + total += len(term) + } + return total +} + +// validateSearchTerms reports whether one capability's discovery phrases stay in bounds. +func validateSearchTerms(terms []string) error { + if len(terms) > maxSearchTermPhrases { + return fmt.Errorf("search terms exceed %d phrases", maxSearchTermPhrases) + } + total := 0 + for index, term := range terms { + if term != strings.TrimSpace(term) { + return fmt.Errorf("search term %d must not have surrounding whitespace", index) + } + if term == "" { + return fmt.Errorf("search term %d must not be empty", index) + } + total += len(term) + } + if total > maxSearchTermBytes { + return fmt.Errorf("search terms exceed %d bytes", maxSearchTermBytes) + } + return nil +} diff --git a/internal/catalog/search_test.go b/internal/catalog/search_test.go new file mode 100644 index 0000000..674b96d --- /dev/null +++ b/internal/catalog/search_test.go @@ -0,0 +1,385 @@ +package catalog + +import ( + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestSearchRanksTokensAndHonorsBounds proves ranked discovery, eligibility, and truncation. +func TestSearchRanksTokensAndHonorsBounds(t *testing.T) { + tests := []struct { + name string + registrations []Registration + options Options + query string + wantNames []string + wantTruncated bool + wantEmpty bool + }{ + { + name: "sql does not match mysql infix and ranks snowflake first", + registrations: []Registration{ + validRegistration("cap.mysql.users", "mysql.users.list", "List MySQL users"), + validRegistration("cap.mysql.queries", "mysql.queries.execute", "Execute a MySQL query"), + validRegistration("cap.snowflake", "snowflake.queries.execute", "Execute a Snowflake SQL query"), + }, + options: testOptions(), + query: "sql", + wantNames: []string{"snowflake.queries.execute"}, + }, + { + name: "mysql query matches mysql without emitting sql tokens", + registrations: []Registration{ + validRegistration("cap.mysql.users", "mysql.users.list", "List MySQL users"), + validRegistration("cap.snowflake", "snowflake.queries.execute", "Execute a Snowflake SQL query"), + }, + options: testOptions(), + query: "mysql", + wantNames: []string{"mysql.users.list"}, + }, + { + name: "acronyms split before a camel tail", + registrations: []Registration{ + validRegistration("cap.xml", "records.parseXMLFile", "Parse an XML file"), + }, + options: testOptions(), + query: "xml file", + wantNames: []string{"records.parseXMLFile"}, + }, + { + name: "GitHub stays one token", + registrations: []Registration{ + validRegistration("cap.review", "github.pulls.createReview", "Create a GitHub pull review"), + }, + options: testOptions(), + query: "hub", + wantEmpty: true, + }, + { + name: "compound task query matches without a contiguous phrase", + registrations: []Registration{ + validRegistration("cap.create", "github.issues.create", "Create a repository issue"), + validRegistration("cap.merge", "github.pulls.merge", "Merge a pull request"), + }, + options: testOptions(), + query: "create github issue", + wantNames: []string{"github.issues.create"}, + }, + { + name: "camel and dotted names tokenize on boundaries", + registrations: []Registration{ + validRegistration("cap.review", "github.pulls.createReview", "Create a pull review"), + }, + options: testOptions(), + query: "create review", + wantNames: []string{"github.pulls.createReview"}, + }, + { + name: "prefix tokens match longer document tokens", + registrations: []Registration{ + validRegistration("cap.review", "github.pulls.createReview", "Create a pull review"), + }, + options: testOptions(), + query: "crea", + wantNames: []string{"github.pulls.createReview"}, + }, + { + name: "common exact list beats rare prefix listener in the same field", + registrations: []Registration{ + withSummary(validRegistration("cap.alpha", "records.alpha", "Alpha record"), "List records"), + withSummary(validRegistration("cap.zeta", "records.zeta", "Zeta record"), "List listener records"), + }, + options: testOptions(), + query: "list", + wantNames: []string{"records.alpha", "records.zeta"}, + }, + { + name: "search terms are required for otherwise undiscoverable vocabulary", + registrations: []Registration{ + withSearchTerms( + validRegistration("cap.create", "github.issues.create", "Create a repository issue"), + "open ticket", + "file bug report", + ), + validRegistration("cap.merge", "github.pulls.merge", "Merge a pull request"), + }, + options: testOptions(), + query: "ticket", + wantNames: []string{"github.issues.create"}, + }, + { + name: "name outranks search terms, summary, and description", + registrations: []Registration{ + withDescription( + validRegistration("cap.desc", "other.describe", "Unrelated summary"), + "Contains the widget token in description.", + ), + withSummary( + validRegistration("cap.summary", "other.summarize", "Handles a widget"), + "Handles a widget", + ), + withSearchTerms(validRegistration("cap.terms", "other.terms", "Unrelated summary"), "widget"), + validRegistration("cap.name", "widget.alpha", "Unrelated summary"), + }, + options: testOptions(), + query: "widget", + wantNames: []string{"widget.alpha", "other.terms", "other.summarize", "other.describe"}, + }, + { + name: "unrelated two-token queries stay empty", + registrations: []Registration{ + validRegistration("cap.deploy", "apps.deploy", "Deploy an application"), + }, + options: testOptions(), + query: "deploy rocket", + wantEmpty: true, + }, + { + name: "exact normalized name is first even when another document scores", + registrations: []Registration{ + validRegistration("cap.generic", "records.lookup", "Lookup records"), + validRegistration("cap.exact", "records.alpha", "Alpha record"), + }, + options: testOptions(), + query: "records.alpha", + wantNames: []string{"records.alpha"}, + }, + { + name: "equal scores sort by exact dotted name", + registrations: []Registration{ + validRegistration("cap.zeta", "records.zeta", "Shared token"), + validRegistration("cap.alpha", "records.alpha", "Shared token"), + }, + options: testOptions(), + query: "shared", + wantNames: []string{"records.alpha", "records.zeta"}, + }, + { + name: "disabled metadata cannot match or occupy truncation slots", + registrations: []Registration{ + validRegistration("cap.alpha", "records.alpha", "Alpha record"), + validRegistration("cap.beta", "records.beta", "Beta record"), + validRegistration("cap.disabled", "records.disabled", "Disabled record"), + }, + options: testOptions("cap.disabled"), + query: "record", + wantNames: []string{"records.alpha", "records.beta"}, + wantTruncated: false, + }, + { + name: "connector and blank queries return non-nil empty results", + registrations: []Registration{ + validRegistration("cap.alpha", "records.alpha", "Alpha record"), + }, + options: testOptions(), + query: "the and of", + wantEmpty: true, + }, + { + name: "count limit sets truncated", + registrations: []Registration{ + validRegistration("cap.zeta", "records.zeta", "Alpha Zeta record"), + validRegistration("cap.beta", "records.beta", "Alpha Beta record"), + validRegistration("cap.alpha", "records.alpha", "Alpha primary record"), + }, + options: func() Options { + options := testOptions() + options.MaxSearchResults = 2 + return options + }(), + query: "alpha", + wantNames: []string{"records.alpha", "records.beta"}, + wantTruncated: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + built, err := Build(tt.registrations, tt.options) + require.NoError(t, err) + + response, err := built.Search(tt.query) + + require.NoError(t, err) + require.NotNil(t, response.Results) + if tt.wantEmpty { + assert.Empty(t, response.Results) + assert.False(t, response.Truncated) + return + } + require.Equal(t, tt.wantNames, searchNames(response)) + assert.Equal(t, tt.wantTruncated, response.Truncated) + }) + } +} + +// TestSearchOmitsOversizedTrailingResults proves byte packing truncates without reordering. +func TestSearchOmitsOversizedTrailingResults(t *testing.T) { + summary := "token " + strings.Repeat("x", 40_000) + built, err := Build([]Registration{ + withSummary(validRegistration("cap.alpha", "records.alpha", "token"), summary), + withSummary(validRegistration("cap.beta", "records.beta", "token"), summary), + }, testOptions()) + require.NoError(t, err) + + response, err := built.Search("token") + + require.NoError(t, err) + require.NotNil(t, response.Results) + require.Equal(t, []string{"records.alpha"}, searchNames(response)) + assert.True(t, response.Truncated) +} + +// TestSearchRejectsQueryTokenOverflow proves distinct token excess uses the search-limit sentinel. +func TestSearchRejectsQueryTokenOverflow(t *testing.T) { + options := testOptions() + options.MaxSearchQueryBytes = 256 + built, err := Build([]Registration{ + validRegistration("cap.alpha", "records.alpha", "Alpha record"), + }, options) + require.NoError(t, err) + + _, err = built.Search( + "alpha bravo charlie delta echo foxtrot golf hotel india juliet kilo lima mike november oscar papa quebec", + ) + + require.Error(t, err) + require.ErrorIs(t, err, ErrSearchQueryLimit) + assert.NotContains(t, err.Error(), "bravo") +} + +// TestDescribeIgnoresSearchVocabulary proves exact lookup is unchanged by discovery metadata. +func TestDescribeIgnoresSearchVocabulary(t *testing.T) { + built, err := Build([]Registration{ + withSearchTerms(validRegistration("cap.alpha", "records.alpha", "Alpha record"), "open ticket"), + }, testOptions()) + require.NoError(t, err) + + _, foundTerm := built.Describe("open ticket") + _, foundCase := built.Describe("Records.alpha") + _, foundPrefix := built.Describe("records.alp") + description, found := built.Describe("records.alpha") + + assert.False(t, foundTerm) + assert.False(t, foundCase) + assert.False(t, foundPrefix) + require.True(t, found) + assert.Equal(t, "records.alpha", description.Name) + assert.Equal(t, "Alpha record", description.Summary) +} + +// TestBuildRejectsSearchMetadataBounds proves registration, term, and result-size ceilings fail closed. +func TestBuildRejectsSearchMetadataBounds(t *testing.T) { + valid := validRegistration("cap.one", "records.one", "First record") + tests := []struct { + name string + registrations []Registration + }{ + { + name: "too many search-term phrases", + registrations: []Registration{ + withSearchTerms(valid, makePhrases(maxSearchTermPhrases+1)...), + }, + }, + { + name: "search-term aggregate bytes", + registrations: []Registration{ + withSearchTerms(valid, strings.Repeat("token", maxSearchTermBytes)), + }, + }, + { + name: "blank search term", + registrations: []Registration{ + withSearchTerms(valid, "open ticket", ""), + }, + }, + { + name: "surrounding whitespace search term", + registrations: []Registration{ + withSearchTerms(valid, " open ticket"), + }, + }, + { + name: "single result exceeds response cap", + registrations: []Registration{ + withSummary(valid, strings.Repeat("x", maxSearchResponseBytes)), + }, + }, + { + name: "too many registrations", + registrations: overflowRegistrations(maxSupportedRegistrations + 1), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := Build(tt.registrations, testOptions()) + + require.Error(t, err) + require.ErrorIs(t, err, ErrInvalidRegistration) + }) + } +} + +// TestDisabledSearchTermsDoNotAffectRanking proves filtered vocabulary cannot change enabled order. +func TestDisabledSearchTermsDoNotAffectRanking(t *testing.T) { + options := testOptions("cap.disabled") + built, err := Build([]Registration{ + validRegistration("cap.alpha", "records.alpha", "Shared token"), + validRegistration("cap.beta", "records.beta", "Shared token"), + withSearchTerms(validRegistration("cap.disabled", "records.disabled", "Shared token"), "unique disabled term"), + }, options) + require.NoError(t, err) + + disabled, err := built.Search("unique") + require.NoError(t, err) + assert.Empty(t, disabled.Results) + assert.False(t, disabled.Truncated) + + shared, err := built.Search("shared") + require.NoError(t, err) + require.Equal(t, []string{"records.alpha", "records.beta"}, searchNames(shared)) + assert.False(t, shared.Truncated) +} + +// withSearchTerms returns a copy with explicit discovery phrases. +func withSearchTerms(registration Registration, terms ...string) Registration { + registration.SearchTerms = terms + return registration +} + +// searchNames extracts ranked result names. +func searchNames(response SearchResponse) []string { + names := make([]string, len(response.Results)) + for index, result := range response.Results { + names[index] = result.Name + } + return names +} + +// makePhrases returns count distinct short search-term phrases. +func makePhrases(count int) []string { + phrases := make([]string, count) + for index := range phrases { + phrases[index] = "term" + strings.Repeat("x", index) + } + return phrases +} + +// overflowRegistrations constructs count valid registrations with distinct names. +func overflowRegistrations(count int) []Registration { + registrations := make([]Registration, count) + for index := range registrations { + suffix := strconv.Itoa(index) + registrations[index] = validRegistration( + "cap.item"+suffix, + "records.item"+suffix, + "Item record", + ) + } + return registrations +} diff --git a/limits_test.go b/limits_test.go index e7fe59c..8ab934c 100644 --- a/limits_test.go +++ b/limits_test.go @@ -48,9 +48,11 @@ func TestBuildDefaultsZeroLimitFields(t *testing.T) { partialServer, err := partialBuilder.Build() require.NoError(t, err) - results, err := partialServer.Search("records") + response, err := partialServer.Search("records") require.NoError(t, err) - assert.Len(t, results, 1) + require.NotNil(t, response.Results) + assert.Len(t, response.Results, 1) + assert.True(t, response.Truncated) } // TestLimitsRejectNonPositiveValues proves zero never selects an unlimited budget. diff --git a/mcpserver/e2e_test.go b/mcpserver/e2e_test.go index 65ac594..836eaee 100644 --- a/mcpserver/e2e_test.go +++ b/mcpserver/e2e_test.go @@ -322,6 +322,10 @@ func TestActualMCPSecureLoop(t *testing.T) { session, err := client.Connect(t.Context(), clientTransport, nil) require.NoError(t, err) t.Cleanup(func() { _ = session.Close() }) + initialized := session.InitializeResult() + require.NotNil(t, initialized) + require.NotNil(t, initialized.ServerInfo) + assert.Equal(t, "2", initialized.ServerInfo.Version) listed, err := session.ListTools(t.Context(), &mcp.ListToolsParams{}) require.NoError(t, err) @@ -342,14 +346,15 @@ func TestActualMCPSecureLoop(t *testing.T) { require.NoError(t, err) assertSuccessfulTool(t, searched) assertNoCanary(t, searched) - searchResults := decodeStructured[[]codemode.SearchResult](t, searched) - require.Len(t, searchResults, 1) - assert.Equal(t, "records.lookup", searchResults[0].Name) - assert.Equal(t, "records.lookup(*, key: str, limit: int | None)", searchResults[0].Signature) - assert.Equal(t, "Look up one record by key.", searchResults[0].Summary) + searchResponse := decodeStructured[codemode.SearchResponse](t, searched) + require.Len(t, searchResponse.Results, 1) + assert.False(t, searchResponse.Truncated) + assert.Equal(t, "records.lookup", searchResponse.Results[0].Name) + assert.Equal(t, "records.lookup(*, key: str, limit: int | None)", searchResponse.Results[0].Signature) + assert.Equal(t, "Look up one record by key.", searchResponse.Results[0].Summary) assertDiscoveryOmitsGoTypeNames(t, searched, "lookupResult", "StatusResult") - requireNonNullJSONArray(t, searched.StructuredContent) - requireJSONTextMirror(t, searched, searchResults) + requireNonNullSearchResults(t, searched) + requireJSONTextMirror(t, searched, searchResponse) described, err := session.CallTool(t.Context(), &mcp.CallToolParams{ Name: "describe_api", @@ -385,13 +390,14 @@ func TestActualMCPSecureLoop(t *testing.T) { require.NoError(t, err) assertSuccessfulTool(t, statusSearch) assertNoCanary(t, statusSearch) - statusResults := decodeStructured[[]codemode.SearchResult](t, statusSearch) - require.Len(t, statusResults, 1) - assert.Equal(t, "health.status", statusResults[0].Name) - assert.Equal(t, "health.status()", statusResults[0].Signature) + statusResponse := decodeStructured[codemode.SearchResponse](t, statusSearch) + require.Len(t, statusResponse.Results, 1) + assert.False(t, statusResponse.Truncated) + assert.Equal(t, "health.status", statusResponse.Results[0].Name) + assert.Equal(t, "health.status()", statusResponse.Results[0].Signature) assertDiscoveryOmitsGoTypeNames(t, statusSearch, "lookupResult", "StatusResult") - requireNonNullJSONArray(t, statusSearch.StructuredContent) - requireJSONTextMirror(t, statusSearch, statusResults) + requireNonNullSearchResults(t, statusSearch) + requireJSONTextMirror(t, statusSearch, statusResponse) statusDescribed, err := session.CallTool(t.Context(), &mcp.CallToolParams{ Name: "describe_api", @@ -417,7 +423,7 @@ func TestActualMCPSecureLoop(t *testing.T) { require.NoError(t, err) assertSuccessfulTool(t, emptySearch) assertNoCanary(t, emptySearch) - requireSuccessfulStructuredValue(t, emptySearch, []any{}) + requireSuccessfulStructuredValue(t, emptySearch, map[string]any{"results": []any{}, "truncated": false}) hidden, err := session.CallTool(t.Context(), &mcp.CallToolParams{ Name: "describe_api", @@ -529,6 +535,10 @@ func TestActualMCPCompositeProgram(t *testing.T) { session, err := client.Connect(t.Context(), clientTransport, nil) require.NoError(t, err) t.Cleanup(func() { _ = session.Close() }) + initialized := session.InitializeResult() + require.NotNil(t, initialized) + require.NotNil(t, initialized.ServerInfo) + assert.Equal(t, "2", initialized.ServerInfo.Version) const ( compositeSignature = "records.search(*, count: int, active: bool, score: float, label: str | None)" @@ -541,11 +551,12 @@ func TestActualMCPCompositeProgram(t *testing.T) { }) 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) + searchResponse := decodeStructured[codemode.SearchResponse](t, searched) + require.Len(t, searchResponse.Results, 1) + assert.False(t, searchResponse.Truncated) + assert.Equal(t, "records.search", searchResponse.Results[0].Name) + assert.Equal(t, compositeSignature, searchResponse.Results[0].Signature) + assert.Equal(t, "Search records and return nested items.", searchResponse.Results[0].Summary) assertDiscoveryOmitsGoTypeNames(t, searched, "NamedItem", "searchOutput") described, err := session.CallTool(t.Context(), &mcp.CallToolParams{ @@ -766,3 +777,13 @@ func requireNonNullDescribeFieldArrays(t *testing.T, result *mcp.CallToolResult) requireNonNullJSONArray(t, object["input"]) requireNonNullJSONArray(t, object["output"]) } + +// requireNonNullSearchResults requires search structured content to carry a +// non-null results array and a Boolean truncated flag. +func requireNonNullSearchResults(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + object := requireJSONObject(t, result.StructuredContent) + requireNonNullJSONArray(t, object["results"]) + _, ok := object["truncated"].(bool) + require.True(t, ok, "truncated must be a Boolean") +} diff --git a/mcpserver/mocks/service.go b/mcpserver/mocks/service.go index 5b15918..f9ff27f 100644 --- a/mcpserver/mocks/service.go +++ b/mcpserver/mocks/service.go @@ -174,24 +174,22 @@ func (_c *MockService_Execute_Call) RunAndReturn(run func(ctx context.Context, s } // Search provides a mock function for the type MockService -func (_mock *MockService) Search(query string) ([]codemode.SearchResult, error) { +func (_mock *MockService) Search(query string) (codemode.SearchResponse, error) { ret := _mock.Called(query) if len(ret) == 0 { panic("no return value specified for Search") } - var r0 []codemode.SearchResult + var r0 codemode.SearchResponse var r1 error - if returnFunc, ok := ret.Get(0).(func(string) ([]codemode.SearchResult, error)); ok { + if returnFunc, ok := ret.Get(0).(func(string) (codemode.SearchResponse, error)); ok { return returnFunc(query) } - if returnFunc, ok := ret.Get(0).(func(string) []codemode.SearchResult); ok { + if returnFunc, ok := ret.Get(0).(func(string) codemode.SearchResponse); ok { r0 = returnFunc(query) } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).([]codemode.SearchResult) - } + r0 = ret.Get(0).(codemode.SearchResponse) } if returnFunc, ok := ret.Get(1).(func(string) error); ok { r1 = returnFunc(query) @@ -225,12 +223,12 @@ func (_c *MockService_Search_Call) Run(run func(query string)) *MockService_Sear return _c } -func (_c *MockService_Search_Call) Return(vs []codemode.SearchResult, err error) *MockService_Search_Call { - _c.Call.Return(vs, err) +func (_c *MockService_Search_Call) Return(v codemode.SearchResponse, err error) *MockService_Search_Call { + _c.Call.Return(v, err) return _c } -func (_c *MockService_Search_Call) RunAndReturn(run func(query string) ([]codemode.SearchResult, error)) *MockService_Search_Call { +func (_c *MockService_Search_Call) RunAndReturn(run func(query string) (codemode.SearchResponse, error)) *MockService_Search_Call { _c.Call.Return(run) return _c } diff --git a/mcpserver/server.go b/mcpserver/server.go index 15df710..ce36deb 100644 --- a/mcpserver/server.go +++ b/mcpserver/server.go @@ -67,13 +67,16 @@ func New(service Service, resolver InvocationResolver) (*mcp.Server, error) { return nil, fmt.Errorf("%w: invocation resolver is required", codemode.ErrInvalidRegistration) } - searchOutputSchema, schemaErr := jsonschema.For[[]codemode.SearchResult](nil) + searchOutputSchema, schemaErr := jsonschema.For[codemode.SearchResponse](nil) if schemaErr != nil { return nil, fmt.Errorf("%w: search_api output schema: %w", codemode.ErrInvalidRegistration, schemaErr) } - schemaErr = requireNonNullArray(searchOutputSchema) + if searchOutputSchema == nil || searchOutputSchema.Properties == nil { + return nil, fmt.Errorf("%w: search_api output schema: missing properties", codemode.ErrInvalidRegistration) + } + schemaErr = requireNonNullArray(searchOutputSchema.Properties["results"]) if schemaErr != nil { - return nil, fmt.Errorf("%w: search_api output schema: %w", codemode.ErrInvalidRegistration, schemaErr) + return nil, fmt.Errorf("%w: search_api output schema: results: %w", codemode.ErrInvalidRegistration, schemaErr) } schemaErr = requireResolvedSchema(searchOutputSchema) if schemaErr != nil { @@ -110,10 +113,10 @@ func New(service Service, resolver InvocationResolver) (*mcp.Server, error) { } bound := &adapter{service: service, resolver: resolver} - server := mcp.NewServer(&mcp.Implementation{Name: "codemode", Version: "1"}, nil) + server := mcp.NewServer(&mcp.Implementation{Name: "codemode", Version: "2"}, nil) mcp.AddTool(server, &mcp.Tool{ Name: "search_api", - Description: "Search enabled names and summaries with a short literal substring. Retry an empty result with a shorter term.", + Description: "Search enabled capabilities using task, resource, or exact-name vocabulary. Results are relevance-ranked. Pass the exact returned name to describe_api. If truncated is true and no result fits, submit a more specific task/resource query.", OutputSchema: searchOutputSchema, }, bound.search) mcp.AddTool(server, &mcp.Tool{ @@ -134,15 +137,15 @@ func (bound *adapter) search( ctx context.Context, _ *mcp.CallToolRequest, input searchInput, -) (*mcp.CallToolResult, []codemode.SearchResult, error) { - outcome := runToolOperation(func() ([]codemode.SearchResult, error) { +) (*mcp.CallToolResult, codemode.SearchResponse, error) { + outcome := runToolOperation(func() (codemode.SearchResponse, error) { if _, err := resolveSubject(ctx, bound.resolver); err != nil { - return nil, err + return codemode.SearchResponse{}, err } return bound.service.Search(input.Query) }) if outcome.err == nil { - outcome.value = nonNilSlice(outcome.value) + outcome.value.Results = nonNilSlice(outcome.value.Results) } return nil, outcome.value, outcome.err } diff --git a/mcpserver/server_test.go b/mcpserver/server_test.go index 49cdfbc..982af61 100644 --- a/mcpserver/server_test.go +++ b/mcpserver/server_test.go @@ -85,8 +85,14 @@ func TestNewRegistersExactlyThreeTools(t *testing.T) { assertOutput func(*testing.T, map[string]any) }{ { - name: "search_api", - cues: []string{"short literal substring", "shorter term"}, + name: "search_api", + cues: []string{ + "task, resource, or exact-name vocabulary", + "relevance-ranked", + "exact returned name to describe_api", + "truncated is true", + "more specific task/resource query", + }, assertOutput: requireSearchAPIOutputSchema, }, { @@ -170,10 +176,12 @@ func TestToolsResolveSubjectBeforeServiceWork(t *testing.T) { events = append(events, "resolve") return authz.Subject{ID: "subject-1"}, nil }).Times(3) - service.EXPECT().Search("lookup").RunAndReturn(func(string) ([]codemode.SearchResult, error) { + service.EXPECT().Search("lookup").RunAndReturn(func(string) (codemode.SearchResponse, error) { events = append(events, "search") - return []codemode.SearchResult{ - {Name: "records.lookup", Signature: "records.lookup()", Summary: "lookup"}, + return codemode.SearchResponse{ + Results: []codemode.SearchResult{ + {Name: "records.lookup", Signature: "records.lookup()", Summary: "lookup"}, + }, }, nil }).Once() service.EXPECT(). @@ -230,7 +238,7 @@ func TestToolsIgnoreUntrustedClientMetadata(t *testing.T) { service := mocks.NewMockService(t) resolver := mocks.NewMockInvocationResolver(t) resolver.EXPECT().Resolve(mock.Anything).Return(authz.Subject{ID: "subject-1"}, nil).Once() - service.EXPECT().Search("lookup").Return([]codemode.SearchResult{}, nil).Once() + service.EXPECT().Search("lookup").Return(codemode.SearchResponse{Results: []codemode.SearchResult{}}, nil).Once() session := newTestSession(t, service, resolver) result, err := session.client.CallTool(t.Context(), &mcp.CallToolParams{ @@ -241,7 +249,7 @@ func TestToolsIgnoreUntrustedClientMetadata(t *testing.T) { require.NoError(t, err) require.False(t, result.IsError) - assert.Equal(t, []any{}, result.StructuredContent) + assert.Equal(t, map[string]any{"results": []any{}, "truncated": false}, result.StructuredContent) } // TestToolsSerializeSuccessfulEmptySlicesAsArrays proves successful nil and empty @@ -268,34 +276,64 @@ func TestToolsSerializeSuccessfulEmptySlicesAsArrays(t *testing.T) { tool: "search_api", arguments: map[string]any{"query": "lookup"}, configure: func(service *mocks.MockService) { - service.EXPECT().Search("lookup").Return(nil, nil).Once() + service.EXPECT().Search("lookup").Return(codemode.SearchResponse{}, nil).Once() }, - want: []any{}, + want: map[string]any{"results": []any{}, "truncated": false}, }, { name: "search empty results", tool: "search_api", arguments: map[string]any{"query": "lookup"}, configure: func(service *mocks.MockService) { - service.EXPECT().Search("lookup").Return([]codemode.SearchResult{}, nil).Once() + service.EXPECT().Search("lookup").Return(codemode.SearchResponse{ + Results: []codemode.SearchResult{}, + }, nil).Once() }, - want: []any{}, + want: map[string]any{"results": []any{}, "truncated": false}, }, { name: "search populated results", tool: "search_api", arguments: map[string]any{"query": "lookup"}, configure: func(service *mocks.MockService) { - service.EXPECT().Search("lookup").Return([]codemode.SearchResult{ - {Name: "records.lookup", Signature: "records.lookup()", Summary: "lookup"}, + service.EXPECT().Search("lookup").Return(codemode.SearchResponse{ + Results: []codemode.SearchResult{ + {Name: "records.lookup", Signature: "records.lookup()", Summary: "lookup"}, + }, }, nil).Once() }, - want: []any{ - map[string]any{ - "name": "records.lookup", - "signature": "records.lookup()", - "summary": "lookup", + want: map[string]any{ + "results": []any{ + map[string]any{ + "name": "records.lookup", + "signature": "records.lookup()", + "summary": "lookup", + }, }, + "truncated": false, + }, + }, + { + name: "search truncated results", + tool: "search_api", + arguments: map[string]any{"query": "lookup"}, + configure: func(service *mocks.MockService) { + service.EXPECT().Search("lookup").Return(codemode.SearchResponse{ + Results: []codemode.SearchResult{ + {Name: "records.lookup", Signature: "records.lookup()", Summary: "lookup"}, + }, + Truncated: true, + }, nil).Once() + }, + want: map[string]any{ + "results": []any{ + map[string]any{ + "name": "records.lookup", + "signature": "records.lookup()", + "summary": "lookup", + }, + }, + "truncated": true, }, }, { @@ -468,7 +506,7 @@ func TestToolsProjectStableServiceErrors(t *testing.T) { configure: func(service *mocks.MockService) { service.EXPECT(). Search("oversized"). - Return(nil, fmt.Errorf("trusted budget: %w", codemode.ErrResourceLimit)). + Return(codemode.SearchResponse{}, fmt.Errorf("trusted budget: %w", codemode.ErrResourceLimit)). Once() }, want: codemode.ErrResourceLimit.Error(), @@ -514,7 +552,10 @@ func TestToolsProjectStableServiceErrors(t *testing.T) { tool: "search_api", arguments: map[string]any{"query": "lookup"}, configure: func(service *mocks.MockService) { - service.EXPECT().Search("lookup").Return(nil, errors.New("trusted stack dump")).Once() + service.EXPECT(). + Search("lookup"). + Return(codemode.SearchResponse{}, errors.New("trusted stack dump")). + Once() }, want: codemode.ErrInternal.Error(), }, @@ -607,7 +648,7 @@ func TestToolsSanitizePanics(t *testing.T) { resolver.EXPECT().Resolve(mock.Anything).Return(authz.Subject{ID: "subject-1"}, nil).Once() service.EXPECT().Search("lookup").Run(func(string) { panic("trusted service panic") - }).Return(nil, nil).Once() + }).Return(codemode.SearchResponse{}, nil).Once() }, }, } @@ -654,6 +695,10 @@ func newTestSession(t *testing.T, service mcpserver.Service, resolver mcpserver. t.Cleanup(func() { _ = clientSession.Close() }) + initialized := clientSession.InitializeResult() + require.NotNil(t, initialized) + require.NotNil(t, initialized.ServerInfo) + assert.Equal(t, "2", initialized.ServerInfo.Version) return &testSession{client: clientSession} } @@ -691,11 +736,17 @@ func listedOutputSchema(t *testing.T, tools []*mcp.Tool, name string) map[string return requireJSONObject(t, listedTool(t, tools, name).OutputSchema) } -// requireSearchAPIOutputSchema requires search_api to advertise a non-null SearchResult array. +// requireSearchAPIOutputSchema requires search_api to advertise an object whose +// required results array is non-null and whose truncated flag is Boolean. func requireSearchAPIOutputSchema(t *testing.T, schema map[string]any) { t.Helper() - requireNonNullJSONType(t, schema, "array") - requireSearchResultItemSchema(t, schema["items"]) + requireNonNullJSONType(t, schema, "object") + requireRequiredNames(t, schema, "results", "truncated") + properties := requireJSONObject(t, schema["properties"]) + results := requireJSONObject(t, properties["results"]) + requireNonNullJSONType(t, results, "array") + requireSearchResultItemSchema(t, results["items"]) + requireNonNullJSONType(t, requireJSONObject(t, properties["truncated"]), "boolean") } // requireDescribeAPIOutputSchema requires describe_api to advertise an object whose diff --git a/mcpserver/service.go b/mcpserver/service.go index a129a5a..29e4efd 100644 --- a/mcpserver/service.go +++ b/mcpserver/service.go @@ -12,8 +12,8 @@ import ( // The root *codemode.Server implements this port. The adapter does not re-enforce // catalog bounds, hidden-capability filtering, or execution restrictions. type Service interface { - // Search returns a bounded name-sorted scan of enabled capability names and summaries. - Search(query string) ([]codemode.SearchResult, error) + // Search returns a bounded relevance-ranked scan of enabled capabilities. + Search(query string) (codemode.SearchResponse, error) // Describe returns one exact enabled capability description or a not-found error. Describe(name codemode.CapabilityName) (codemode.Description, error) diff --git a/server.go b/server.go index 70f6fd0..0fd719c 100644 --- a/server.go +++ b/server.go @@ -17,6 +17,9 @@ type Program string // SearchResult is one compact enabled-capability discovery record. type SearchResult = catalog.SearchResult +// SearchResponse is one bounded ranked discovery result set. +type SearchResponse = catalog.SearchResponse + // Description is one exact enabled-capability description and supported binding shape. type Description = catalog.Description @@ -84,19 +87,19 @@ func capabilityBindings(capabilityCatalog *catalog.Catalog) []execution.Capabili return bindings } -// Search returns a bounded name-sorted scan of enabled capability names and summaries. -func (server *Server) Search(query string) ([]SearchResult, error) { +// Search returns a bounded relevance-ranked scan of enabled capabilities. +func (server *Server) Search(query string) (SearchResponse, error) { if server == nil || server.catalog == nil { - return nil, ErrInternal + return SearchResponse{}, ErrInternal } - results, err := server.catalog.Search(query) + response, err := server.catalog.Search(query) if err != nil { if errors.Is(err, catalog.ErrSearchQueryLimit) { - return nil, ErrResourceLimit + return SearchResponse{}, ErrResourceLimit } - return nil, ErrInternal + return SearchResponse{}, ErrInternal } - return results, nil + return response, nil } // Describe returns one exact enabled capability description or ErrNotFound. diff --git a/server_test.go b/server_test.go index 975823b..f0d4190 100644 --- a/server_test.go +++ b/server_test.go @@ -31,10 +31,12 @@ func TestServerSearchAndDescribeExposeOnlyEnabledCapabilities(t *testing.T) { server, err := builder.Build() require.NoError(t, err) - results, err := server.Search("record") + response, err := server.Search("record") require.NoError(t, err) - require.Len(t, results, 1) - assert.Equal(t, "records.alpha", results[0].Name) + require.NotNil(t, response.Results) + require.Len(t, response.Results, 1) + assert.Equal(t, "records.alpha", response.Results[0].Name) + assert.False(t, response.Truncated) description, err := server.Describe("records.alpha") require.NoError(t, err) assert.Equal(t, "records.alpha", description.Name) @@ -55,9 +57,11 @@ func TestCapabilityIDDefaultsToName(t *testing.T) { server, err := builder.Build() require.NoError(t, err) - results, err := server.Search("") + response, err := server.Search("") require.NoError(t, err) - assert.Empty(t, results) + require.NotNil(t, response.Results) + assert.Empty(t, response.Results) + assert.False(t, response.Truncated) }) t.Run("authorization", func(t *testing.T) { @@ -89,16 +93,29 @@ def main(): // TestServerSearchProjectsQueryLimits proves internal discovery details do not escape the public taxonomy. func TestServerSearchProjectsQueryLimits(t *testing.T) { - limits := codemode.DefaultLimits() - limits.MaxSearchQueryBytes = 4 - builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: limits}) - server, err := builder.Build() - require.NoError(t, err) + t.Run("raw query bytes", func(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxSearchQueryBytes = 4 + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: limits}) + server, err := builder.Build() + require.NoError(t, err) - _, err = server.Search("oversized") + _, err = server.Search("oversized") - require.ErrorIs(t, err, codemode.ErrResourceLimit) - assert.Equal(t, codemode.ErrResourceLimit.Error(), err.Error()) + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit.Error(), err.Error()) + }) + + t.Run("distinct query tokens", func(t *testing.T) { + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll()}) + server, err := builder.Build() + require.NoError(t, err) + + _, err = server.Search("t0 t1 t2 t3 t4 t5 t6 t7 t8 t9 t10 t11 t12 t13 t14 t15 t16") + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit.Error(), err.Error()) + }) } // TestServerExecuteAuthorizesCanonicalArgumentsBeforeDispatch proves the complete native-call ordering.