You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@maka.apache.org by wy...@apache.org on 2026/09/22 14:05:09 UTC
(maka) branch feat/runtime-host-rust updated: feat(capabilities): scope client publications to sessions
This is an automated email from the ASF dual-hosted git repository.
M4n5ter pushed a commit to branch feat/runtime-host-rust
in repository https://gitbox.apache.org/repos/asf/maka.git
The following commit(s) were added to refs/heads/feat/runtime-host-rust by this push:
new 2a206aec49 feat(capabilities): scope client publications to sessions
2a206aec49 is described below
commit 2a206aec4948d4391e6efbcd7d7aa95e3d37d65c
Author: Wang <m4...@gmail.com>
AuthorDate: Tue Sep 22 22:04:58 2026 +0800
feat(capabilities): scope client publications to sessions
Preserve pinned generations until explicit scope retirement, enforce session visibility across native and plugin calls, and add exact MCP admission with bounded progress delivery.
Generated-by: Codex
---
apps/desktop/src/main/runtime-host-client.ts | 9 +-
crates/client-capability/src/broker.rs | 4 +-
crates/client-capability/src/broker/dispatch.rs | 5 +
crates/client-capability/src/broker/state.rs | 5 +-
crates/client-capability/src/managed.rs | 17 +-
crates/client-capability/src/registry.rs | 153 +++----------
crates/client-capability/src/registry/bindings.rs | 73 +++++-
.../src/registry/bindings/restoration.rs | 55 +++--
.../src/registry/bindings/selection.rs | 49 ++--
.../src/registry/bindings/snapshot.rs | 13 +-
.../src/{registry.rs => registry/mutation.rs} | 248 ++++++++-------------
.../client-capability/src/registry/publication.rs | 38 +++-
crates/client-capability/tests/bindings.rs | 155 +++++++++++++
crates/client-capability/tests/broker.rs | 1 +
crates/client-capability/tests/managed_target.rs | 28 +++
crates/client-capability/tests/support/forms.rs | 1 +
crates/protocol/src/capability/manifest.rs | 35 ++-
crates/protocol/tests/capability_contract.rs | 19 ++
crates/protocol/tests/interaction_contract.rs | 8 +
.../runtime-host/src/execution/plugins/scopes.rs | 7 +-
crates/runtime-host/src/server/capabilities.rs | 20 +-
crates/runtime-host/src/server/sessions.rs | 7 +
.../tests/integration/plugin_authorization.rs | 8 +
crates/runtime/src/capability.rs | 12 +-
crates/runtime/src/capability/composition.rs | 12 +-
crates/runtime/src/capability/contract.rs | 1 +
crates/runtime/src/capability/frame.rs | 11 +
crates/runtime/src/interaction.rs | 1 +
crates/runtime/src/interaction/validation.rs | 2 +-
docs/rust-runtime.md | 7 +
docs/rust-runtime.zh-CN.md | 5 +
packages/cli/src/mcp-capability-provider.ts | 7 +
packages/cli/src/runtime-host-onboarding.ts | 2 +-
packages/cli/src/tui-mcp-remote-publication.ts | 16 +-
packages/core/src/client-capability-grant.ts | 6 +-
packages/mcp/src/__tests__/modern-manager.test.ts | 77 ++++++-
packages/mcp/src/__tests__/oauth.test.ts | 6 +
packages/mcp/src/index.ts | 25 +++
packages/mcp/src/oauth.ts | 4 +-
packages/mcp/src/tool-progress.ts | 53 +++++
.../__tests__/capability-provider-service.test.ts | 19 +-
.../__tests__/client-capability-channel.test.ts | 73 +++++-
.../src/client/client-capability-channel.ts | 105 ++++++---
.../runtime-host/src/client/client-capability.ts | 5 +
packages/runtime-host/src/client/connection.ts | 25 ++-
packages/runtime-host/src/client/index.ts | 5 +-
.../src/client/reconnecting-connection.ts | 13 +-
.../runtime-host/src/protocol/client-capability.ts | 25 ++-
packages/ui/src/client-capability-prompt.tsx | 1 +
tests/fixtures/client-capability-host.mjs | 34 ++-
tests/fixtures/client-pairing-access.mjs | 10 +-
tests/fixtures/client-provider-access.mjs | 10 +-
tests/fixtures/client-scheduler-entry.mjs | 6 +-
53 files changed, 1101 insertions(+), 435 deletions(-)
diff --git a/apps/desktop/src/main/runtime-host-client.ts b/apps/desktop/src/main/runtime-host-client.ts
index 4376c6a915..4246652a7a 100644
--- a/apps/desktop/src/main/runtime-host-client.ts
+++ b/apps/desktop/src/main/runtime-host-client.ts
@@ -41,6 +41,7 @@ import {
import type { PricingConfig } from "@maka/core/usage-stats/types";
import {
type ClientCapabilityProvider,
+ type ClientCapabilityRegistrationOptions,
type DecodedSessionTranscriptPage,
type DirectRequestOperationKey,
type RuntimeHostConnection,
@@ -1511,17 +1512,17 @@ export class DesktopRuntimeHostClient {
replaceClientCapabilities(
provider: ClientCapabilityProvider,
- timeoutMs?: number,
+ options?: ClientCapabilityRegistrationOptions,
): Promise<ClientCapabilityReplaceResult> {
this.#assertOpen();
- return this.connection.replaceClientCapabilities(provider, timeoutMs);
+ return this.connection.replaceClientCapabilities(provider, options);
}
unregisterClientCapabilities(
- timeoutMs?: number,
+ options?: ClientCapabilityRegistrationOptions,
): Promise<ClientCapabilityUnregisterResult> {
this.#assertOpen();
- return this.connection.unregisterClientCapabilities(timeoutMs);
+ return this.connection.unregisterClientCapabilities(options);
}
async openSession(sessionId: string): Promise<DesktopRuntimeHostSession> {
diff --git a/crates/client-capability/src/broker.rs b/crates/client-capability/src/broker.rs
index 5e9d3c0982..da30c9fc4d 100644
--- a/crates/client-capability/src/broker.rs
+++ b/crates/client-capability/src/broker.rs
@@ -111,7 +111,7 @@ impl Broker {
&& call.pending_terminal.is_none()
&& !call.cancellation.is_cancelled()
&& !endpoint.closed().is_cancelled()
- && !endpoint.invocations().is_cancelled()
+ && call.registration.available()
&& call.tool.as_ref().is_some_and(|tool| {
tool.session_id == session
&& tool.turn_id == turn
@@ -158,7 +158,7 @@ impl Broker {
serde_json::to_value(&frame).map_err(|_| CallError::Invalid("JSON encoding"))?;
decode_host_frame(&encoded).map_err(|_| CallError::Invalid("outbound wire boundary"))?;
let endpoint = registration.endpoint().clone();
- let provider = endpoint.invocations();
+ let provider = registration.invocations();
let mut state = self.inner.active.lock().unwrap_or_else(|e| e.into_inner());
if self.inner.draining.is_cancelled() || cancellation.is_cancelled() {
return Err(CallError::Cancelled);
diff --git a/crates/client-capability/src/broker/dispatch.rs b/crates/client-capability/src/broker/dispatch.rs
index cc1bf815b9..92a7023a8e 100644
--- a/crates/client-capability/src/broker/dispatch.rs
+++ b/crates/client-capability/src/broker/dispatch.rs
@@ -87,6 +87,11 @@ impl Broker {
timeout: Duration,
cancellation: CancellationToken,
) -> Result<PendingCall, CallError> {
+ if !registration.visible_to(call.source.session_id()) {
+ return Err(CallError::Invalid(
+ "tool publication belongs to another Session",
+ ));
+ }
let offer = registration
.manifest()
.offers
diff --git a/crates/client-capability/src/broker/state.rs b/crates/client-capability/src/broker/state.rs
index f9f5337e11..0c955a229e 100644
--- a/crates/client-capability/src/broker/state.rs
+++ b/crates/client-capability/src/broker/state.rs
@@ -169,7 +169,7 @@ impl Inner {
return;
}
let endpoint = invocation.registration.endpoint();
- if endpoint.closed().is_cancelled() || endpoint.invocations().is_cancelled() {
+ if !invocation.registration.available() {
Self::settle(&mut state, id, Err(CallError::CapabilityLost), false);
return;
}
@@ -209,8 +209,7 @@ impl Inner {
let Some(invocation) = state.calls.get(id) else {
return true;
};
- let endpoint = invocation.registration.endpoint();
- if endpoint.closed().is_cancelled() || endpoint.invocations().is_cancelled() {
+ if !invocation.registration.available() {
Self::stop_locked(state, id, CallError::CapabilityLost, false);
true
} else if invocation.cancellation.is_cancelled() || self.draining.is_cancelled() {
diff --git a/crates/client-capability/src/managed.rs b/crates/client-capability/src/managed.rs
index ecd6f91d0d..9eaef127e4 100644
--- a/crates/client-capability/src/managed.rs
+++ b/crates/client-capability/src/managed.rs
@@ -18,7 +18,7 @@
*/
use maka_runtime::{
- capability::{AdmissionEvidence, ToolDescriptor},
+ capability::{Admission, AdmissionEvidence, Offer, ToolDescriptor},
interaction::{GrantCapability, GrantScope},
};
@@ -39,10 +39,23 @@ pub enum ManagedAdmissionError {
}
pub(crate) fn scope(
- offer_id: &str,
+ offer: &Offer,
tool: &ToolDescriptor,
evidence: &AdmissionEvidence,
) -> Result<Option<(GrantCapability, GrantScope)>, ManagedAdmissionError> {
+ if offer.admission == Some(Admission::Mcp) {
+ if !matches!(evidence, AdmissionEvidence::None) {
+ return Err(ManagedAdmissionError::InvalidEvidence);
+ }
+ return Ok(Some((
+ GrantCapability::Mcp,
+ GrantScope::McpTool {
+ server_id: tool.server_id.clone(),
+ tool_name: tool.name.clone(),
+ },
+ )));
+ }
+ let offer_id = offer.offer_id.as_str();
if offer_id == "desktop_settings"
&& tool.server_id == "desktop_settings"
&& matches!(
diff --git a/crates/client-capability/src/registry.rs b/crates/client-capability/src/registry.rs
index fd3d208794..910751ea54 100644
--- a/crates/client-capability/src/registry.rs
+++ b/crates/client-capability/src/registry.rs
@@ -18,14 +18,11 @@
*/
use crate::{Endpoint, Error, Identity, PrincipalKind};
-use maka_runtime::capability::{Affinity, HostPathAccess, Manifest, RegistrationResult};
-use std::{
- collections::{HashMap, HashSet},
- sync::Arc,
-};
+use std::{collections::HashMap, sync::Arc};
use uuid::Uuid;
mod bindings;
+mod mutation;
mod publication;
pub use bindings::{
BindingError, BindingMode, PreparedBindings, RestoredBindings, Snapshot, SnapshotOffer,
@@ -37,6 +34,7 @@ struct Provider {
identity: Arc<Identity>,
residency: Residency,
current: Option<Arc<Registration>>,
+ scoped: HashMap<String, Arc<Registration>>,
active: Option<Uuid>,
}
@@ -53,6 +51,7 @@ pub struct Registry {
providers: HashMap<String, Provider>,
connections: HashMap<Uuid, Connection>,
revision: u64,
+ retirement_revision: u64,
draining: bool,
sessions: HashMap<String, bindings::SessionBindings>,
}
@@ -96,6 +95,7 @@ impl Registry {
identity: Arc::new(identity),
residency: Arc::default(),
current: None,
+ scoped: HashMap::new(),
active: None,
},
);
@@ -111,122 +111,6 @@ impl Registry {
Ok(provider_id)
}
- /// Input must first pass the protocol manifest decoder.
- pub fn replace(
- &mut self,
- connection_id: Uuid,
- manifest: Manifest,
- ) -> Result<RegistrationResult, Error> {
- if self.draining {
- return Err(Error::Draining);
- }
- let connection = self
- .connections
- .get(&connection_id)
- .ok_or(Error::Unavailable)?;
- if connection.superseded || connection.endpoint.closed().is_cancelled() {
- return Err(Error::Invalid("connection has been superseded or closed"));
- }
- let provider = self
- .providers
- .get_mut(&connection.provider_id)
- .expect("attached provider");
- if provider.identity.principal_kind == PrincipalKind::CapabilityProvider
- && (manifest.services.as_ref().is_some_and(|s| !s.is_empty())
- || manifest.offers.iter().any(|offer| {
- offer.affinity != Affinity::Session
- || offer.host_path_access != HostPathAccess::None
- }))
- {
- return Err(Error::Invalid(
- "trusted providers may publish only path-independent session tools",
- ));
- }
- let revision = self
- .revision
- .checked_add(1)
- .filter(|r| *r <= 9_007_199_254_740_991)
- .ok_or(Error::Invalid("revision exhausted"))?;
- let mut proxy_names = HashSet::new();
- for tool in manifest.offers.iter().flat_map(|offer| &offer.tools) {
- if !proxy_names.insert(crate::proxy_tool_name(&tool.server_id, &tool.name)) {
- return Err(Error::Invalid("model proxy tool name collision"));
- }
- }
- let mut resident = provider.residency.lock().unwrap_or_else(|e| e.into_inner());
- if !resident.insert(manifest.registration_id.clone()) {
- return Err(Error::Invalid("registration identity is still resident"));
- }
- drop(resident);
- let registration = Arc::new(Registration {
- provider_id: connection.provider_id.clone(),
- connection_id,
- identity: provider.identity.clone(),
- contracts: manifest
- .offers
- .iter()
- .enumerate()
- .map(|(index, offer)| (crate::ContractId::of(offer), index))
- .collect(),
- manifest,
- endpoint: connection.endpoint.clone(),
- residency: provider.residency.clone(),
- });
- if let Some(old) = provider.active.filter(|old| *old != connection_id)
- && let Some(connection) = self.connections.get_mut(&old)
- {
- connection.superseded = true;
- connection.endpoint.supersede();
- }
- provider.active = Some(connection_id);
- let result = RegistrationResult {
- registration_id: registration.manifest.registration_id.clone(),
- revision,
- };
- let previous = provider.current.replace(registration.clone());
- self.revision = revision;
- self.publication_changed(previous.as_deref(), ®istration);
- Ok(result)
- }
-
- pub fn unregister(
- &mut self,
- connection_id: Uuid,
- registration_id: &str,
- ) -> Result<RegistrationResult, Error> {
- let connection = self
- .connections
- .get(&connection_id)
- .ok_or(Error::Unavailable)?;
- let provider = self
- .providers
- .get_mut(&connection.provider_id)
- .expect("attached provider");
- if provider.active != Some(connection_id)
- || provider
- .current
- .as_ref()
- .is_none_or(|r| r.manifest.registration_id != registration_id)
- {
- return Err(Error::Invalid("registration is not current"));
- }
- let revision = self
- .revision
- .checked_add(1)
- .filter(|r| *r <= 9_007_199_254_740_991)
- .ok_or(Error::Invalid("revision exhausted"))?;
- let previous = provider
- .current
- .take()
- .expect("validated current registration");
- self.revision = revision;
- self.publication_removed(&previous);
- Ok(RegistrationResult {
- registration_id: registration_id.into(),
- revision,
- })
- }
-
/// Disconnection marks session bindings Lost, unlike explicit unregister's
/// removal. Inactive socket teardown cannot clear its replacement.
pub fn detach(&mut self, id: Uuid) {
@@ -234,14 +118,26 @@ impl Registry {
return;
};
connection.endpoint.close();
- if let Some(provider) = self.providers.get_mut(&connection.provider_id)
- && provider.active == Some(id)
- {
- provider.active = None;
- if provider.current.take().is_some() {
- self.revision += 1;
+ let mut lost = Vec::new();
+ if let Some(provider) = self.providers.get_mut(&connection.provider_id) {
+ if provider.active == Some(id) {
+ provider.active = None;
+ if let Some(r) = provider.current.take() {
+ lost.push(r);
+ }
}
- self.provider_lost(&connection.provider_id);
+ provider.scoped.retain(|_, r| {
+ if r.connection_id == id {
+ lost.push(r.clone());
+ false
+ } else {
+ true
+ }
+ });
+ }
+ for r in lost {
+ self.revision += 1;
+ self.provider_lost(&r);
}
self.prune();
}
@@ -259,6 +155,7 @@ impl Registry {
fn prune(&mut self) {
self.providers.retain(|id, provider| {
provider.current.is_some()
+ || !provider.scoped.is_empty()
|| provider.active.is_some()
|| Arc::strong_count(&provider.identity) > 1
|| self
diff --git a/crates/client-capability/src/registry/bindings.rs b/crates/client-capability/src/registry/bindings.rs
index 43cb106e18..16913808da 100644
--- a/crates/client-capability/src/registry/bindings.rs
+++ b/crates/client-capability/src/registry/bindings.rs
@@ -60,12 +60,26 @@ pub enum BindingError {
struct ProviderRef {
id: String,
identity: Arc<Identity>,
+ session_id: Option<String>,
}
impl ProviderRef {
fn of(registration: &Registration) -> Self {
Self {
id: registration.provider_id().into(),
identity: registration.identity().clone(),
+ session_id: registration.session_id().map(str::to_owned),
+ }
+ }
+
+ fn matches(&self, registration: &Registration) -> bool {
+ self.id == registration.provider_id()
+ && self.session_id.as_deref() == registration.session_id()
+ }
+
+ fn publication(&self) -> maka_runtime::capability::PublicationIdentity {
+ maka_runtime::capability::PublicationIdentity {
+ identity: (*self.identity).clone(),
+ session_id: self.session_id.clone(),
}
}
}
@@ -94,9 +108,12 @@ pub(super) struct SessionBindings {
type Candidates = BTreeMap<ContractId, Vec<Arc<Registration>>>;
impl Registry {
- fn eligible(&self) -> Candidates {
+ fn eligible(&self, session_id: Option<&str>) -> Candidates {
let mut eligible: Candidates = BTreeMap::new();
- for registration in self.published().filter(|r| r.available()) {
+ for registration in self
+ .published()
+ .filter(|r| r.available() && r.visible_to(session_id))
+ {
for (contract, _) in registration.offers() {
eligible
.entry(contract.clone())
@@ -111,7 +128,26 @@ impl Registry {
}
pub fn release_session(&mut self, session_id: &str) {
+ self.retirement_revision += 1;
self.sessions.remove(session_id);
+ for provider in self.providers.values_mut() {
+ // Include old generations still pinned by prepared or admitted work.
+ let resident: Vec<_> = provider
+ .residency
+ .lock()
+ .unwrap_or_else(|e| e.into_inner())
+ .values()
+ .filter_map(std::sync::Weak::upgrade)
+ .collect();
+ for r in resident {
+ if r.session_id() == Some(session_id) {
+ r.invocations.cancel();
+ }
+ }
+ if provider.scoped.remove(session_id).is_some() {
+ self.revision += 1;
+ }
+ }
self.prune();
}
@@ -126,14 +162,18 @@ impl Registry {
&& current.offer(contract).is_none()
};
state.session.retain(|contract, binding| {
- !matches!(binding, Binding::Bound(provider) if provider.id == current.provider_id() && removed(contract))
- });
- state.turn.retain(|contract, provider| {
- provider.id != current.provider_id() || !removed(contract)
+ let provider = binding.provider();
+ !(provider.matches(current)
+ && (removed(contract)
+ || (matches!(binding, Binding::Lost(_))
+ && current.offer(contract).is_none())))
});
+ state
+ .turn
+ .retain(|contract, provider| !provider.matches(current) || !removed(contract));
for (contract, binding) in &mut state.session {
if let Binding::Lost(provider) = binding
- && provider.id == current.provider_id()
+ && provider.matches(current)
&& current.offer(contract).is_some()
{
*binding = Binding::Bound(provider.clone());
@@ -147,25 +187,27 @@ impl Registry {
for state in self.sessions.values_mut() {
state.session.retain(|contract, binding| {
!matches!(binding, Binding::Bound(provider)
- if provider.id == previous.provider_id() && previous.offer(contract).is_some())
+ if provider.matches(previous) && previous.offer(contract).is_some())
});
state.turn.retain(|contract, provider| {
- provider.id != previous.provider_id() || previous.offer(contract).is_none()
+ !provider.matches(previous) || previous.offer(contract).is_none()
});
}
self.prune_sessions();
}
- pub(super) fn provider_lost(&mut self, provider_id: &str) {
+ pub(super) fn provider_lost(&mut self, registration: &Registration) {
for state in self.sessions.values_mut() {
for binding in state.session.values_mut() {
if let Binding::Bound(provider) = binding
- && provider.id == provider_id
+ && provider.matches(registration)
{
*binding = Binding::Lost(provider.clone());
}
}
- state.turn.retain(|_, provider| provider.id != provider_id);
+ state
+ .turn
+ .retain(|_, provider| !provider.matches(registration));
}
self.prune_sessions();
}
@@ -200,6 +242,13 @@ fn choose(
}
}
+fn choose_bound(
+ candidates: &[Arc<Registration>],
+ provider: &ProviderRef,
+) -> Option<Arc<Registration>> {
+ candidates.iter().find(|r| provider.matches(r)).cloned()
+}
+
fn claim_names(
registration: &Registration,
contract: &ContractId,
diff --git a/crates/client-capability/src/registry/bindings/restoration.rs b/crates/client-capability/src/registry/bindings/restoration.rs
index 5440234619..2b6452749c 100644
--- a/crates/client-capability/src/registry/bindings/restoration.rs
+++ b/crates/client-capability/src/registry/bindings/restoration.rs
@@ -21,7 +21,9 @@ use super::{
Binding, BindingError, ProviderRef, SessionBindings, Snapshot, SnapshotOffer, snapshot::Source,
};
use crate::{ContractId, Identity, Registry};
-use maka_runtime::capability::{Affinity, ClientComposition, ClientOffer, PinnedAffinity};
+use maka_runtime::capability::{
+ Affinity, ClientComposition, ClientOffer, PinnedAffinity, PublicationIdentity,
+};
use std::{
collections::{BTreeMap, HashSet},
sync::Arc,
@@ -29,6 +31,7 @@ use std::{
/// Restores only the admitted surface. It never discovers additional tools.
pub struct RestoredBindings {
+ retirement_revision: u64,
session_id: String,
previous: Option<SessionBindings>,
selected: SessionBindings,
@@ -47,13 +50,15 @@ impl Registry {
}
let mut identities = BTreeMap::new();
let mut selected = SessionBindings::default();
- for (contract, identity) in &composition.session_bindings {
- let provider = self.restore_identity(identity, &mut identities)?;
- let available = self.current(&provider.id).is_some_and(|r| {
- r.available()
- && r.offer(contract)
- .is_some_and(|o| o.affinity == Affinity::Session)
- });
+ for (contract, publication) in &composition.session_bindings {
+ let provider = self.restore_publication(session_id, publication, &mut identities)?;
+ let available = self
+ .current_scoped(&provider.id, provider.session_id.as_deref())
+ .is_some_and(|r| {
+ r.available()
+ && r.offer(contract)
+ .is_some_and(|o| o.affinity == Affinity::Session)
+ });
selected.session.insert(
contract.clone(),
if available {
@@ -71,11 +76,12 @@ impl Registry {
ClientOffer::Pinned {
contract,
affinity,
- identity,
+ publication,
} => {
- let provider = self.restore_identity(identity, &mut identities)?;
+ let provider =
+ self.restore_publication(session_id, publication, &mut identities)?;
let registration = self
- .current(&provider.id)
+ .current_scoped(&provider.id, provider.session_id.as_deref())
.filter(|r| r.available())
.ok_or(BindingError::Lost)?;
let offer = registration.offer(contract).ok_or(BindingError::Lost)?;
@@ -132,6 +138,7 @@ impl Registry {
let snapshot = Snapshot { offers };
Ok((
RestoredBindings {
+ retirement_revision: self.retirement_revision,
session_id: session_id.into(),
previous: self.sessions.get(session_id).cloned(),
selected,
@@ -151,7 +158,9 @@ impl Registry {
if self.draining {
return Err(BindingError::Draining);
}
- if self.sessions.get(&restored.session_id) != restored.previous.as_ref() {
+ if self.retirement_revision != restored.retirement_revision
+ || self.sessions.get(&restored.session_id) != restored.previous.as_ref()
+ {
return Ok(false);
}
for (id, identity) in &restored.identities {
@@ -170,7 +179,7 @@ impl Registry {
if let Source::Pinned(pinned) = &entry.source
&& (!pinned.available()
|| !self
- .current(pinned.provider_id())
+ .current_scoped(pinned.provider_id(), pinned.session_id())
.is_some_and(|current| Arc::ptr_eq(¤t, pinned)))
{
return Ok(false);
@@ -181,6 +190,7 @@ impl Registry {
identity,
residency: Arc::default(),
current: None,
+ scoped: Default::default(),
active: None,
});
}
@@ -212,6 +222,25 @@ impl Registry {
Ok(ProviderRef {
id,
identity: stored,
+ session_id: None,
})
}
+
+ fn restore_publication(
+ &self,
+ session_id: &str,
+ publication: &PublicationIdentity,
+ identities: &mut BTreeMap<String, Arc<Identity>>,
+ ) -> Result<ProviderRef, BindingError> {
+ if publication
+ .session_id
+ .as_deref()
+ .is_some_and(|id| id != session_id)
+ {
+ return Err(BindingError::InvalidComposition);
+ }
+ let mut provider = self.restore_identity(&publication.identity, identities)?;
+ provider.session_id.clone_from(&publication.session_id);
+ Ok(provider)
+ }
}
diff --git a/crates/client-capability/src/registry/bindings/selection.rs b/crates/client-capability/src/registry/bindings/selection.rs
index 6af7ee3f6c..6a43ee6426 100644
--- a/crates/client-capability/src/registry/bindings/selection.rs
+++ b/crates/client-capability/src/registry/bindings/selection.rs
@@ -18,7 +18,8 @@
*/
use super::{
- Binding, BindingError, BindingMode, ProviderRef, SessionBindings, choose, claim_names,
+ Binding, BindingError, BindingMode, ProviderRef, SessionBindings, choose, choose_bound,
+ claim_names,
};
use crate::{PrincipalKind, Registry};
use maka_runtime::capability::Affinity;
@@ -28,6 +29,7 @@ use uuid::Uuid;
/// A candidate owns no Session binding. Commit verifies the same selection and
/// publications under the registry lock before making it authoritative.
pub struct PreparedBindings {
+ retirement_revision: u64,
session_id: String,
initiating: Option<Uuid>,
mode: BindingMode,
@@ -51,7 +53,7 @@ fn composition(
session_bindings: selected
.session
.iter()
- .map(|(contract, binding)| (contract.clone(), (*binding.provider().identity).clone()))
+ .map(|(contract, binding)| (contract.clone(), binding.provider().publication()))
.collect(),
offers: snapshot.composition(),
}
@@ -65,10 +67,12 @@ impl Registry {
mode: BindingMode,
) -> Result<(PreparedBindings, super::Snapshot), BindingError> {
let previous = self.sessions.get(session_id).cloned();
- let selected = self.select_bindings(previous.as_ref(), initiating, mode)?;
+ let selected =
+ self.select_bindings(Some(session_id), previous.as_ref(), initiating, mode)?;
let snapshot = self.snapshot_bindings(Some(&selected))?;
Ok((
PreparedBindings {
+ retirement_revision: self.retirement_revision,
session_id: session_id.into(),
initiating,
mode,
@@ -83,10 +87,13 @@ impl Registry {
/// False means stale preparation; no binding was changed.
pub fn commit_bindings(&mut self, prepared: PreparedBindings) -> Result<bool, BindingError> {
- if self.sessions.get(&prepared.session_id) != prepared.previous.as_ref() {
+ if self.retirement_revision != prepared.retirement_revision
+ || self.sessions.get(&prepared.session_id) != prepared.previous.as_ref()
+ {
return Ok(false);
}
let selected = self.select_bindings(
+ Some(&prepared.session_id),
prepared.previous.as_ref(),
prepared.initiating,
prepared.mode,
@@ -122,9 +129,7 @@ impl Registry {
continue;
};
let Some(registration) = self
- .providers
- .get(&provider.id)
- .and_then(|provider| provider.current.as_ref())
+ .current_scoped(&provider.id, provider.session_id.as_deref())
.filter(|r| r.available())
else {
continue;
@@ -179,7 +184,12 @@ impl Registry {
initiating: Option<Uuid>,
mode: BindingMode,
) -> Result<bool, BindingError> {
- let next = self.select_bindings(self.sessions.get(session_id), initiating, mode)?;
+ let next = self.select_bindings(
+ Some(session_id),
+ self.sessions.get(session_id),
+ initiating,
+ mode,
+ )?;
let changed = self.sessions.get(session_id) != Some(&next);
self.sessions.insert(session_id.into(), next);
self.prune_sessions();
@@ -195,12 +205,13 @@ impl Registry {
mode: BindingMode,
) -> Result<super::Snapshot, BindingError> {
let previous = session_id.and_then(|id| self.sessions.get(id));
- let selected = self.select_bindings(previous, initiating, mode)?;
+ let selected = self.select_bindings(session_id, previous, initiating, mode)?;
self.snapshot_bindings(Some(&selected))
}
fn select_bindings(
&self,
+ session_id: Option<&str>,
previous: Option<&SessionBindings>,
initiating: Option<Uuid>,
mode: BindingMode,
@@ -212,7 +223,12 @@ impl Registry {
.and_then(|id| self.connections.get(&id))
.and_then(|connection| self.providers.get(&connection.provider_id));
let direct = initiating
- .and_then(|provider| provider.current.as_ref())
+ .and_then(|provider| {
+ provider
+ .current
+ .as_ref()
+ .or_else(|| session_id.and_then(|id| provider.scoped.get(id)))
+ })
.filter(|r| r.available());
let first_associated = {
let mut associated = self.published().filter(|r| {
@@ -223,6 +239,7 @@ impl Registry {
identity.credential_bound_client_instance_id.as_deref()
== Some(identity.client_instance_id.as_str())
&& r.available()
+ && r.visible_to(session_id)
&& r.identity().trusted()
&& r.identity().capability_owner.as_ref().is_some_and(|owner| {
owner.principal_id == identity.principal_id
@@ -230,11 +247,16 @@ impl Registry {
})
});
let first_associated = associated.next();
- if direct.is_none() && first_associated.is_some() && associated.next().is_some() {
+ if direct.is_none()
+ && first_associated
+ .is_some_and(|first| associated.any(|r| r.provider_id() != first.provider_id()))
+ {
return Err(BindingError::Ambiguous);
}
first_associated
};
+ let first_associated =
+ first_associated.map(|r| self.providers[&r.provider_id].current.as_ref().unwrap_or(r));
let selected = direct.or(first_associated);
let selector = selected.map(|r| ProviderRef::of(r)).or_else(|| {
initiating
@@ -242,6 +264,7 @@ impl Registry {
.map(|p| ProviderRef {
id: p.identity.provider_id(),
identity: p.identity.clone(),
+ session_id: None,
})
});
let mut next = SessionBindings {
@@ -258,7 +281,7 @@ impl Registry {
}),
..SessionBindings::default()
};
- let eligible = self.eligible();
+ let eligible = self.eligible(session_id);
let contracts: BTreeSet<_> = previous
.into_iter()
.flat_map(|s| s.session.keys())
@@ -278,7 +301,7 @@ impl Registry {
.map(Vec::as_slice)
.unwrap_or_default();
let candidate = if let Some(prior) = prior {
- match choose(candidates, Some(prior.provider()))? {
+ match choose_bound(candidates, prior.provider()) {
Some(candidate) => Some(candidate),
None if mode == BindingMode::Degrade => {
next.session
diff --git a/crates/client-capability/src/registry/bindings/snapshot.rs b/crates/client-capability/src/registry/bindings/snapshot.rs
index f8416354c4..8f5fd758b5 100644
--- a/crates/client-capability/src/registry/bindings/snapshot.rs
+++ b/crates/client-capability/src/registry/bindings/snapshot.rs
@@ -41,7 +41,7 @@ impl Snapshot {
Affinity::Turn => PinnedAffinity::Turn,
Affinity::Call => unreachable!("Call offers are not pinned"),
},
- identity: (**registration.identity()).clone(),
+ publication: ProviderRef::of(registration).publication(),
},
Source::Call { offer, selector } => ClientOffer::Call {
offer: offer.clone(),
@@ -132,8 +132,7 @@ impl SnapshotOffer {
.iter()
.find(|tool| tool.server_id == server_id && tool.name == tool_name)
.ok_or(ManagedAdmissionError::UnknownTool)?;
- let Some((capability, scope)) = crate::managed::scope(&offer.offer_id, tool, evidence)?
- else {
+ let Some((capability, scope)) = crate::managed::scope(offer, tool, evidence)? else {
return Ok(None);
};
Ok(Some(GrantTarget {
@@ -157,7 +156,7 @@ impl SnapshotOffer {
}
}
Source::Call { selector, .. } => {
- let eligible = registry.eligible();
+ let eligible = registry.eligible(None);
choose(
eligible
.get(&self.contract)
@@ -190,7 +189,7 @@ impl Registry {
continue;
};
let registration = self
- .current(&provider.id)
+ .current_scoped(&provider.id, provider.session_id.as_deref())
.filter(|r| r.available() && r.offer(contract).is_some())
.ok_or(BindingError::Lost)?;
if !claim_names(®istration, contract, &mut names) {
@@ -203,7 +202,7 @@ impl Registry {
}
for (contract, provider) in &state.turn {
let Some(registration) = self
- .current(&provider.id)
+ .current_scoped(&provider.id, provider.session_id.as_deref())
.filter(|r| r.available() && r.offer(contract).is_some())
else {
continue;
@@ -216,7 +215,7 @@ impl Registry {
}
}
}
- for (contract, candidates) in self.eligible() {
+ for (contract, candidates) in self.eligible(None) {
let selector = state.and_then(|s| s.initiating.as_ref());
let candidate = if selector.is_some() {
choose(&candidates, selector)?
diff --git a/crates/client-capability/src/registry.rs b/crates/client-capability/src/registry/mutation.rs
similarity index 51%
copy from crates/client-capability/src/registry.rs
copy to crates/client-capability/src/registry/mutation.rs
index fd3d208794..4f8da1de3e 100644
--- a/crates/client-capability/src/registry.rs
+++ b/crates/client-capability/src/registry/mutation.rs
@@ -17,100 +17,13 @@
* under the License.
*/
-use crate::{Endpoint, Error, Identity, PrincipalKind};
+use super::{Registration, Registry};
+use crate::{Error, PrincipalKind};
use maka_runtime::capability::{Affinity, HostPathAccess, Manifest, RegistrationResult};
-use std::{
- collections::{HashMap, HashSet},
- sync::Arc,
-};
+use std::{collections::HashSet, sync::Arc};
use uuid::Uuid;
-mod bindings;
-mod publication;
-pub use bindings::{
- BindingError, BindingMode, PreparedBindings, RestoredBindings, Snapshot, SnapshotOffer,
-};
-pub use publication::Registration;
-use publication::Residency;
-
-struct Provider {
- identity: Arc<Identity>,
- residency: Residency,
- current: Option<Arc<Registration>>,
- active: Option<Uuid>,
-}
-
-struct Connection {
- provider_id: String,
- endpoint: Endpoint,
- superseded: bool,
-}
-
-/// Mutations are synchronous so composition can serialize registry and binding
-/// changes under one admission gate without holding a lock across I/O.
-#[derive(Default)]
-pub struct Registry {
- providers: HashMap<String, Provider>,
- connections: HashMap<Uuid, Connection>,
- revision: u64,
- draining: bool,
- sessions: HashMap<String, bindings::SessionBindings>,
-}
-
impl Registry {
- pub fn revision(&self) -> u64 {
- self.revision
- }
-
- pub fn attach(
- &mut self,
- id: Uuid,
- identity: Identity,
- endpoint: Endpoint,
- ) -> Result<String, Error> {
- if self.draining {
- return Err(Error::Draining);
- }
- if self.connections.contains_key(&id) || endpoint.closed().is_cancelled() {
- return Err(Error::Invalid("connection is already attached or closed"));
- }
- if identity.capability_owner.is_some()
- && identity.principal_kind != PrincipalKind::CapabilityProvider
- {
- return Err(Error::Invalid(
- "only capability providers may declare an owner",
- ));
- }
- self.prune();
- let provider_id = identity.provider_id();
- if let Some(provider) = self.providers.get(&provider_id) {
- if *provider.identity != identity {
- return Err(Error::Invalid(
- "provider authority changed across connections",
- ));
- }
- } else {
- self.providers.insert(
- provider_id.clone(),
- Provider {
- identity: Arc::new(identity),
- residency: Arc::default(),
- current: None,
- active: None,
- },
- );
- }
- self.connections.insert(
- id,
- Connection {
- provider_id: provider_id.clone(),
- endpoint,
- superseded: false,
- },
- );
- Ok(provider_id)
- }
-
/// Input must first pass the protocol manifest decoder.
pub fn replace(
&mut self,
@@ -124,13 +37,29 @@ impl Registry {
.connections
.get(&connection_id)
.ok_or(Error::Unavailable)?;
- if connection.superseded || connection.endpoint.closed().is_cancelled() {
+ if (connection.superseded && manifest.session_id.is_none())
+ || connection.endpoint.closed().is_cancelled()
+ {
return Err(Error::Invalid("connection has been superseded or closed"));
}
let provider = self
.providers
.get_mut(&connection.provider_id)
.expect("attached provider");
+ if let Some(session_id) = &manifest.session_id {
+ if manifest.services.as_ref().is_some_and(|s| !s.is_empty())
+ || manifest.offers.iter().any(|o| {
+ o.affinity != Affinity::Session || o.host_path_access != HostPathAccess::None
+ })
+ {
+ return Err(Error::Invalid(
+ "Session publications support only path-independent Session tools",
+ ));
+ }
+ if !provider.scoped.contains_key(session_id) && provider.scoped.len() >= 128 {
+ return Err(Error::Invalid("too many Session publications"));
+ }
+ }
if provider.identity.principal_kind == PrincipalKind::CapabilityProvider
&& (manifest.services.as_ref().is_some_and(|s| !s.is_empty())
|| manifest.offers.iter().any(|offer| {
@@ -153,11 +82,31 @@ impl Registry {
return Err(Error::Invalid("model proxy tool name collision"));
}
}
+ let overlaps = |r: &Registration| {
+ r.manifest
+ .offers
+ .iter()
+ .flat_map(|o| &o.tools)
+ .any(|t| proxy_names.contains(&crate::proxy_tool_name(&t.server_id, &t.name)))
+ };
+ if if manifest.session_id.is_some() {
+ provider.current.as_ref().is_some_and(|r| overlaps(r))
+ } else {
+ provider.scoped.values().any(|r| overlaps(r))
+ } {
+ return Err(Error::Invalid(
+ "global and Session publications expose overlapping tools",
+ ));
+ }
let mut resident = provider.residency.lock().unwrap_or_else(|e| e.into_inner());
- if !resident.insert(manifest.registration_id.clone()) {
+ if resident.contains_key(&manifest.registration_id) {
return Err(Error::Invalid("registration identity is still resident"));
}
- drop(resident);
+ let invocations = if manifest.session_id.is_some() {
+ connection.endpoint.closed().child_token()
+ } else {
+ connection.endpoint.invocations().child_token()
+ };
let registration = Arc::new(Registration {
provider_id: connection.provider_id.clone(),
connection_id,
@@ -171,19 +120,45 @@ impl Registry {
manifest,
endpoint: connection.endpoint.clone(),
residency: provider.residency.clone(),
+ invocations,
});
- if let Some(old) = provider.active.filter(|old| *old != connection_id)
+ resident.insert(
+ registration.manifest.registration_id.clone(),
+ Arc::downgrade(®istration),
+ );
+ drop(resident);
+ if registration.session_id().is_none()
+ && let Some(old) = provider.active.filter(|old| *old != connection_id)
&& let Some(connection) = self.connections.get_mut(&old)
{
connection.superseded = true;
connection.endpoint.supersede();
}
- provider.active = Some(connection_id);
let result = RegistrationResult {
registration_id: registration.manifest.registration_id.clone(),
revision,
};
- let previous = provider.current.replace(registration.clone());
+ let previous = match registration.session_id() {
+ Some(id) => provider.scoped.insert(id.into(), registration.clone()),
+ None => {
+ provider.active = Some(connection_id);
+ provider.current.replace(registration.clone())
+ }
+ };
+ // Explicit withdrawal can leave pinned generations without a current
+ // slot. A new connection still takes over their publication scope.
+ let resident: Vec<_> = provider
+ .residency
+ .lock()
+ .unwrap_or_else(|e| e.into_inner())
+ .values()
+ .filter_map(std::sync::Weak::upgrade)
+ .collect();
+ for r in resident {
+ if r.session_id() == registration.session_id() && r.connection_id != connection_id {
+ r.invocations.cancel();
+ }
+ }
self.revision = revision;
self.publication_changed(previous.as_deref(), ®istration);
Ok(result)
@@ -202,23 +177,28 @@ impl Registry {
.providers
.get_mut(&connection.provider_id)
.expect("attached provider");
- if provider.active != Some(connection_id)
- || provider
- .current
- .as_ref()
- .is_none_or(|r| r.manifest.registration_id != registration_id)
- {
- return Err(Error::Invalid("registration is not current"));
- }
+ let previous = provider
+ .current
+ .iter()
+ .chain(provider.scoped.values())
+ .find(|r| {
+ r.connection_id == connection_id && r.manifest.registration_id == registration_id
+ })
+ .cloned()
+ .ok_or(Error::Invalid("registration is not current"))?;
let revision = self
.revision
.checked_add(1)
.filter(|r| *r <= 9_007_199_254_740_991)
.ok_or(Error::Invalid("revision exhausted"))?;
- let previous = provider
- .current
- .take()
- .expect("validated current registration");
+ match previous.session_id() {
+ Some(id) => {
+ provider.scoped.remove(id);
+ }
+ None => {
+ provider.current.take();
+ }
+ }
self.revision = revision;
self.publication_removed(&previous);
Ok(RegistrationResult {
@@ -226,56 +206,4 @@ impl Registry {
revision,
})
}
-
- /// Disconnection marks session bindings Lost, unlike explicit unregister's
- /// removal. Inactive socket teardown cannot clear its replacement.
- pub fn detach(&mut self, id: Uuid) {
- let Some(connection) = self.connections.remove(&id) else {
- return;
- };
- connection.endpoint.close();
- if let Some(provider) = self.providers.get_mut(&connection.provider_id)
- && provider.active == Some(id)
- {
- provider.active = None;
- if provider.current.take().is_some() {
- self.revision += 1;
- }
- self.provider_lost(&connection.provider_id);
- }
- self.prune();
- }
-
- pub fn begin_drain(&mut self) {
- self.draining = true;
- for connection in self.connections.values() {
- connection.endpoint.close();
- }
- self.connections.clear();
- self.providers.clear();
- self.sessions.clear();
- }
-
- fn prune(&mut self) {
- self.providers.retain(|id, provider| {
- provider.current.is_some()
- || provider.active.is_some()
- || Arc::strong_count(&provider.identity) > 1
- || self
- .connections
- .values()
- .any(|connection| &connection.provider_id == id)
- || !provider
- .residency
- .lock()
- .unwrap_or_else(|e| e.into_inner())
- .is_empty()
- });
- }
-}
-
-impl Drop for Registry {
- fn drop(&mut self) {
- self.begin_drain();
- }
}
diff --git a/crates/client-capability/src/registry/publication.rs b/crates/client-capability/src/registry/publication.rs
index d9c195dcef..bd51d5bdef 100644
--- a/crates/client-capability/src/registry/publication.rs
+++ b/crates/client-capability/src/registry/publication.rs
@@ -20,12 +20,13 @@
use crate::{ContractId, Endpoint, Identity};
use maka_runtime::capability::{HostFrame, Manifest, Offer};
use std::{
- collections::{BTreeMap, HashSet},
- sync::{Arc, Mutex},
+ collections::{BTreeMap, HashMap},
+ sync::{Arc, Mutex, Weak},
};
+use tokio_util::sync::CancellationToken;
use uuid::Uuid;
-pub(super) type Residency = Arc<Mutex<HashSet<String>>>;
+pub(super) type Residency = Arc<Mutex<HashMap<String, Weak<Registration>>>>;
/// An immutable publication. Runs and invocations pin this exact allocation.
pub struct Registration {
@@ -36,6 +37,7 @@ pub struct Registration {
pub(super) contracts: BTreeMap<ContractId, usize>,
pub(super) endpoint: Endpoint,
pub(super) residency: Residency,
+ pub(super) invocations: CancellationToken,
}
impl Registration {
@@ -64,7 +66,19 @@ impl Registration {
}
pub fn available(&self) -> bool {
- !self.endpoint.closed().is_cancelled() && !self.endpoint.invocations().is_cancelled()
+ !self.invocations.is_cancelled()
+ }
+
+ pub fn invocations(&self) -> CancellationToken {
+ self.invocations.clone()
+ }
+
+ pub fn session_id(&self) -> Option<&str> {
+ self.manifest.session_id.as_deref()
+ }
+
+ pub fn visible_to(&self, session_id: Option<&str>) -> bool {
+ self.session_id().is_none() || self.session_id() == session_id
}
pub fn endpoint(&self) -> &Endpoint {
@@ -86,7 +100,19 @@ impl Drop for Registration {
impl super::Registry {
pub fn current(&self, provider_id: &str) -> Option<Arc<Registration>> {
- self.providers.get(provider_id)?.current.clone()
+ self.current_scoped(provider_id, None)
+ }
+
+ pub(super) fn current_scoped(
+ &self,
+ provider_id: &str,
+ session_id: Option<&str>,
+ ) -> Option<Arc<Registration>> {
+ let provider = self.providers.get(provider_id)?;
+ match session_id {
+ Some(id) => provider.scoped.get(id).cloned(),
+ None => provider.current.clone(),
+ }
}
pub fn current_for_connection(&self, id: Uuid) -> Option<Arc<Registration>> {
@@ -99,6 +125,6 @@ impl super::Registry {
pub fn published(&self) -> impl Iterator<Item = &Arc<Registration>> {
self.providers
.values()
- .filter_map(|provider| provider.current.as_ref())
+ .flat_map(|provider| provider.current.iter().chain(provider.scoped.values()))
}
}
diff --git a/crates/client-capability/tests/bindings.rs b/crates/client-capability/tests/bindings.rs
index d88296905c..f4fd25ce42 100644
--- a/crates/client-capability/tests/bindings.rs
+++ b/crates/client-capability/tests/bindings.rs
@@ -51,6 +51,161 @@ fn attach(
(id, provider, out)
}
+#[tokio::test]
+async fn scoped_publications_isolate_selection_restore_and_retire_all_pinned_generations() {
+ let mut registry = Registry::default();
+ let (one, _, _out) = attach(&mut registry, identity("desktop"));
+ let scoped = |id: &str, session: &str| {
+ let mut m = manifest(id, &["session"]);
+ m.session_id = Some(session.into());
+ m
+ };
+ registry.replace(one, scoped("a1", "a")).unwrap();
+ registry.replace(one, scoped("b1", "b")).unwrap();
+ assert!(
+ registry
+ .replace(one, manifest("collision", &["session"]))
+ .is_err()
+ );
+ assert!(
+ registry
+ .preview_bindings(None, Some(one), BindingMode::Strict)
+ .unwrap()
+ .offers()
+ .is_empty()
+ );
+ let (prepared, a1) = registry
+ .prepare_bindings("a", Some(one), BindingMode::Strict)
+ .unwrap();
+ let composition = prepared.composition();
+ assert!(registry.commit_bindings(prepared).unwrap());
+ registry
+ .bind_session("b", Some(one), BindingMode::Strict)
+ .unwrap();
+ let b = registry.snapshot("b").unwrap();
+ let broker = maka_client_capability::broker::Broker::default();
+ let call = |session: &str| maka_client_capability::broker::ToolCall {
+ offer_id: "session".into(),
+ server_id: "session".into(),
+ tool_name: "effect".into(),
+ arguments: Default::default(),
+ source: maka_runtime::capability::CallSource::Agent {
+ session_id: session.into(),
+ turn_id: "turn".into(),
+ },
+ tool_call_id: "tool".into(),
+ cwd: "/not-forwarded".into(),
+ };
+ assert!(
+ broker
+ .prepare_tool(
+ b.offers()[0].resolve(®istry).unwrap(),
+ call("a"),
+ std::time::Duration::from_secs(1),
+ Default::default()
+ )
+ .is_err()
+ );
+ let pending = broker
+ .prepare_tool(
+ b.offers()[0].resolve(®istry).unwrap(),
+ call("b"),
+ std::time::Duration::from_secs(10),
+ Default::default(),
+ )
+ .unwrap();
+ assert_eq!(a1.offers().len(), 1);
+ assert_eq!(
+ a1.offers()[0]
+ .resolve(®istry)
+ .unwrap()
+ .manifest()
+ .registration_id,
+ "a1"
+ );
+ assert_eq!(
+ b.offers()[0]
+ .resolve(®istry)
+ .unwrap()
+ .manifest()
+ .registration_id,
+ "b1"
+ );
+ assert!(matches!(
+ registry.restore_bindings("b", &composition),
+ Err(BindingError::InvalidComposition)
+ ));
+ registry.replace(one, scoped("a2", "a")).unwrap();
+ let a2 = registry.snapshot("a").unwrap();
+ assert!(
+ a1.offers()[0].resolve(®istry).is_ok(),
+ "same connection replacement keeps frozen calls"
+ );
+ let (two, _, _out2) = attach(&mut registry, identity("desktop"));
+ registry.unregister(one, "a2").unwrap();
+ assert!(a2.offers()[0].resolve(®istry).is_ok());
+ registry.replace(two, scoped("a3", "a")).unwrap();
+ assert!(a1.offers()[0].resolve(®istry).is_err());
+ assert!(a2.offers()[0].resolve(®istry).is_err());
+ assert!(
+ b.offers()[0].resolve(®istry).is_ok(),
+ "another Session keeps its original connection"
+ );
+ let (prepared, _) = registry.restore_bindings("a", &composition).unwrap();
+ assert!(registry.commit_restored_bindings(prepared).unwrap());
+ let a3 = registry.snapshot("a").unwrap();
+ registry
+ .replace(one, manifest("global", &["turn"]))
+ .unwrap();
+ registry
+ .replace(two, manifest("global-reconnected", &["turn"]))
+ .unwrap();
+ assert!(
+ b.offers()[0].resolve(®istry).is_ok(),
+ "global supersession must not cancel scoped calls"
+ );
+ registry.replace(one, scoped("b2", "b")).unwrap();
+ registry.detach(two);
+ assert!(a3.offers()[0].resolve(®istry).is_err());
+ assert!(
+ registry
+ .bind_session("a", None, BindingMode::Strict)
+ .is_err()
+ );
+ let mut empty = scoped("a-empty", "a");
+ empty.offers.clear();
+ registry.replace(one, empty).unwrap();
+ registry
+ .bind_session("a", None, BindingMode::Strict)
+ .unwrap();
+ assert!(
+ registry.snapshot("a").unwrap().offers().is_empty(),
+ "empty scoped publication clears lost contracts"
+ );
+ let (prepared, _) = registry
+ .prepare_bindings("b", Some(one), BindingMode::Strict)
+ .unwrap();
+ registry.release_session("b");
+ assert!(!registry.commit_bindings(prepared).unwrap());
+ assert!(
+ b.offers()[0].resolve(®istry).is_err(),
+ "retirement reaches pinned prior generations"
+ );
+ assert!(matches!(
+ pending.accepted().await,
+ Err(maka_client_capability::broker::CallError::CapabilityLost)
+ ));
+ broker.shutdown().await;
+ let (prepared, _) = registry
+ .prepare_bindings("new", None, BindingMode::Strict)
+ .unwrap();
+ registry.release_session("new");
+ assert!(
+ !registry.commit_bindings(prepared).unwrap(),
+ "retirement also invalidates an empty preview"
+ );
+}
+
#[test]
fn restart_restores_identity_ceiling_without_discovery_or_call_pinning() {
for initiating in [false, true] {
diff --git a/crates/client-capability/tests/broker.rs b/crates/client-capability/tests/broker.rs
index 7a407a2ea8..e2e0e81d38 100644
--- a/crates/client-capability/tests/broker.rs
+++ b/crates/client-capability/tests/broker.rs
@@ -56,6 +56,7 @@ fn setup(
.replace(
connection,
Manifest {
+ session_id: None,
registration_id: "r".into(),
offers: vec![],
services: Some(vec![ServiceOffer {
diff --git a/crates/client-capability/tests/managed_target.rs b/crates/client-capability/tests/managed_target.rs
index 3da4ad6eec..4cead41cec 100644
--- a/crates/client-capability/tests/managed_target.rs
+++ b/crates/client-capability/tests/managed_target.rs
@@ -237,7 +237,35 @@ fn managed_policy_is_exact_and_requires_trusted_resolved_publication() {
);
}
}
+ for kind in [PrincipalKind::LocalOwner, PrincipalKind::RemoteOwner] {
+ let mut publication = manifest("acp_mcp", "server", &["read"], "session");
+ publication.session_id = Some("s".into());
+ publication.offers[0].admission = Some(maka_runtime::capability::Admission::Mcp);
+ let (registry, _, _) = registry(publication, kind);
+ let snapshot = registry.snapshot("s").unwrap();
+ let offer = &snapshot.offers()[0];
+ let target = offer.managed_target(
+ &offer.resolve(®istry).unwrap(),
+ "server",
+ "read",
+ &AdmissionEvidence::None,
+ );
+ if kind == PrincipalKind::RemoteOwner {
+ assert_eq!(target, Err(Error::UntrustedProvider));
+ } else {
+ let target = target.unwrap().unwrap();
+ assert_eq!(target.capability, GrantCapability::Mcp);
+ assert_eq!(
+ target.scope,
+ GrantScope::McpTool {
+ server_id: "server".into(),
+ tool_name: "read".into()
+ }
+ );
+ }
+ }
for (id, server, name) in [
+ ("acp_mcp", "real_server", "read"),
("desktop_browser", "desktop_browser", "browser_unknown"),
("desktop_browser", "impostor", "browser_navigate"),
("impostor", "desktop_browser", "browser_navigate"),
diff --git a/crates/client-capability/tests/support/forms.rs b/crates/client-capability/tests/support/forms.rs
index 71579183af..8cdc88c3ad 100644
--- a/crates/client-capability/tests/support/forms.rs
+++ b/crates/client-capability/tests/support/forms.rs
@@ -59,6 +59,7 @@ pub fn setup(
.replace(
connection,
Manifest {
+ session_id: None,
registration_id: "r".into(),
offers: vec![],
services: Some(vec![ServiceOffer {
diff --git a/crates/protocol/src/capability/manifest.rs b/crates/protocol/src/capability/manifest.rs
index d3a0654dd6..4fbc18ee3f 100644
--- a/crates/protocol/src/capability/manifest.rs
+++ b/crates/protocol/src/capability/manifest.rs
@@ -31,9 +31,20 @@ use std::collections::HashSet;
pub fn decode_replace_input(value: &Value) -> Result<Manifest> {
let frame = record(value, "Client Capability replacement")?;
- shaped(frame, &["registrationId", "offers"], &["services"])?;
+ shaped(
+ frame,
+ &["registrationId", "offers"],
+ &["services", "sessionId"],
+ )?;
+ let session_id = frame
+ .get("sessionId")
+ .map(|value| entity(value, "sessionId"))
+ .transpose()?;
let values = array(&frame["offers"], MAX_OFFERS)?;
let offers = values.iter().map(offer).collect::<Result<Vec<_>>>()?;
+ if session_id.is_none() && offers.iter().any(|o| o.admission.is_some()) {
+ return Err(invalid("MCP admission requires a target Session"));
+ }
// The source treats explicit null as an empty services array, preserving
// field presence in its normalized value; omission remains omission.
let services = match frame.get("services") {
@@ -46,9 +57,19 @@ pub fn decode_replace_input(value: &Value) -> Result<Manifest> {
.collect::<Result<Vec<_>>>()?,
),
};
- if offers.is_empty() && services.as_ref().is_none_or(Vec::is_empty) {
+ if session_id.is_none() && offers.is_empty() && services.as_ref().is_none_or(Vec::is_empty) {
return Err(invalid("Client Capability registration is empty"));
}
+ if session_id.is_some()
+ && (services.as_ref().is_some_and(|s| !s.is_empty())
+ || offers.iter().any(|o| {
+ o.affinity != Affinity::Session || o.host_path_access != HostPathAccess::None
+ }))
+ {
+ return Err(invalid(
+ "Session publications support only path-independent Session tools",
+ ));
+ }
let mut offer_ids = HashSet::new();
let mut tool_ids = HashSet::new();
for offer in &offers {
@@ -72,6 +93,7 @@ pub fn decode_replace_input(value: &Value) -> Result<Manifest> {
}
let manifest = Manifest {
registration_id: entity(&frame["registrationId"], "registrationId")?,
+ session_id,
offers,
services,
};
@@ -108,7 +130,7 @@ fn offer(value: &Value) -> Result<Offer> {
"label",
"tools",
],
- &["description"],
+ &["description", "admission"],
)?;
let tools = array(&frame["tools"], MAX_TOOLS_PER_OFFER)?;
if tools.is_empty() {
@@ -128,6 +150,13 @@ fn offer(value: &Value) -> Result<Offer> {
Some("cwd") => HostPathAccess::Cwd,
_ => return Err(invalid("Invalid Client Capability host path access")),
},
+ admission: frame
+ .get("admission")
+ .map(|value| match value.as_str() {
+ Some("mcp") => Ok(Admission::Mcp),
+ _ => Err(invalid("Invalid Client Capability admission")),
+ })
+ .transpose()?,
label: string(&frame["label"], "label", 128)?,
description: optional_string(frame, "description", 1024)?,
tools: tools.iter().map(tool).collect::<Result<_>>()?,
diff --git a/crates/protocol/tests/capability_contract.rs b/crates/protocol/tests/capability_contract.rs
index a6e9c2c811..1f33bc495e 100644
--- a/crates/protocol/tests/capability_contract.rs
+++ b/crates/protocol/tests/capability_contract.rs
@@ -71,6 +71,25 @@ fn manifests_normalize_and_reject_the_same_boundaries_as_the_current_client() {
cases.push(json!({"kind":kind,"input":input,"expected":expected}));
};
add("manifest", manifest());
+ let mut scoped = manifest();
+ scoped["sessionId"] = json!("session-a");
+ add("manifest", scoped.clone()); // Services cannot be scoped.
+ scoped.as_object_mut().unwrap().remove("services");
+ add("manifest", scoped.clone());
+ scoped["offers"][0]["admission"] = json!("mcp");
+ add("manifest", scoped.clone());
+ let mut global_mcp = scoped.clone();
+ global_mcp.as_object_mut().unwrap().remove("sessionId");
+ add("manifest", global_mcp);
+ scoped["offers"][0]["hostPathAccess"] = json!("cwd");
+ add("manifest", scoped.clone());
+ scoped["offers"][0]["hostPathAccess"] = json!("none");
+ scoped["offers"][0]["affinity"] = json!("turn");
+ add("manifest", scoped.clone());
+ scoped["offers"] = json!([]);
+ add("manifest", scoped.clone());
+ scoped["sessionId"] = json!("invalid.id");
+ add("manifest", scoped);
let mut absent = manifest();
absent.as_object_mut().unwrap().remove("services");
add("manifest", absent);
diff --git a/crates/protocol/tests/interaction_contract.rs b/crates/protocol/tests/interaction_contract.rs
index 6bc3a11217..01a100cbba 100644
--- a/crates/protocol/tests/interaction_contract.rs
+++ b/crates/protocol/tests/interaction_contract.rs
@@ -112,6 +112,14 @@ fn original_typescript_agrees_on_grants_lifecycle_identity_and_bounds() {
add("request", value);
}
for (capability, scope) in [
+ (
+ "mcp",
+ json!({"kind":"mcp_tool","serverId":"server","toolName":"read"}),
+ ),
+ (
+ "mcp",
+ json!({"kind":"browser_origin","origin":"https://example.com"}),
+ ),
("computer_use", json!({"kind":"capability"})),
(
"computer_use",
diff --git a/crates/runtime-host/src/execution/plugins/scopes.rs b/crates/runtime-host/src/execution/plugins/scopes.rs
index cfd7a6325c..62c3107f48 100644
--- a/crates/runtime-host/src/execution/plugins/scopes.rs
+++ b/crates/runtime-host/src/execution/plugins/scopes.rs
@@ -254,7 +254,11 @@ impl Executions {
scope: &Scope,
capability: Capability,
) -> Result<Vec<std::sync::Arc<maka_client_capability::Registration>>, Error> {
- self.plugin_resource_boundary(scope, capability).await?;
+ let boundary = self.plugin_resource_boundary(scope, capability).await?;
+ let session_id = match &boundary {
+ Boundary::Session { boundary, .. } => Some(boundary.session_id.as_str()),
+ _ => None,
+ };
let evidence = self
.plugin_calls
.evidence::<Evidence>(scope)
@@ -283,6 +287,7 @@ impl Executions {
.filter(|registration| {
let provider = registration.identity();
registration.available()
+ && registration.visible_to(session_id)
&& (**provider == identity
|| (provider.trusted()
&& identity.credential_bound_client_instance_id.as_deref()
diff --git a/crates/runtime-host/src/server/capabilities.rs b/crates/runtime-host/src/server/capabilities.rs
index 6c9f3bd76f..05197688ba 100644
--- a/crates/runtime-host/src/server/capabilities.rs
+++ b/crates/runtime-host/src/server/capabilities.rs
@@ -232,6 +232,24 @@ pub(super) async fn execute(
message: "Host is draining".into(),
}));
}
+ let manifest = if operation == Operation::ClientCapabilityReplace {
+ let manifest = capability::decode_replace_input(input)?;
+ if let Some(session_id) = &manifest.session_id
+ && host
+ .log
+ .get_session::<crate::session::SessionConfiguration>(session_id)
+ .await?
+ .is_some_and(|session| session.archived)
+ {
+ return Ok(Outcome::failure(OperationError {
+ code: Code::InvalidRequest,
+ message: "Cannot publish capabilities for an archived Session".into(),
+ }));
+ }
+ Some(manifest)
+ } else {
+ None
+ };
let result = {
let mut registry = host
.capabilities
@@ -240,7 +258,7 @@ pub(super) async fn execute(
.unwrap_or_else(|e| e.into_inner());
match operation {
Operation::ClientCapabilityReplace => {
- registry.replace(connection_id, capability::decode_replace_input(input)?)
+ registry.replace(connection_id, manifest.expect("decoded replacement"))
}
Operation::ClientCapabilityUnregister => registry.unregister(
connection_id,
diff --git a/crates/runtime-host/src/server/sessions.rs b/crates/runtime-host/src/server/sessions.rs
index 35d5eff7e9..f411c447ce 100644
--- a/crates/runtime-host/src/server/sessions.rs
+++ b/crates/runtime-host/src/server/sessions.rs
@@ -156,6 +156,13 @@ pub(super) async fn execute(
.await
.map_err(stored)?;
let item = item(record);
+ if input.state == SessionLifecycleState::Archived {
+ host.capabilities
+ .registry
+ .lock()
+ .unwrap_or_else(|e| e.into_inner())
+ .release_session(&input.session_id);
+ }
assert_lifecycle_output_for_input(&input, &item).map_err(invalid)?;
Ok(Output::Item(Box::new(item)))
}
diff --git a/crates/runtime-host/tests/integration/plugin_authorization.rs b/crates/runtime-host/tests/integration/plugin_authorization.rs
index 0d157f6a23..c4951c8c98 100644
--- a/crates/runtime-host/tests/integration/plugin_authorization.rs
+++ b/crates/runtime-host/tests/integration/plugin_authorization.rs
@@ -154,6 +154,10 @@ async fn scenario() {
success(peer.rpc("session.create", json!({"sessionId":"background-session", "workspace":{"kind":"host_path","path":fixture.workspace}, "executorId":"example.background","sandboxMode":"danger-full-access"})).await);
}
let mut publication = super::plugin_clients::publication("desktop", "inspect");
+ let mut scoped = publication.clone();
+ scoped["sessionId"] = json!("background-session");
+ publication["registrationId"] = json!("global-services");
+ publication["offers"] = json!([]);
publication["services"] = json!([{"serviceId":"maka_notifications","version":"1"}]);
let mut turn_only =
super::plugin_clients::publication("desktop-turn", "turn_only")["offers"][0].clone();
@@ -163,6 +167,10 @@ async fn scenario() {
.unwrap()
.push(turn_only);
success(peer.rpc("client.capability.replace", publication).await);
+ success(peer.rpc("client.capability.replace", scoped).await);
+ let mut unrelated = super::plugin_clients::publication("other-session", "unexpected");
+ unrelated["sessionId"] = json!("unrelated-session");
+ success(peer.rpc("client.capability.replace", unrelated).await);
let mut foreign = Peer::new(host.clone(), "another-client").await;
success(
foreign
diff --git a/crates/runtime/src/capability.rs b/crates/runtime/src/capability.rs
index 804f61b450..f7fe993dc7 100644
--- a/crates/runtime/src/capability.rs
+++ b/crates/runtime/src/capability.rs
@@ -25,7 +25,7 @@ mod form_validation;
pub mod json;
pub use contract::ContractId;
mod composition;
-pub use composition::{ClientComposition, ClientOffer, PinnedAffinity};
+pub use composition::{ClientComposition, ClientOffer, PinnedAffinity, PublicationIdentity};
mod identity;
pub use identity::{Identity, PrincipalKind};
mod result;
@@ -44,6 +44,8 @@ use serde_json::{Map, Value};
#[serde(rename_all = "camelCase")]
pub struct Manifest {
pub registration_id: String,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub session_id: Option<String>,
pub offers: Vec<Offer>,
#[serde(skip_serializing_if = "Option::is_none")]
pub services: Option<Vec<ServiceOffer>>,
@@ -56,12 +58,20 @@ pub struct Offer {
pub version: String,
pub affinity: Affinity,
pub host_path_access: HostPathAccess,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub admission: Option<Admission>,
pub label: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub tools: Vec<ToolDescriptor>,
}
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Admission {
+ Mcp,
+}
+
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServiceOffer {
diff --git a/crates/runtime/src/capability/composition.rs b/crates/runtime/src/capability/composition.rs
index d4c2e958e9..badf0a7c88 100644
--- a/crates/runtime/src/capability/composition.rs
+++ b/crates/runtime/src/capability/composition.rs
@@ -26,7 +26,7 @@ use std::collections::BTreeMap;
#[serde(deny_unknown_fields)]
pub struct ClientComposition {
/// Includes unavailable Session owners, which still constrain future binding.
- pub session_bindings: BTreeMap<ContractId, Identity>,
+ pub session_bindings: BTreeMap<ContractId, PublicationIdentity>,
pub offers: Vec<ClientOffer>,
}
@@ -36,7 +36,7 @@ pub enum ClientOffer {
Pinned {
contract: ContractId,
affinity: PinnedAffinity,
- identity: Identity,
+ publication: PublicationIdentity,
},
Call {
offer: Offer,
@@ -44,6 +44,14 @@ pub enum ClientOffer {
},
}
+/// Authentication authority and publication scope are independent identities.
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct PublicationIdentity {
+ pub identity: Identity,
+ pub session_id: Option<String>,
+}
+
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PinnedAffinity {
diff --git a/crates/runtime/src/capability/contract.rs b/crates/runtime/src/capability/contract.rs
index 0eb3426367..2df69c4de4 100644
--- a/crates/runtime/src/capability/contract.rs
+++ b/crates/runtime/src/capability/contract.rs
@@ -41,6 +41,7 @@ impl ContractId {
"version": offer.version,
"affinity": offer.affinity,
"hostPathAccess": offer.host_path_access,
+ "admission": offer.admission,
"tools": tools,
});
let mut canonical = String::new();
diff --git a/crates/runtime/src/capability/frame.rs b/crates/runtime/src/capability/frame.rs
index 2a08d27bac..599b3837af 100644
--- a/crates/runtime/src/capability/frame.rs
+++ b/crates/runtime/src/capability/frame.rs
@@ -51,6 +51,17 @@ pub enum CallSource {
},
}
+impl CallSource {
+ pub fn session_id(&self) -> Option<&str> {
+ match self {
+ Self::Agent { session_id, .. } => Some(session_id),
+ Self::Remote { session_id, .. } | Self::Background { session_id, .. } => {
+ session_id.as_deref()
+ }
+ }
+ }
+}
+
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all_fields = "camelCase")]
pub enum ClientFrame {
diff --git a/crates/runtime/src/interaction.rs b/crates/runtime/src/interaction.rs
index b1d19fb14d..56e663463c 100644
--- a/crates/runtime/src/interaction.rs
+++ b/crates/runtime/src/interaction.rs
@@ -33,6 +33,7 @@ pub enum GrantCapability {
Browser,
ComputerUse,
DesktopMcp,
+ Mcp,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
diff --git a/crates/runtime/src/interaction/validation.rs b/crates/runtime/src/interaction/validation.rs
index 6093052fee..83b64e9e85 100644
--- a/crates/runtime/src/interaction/validation.rs
+++ b/crates/runtime/src/interaction/validation.rs
@@ -56,7 +56,7 @@ impl GrantTarget {
}
(GrantCapability::ComputerUse, GrantScope::Capability {}) => {}
(
- GrantCapability::DesktopMcp,
+ GrantCapability::DesktopMcp | GrantCapability::Mcp,
GrantScope::McpTool {
server_id,
tool_name,
diff --git a/docs/rust-runtime.md b/docs/rust-runtime.md
index 5a471320f4..a6150cc3eb 100644
--- a/docs/rust-runtime.md
+++ b/docs/rust-runtime.md
@@ -329,6 +329,13 @@ The runtime core has no V8 or SQLite dependency. SQLx migrations own persistent
schema changes. Client Capability registration and reverse-call ownership live in
`client-capability`; the host composes them with execution.
+Client capabilities have a connection-wide slot and bounded per-Session slots.
+Session publications contain only path-independent Session-affinity tools; frozen
+bindings retain their publication scope across reconnects. Archiving a Session
+retires all its publication generations and rejects later registration until reopened.
+MCP admission derives an exact tool grant from a trusted, frozen publication;
+declaring MCP admission does not grant provider trust.
+
## Development
`node scripts/rust/release-cli.mjs --source <source.tar.gz> --keys <KEYS>
diff --git a/docs/rust-runtime.zh-CN.md b/docs/rust-runtime.zh-CN.md
index f8b2d58df4..c8f6245d54 100644
--- a/docs/rust-runtime.zh-CN.md
+++ b/docs/rust-runtime.zh-CN.md
@@ -255,6 +255,11 @@ SQLx 管理的可重建文本投影避免反复解析大型 JSON 结果。
Runtime core 不依赖 V8 或 SQLite；持久 schema 由 SQLx migration 管理。
Client Capability 注册与反向调用所有权位于 `client-capability`，Host 负责组合执行。
+客户端能力分为连接级发布和有数量上限的 Session 级发布。后者仅允许不访问 Host 路径的
+Session-affinity 工具；冻结绑定在重连后仍保留发布作用域。归档 Session 会退休其全部
+发布代次，重新打开前拒绝再次注册。MCP 准入从可信的冻结发布推导精确工具授权；
+声明 MCP 准入方式不赋予提供者信任。
+
## 开发验证
`node scripts/rust/release-cli.mjs --source <源码.tar.gz> --keys <KEYS>
diff --git a/packages/cli/src/mcp-capability-provider.ts b/packages/cli/src/mcp-capability-provider.ts
index ca0849514d..bb15f54aa4 100644
--- a/packages/cli/src/mcp-capability-provider.ts
+++ b/packages/cli/src/mcp-capability-provider.ts
@@ -110,6 +110,13 @@ export function createMcpCapabilityProvider(
await manager.callTool(binding, frame.arguments, {
signal: options.signal,
requestInteraction: options.requestInteraction,
+ ...(options.progress
+ ? {
+ onProgress: (current, total) => {
+ options.progress?.(current, total);
+ },
+ }
+ : {}),
}),
);
},
diff --git a/packages/cli/src/runtime-host-onboarding.ts b/packages/cli/src/runtime-host-onboarding.ts
index 246c9067d3..9e59bd0181 100644
--- a/packages/cli/src/runtime-host-onboarding.ts
+++ b/packages/cli/src/runtime-host-onboarding.ts
@@ -261,7 +261,7 @@ async function runOAuthLogin(
const timeoutMs = remainingTimeout();
try {
await waitForOAuthOperation(
- connection.replaceClientCapabilities(provider, timeoutMs),
+ connection.replaceClientCapabilities(provider, { timeoutMs }),
timeoutMs,
observation.signal,
);
diff --git a/packages/cli/src/tui-mcp-remote-publication.ts b/packages/cli/src/tui-mcp-remote-publication.ts
index b0dec0b88e..2c50725aa6 100644
--- a/packages/cli/src/tui-mcp-remote-publication.ts
+++ b/packages/cli/src/tui-mcp-remote-publication.ts
@@ -40,7 +40,10 @@ import {
type RuntimeHostRemoteProfileIncarnation,
type RuntimeHostReconnectingConnection,
} from '@maka/runtime-host/client';
-import type { ClientCapabilityProvider } from '@maka/runtime-host/client';
+import type {
+ ClientCapabilityProvider,
+ ClientCapabilityRegistrationOptions,
+} from '@maka/runtime-host/client';
import {
tryAcquireFileLifetimeOwner,
type FileLifetimeOwner,
@@ -174,7 +177,10 @@ class RemoteTuiMcpPublicationTarget implements TuiMcpPublicationTarget {
});
}
- replaceClientCapabilities(provider: ClientCapabilityProvider, timeoutMs?: number) {
+ replaceClientCapabilities(
+ provider: ClientCapabilityProvider,
+ options?: ClientCapabilityRegistrationOptions,
+ ) {
return this.#serialize(async () => {
if (this.#closed) throw new Error('Remote MCP publication is closed');
let result:
@@ -183,7 +189,7 @@ class RemoteTuiMcpPublicationTarget implements TuiMcpPublicationTarget {
const committed = await this.#deps.profiles.mutateRemoteProfileIfCurrent(
this.#profileTarget(),
async () => {
- result = await this.#requireConnection().replaceClientCapabilities(provider, timeoutMs);
+ result = await this.#requireConnection().replaceClientCapabilities(provider, options);
},
);
if (!committed) {
@@ -198,8 +204,8 @@ class RemoteTuiMcpPublicationTarget implements TuiMcpPublicationTarget {
});
}
- unregisterClientCapabilities(timeoutMs?: number) {
- return this.#requireConnection().unregisterClientCapabilities(timeoutMs);
+ unregisterClientCapabilities(options?: ClientCapabilityRegistrationOptions) {
+ return this.#requireConnection().unregisterClientCapabilities(options);
}
subscribeConnectionAvailability(
diff --git a/packages/core/src/client-capability-grant.ts b/packages/core/src/client-capability-grant.ts
index d2ffca1f0f..327ba51911 100644
--- a/packages/core/src/client-capability-grant.ts
+++ b/packages/core/src/client-capability-grant.ts
@@ -21,7 +21,7 @@ import { defineObjectShape, hasExactShape } from './record-schema.js';
const SAFE_ID = /^[A-Za-z0-9_-]{1,128}$/;
-export type ClientCapabilityGrantCapability = 'browser' | 'computer_use' | 'desktop_mcp';
+export type ClientCapabilityGrantCapability = 'browser' | 'computer_use' | 'desktop_mcp' | 'mcp';
export type ClientCapabilityGrantScope =
| { readonly kind: 'browser_origin'; readonly origin: string }
@@ -89,14 +89,14 @@ export function decodeClientCapabilityGrantTarget(value: unknown): ClientCapabil
const record = plainRecord(value, 'Client Capability Grant target');
const capability = oneOf(
record.capability,
- ['browser', 'computer_use', 'desktop_mcp'] as const,
+ ['browser', 'computer_use', 'desktop_mcp', 'mcp'] as const,
'capability',
);
const scope = decodeClientCapabilityGrantScope(record.scope);
if (
(capability === 'browser' && scope.kind !== 'browser_origin') ||
(capability === 'computer_use' && scope.kind !== 'capability') ||
- (capability === 'desktop_mcp' && scope.kind !== 'mcp_tool')
+ ((capability === 'desktop_mcp' || capability === 'mcp') && scope.kind !== 'mcp_tool')
) {
throw new Error('Client Capability Session Grant scope does not match capability');
}
diff --git a/packages/mcp/src/__tests__/modern-manager.test.ts b/packages/mcp/src/__tests__/modern-manager.test.ts
index 9554e5d98e..f4da0e5aa3 100644
--- a/packages/mcp/src/__tests__/modern-manager.test.ts
+++ b/packages/mcp/src/__tests__/modern-manager.test.ts
@@ -121,6 +121,39 @@ describe('McpClientManager modern Streamable HTTP E2E', () => {
}
});
+ test('tools/call progress is bounded, opt-in and cannot change the result', async () => {
+ const fixture = await createModernRemoteFixture();
+ const manager = createManager();
+ await manager.sync(modernConfig(fixture.url, 'auto'));
+ for (const key of ['onProgress', 'emitProgress'] as const) {
+ const seen: Array<[number, number]> = [];
+ const result = await manager.callTool(
+ bindingFor(manager, 'echo'),
+ { value: 'go' },
+ {
+ [key]: (current: number, total: number) => {
+ seen.push([current, total]);
+ throw new Error('advisory listener failed');
+ },
+ },
+ );
+ assert.deepEqual(seen, [
+ [1, 3],
+ [3, 3],
+ ]);
+ assert.deepEqual(result, {
+ content: [{ type: 'text', text: 'go' }],
+ structuredContent: undefined,
+ });
+ }
+ await manager.callTool(bindingFor(manager, 'echo'), { value: 'silent' });
+ assert.equal(fixture.toolCalls.get('echo'), 3);
+ assert.equal(fixture.toolCallProgressTokens.length, 3);
+ assert.notEqual(fixture.toolCallProgressTokens[0], undefined);
+ assert.notEqual(fixture.toolCallProgressTokens[1], undefined);
+ assert.equal(fixture.toolCallProgressTokens[2], undefined);
+ });
+
test('surfaces input_required without silently retrying the tool', async () => {
const fixture = await createModernRemoteFixture();
const manager = createManager();
@@ -228,6 +261,7 @@ interface ModernRemoteFixture {
protocolMethods: string[];
toolCalls: Map<string, number>;
toolCallHeaders: Map<string, string | undefined>;
+ toolCallProgressTokens: unknown[];
close(): Promise<void>;
}
@@ -236,6 +270,7 @@ async function createModernRemoteFixture(
): Promise<ModernRemoteFixture> {
const advertiseTools = options.advertiseTools !== false;
const protocolMethods: string[] = [];
+ const toolCallProgressTokens: unknown[] = [];
const toolCalls = new Map<string, number>();
const toolCallHeaders = new Map<string, string | undefined>();
const errors: Error[] = [];
@@ -278,9 +313,26 @@ async function createModernRemoteFixture(
}
return { tools };
});
- server.setRequestHandler('tools/call', async ({ params }) => {
+ server.setRequestHandler('tools/call', async ({ params }, extra) => {
toolCalls.set(params.name, (toolCalls.get(params.name) ?? 0) + 1);
const args = params.arguments ?? {};
+ const progressToken = extra.mcpReq._meta?.progressToken;
+ if (progressToken !== undefined) {
+ for (const [progress, total] of [
+ [1, 3],
+ [1, 3],
+ [0, 3],
+ [2, 4],
+ [2, 1025],
+ [1.5, 3],
+ [3, 3],
+ ]) {
+ await extra.mcpReq.notify({
+ method: 'notifications/progress',
+ params: { progressToken, progress, total },
+ });
+ }
+ }
if (params.name === 'needs-input') {
return inputRequired({ requestState: 'awaiting-confirmation' });
}
@@ -324,6 +376,7 @@ async function createModernRemoteFixture(
for (const toolName of readToolCallNames(body)) {
toolCallHeaders.set(toolName, headerValue(req, 'mcp-param-shard'));
}
+ toolCallProgressTokens.push(...readToolCallProgressTokens(body));
await nodeHandler(req, res, body);
} catch (error) {
errors.push(error instanceof Error ? error : new Error(String(error)));
@@ -342,6 +395,7 @@ async function createModernRemoteFixture(
protocolMethods,
toolCalls,
toolCallHeaders,
+ toolCallProgressTokens,
close: async () => {
await handler.close();
await new Promise<void>((resolve, reject) =>
@@ -385,6 +439,27 @@ function readProtocolMethods(body: unknown): string[] {
});
}
+function readToolCallProgressTokens(body: unknown): unknown[] {
+ return (Array.isArray(body) ? body : [body]).flatMap((message) => {
+ if (
+ typeof message === 'object' &&
+ message !== null &&
+ 'method' in message &&
+ message.method === 'tools/call' &&
+ 'params' in message &&
+ typeof message.params === 'object' &&
+ message.params !== null
+ ) {
+ const meta =
+ '_meta' in message.params && typeof message.params._meta === 'object'
+ ? message.params._meta
+ : undefined;
+ return [meta && 'progressToken' in meta && meta !== null ? meta.progressToken : undefined];
+ }
+ return [];
+ });
+}
+
function readToolCallNames(body: unknown): string[] {
return (Array.isArray(body) ? body : [body]).flatMap((message) => {
if (
diff --git a/packages/mcp/src/__tests__/oauth.test.ts b/packages/mcp/src/__tests__/oauth.test.ts
index c25b265b65..8692c29d4d 100644
--- a/packages/mcp/src/__tests__/oauth.test.ts
+++ b/packages/mcp/src/__tests__/oauth.test.ts
@@ -80,6 +80,12 @@ describe('McpClientManager OAuth E2E', () => {
assert.ok(authorizationUrl.searchParams.get('code_challenge'));
// Dynamic registration ran before the redirect.
assert.ok(fixture.registrations.length >= 1);
+ const registration = fixture.registrations[0];
+ assert.ok(registration && typeof registration === 'object');
+ const body = registration as Record<string, unknown>;
+ assert.equal(body.client_uri, 'https://maka.apache.org/en/');
+ assert.equal(body.software_id, 'maka');
+ assert.equal(body.client_name, 'maka');
// Consent disclosure material: the resolved issuer, the scope the round
// requests, and the round's state travel back to the caller so a UI can
// show what is being granted before a browser opens.
diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts
index 4d7b35d773..c66dd76787 100644
--- a/packages/mcp/src/index.ts
+++ b/packages/mcp/src/index.ts
@@ -84,6 +84,7 @@ import { createMcpToolBinding, parseMcpToolBinding } from './tool-binding.js';
import { McpToolCallError, normalizeToolCallError } from './tool-call-error.js';
import { discoverMcpTools, type McpDiscoveredTool } from './tool-discovery.js';
import { McpToolCallPreparer, type McpToolCallPreparationState } from './tool-output-validation.js';
+import { mapMcpToolProgress } from './tool-progress.js';
import { McpCredentialCoordinator } from './credential-coordinator.js';
import {
assertTransportSecurity,
@@ -804,6 +805,9 @@ export class McpClientManager {
options: {
signal?: AbortSignal;
timeoutMs?: number;
+ onProgress?: (current: number, total: number) => void;
+ /** McpToolProvider name; Desktop in-process tools pass this instead of onProgress. */
+ emitProgress?: (current: number, total: number) => void;
requestInteraction?: (
form: InteractionFormInput,
options?: { cancellationSignal?: AbortSignal },
@@ -901,6 +905,26 @@ export class McpClientManager {
)
throw new McpToolCallError(serverId, toolName, 'tool binding is stale');
};
+ const progressListener = options.onProgress ?? options.emitProgress;
+ let forwardedProgressTotal: number | undefined;
+ let forwardedProgressCurrent = -1;
+ const forwardProgress = (progress: unknown): void => {
+ const mapped = mapMcpToolProgress(progress);
+ if (!mapped) return;
+ if (
+ (forwardedProgressTotal !== undefined && mapped.total !== forwardedProgressTotal) ||
+ mapped.current <= forwardedProgressCurrent
+ ) {
+ return;
+ }
+ forwardedProgressTotal ??= mapped.total;
+ forwardedProgressCurrent = mapped.current;
+ try {
+ progressListener?.(mapped.current, mapped.total);
+ } catch {
+ // Progress is advisory: a listener failure must not fail the tool.
+ }
+ };
try {
const originalArguments = requestInteraction ? structuredClone(args) : args;
let continuation: { inputResponses?: Record<string, ElicitResult>; requestState?: string } =
@@ -924,6 +948,7 @@ export class McpClientManager {
signal,
timeout: options.timeoutMs ?? this.timeouts.callToolMs,
toolDefinition: structuredClone(preparation.value.definitionForSdk),
+ ...(progressListener ? { onprogress: forwardProgress } : {}),
...(requestInteraction ? { allowInputRequired: true } : {}),
},
);
diff --git a/packages/mcp/src/oauth.ts b/packages/mcp/src/oauth.ts
index 9ca8834d82..a8e3aad8e0 100644
--- a/packages/mcp/src/oauth.ts
+++ b/packages/mcp/src/oauth.ts
@@ -177,8 +177,8 @@ export class McpOAuthProvider implements OAuthClientProvider {
get clientMetadata(): OAuthClientMetadata {
return {
client_name: this.options.clientName,
- client_uri: 'https://github.com/maka-agent/maka-agent',
- software_id: 'maka-desktop',
+ client_uri: 'https://maka.apache.org/en/',
+ software_id: this.options.clientName,
software_version: this.options.clientVersion,
redirect_uris: this.options.interactive ? [this.options.interactive.redirectUrl] : [],
grant_types: ['authorization_code', 'refresh_token'],
diff --git a/packages/mcp/src/tool-progress.ts b/packages/mcp/src/tool-progress.ts
new file mode 100644
index 0000000000..fb9848f876
--- /dev/null
+++ b/packages/mcp/src/tool-progress.ts
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export interface McpToolProgress {
+ readonly current: number;
+ readonly total: number;
+}
+
+export const MCP_TOOL_PROGRESS_MAX_TOTAL = 1_024;
+
+/**
+ * Map an SDK `notifications/progress` payload onto the Host progress pair.
+ * Only finite step counts accepted by the shared tool-progress codec.
+ * Missing `total`, fractions, and inverted ranges are
+ * dropped so a noisy server cannot fail the tool call.
+ */
+export function mapMcpToolProgress(value: unknown): McpToolProgress | undefined {
+ if (!isRecord(value)) return undefined;
+ const current = value.progress;
+ const total = value.total;
+ if (typeof current !== 'number' || typeof total !== 'number') return undefined;
+ if (
+ !Number.isSafeInteger(current) ||
+ !Number.isSafeInteger(total) ||
+ current < 0 ||
+ total < 1 ||
+ total > MCP_TOOL_PROGRESS_MAX_TOTAL ||
+ current > total
+ ) {
+ return undefined;
+ }
+ return { current, total };
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+ return typeof value === 'object' && value !== null && !Array.isArray(value);
+}
diff --git a/packages/runtime-host/src/__tests__/capability-provider-service.test.ts b/packages/runtime-host/src/__tests__/capability-provider-service.test.ts
index 7c3854910f..ccd978a525 100644
--- a/packages/runtime-host/src/__tests__/capability-provider-service.test.ts
+++ b/packages/runtime-host/src/__tests__/capability-provider-service.test.ts
@@ -21,7 +21,10 @@ import { deferred } from '@maka/core/test-only/async-primitives';
import assert from 'node:assert/strict';
import { test } from 'node:test';
import { ClientCapabilityChannel } from '../client/client-capability-channel.js';
-import type { ClientCapabilityProvider } from '../client/client-capability.js';
+import type {
+ ClientCapabilityProvider,
+ ClientCapabilityRegistrationOptions,
+} from '../client/client-capability.js';
import { startRuntimeHostCapabilityProviderService } from '../client/capability-provider-service.js';
import type { ClientCapabilityReplaceInput } from '../protocol/index.js';
import { waitFor as pollFor } from '@maka/core/test-only/async-primitives';
@@ -97,12 +100,18 @@ class FakeConnection {
readonly closed = this.#closed.promise;
- replaceClientCapabilities(provider: ClientCapabilityProvider, timeoutMs = 1_000) {
- return this.#channel.replace(provider, timeoutMs);
+ replaceClientCapabilities(
+ provider: ClientCapabilityProvider,
+ { timeoutMs = 1_000, sessionId }: ClientCapabilityRegistrationOptions = {},
+ ) {
+ return this.#channel.replace(provider, timeoutMs, sessionId);
}
- unregisterClientCapabilities(timeoutMs = 1_000) {
- return this.#channel.unregister(timeoutMs);
+ unregisterClientCapabilities({
+ timeoutMs = 1_000,
+ sessionId,
+ }: ClientCapabilityRegistrationOptions = {}) {
+ return this.#channel.unregister(timeoutMs, sessionId);
}
async close(): Promise<void> {
diff --git a/packages/runtime-host/src/__tests__/client-capability-channel.test.ts b/packages/runtime-host/src/__tests__/client-capability-channel.test.ts
index 339ba6094a..923db63a79 100644
--- a/packages/runtime-host/src/__tests__/client-capability-channel.test.ts
+++ b/packages/runtime-host/src/__tests__/client-capability-channel.test.ts
@@ -22,6 +22,74 @@ import { test } from 'node:test';
import { ClientCapabilityChannel } from '../client/client-capability-channel.js';
import type { ClientCapabilityProvider } from '../client/client-capability.js';
+test('Session slots retire independently and reject calls from another Session', async () => {
+ const closed: string[] = [];
+ const writes: unknown[] = [];
+ const removed: string[] = [];
+ const channel = new ClientCapabilityChannel({
+ write: async (frame) => {
+ writes.push(frame);
+ },
+ replace: async (input) => ({ registrationId: input.registrationId, revision: 1 }),
+ unregister: async (input) => {
+ removed.push(input.registrationId);
+ return { registrationId: input.registrationId, revision: 2 };
+ },
+ onFailure: (error) => {
+ throw error;
+ },
+ });
+ const provider = (id: string): ClientCapabilityProvider => ({
+ offers: () => [
+ {
+ offerId: 'same',
+ version: '1',
+ affinity: 'session',
+ hostPathAccess: 'none',
+ label: 'same',
+ tools: [{ serverId: 'mcp', name: 'read', inputSchema: { type: 'object' } }],
+ },
+ ],
+ call: async () => {
+ throw new Error('cross-Session call reached provider');
+ },
+ close: () => {
+ closed.push(id);
+ },
+ });
+ const [a, b] = await Promise.all([
+ channel.replace(provider('a'), 1000, 'a'),
+ channel.replace(provider('b'), 1000, 'b'),
+ ]);
+ channel.accept({
+ kind: 'client.capability.call',
+ invocationId: 'cross-session',
+ registrationId: a.registrationId,
+ offerId: 'same',
+ serverId: 'mcp',
+ toolName: 'read',
+ arguments: {},
+ source: { kind: 'agent', sessionId: 'b', turnId: 'turn' },
+ toolCallId: 'call',
+ });
+ assert.equal((writes[0] as { kind: string }).kind, 'client.capability.rejected');
+ channel.accept({
+ kind: 'client.capability.registration_release',
+ registrationId: a.registrationId,
+ });
+ assert.deepEqual(closed, ['a']);
+ await assert.rejects(channel.unregister(1000, 'a'), /No Client Capability/);
+ await channel.unregister(1000, 'b');
+ assert.deepEqual(removed, [b.registrationId]);
+ channel.accept({
+ kind: 'client.capability.registration_release',
+ registrationId: b.registrationId,
+ });
+ assert.deepEqual(closed, ['a', 'b']);
+ channel.close(new Error('done'));
+ assert.deepEqual(closed, ['a', 'b']);
+});
+
test('Client Capability channel closes a provider after its final registration is released', async () => {
let closeCalls = 0;
const replacements: string[] = [];
@@ -233,6 +301,9 @@ test('Client Capability channel forwards admitted tool progress before the resul
await options.accept({ kind: 'none' });
options.progress?.(1, 3);
options.progress?.(2, 3);
+ options.progress?.(3, 3);
+ options.progress?.(2, 3);
+ options.progress?.(1, 1_025);
return { content: [] };
},
};
@@ -288,7 +359,7 @@ test('Client Capability channel forwards admitted tool progress before the resul
{
kind: 'client.capability.progress',
invocationId: 'progress-invocation',
- current: 2,
+ current: 3,
total: 3,
},
{
diff --git a/packages/runtime-host/src/client/client-capability-channel.ts b/packages/runtime-host/src/client/client-capability-channel.ts
index 4467f961ae..7a9f3ed22b 100644
--- a/packages/runtime-host/src/client/client-capability-channel.ts
+++ b/packages/runtime-host/src/client/client-capability-channel.ts
@@ -20,6 +20,7 @@
import { randomUUID } from 'node:crypto';
import type { InteractionFormInput, InteractionFormResult } from '@maka/core/interaction';
import {
+ CLIENT_CAPABILITY_MAX_PROGRESS_TOTAL,
CLIENT_CAPABILITY_MAX_RESULT_BYTES,
CLIENT_CAPABILITY_RESULT_CHUNK_MAX_BYTES,
decodeClientCapabilityClientFrame,
@@ -39,6 +40,7 @@ import type { ClientCapabilityProvider } from './client-capability.js';
interface ClientCapabilityRegistration {
readonly registrationId: string;
+ readonly sessionId: string | undefined;
readonly provider: ClientCapabilityProvider;
readonly offers: ReturnType<ClientCapabilityProvider['offers']>;
readonly services: NonNullable<ReturnType<NonNullable<ClientCapabilityProvider['services']>>>;
@@ -48,6 +50,7 @@ interface ClientCapabilityInvocation {
readonly controller: AbortController;
admission?: ClientCapabilityAdmission;
interaction?: ClientCapabilityPendingInteraction;
+ progress?: { current: number; total: number };
released: boolean;
}
@@ -90,8 +93,8 @@ export class ClientCapabilityChannel {
readonly #registrations = new Map<string, ClientCapabilityRegistration>();
readonly #invocations = new Map<string, ClientCapabilityInvocation>();
readonly #releasedRegistrationIds = new Set<string>();
- #currentRegistrationId: string | undefined;
- #mutationPending = false;
+ readonly #currentRegistrations = new Map<string | undefined, string>();
+ readonly #mutations = new Set<string | undefined>();
#closedError: Error | undefined;
constructor(options: ClientCapabilityChannelOptions) {
@@ -101,23 +104,26 @@ export class ClientCapabilityChannel {
async replace(
provider: ClientCapabilityProvider,
timeoutMs: number,
+ sessionId?: string,
): Promise<ClientCapabilityReplaceResult> {
this.#assertOpen();
- if (this.#mutationPending) {
+ if (this.#mutations.has(sessionId)) {
throw new Error('A Client Capability registration mutation is already pending');
}
- this.#mutationPending = true;
+ this.#mutations.add(sessionId);
const registrationId = randomUUID();
let registration: ClientCapabilityRegistration | undefined;
try {
const services = provider.services?.() ?? [];
const canonical = decodeClientCapabilityReplaceInput({
registrationId,
+ ...(sessionId === undefined ? {} : { sessionId }),
offers: provider.offers(),
...(services.length === 0 ? {} : { services }),
});
registration = {
registrationId,
+ sessionId,
provider,
offers: canonical.offers,
services: canonical.services ?? [],
@@ -127,39 +133,44 @@ export class ClientCapabilityChannel {
if (result.registrationId !== registrationId) {
throw new Error('Runtime Host replaced a different Client Capability registration');
}
- this.#currentRegistrationId = registrationId;
+ this.#currentRegistrations.set(sessionId, registrationId);
this.#collectReleasedRegistrations();
return result;
} catch (error) {
- if (registration && this.#currentRegistrationId !== registrationId) {
+ if (registration && this.#currentRegistrations.get(sessionId) !== registrationId) {
this.#registrations.delete(registrationId);
}
throw error;
} finally {
- this.#mutationPending = false;
+ this.#mutations.delete(sessionId);
+ this.#collectReleasedRegistrations();
}
}
- async unregister(timeoutMs: number): Promise<ClientCapabilityUnregisterResult> {
+ async unregister(
+ timeoutMs: number,
+ sessionId?: string,
+ ): Promise<ClientCapabilityUnregisterResult> {
this.#assertOpen();
- if (this.#mutationPending) {
+ if (this.#mutations.has(sessionId)) {
throw new Error('A Client Capability registration mutation is already pending');
}
- const registrationId = this.#currentRegistrationId;
+ const registrationId = this.#currentRegistrations.get(sessionId);
if (!registrationId) throw new Error('No Client Capability registration is active');
- this.#mutationPending = true;
+ this.#mutations.add(sessionId);
try {
const result = await this.#options.unregister({ registrationId }, timeoutMs);
if (result.registrationId !== registrationId) {
throw new Error('Runtime Host unregistered a different Client Capability registration');
}
- if (this.#currentRegistrationId === registrationId) {
- this.#currentRegistrationId = undefined;
+ if (this.#currentRegistrations.get(sessionId) === registrationId) {
+ this.#currentRegistrations.delete(sessionId);
}
this.#collectReleasedRegistrations();
return result;
} finally {
- this.#mutationPending = false;
+ this.#mutations.delete(sessionId);
+ this.#collectReleasedRegistrations();
}
}
@@ -236,7 +247,7 @@ export class ClientCapabilityChannel {
);
this.#registrations.clear();
this.#releasedRegistrationIds.clear();
- this.#currentRegistrationId = undefined;
+ this.#currentRegistrations.clear();
for (const provider of providers) this.#closeProvider(provider);
}
@@ -249,7 +260,13 @@ export class ClientCapabilityChannel {
const offered = offer?.tools.some(
(tool) => tool.serverId === frame.serverId && tool.name === frame.toolName,
);
- if (!registration || !offer || !offered || !registration.provider.call) {
+ if (
+ !registration ||
+ !offer ||
+ !offered ||
+ !registration.provider.call ||
+ (registration.sessionId !== undefined && registration.sessionId !== frame.source.sessionId)
+ ) {
void this.#options
.write({
kind: 'client.capability.rejected',
@@ -340,6 +357,38 @@ export class ClientCapabilityChannel {
});
return accepting;
};
+ let pendingProgress: { current: number; total: number } | undefined;
+ let progressWriting = false;
+ let progressWrite: Promise<void> = Promise.resolve();
+ const scheduleProgressWrite = (): void => {
+ if (progressWriting) return;
+ progressWriting = true;
+ progressWrite = (async () => {
+ while (pendingProgress) {
+ const next = pendingProgress;
+ pendingProgress = undefined;
+ try {
+ await this.#options.write({
+ kind: 'client.capability.progress',
+ invocationId,
+ current: next.current,
+ total: next.total,
+ });
+ } catch {
+ // Progress is advisory: a failed write must not fail the invocation.
+ return;
+ }
+ }
+ })().finally(() => {
+ progressWriting = false;
+ });
+ };
+ const flushProgress = async (): Promise<void> => {
+ do {
+ scheduleProgressWrite();
+ await progressWrite;
+ } while (pendingProgress);
+ };
try {
const progress = (current: number, total: number): void => {
if (
@@ -349,18 +398,18 @@ export class ClientCapabilityChannel {
!Number.isInteger(total) ||
current < 0 ||
total < 1 ||
+ total > CLIENT_CAPABILITY_MAX_PROGRESS_TOTAL ||
current > total
) {
return;
}
- void this.#options
- .write({
- kind: 'client.capability.progress',
- invocationId,
- current,
- total,
- })
- .catch((error: unknown) => this.#options.onFailure(asError(error)));
+ const previous = invocation.progress;
+ if (previous && (previous.total !== total || current <= previous.current)) {
+ return;
+ }
+ invocation.progress = { current, total };
+ pendingProgress = { current, total };
+ scheduleProgressWrite();
};
const requestInteraction = async (
request: InteractionFormInput,
@@ -409,6 +458,7 @@ export class ClientCapabilityChannel {
throw new Error('Client Capability provider returned with a pending interaction');
}
await accept({ kind: 'none' });
+ await flushProgress();
await this.#sendResult(invocationId, result, invocation);
} catch (error) {
if (invocation.released) return;
@@ -464,10 +514,13 @@ export class ClientCapabilityChannel {
#collectReleasedRegistrations(): void {
for (const registrationId of this.#releasedRegistrationIds) {
- if (registrationId === this.#currentRegistrationId) continue;
- this.#releasedRegistrationIds.delete(registrationId);
const registration = this.#registrations.get(registrationId);
+ if (registration && this.#mutations.has(registration.sessionId)) continue;
+ this.#releasedRegistrationIds.delete(registrationId);
if (!registration) continue;
+ if (this.#currentRegistrations.get(registration.sessionId) === registrationId) {
+ this.#currentRegistrations.delete(registration.sessionId);
+ }
this.#registrations.delete(registrationId);
if (
![...this.#registrations.values()].some(
diff --git a/packages/runtime-host/src/client/client-capability.ts b/packages/runtime-host/src/client/client-capability.ts
index f0a4946333..3e0f58085c 100644
--- a/packages/runtime-host/src/client/client-capability.ts
+++ b/packages/runtime-host/src/client/client-capability.ts
@@ -28,6 +28,11 @@ import type {
} from '../protocol/index.js';
/** A Client-owned open-world capability provider registered on one Host connection. */
+export interface ClientCapabilityRegistrationOptions {
+ readonly sessionId?: string;
+ readonly timeoutMs?: number;
+}
+
export interface ClientCapabilityProvider {
offers(): readonly ClientCapabilityOffer[];
services?(): readonly ClientCapabilityServiceOffer[];
diff --git a/packages/runtime-host/src/client/connection.ts b/packages/runtime-host/src/client/connection.ts
index 91a50596d3..22dda89409 100644
--- a/packages/runtime-host/src/client/connection.ts
+++ b/packages/runtime-host/src/client/connection.ts
@@ -80,7 +80,10 @@ import {
type RuntimeHostSessionSubscription,
} from './session-subscription.js';
import { ClientCapabilityChannel } from './client-capability-channel.js';
-import type { ClientCapabilityProvider } from './client-capability.js';
+import type {
+ ClientCapabilityProvider,
+ ClientCapabilityRegistrationOptions,
+} from './client-capability.js';
import {
readRuntimeHostProcessIdentity,
type RuntimeHostProcessIdentity,
@@ -277,9 +280,11 @@ export interface RuntimeHostConnection {
close(): Promise<void>;
replaceClientCapabilities(
provider: ClientCapabilityProvider,
- timeoutMs?: number,
+ options?: ClientCapabilityRegistrationOptions,
): Promise<ClientCapabilityReplaceResult>;
- unregisterClientCapabilities(timeoutMs?: number): Promise<ClientCapabilityUnregisterResult>;
+ unregisterClientCapabilities(
+ options?: ClientCapabilityRegistrationOptions,
+ ): Promise<ClientCapabilityUnregisterResult>;
subscribeConfigurationChanges(listener: (revision: number) => void): () => void;
subscribeConnectionCatalogChanges(listener: (revision: number) => void): () => void;
subscribeProjectCatalogChanges(listener: (revision: number) => void): () => void;
@@ -700,15 +705,19 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection {
async replaceClientCapabilities(
provider: ClientCapabilityProvider,
- timeoutMs = DEFAULT_HANDSHAKE_TIMEOUT_MS,
+ {
+ timeoutMs = DEFAULT_HANDSHAKE_TIMEOUT_MS,
+ sessionId,
+ }: ClientCapabilityRegistrationOptions = {},
): Promise<ClientCapabilityReplaceResult> {
- return this.#clientCapabilities.replace(provider, timeoutMs);
+ return this.#clientCapabilities.replace(provider, timeoutMs, sessionId);
}
- async unregisterClientCapabilities(
+ async unregisterClientCapabilities({
timeoutMs = DEFAULT_HANDSHAKE_TIMEOUT_MS,
- ): Promise<ClientCapabilityUnregisterResult> {
- return this.#clientCapabilities.unregister(timeoutMs);
+ sessionId,
+ }: ClientCapabilityRegistrationOptions = {}): Promise<ClientCapabilityUnregisterResult> {
+ return this.#clientCapabilities.unregister(timeoutMs, sessionId);
}
subscribeConfigurationChanges(listener: (revision: number) => void): () => void {
diff --git a/packages/runtime-host/src/client/index.ts b/packages/runtime-host/src/client/index.ts
index a52c1c6921..a756090b31 100644
--- a/packages/runtime-host/src/client/index.ts
+++ b/packages/runtime-host/src/client/index.ts
@@ -168,7 +168,10 @@ export {
type RuntimeHostCandidateLaunchBarrier,
} from './candidate-launch-barrier.js';
export { runHostedExecution, type RunHostedExecutionInput } from './hosted-execution.js';
-export { type ClientCapabilityProvider } from './client-capability.js';
+export {
+ type ClientCapabilityProvider,
+ type ClientCapabilityRegistrationOptions,
+} from './client-capability.js';
export {
readRuntimeHostAgentGraphEpochs,
type AgentGraphEpochDirectory,
diff --git a/packages/runtime-host/src/client/reconnecting-connection.ts b/packages/runtime-host/src/client/reconnecting-connection.ts
index 412c694be2..a5e936455f 100644
--- a/packages/runtime-host/src/client/reconnecting-connection.ts
+++ b/packages/runtime-host/src/client/reconnecting-connection.ts
@@ -29,7 +29,10 @@ import {
type ScheduledTaskChangedFrame,
type SubscriptionOpenInput,
} from '../protocol/index.js';
-import type { ClientCapabilityProvider } from './client-capability.js';
+import type {
+ ClientCapabilityProvider,
+ ClientCapabilityRegistrationOptions,
+} from './client-capability.js';
import {
RuntimeHostOperationError,
RuntimeHostRequestInterruptedError,
@@ -226,17 +229,17 @@ class RuntimeHostReconnectingConnectionImpl implements RuntimeHostReconnectingCo
async replaceClientCapabilities(
provider: ClientCapabilityProvider,
- timeoutMs?: number,
+ options?: ClientCapabilityRegistrationOptions,
): Promise<ClientCapabilityReplaceResult> {
const connection = this.#requireCurrent('client.capability.replace');
- return connection.replaceClientCapabilities(provider, timeoutMs);
+ return connection.replaceClientCapabilities(provider, options);
}
async unregisterClientCapabilities(
- timeoutMs?: number,
+ options?: ClientCapabilityRegistrationOptions,
): Promise<ClientCapabilityUnregisterResult> {
const connection = this.#requireCurrent('client.capability.unregister');
- return connection.unregisterClientCapabilities(timeoutMs);
+ return connection.unregisterClientCapabilities(options);
}
subscribeConfigurationChanges(listener: (revision: number) => void): () => void {
diff --git a/packages/runtime-host/src/protocol/client-capability.ts b/packages/runtime-host/src/protocol/client-capability.ts
index 4ed7b45007..2d040fedb2 100644
--- a/packages/runtime-host/src/protocol/client-capability.ts
+++ b/packages/runtime-host/src/protocol/client-capability.ts
@@ -111,6 +111,7 @@ export interface ClientCapabilityOffer {
readonly version: string;
readonly affinity: ClientCapabilityAffinity;
readonly hostPathAccess: ClientCapabilityHostPathAccess;
+ readonly admission?: 'mcp';
readonly label: string;
readonly description?: string;
readonly tools: readonly ClientCapabilityToolDescriptor[];
@@ -123,6 +124,7 @@ export interface ClientCapabilityServiceOffer {
export interface ClientCapabilityReplaceInput {
readonly registrationId: string;
+ readonly sessionId?: string;
readonly offers: readonly ClientCapabilityOffer[];
readonly services?: readonly ClientCapabilityServiceOffer[];
}
@@ -305,7 +307,7 @@ export function decodeClientCapabilityReplaceInput(value: unknown): ClientCapabi
record,
'Client Capability replacement',
['registrationId', 'offers'],
- ['services'],
+ ['services', 'sessionId'],
);
if (!Array.isArray(record.offers) || record.offers.length > CLIENT_CAPABILITY_MAX_OFFERS) {
throw invalidProtocolFrame('Invalid Client Capability offers');
@@ -314,11 +316,23 @@ export function decodeClientCapabilityReplaceInput(value: unknown): ClientCapabi
if (!Array.isArray(serviceValues) || serviceValues.length > CLIENT_CAPABILITY_MAX_SERVICES) {
throw invalidProtocolFrame('Invalid Client Capability services');
}
- if (record.offers.length === 0 && serviceValues.length === 0) {
+ const sessionId =
+ record.sessionId === undefined ? undefined : requireEntityId(record.sessionId, 'sessionId');
+ if (sessionId === undefined && record.offers.length === 0 && serviceValues.length === 0) {
throw invalidProtocolFrame('Client Capability registration is empty');
}
const offers = record.offers.map((offer) => decodeClientCapabilityOffer(offer));
+ if (sessionId === undefined && offers.some((offer) => offer.admission === 'mcp')) {
+ throw invalidProtocolFrame('MCP admission requires a target Session');
+ }
const services = serviceValues.map((service) => decodeClientCapabilityServiceOffer(service));
+ if (
+ sessionId !== undefined &&
+ (services.length !== 0 ||
+ offers.some((offer) => offer.affinity !== 'session' || offer.hostPathAccess !== 'none'))
+ ) {
+ throw invalidProtocolFrame('Session publications support only path-independent Session tools');
+ }
const offerIds = new Set<string>();
const serviceContracts = new Set<string>();
const toolIdentities = new Set<string>();
@@ -349,6 +363,7 @@ export function decodeClientCapabilityReplaceInput(value: unknown): ClientCapabi
}
const decoded = {
registrationId: requireEntityId(record.registrationId, 'registrationId'),
+ ...(sessionId === undefined ? {} : { sessionId }),
offers,
...(record.services === undefined ? {} : { services }),
};
@@ -741,7 +756,7 @@ function decodeClientCapabilityOffer(value: unknown): ClientCapabilityOffer {
record,
'Client Capability offer',
['offerId', 'version', 'affinity', 'hostPathAccess', 'label', 'tools'],
- ['description'],
+ ['description', 'admission'],
);
if (
!Array.isArray(record.tools) ||
@@ -750,7 +765,11 @@ function decodeClientCapabilityOffer(value: unknown): ClientCapabilityOffer {
) {
throw invalidProtocolFrame('Invalid Client Capability offer tools');
}
+ if (record.admission !== undefined && record.admission !== 'mcp') {
+ throw invalidProtocolFrame('Invalid Client Capability admission');
+ }
return {
+ ...(record.admission === 'mcp' ? { admission: 'mcp' as const } : {}),
offerId: requireEntityId(record.offerId, 'offerId'),
version: requireString(record.version, 'version', 64),
affinity: decodeClientCapabilityAffinity(record.affinity),
diff --git a/packages/ui/src/client-capability-prompt.tsx b/packages/ui/src/client-capability-prompt.tsx
index 0e54cc1603..f41dddd645 100644
--- a/packages/ui/src/client-capability-prompt.tsx
+++ b/packages/ui/src/client-capability-prompt.tsx
@@ -106,6 +106,7 @@ function clientCapabilityLabel(
case 'computer_use':
return copy.computerUse;
case 'desktop_mcp':
+ case 'mcp':
if (request.scope.kind !== 'mcp_tool') break;
return copy.desktopMcp(request.scope.serverId, request.scope.toolName);
}
diff --git a/tests/fixtures/client-capability-host.mjs b/tests/fixtures/client-capability-host.mjs
index 2a15cb1eeb..3bac53ae46 100644
--- a/tests/fixtures/client-capability-host.mjs
+++ b/tests/fixtures/client-capability-host.mjs
@@ -42,6 +42,10 @@ export async function verifyCapabilityHost(connection, workspace, reopened, open
assert.equal(JSON.stringify(await transcript(connection)), await readFile(snapshot, 'utf8'));
for (let index = 1; index <= 3; index++)
assert.equal(await readFile(join(workspace, `effect-${index}`), 'utf8'), `${index}`);
+ await assert.rejects(
+ connection.replaceClientCapabilities({ offers: () => [] }, { sessionId }),
+ (error) => error.code === 'invalid_request',
+ );
console.log('original-client-capability-host-reopened');
return;
}
@@ -94,7 +98,9 @@ export async function verifyCapabilityHost(connection, workspace, reopened, open
if (frame.offerId === 'none') assert(!Object.hasOwn(frame, 'cwd'));
else assert.equal(frame.cwd, await realpath(workspace));
if (index === 1) {
- const replacement = await connection.replaceClientCapabilities(provider(2), 3000);
+ const replacement = await connection.replaceClientCapabilities(provider(2), {
+ timeoutMs: 3000,
+ });
registrations.set(2, replacement.registrationId);
assert.notEqual(replacement.registrationId, registrations.get(1));
assert.deepEqual(closed, [], 'old snapshot must keep its provider alive');
@@ -191,7 +197,7 @@ export async function verifyCapabilityHost(connection, workspace, reopened, open
});
registrations.set(
1,
- (await connection.replaceClientCapabilities(provider(1), 3000)).registrationId,
+ (await connection.replaceClientCapabilities(provider(1), { timeoutMs: 3000 })).registrationId,
);
for (const turnId of ['capability-first', 'capability-second']) {
const live = await watchSession(connection, sessionId, { kind: 'tail', maxBytes: 2 });
@@ -209,7 +215,7 @@ export async function verifyCapabilityHost(connection, workspace, reopened, open
await live.close();
}
}
- await connection.unregisterClientCapabilities(3000);
+ await connection.unregisterClientCapabilities({ timeoutMs: 3000 });
for (let attempt = 0; attempt < 100 && closed.length !== 2; attempt++) await delay(5);
assert.deepEqual(closed.sort(), [1, 2]);
model.verify();
@@ -235,6 +241,28 @@ export async function verifyCapabilityHost(connection, workspace, reopened, open
assert.deepEqual(result.content, { kind: 'json', value: capabilityResult(index + 1) });
}
await writeFile(snapshot, JSON.stringify(rows));
+ const retired = [];
+ await Promise.all(
+ [sessionId, 'not-created-yet'].map((id) =>
+ connection.replaceClientCapabilities(
+ {
+ offers: () => [],
+ close: () => {
+ retired.push(id);
+ },
+ },
+ { sessionId: id },
+ ),
+ ),
+ );
+ await request('session.lifecycle.set', { sessionId, state: 'archived' });
+ for (let attempt = 0; attempt < 100 && retired.length === 0; attempt++) await delay(5);
+ assert.deepEqual(retired, [sessionId], 'archive releases only its scoped publication');
+ await assert.rejects(
+ connection.replaceClientCapabilities({ offers: () => [] }, { sessionId }),
+ (error) => error.code === 'invalid_request',
+ );
+ await connection.unregisterClientCapabilities({ sessionId: 'not-created-yet' });
console.log('original-client-capability-host');
} finally {
await model.close();
diff --git a/tests/fixtures/client-pairing-access.mjs b/tests/fixtures/client-pairing-access.mjs
index 8505e5bb3f..a94747de57 100644
--- a/tests/fixtures/client-pairing-access.mjs
+++ b/tests/fixtures/client-pairing-access.mjs
@@ -127,7 +127,7 @@ export async function verifyPairingAccess(local, handshake, url, control, bounde
const loser = await ready(candidate, 'pairing-losing-client');
await assert.rejects(query(winner), { code: 'unauthorized' }, 'pending winner catalog denied');
await assert.rejects(
- winner.replaceClientCapabilities(capability, 3000),
+ winner.replaceClientCapabilities(capability, { timeoutMs: 3000 }),
{ code: 'unauthorized' },
'pending winner capability publication denied',
);
@@ -173,15 +173,15 @@ export async function verifyPairingAccess(local, handshake, url, control, bounde
{ code: 'unauthorized' },
'associated provider catalog denied',
);
- await providerPeer.replaceClientCapabilities(capability, 3000);
- await providerPeer.unregisterClientCapabilities(3000);
+ await providerPeer.replaceClientCapabilities(capability, { timeoutMs: 3000 });
+ await providerPeer.unregisterClientCapabilities({ timeoutMs: 3000 });
await revoke(candidate);
await bounded(bound.closed, 'bound owner revocation');
assert.equal((await providerPeer.status(3000)).state, 'ready');
await providerPeer.close();
const survivingProvider = await ready(provider, 'pairing-provider-after-owner-revoke');
- await survivingProvider.replaceClientCapabilities(capability, 3000);
- await survivingProvider.unregisterClientCapabilities(3000);
+ await survivingProvider.replaceClientCapabilities(capability, { timeoutMs: 3000 });
+ await survivingProvider.unregisterClientCapabilities({ timeoutMs: 3000 });
await assert.rejects(request(local, 'access.credential.issue', associated(candidate)), {
code: 'invalid_request',
});
diff --git a/tests/fixtures/client-provider-access.mjs b/tests/fixtures/client-provider-access.mjs
index e40176aaaf..b2292f1c7e 100644
--- a/tests/fixtures/client-provider-access.mjs
+++ b/tests/fixtures/client-provider-access.mjs
@@ -140,7 +140,7 @@ export async function verifyProviderAccess(local, handshake, url, control, bound
}
const first = provider();
- const registered = await peer.replaceClientCapabilities(first.value, 3000);
+ const registered = await peer.replaceClientCapabilities(first.value, { timeoutMs: 3000 });
assert.equal(typeof registered.registrationId, 'string');
assert(registered.registrationId.length > 0);
for (const [invalid, code] of [
@@ -149,13 +149,15 @@ export async function verifyProviderAccess(local, handshake, url, control, bound
[provider({ hostPathAccess: 'cwd' }), 'unauthorized'],
]) {
try {
- await assert.rejects(peer.replaceClientCapabilities(invalid.value, 3000), { code });
+ await assert.rejects(peer.replaceClientCapabilities(invalid.value, { timeoutMs: 3000 }), {
+ code,
+ });
} finally {
invalid.value.close();
}
}
// Rejections preserve the active registration and consume no registry revisions.
- const unregistered = await peer.unregisterClientCapabilities(3000);
+ const unregistered = await peer.unregisterClientCapabilities({ timeoutMs: 3000 });
assert.deepEqual(unregistered, {
registrationId: registered.registrationId,
revision: registered.revision + 1,
@@ -164,7 +166,7 @@ export async function verifyProviderAccess(local, handshake, url, control, bound
assert.equal(first.closeCount(), 1);
const active = provider();
- const replaced = await peer.replaceClientCapabilities(active.value, 3000);
+ const replaced = await peer.replaceClientCapabilities(active.value, { timeoutMs: 3000 });
assert.notEqual(replaced.registrationId, registered.registrationId);
assert.equal(replaced.revision, unregistered.revision + 1);
assert.deepEqual(
diff --git a/tests/fixtures/client-scheduler-entry.mjs b/tests/fixtures/client-scheduler-entry.mjs
index e5bd8aebf9..23c0a56c64 100644
--- a/tests/fixtures/client-scheduler-entry.mjs
+++ b/tests/fixtures/client-scheduler-entry.mjs
@@ -122,7 +122,7 @@ try {
},
close() {},
};
- await connection.replaceClientCapabilities(provider, 3000);
+ await connection.replaceClientCapabilities(provider, { timeoutMs: 3000 });
const { grant } = await request('plugin.authorization', {
client,
scope: 'profile',
@@ -219,7 +219,7 @@ try {
assert.equal(withdrawn.status, 'paused');
assert.equal(withdrawn.fireCount, 0);
assert.equal(calls.length, 1);
- await connection.unregisterClientCapabilities(3000);
+ await connection.unregisterClientCapabilities({ timeoutMs: 3000 });
for (const waitingId of ids.slice(1, 4)) {
await mutate({ kind: 'trigger_now', taskId: waitingId });
for (let attempt = 0; attempt < 300; attempt++) {
@@ -231,7 +231,7 @@ try {
await mutate({ kind: 'pause', taskId: ids[1] });
await mutate({ kind: 'snooze', taskId: ids[2], delayMs: 3600000 });
await mutate({ kind: 'delete', taskId: ids[3] });
- await connection.replaceClientCapabilities(provider, 3000);
+ await connection.replaceClientCapabilities(provider, { timeoutMs: 3000 });
await delay(350);
assert.equal(calls.length, 1, 'cancelled waiting notifications must not be admitted');
await request('plugin.composition.apply', {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@maka.apache.org
For additional commands, e-mail: commits-help@maka.apache.org