feat(mcp): support client-side sampling with audio - #1972
Merged
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.
…guessed The sampling model call was `generateText`. Non-streaming makes "has this hung?" unanswerable from inside the call, so the only available defence was a total wall-clock bound — and the number filling that role, 120 s, has no derivation, as its own comment says. Adding the progress keepalive made that worse rather than better: before it, a peer on the SDK's 60 s default gave up first and our 120 s almost never fired; after it, a peer that opted into `resetTimeoutOnProgress` never fires at all, which promoted an undetermined number to the only thing between a hung provider and forever. Switch to `streamText` and use the stream as evidence: - A stall bound (45 s, injectable) fails with `data.phase: "stall"` as soon as the model produces nothing for that long, covering both the wait for the first chunk and every gap after it; each chunk resets it. `data.chunks` and `data.characters` ride along, so `chunks: 0` distinguishes a provider that never produced anything from one that started and died. - The 120 s bound keeps its value and loses its job: it is now the backstop for a stream that trickles without ever finishing, which a stall detector cannot see. The value stays undetermined and is documented as such; only its role narrowed. - Liveness notifications now report observed output instead of a fixed tick, so a peer can tell "we are alive" from "the model is producing". The model's text is deliberately NOT sent: a request that later fails returns no text, so a prefix on the progress channel would disclose only in the failure case what the contract says was never delivered. `total` stays omitted — streaming still cannot say how much output is coming. The response contract is untouched: one `CreateMessageResult`, text-only, the same four fields and the same error mapping. Counting stream parts naively was wrong and measurement caught it — a provider whose HTTP call never answers still yields a `start` part, which reported "1 chunk" for a dead call, so lifecycle parts are excluded by name. Four tests, each revert-probed, plus the fixture converted to server-sent events — a JSON completion body does not fail loudly against a streaming adapter, it silently yields an empty transcript.
… silence bound Sampling carried three wall-clock bounds — 120 s on the model call, an absolute ceiling in `serve` equal to the sum of the phase bounds, and 30 s on the human approval wait — plus its own 45 s silence bound. None had a precedent in this repo, and each was tighter than the comparable path. Measured what the repo already does and deleted all four accordingly. No total bound. `session/llm.ts` puts none on a real conversation: grep it for `Effect.timeout`, `AbortSignal.timeout` or `Schedule.upTo` and nothing comes back, and it says so in words — retry is "intentionally NOT capped", worst case ~97 min, "bounding per-attempt latency via chunkTimeout is the primary lever for hang-time control". For a streaming call elapsed total is not a health signal; silence is. No approval bound. `permission/index.ts` gives an ordinary interactive ask no timeout at all; only a forwarded ask and a skip-all forced ask are bounded, and sampling's is neither. A TUI prompt waits indefinitely while sampling gave up at 30 s. The silence bound is now the provider's `chunkTimeout` rather than a second number: same per-provider config key, same `DEFAULT_CHUNK_TIMEOUT` default. Our 45 s disagreed with it 10x and in the wrong direction, since that value is tuned to tolerate a real ~5 min silent cold-path first token — so it would have failed calls ordinary chat survives. The detector itself stays ours because it observes output-bearing stream parts rather than HTTP bytes, which catches a fetch that never resolves, a keepalive-only stream, and non-SSE adapters that `wrapSSE` cannot see. `<= 0` disables it, as it already does in provider.ts. `data.phase` therefore has one value, `"stall"`; `"model"`, `"total"` and `"approval"` are gone with their bounds. What is no longer caught is named in code and doc: a stream that trickles forever, the pre-model stretch, and an approval nobody answers — each unbounded on the main path too. Tests: the ceiling test is deleted with its mechanism; the reap, phase and constant tests are rewritten to assert the new intent; new tests prove the config is genuinely what bounds a silent call with nothing injected, that a late-answered approval still succeeds, and that a 0 bound disables rather than fires instantly. `chunkTimeout: 0` turns out to be unrepresentable in config (`PositiveInt`), which is noted in the test that injects it instead.
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).session/llm.tsputs no total bound on a conversation: retry "intentionally NOT capped", stated worst case ~97 min,chunkTimeoutnamed as "the primary lever for hang-time control" — so for a streaming call elapsed total is not a health signal, silence is.permission/index.tsgives an ordinary interactive ask no timeout either; only a forwarded ask and a skip-all forced ask are bounded, and sampling's is neither — so a TUI prompt waits indefinitely while sampling gave up at 30 s.data.phasenow has one value,stall:approval,modelandtotalwent with their bounds, whiledata.serverstill rides every error path as the only discriminator against the SDK's own-32001. What that leaves unbounded is named in code and doc rather than implied — a stream that trickles without ever finishing, the pre-model stretch (model selection, provider init), and an approval nobody answers — each equally unbounded on the main chat path, which is why they are accepted rather than covered by a new guess.notifications/progresswere being sent, so a server on the SDK's 60 sDEFAULT_REQUEST_TIMEOUT_MSECabandoned any long request at 60 s while we were still working: our bound was on the wrong side of the interaction. The model call emits a beat every 15 s — the one number sampling still chooses for itself, and now purely a cadence that keeps the connection and the peer's timer from going idle — but only when the server minted a progress token — the SDK mints one solely when its caller passedonprogress, we read it back off the request's_meta, and with no token we send nothing at all. The check is!== undefined, not truthiness, because the first server-initiated request on a connection has token0. Each beat now carries the chunk and character counts the stream has actually produced, so a peer can tell MiMoCode is alive from the model is producing — a fixed tick cannot. It deliberately does not carry the model's text: a request that later stalls, times out or is cancelled returns no text at all, so a prefix here would disclose only in the failure case what the contract says was never delivered, andonprogressasks for liveness, not content.totalstays omitted — streaming still cannot say how much output is coming, so no percentage exists. Beats reset the peer's timer only, never ours.chunkTimeout, not a number of sampling's own. The call streams, so no-output-at-all is observable: it fails withdata.phase: "stall", the clock covers the wait for the first chunk and every gap after it, every chunk resets it, anddata.chunks/data.charactersseparate never produced anything from started and died. Our own 45 s is gone — the repo already had this exact concept inprovider.ts:wrapSSE, per-provider inmimocode.json, default 8 min, so two numbers disagreed 10x about one fact, and in the wrong direction:chunkTimeoutis tuned to tolerate a real cold-path first token that can be ~5 minutes silent, so 45 s would have failed calls ordinary chat survives. Sampling reads that same key with that same default, and<= 0disables it as it does there. The detector stays ours because of where it observes:wrapSSEbounds gaps between HTTP bytes on an already-resolved SSE response and counts keep-alive comments as activity, whereas this counts stream parts that carry model output — so it also catches a fetch that never resolves, a keep-alive-only stream, and adapters that are not SSE-over-HTTP, and misses nothingwrapSSEcatches. The response contract is unchanged: oneCreateMessageResult, text-only, the same four fields, the same error mapping, with a test asserting the request went out as a stream at the provider's HTTP boundary while the result stayed identical. Worth checking: lifecycle parts are excluded by name, because a provider whose HTTP call never answers still yields astartpart and counting it reported 1 chunk for a stone-dead call.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). The stall bound is the only thing that reaps such a request, which is why it is not gated on the server having asked for progress; it caps the resulting waste rather than avoiding 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.resetTimeoutOnProgressis the requester's flag and defaults tofalsein the SDK, so a server that does not pass it ignores our notifications and still gives up at its own deadline. Sending them is necessary but not sufficient, and that half is not ours to decide.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 bound had to be injected to keep CI fast the test comment says so rather than implying an unbounded proof, and the one constant-only test is labelled in-file as a change-detector rather than counted as coverage.