feat(mcp): support client-side sampling with audio - #1972
Open
wqymi wants to merge 7 commits into
Open
Conversation
MCP servers can now ask MiMoCode to run a model call on their behalf via
`sampling/createMessage`, so a local server never needs its own API key. The
motivating case is the MiMo Cut MCP server: it extracts a 16 kHz mono WAV from a
video and asks MiMoCode to transcribe it using the user's existing `mimo-v2.5`
connection.
MiMoCode declares `capabilities.sampling = {}` during initialize and registers
the SDK's native `setRequestHandler(CreateMessageRequestSchema, ...)`. Work runs
on a fresh root Effect fiber, so a sampling request arriving while we are still
awaiting that server's `tools/call` cannot self-lock. `sampling.tools` and
`sampling.context` stay undeclared, and `tools`/`toolChoice` are rejected per
spec. `includeContext` is never honoured beyond `none`.
Model selection goes through a new Model Capability Registry: MCP publishes no
model list and no modality discovery, so eligibility is derived from the actual
content and decided by capability plus configured credentials FIRST, with
`modelPreferences.hints` only ranking the already-eligible set. A hint can never
widen eligibility. When nothing is compatible the server gets a structured error
naming every rejected model and why; audio is never dropped, downgraded to text,
or sent to a model that cannot accept it.
Default policy is ask, per MCP server, with the prompt showing the server,
target model, content types, audio size and prompt previews. Credentials never
appear in responses, logs or errors.
`sampling-e2e.test.ts` imported `wav()` from `sampling.test.ts`, which made Bun attribute all 20 of that file's tests to the importer: the e2e file reported 36 tests (its own 16 plus 20 pulled in) while `sampling.test.ts` reported none of its own in any combined run. Every test still executed exactly once — the scoped suite is 695 either way — but per-file totals were misleading. A plain helper module keeps the attribution honest.
…e SDK's id-0 cancellation Three gaps that survived green CI, each closed by driving the path rather than asserting the shape. 1. `unknown`-modality fail-closed was asserted but never revert-probed. Collapsing `unknown` into an eligible state (`continue` in place of `rejectionFor`'s unknown branch) left BOTH declaration assertions green — they pin the adapter table, not the decision — and `selectModel`, the function the sampling handler actually calls, had no `unknown` case at all. Only the reason-kind assertion bit. Added: the gate's refusal and its operator-facing wording, a hint that must not promote an unproven model, and an end-to-end refusal through the handler using `@ai-sdk/mistral` (bundled, so offline; undeclared, so `unknown`). Under the probe the E2E returned -32603 from the adapter instead of a clean -32602 refusal. 2. The request-timeout path was implemented and untested; only cancellation was exercised. Added a test where the provider accepts the request and never answers: the request is reaped at the bound, answered with our own timeout error, and — the part that was unasserted — the fiber is removed from the in-flight set, checked synchronously rather than polled. The bound is a parameter on `serve` (default unchanged, production passes nothing) so it can be driven in CI time; widening it to 60s fails the test with the SDK's own "Request timed out" instead. 3. `@modelcontextprotocol/sdk@1.27.1` drops a `notifications/cancelled` whose `requestId` is 0. Pinned with three cases reading the requestId off the wire. The measurement that settles the design question: `_requestMessageId` is per Protocol instance, so the id at risk is the SERVER's outgoing counter — spending an id from the client side leaves it at 0 and the cancellation still does not land. Burning id 0 at connection setup is therefore not a costly option, it is an ineffective one. Behaviour left alone with the 120s timeout as the documented bound; patching `_oncancel` flips cases 1 and 2 and fails here, so an SDK upgrade is visible.
… permission deny win Three review findings on the sampling path. 1. `mcp_sampling` was listed as "Simple action-only" in the bundled permissions reference while the example two paragraphs below it used the glob-map form. The handler passes `patterns: [server]`, so the rule IS glob-keyed. Because the .bundle docs are fed to the model as instructions, the misclassification taught it that a valid config shape is invalid. 2. `Effect.timeoutOption` and `cancelAll` interrupted the handler fiber but did not cancel the provider HTTP call already in flight inside it, so every timed-out or abandoned sampling request leaked a live — billable — model call. Only the SDK's per-request signal reached the provider. The provider now gets the union of that signal and the handler fiber's own, so the 120 s bound and a client teardown both stop the call rather than only stopping our waiting for it. The PR body and reference doc claimed the bound already did this; corrected. 3. `mcp.<server>.sampling: "allow"` skips the approval prompt, so an explicit `permission.mcp_sampling` deny was never evaluated and the model ran anyway — a security control that read as configured and did nothing. `handle` now evaluates the ruleset itself and a deny from either control refuses, matching the precedence the permission service already applies internally. Tests: the timeout and cancelAll cases now assert the abort reached the provider stub, not our own bookkeeping. Reverting the abort composition fails both; reverting the deny check makes the denied request succeed.
…parking the ask
`Permission.ask` raced the approval Deferred against the caller's
abortSignal with `Effect.race`. `Effect.race` resolves with the first
*success* and treats a failure as "not a winner", so it waits on the
loser. A human rejection FAILS the Deferred and the abort side never
settles on its own, so any ask that was passed a signal parked forever
on a rejection. Use `Effect.raceFirst`, which resolves with the first
side to *complete*, success or failure.
Two callers pass a signal: MCP sampling (`extra.signal`, always present)
and the tool-context `ask` in session/prompt.ts, which every tool's
`ctx.ask` funnels through. Sampling's `-1` "declined" error was therefore
unreachable in practice — a rejected sampling prompt produced no answer
at all and the request sat in the in-flight set until the timeout bound
reaped it.
The forwarded-ask timeout had the same misuse: a denial failed the
Deferred, counted as no winner, and the caller waited out the whole
FORWARD_DENY_TIMEOUT_MS before seeing the rejection it already had.
Interruption still composes. Measured on effect@4.0.0-beta.48:
`race(failed Deferred, never)` never settles while `raceFirst` yields the
error, and interrupting a `raceFirst` fiber exits with a cause for which
`Cause.hasInterrupts` is true rather than a plain failure — which is why
this is not done by wrapping a side in `Effect.exit`. The abort
listener's cleanup still runs on interrupt, so the no-leak guard in
test/permission/abort.test.ts continues to hold.
Tests, in test/mcp/sampling-e2e.test.ts:
- a genuine `reply({ reply: "reject" })` must answer the server with the
declined error (`-1` + "declined" + `data.server`, since our
request-timeout error and the SDK's own both use -32001 and both carry
`data.timeout`) and drain `inFlightCount` back to 0. Reverting the fix
fails it in 8.6s with code -32001: the bound reaped the parked ask.
The bound is injected so the pre-fix run fails fast instead of hanging;
the test comment says so rather than implying an unbounded proof.
- the mirror case: an abort fired while the prompt is pending must be
abandoned promptly, not at the bound. Labelled in-file as a regression
guard, not a reproducer, because it passes before the fix too — an
abort failed BOTH sides of the race, which is the asymmetry.
Also replaces the stale "NOT COVERED HERE" note that documented this
defect as out of scope.
…mer alive One wall-clock bound wrapped the human approval wait and the model call together, and nothing kept the counterparty's own timer alive. - Emit periodic notifications/progress during the model call, but only when the server minted a progress token. Named liveness, not progress: the call is generateText, so no completion fraction exists and `total` is omitted. Necessary, not sufficient — resetTimeoutOnProgress is the requester's choice and defaults to false. - Bound the approval wait and the model call separately, so a slow human no longer eats the model's budget, and keep their sum as an absolute ceiling over the stretch neither covers. Heartbeats reset the peer's timer only, never ours, so a hung provider is still cut off. - Say which phase expired, in the message and in data.phase. data.server stays on every path; it is the only discriminator against the SDK's own -32001. - Record that DEFAULT_SAMPLING_TIMEOUT has no derivation and exceeds the SDK's 60 s default, and that the 30 s approval default is a placeholder awaiting a product decision.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this adds
MCP client-side sampling (
sampling/createMessage), so an MCP server can ask MiMoCode to run a model on its behalf and never needs its own API key. Driving case: MiMo Cut hands us a 16 kHz mono WAV and gets a transcript back through the user's already-configured provider.Two parts:
src/mcp/sampling.ts— declarescapabilities.sampling = {}, serves the request, enforces permission, converts content, calls the model, tracks in-flight work for shutdown.src/provider/capability-registry.ts— decides which model may serve a request, given the modalities that request actually needs.Requests accept text, image and audio. Responses are text only.
Design decisions a reviewer should check
unknowncapability fails closed.Supportis tri-state (supported/unsupported/unknown); an unknown modality is refused, not attempted. The consumer (selectModel) is what has the assertions — a positive test on the declaration table would not have covered the decision that reads it.modelPreferences.hintsreorder the eligible set and can never widen it. Ineligible stays ineligible; an empty eligible set returns a structured error naming the per-model reasons and the derived requirements.capabilities.input.denyoutrankssampling: "allow". Enforced inhandlebefore content validation and model selection, because theallowfast path is precisely what skipsask— so that is the only point where both controls are visible.ask, per MCP server, glob-keyed by server name.raceFirstis why.Permission.askraced the approvalDeferredagainst the caller'sabortSignalwithEffect.race, which resolves on the first success and treats a failure as "not a winner" — so it waited on the loser. A rejection fails theDeferredand the abort side never settles on its own, so the ask parked forever. Sampling always passesextra.signal, which made the-1"declined" error unreachable in practice: a rejected prompt produced no answer at all and the request sat in the in-flight set until the timeout bound reaped it as-32001. NowEffect.raceFirst(first side to complete, success or failure). Two things to check: this is deliberately not done by wrapping a side inEffect.exit, because interruption must stay interruption — measured, an interruptedraceFirstexits with a cause for whichCause.hasInterruptsis true — and the same misuse on the forwarded-ask timeout made a denial wait out the whole 5-minute bound, so it is fixed there too. The defect reaches everyPermission.askcaller that passes a signal: sampling, plus the tool-contextaskinsession/prompt.tsthat every tool'sctx.askfunnels through.-32001 Request timed out).Known limitations
CreateMessageResult; we produce text.sampling.toolsandsampling.contextare deliberately not declared. Not declaringcontextis what keeps session context from being sent to MCP servers;includeContextdefaults tonone.0.@modelcontextprotocol/sdkprotocol.js:170opens_oncancelwithif (!notification.params.requestId) return, so the first server-initiated request of a connection cannot be cancelled by the server. Present identically in 1.28.0 and 1.30.0, so upgrading does not help, and there is no client-side workaround (_requestMessageIdis per-instance and counts only requests that instance sends). Our request timeout caps the resulting waste; it does not avoid it. Three wire-level tests pin the upstream behaviour so a future SDK fix surfaces as a test change rather than silently.test/provider/capability-registry-wire.test.ts, not asserted in prose: an adapter behaviour change fails CI rather than making the registry lie.Operator-facing contract — permission forms, error-code mapping, size limits, and the cancellation residual — lives in
src/skill/builtin/.bundle/mimocode-docs/reference/mcp-sampling.md.Verification
Every behavioural test here was revert-probed; where a path genuinely cannot be exercised in CI (the 120 s production default) the test says so in-file rather than implying coverage.