fix(bedrock-kb-retrieval-mcp-server): support managed knowledge bases and add agentic retrieval

[ankdav-amazon]

Fixes

Summary

Amazon Bedrock managed knowledge bases (knowledgeBaseConfiguration.type == "MANAGED")
do not work with this server today, and there is no way to get a generated answer from one.
This makes both work.

Changes

1. QueryKnowledgeBases failed on every call against a managed knowledge base

retrieval.py always sent retrievalConfiguration={'vectorSearchConfiguration': {...}},
which Retrieve rejects for a managed knowledge base:

ValidationException: Incompatible configuration: vectorSearchConfiguration is not
supported for managed knowledge bases. Use managedSearchConfiguration instead.

ListKnowledgeBases discovers such a knowledge base happily, so it looks usable and then
fails on every query.

There is a quieter second half. Managed knowledge bases expose data-source identity as
_data_source_id, not the x-amz-bedrock-kb-* reserved keys. Filtering a managed
knowledge base on the vector key is accepted by the API and matches nothing, so
data_source_ids returns an empty result set with no error and the caller concludes the
data source is empty. Fixing only the configuration key would leave data_source_ids
silently broken.

• Detect the knowledge base type and send the matching search configuration. The type comes
  from the get_knowledge_base call discover_knowledge_bases already makes for the ARN,
  so the common path costs no extra API call, and it is cached per knowledge base id.
• Filter on the metadata key appropriate to the type.
• Nest rerankingConfiguration under whichever search configuration is in use.
• Degrade gracefully when the type cannot be determined (for example a caller without
  bedrock:GetKnowledgeBase): retry once with the other shape if the API rejects it by
  name, then cache what that retry proved so later calls go direct. Errors unrelated to
  configuration shape are not retried, and neither is a mismatch on a knowledge base whose
  type was already known.
• ListKnowledgeBases reports each knowledge base's type, so a client can tell which
  tool applies.
• Raise the boto3 floor to >=1.43.0. managedSearchConfiguration first appears in the
  Retrieve input model there; 1.42.50 does not have it, and the previously declared floor
  resolves to a version that fails client-side parameter validation before a request is sent.

Vector knowledge base behaviour is unchanged.

2. New tool: AgenticQueryKnowledgeBases

RetrieveAndGenerate rejects managed knowledge bases outright
(ValidationException: This operation is not supported for managed knowledge bases.), so
AgenticRetrieveStream is the only way to get a generated answer from one. This server did
not expose it.

AgenticRetrieveStream is a streaming operation whereas an MCP tool call returns a single
result, so the stream is consumed server-side and returned as one aggregate. Observed shape
for a default call: 4 trace events, ~440 responseEvent answer chunks, and a single
terminal result event. The terminal event already carries the complete answer and
citations, so the streamed chunks are only a fallback. Raw events are deliberately not
forwarded — there are hundreds per call and they would crowd out the content the model
actually needs.

Returns a single JSON object: results always; answer and citations when
generate_response is true; trace when include_trace is true; warnings, failures
and nextToken only when non-empty.

• Managed knowledge bases only — verified: a vector knowledge base returns
  Knowledge Base with id ... is not supported for Agent.... The type is checked up front
  and reported actionably, reusing the detection above.
• Data-source filtering goes through retrievalOverrides with the managed metadata key.
• Several knowledge base ids become several retrievers in one call.
• Accepts next_token, so the nextToken it surfaces is actually usable.
• Modelled error events inside the stream are raised rather than silently dropped.
• byteContent is a Blob and not JSON serialisable; its presence is flagged rather than
  embedded.

3. ACL-protected content was unreachable

Neither tool exposed userContext, so content in data sources with per-document ACLs
(SharePoint, OneDrive, Confluence) could not be retrieved. Agentic retrieval degraded
partially rather than loudly: its full-document expansion step failed with
UserContext is required for ACL-aware data sources while still returning a partial result.

Both tools now take an optional user_id.

4. Reranking availability is per model, not per region

A single flat region allowlist let reranking_model_name='AMAZON' through in us-east-1,
where that model is not offered, so the request failed at the API with an opaque error
instead of failing fast. Verified with ListFoundationModels across 12 regions:
amazon.rerank-v1:0 is offered in us-west-2, eu-central-1, ca-central-1 and
ap-northeast-1 but not us-east-1; cohere.rerank-v3-5:0 is offered in all five.
Validation is now per (region, model) and names the supported regions.

User experience

Before — with a managed knowledge base:

• QueryKnowledgeBases fails 100% of the time with a ValidationException
• No way to get a generated answer
• ACL-protected content is unreachable
• AMAZON reranking in us-east-1 fails with an opaque error

After:

• QueryKnowledgeBases works, and data_source_ids actually scopes results
• AgenticQueryKnowledgeBases returns a cited answer, or results only via
  generate_response=false
• user_id reaches ACL-protected content, filtered to that user's authorised subset
• Unavailable (region, model) reranking pairs fail immediately with the supported list
• ListKnowledgeBases reports type, so a client can pick the right tool

Testing

75 unit tests pass; ruff check and ruff format clean.

Unit tests alone were not treated as sufficient: mocking the runtime client hides both
Retrieve bugs, since neither client-side parameter validation nor the server's filter-key
semantics are exercised. So everything was also verified against live managed and vector
knowledge bases in us-west-2, both directly and end to end over MCP JSON-RPC.

End to end over MCP — 10/10: tools registered; ListKnowledgeBases reports type for
both kinds; Retrieve on managed (previously always failed); Retrieve on vector (unchanged);
Retrieve data_source_ids on managed returns results (previously 0, silently); agentic
answer with 13 citations over 10 results; generate_response=false honouring
number_of_results; include_trace; agentic data_source_ids scoped to exactly the
requested source; vector knowledge base refused with an actionable message.

Input-combination matrix — 34/34, covering Retrieve × knowledge base type × reranking
model, data-source filters (valid / multiple / non-existent), number_of_results bounds
(1 / 10 / 100 / 101 / 0), type-detection modes (management client present, absent, bogus id),
and agentic × generate_response × include_trace × iteration caps × retriever sets
(one, two, managed+vector, vector only, empty, bogus).

Specific behaviours confirmed live rather than only mocked:

• Reranking under managedSearchConfiguration succeeds with both AMAZON and COHERE
  on a managed knowledge base.
• The filter-key fix genuinely discriminates: filtering to each of three data sources in
  turn returned results only from that source (20 / 20 / 12), a non-existent id returned 0,
  and the vector key returned 0 against the same knowledge base that returns 20 unfiltered.
• The fallback retry recovers and then stops paying for it: with no management client the
  vector shape is attempted, rejected, retried as managed, and the learned type cached.
• user_id resolves the ACL failure: without it the run reports
  UserContext is required for ACL-aware data sources; with it the failure clears and
  results narrow to the user's authorised subset.
• Multiple knowledge bases in one call returned results spanning both.

Not reproducible in my environment, and therefore only unit tested: warning/failure
surfacing, binary-content flagging, in-stream modelled error events, and nextToken
pagination (the service never returned a token for my corpus).

Checklist

If your change doesn't seem to apply, please leave them unchecked.

• I have reviewed the contributing guidelines
• I have performed a self-review of this change
• Changes have been tested
• Changes are documented

Is this a breaking change? (Y/N) N

New parameters are optional and appended after existing ones, so positional callers are
unaffected and omitting them preserves current behaviour. Vector knowledge base behaviour is
unchanged. The reranking region check now refuses an unavailable (region, model) pair that
previously passed the check and then failed at the API — a clearer failure for a request that
could not have succeeded.

RFC issue number: n/a — this is a fix plus a tool on an existing server, not a new server.

Checklist:

• Migration process documented
• Implement warnings (if it can live side by side)

Notes for reviewers

• memoryConfiguration exists on AgenticRetrieveStream in newer models but is not in the
  public botocore version pinned here, so it is not exposed. Straightforward to add later.
• The agentic tool is registered unconditionally. If you would prefer it behind an
  environment flag like BEDROCK_KB_RERANKING_ENABLED, say so and I will gate it — it
  invokes a foundation model and so costs more per call than QueryKnowledgeBases.

Acknowledgment

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.

[ankdav-amazon]

Additional live validation, beyond the end-to-end table in the description.

Fallback retry path, verified against a live managed knowledge base. The description
noted this was unit tested; it is now also confirmed live. Calling with no management
client leaves the type unknown, so the vector shape is attempted first and recovered from:

Knowledge base <id> rejected the vectorSearchConfiguration shape; retrying as managed.

Results come back on the retry. The same call with data_source_ids also returns
results, which specifically verifies the retry rebuilds the request with the managed
metadata key — had it kept the vector key, the response would have been silently empty.

Multiple knowledge bases in one call. Two managed knowledge bases returned 11 results
spanning both, confirming each id becomes its own retriever end to end rather than only in
the unit test.

max_agent_iterations verified through the tool against a live knowledge base.

Reranking — clarifying the caveat in the description. In my test account reranking
fails with AccessDeniedException for both knowledge base types. The vector path is
unmodified by this PR and fails identically, with the cause visible in the error (the
knowledge base execution role lacks the rerank permission). So this is an account-level
model-access limitation rather than a regression here. Worth noting the managed request
passes service validation and only fails at authorization, which indicates the
rerankingConfiguration nesting is structurally accepted — but I still cannot demonstrate
a successful managed rerank, so a reviewer with rerank model access confirming that would
be valuable.

Remaining mock-only coverage, for transparency: a successful explicit
rerankingConfiguration, warning/failure surfacing, binary-content flagging, in-stream
modelled error events, and nextToken pagination. Each of these is unit tested; none were
reproducible on demand in my environment.