tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat(goose-acp): enable parallel sessions with isolated agent state (#6392)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-01-15 06:19:17 +08:00
committed by GitHub
parent fb0eca2c36
commit 7d4a6bd8ff
86 changed files with 2594 additions and 1938 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/mcp_client.rs
+87 -53
View File
@@ -3,7 +3,7 @@ use crate::agents::types::SharedProvider;
use crate::session_context::SESSION_ID_HEADER;
use rmcp::model::{
Content, CreateElicitationRequestParam, CreateElicitationResult, ElicitationAction, ErrorCode,
JsonObject,
Extensions, JsonObject, Meta,
};
/// MCP client implementation for Goose
use rmcp::{
@@ -37,20 +37,25 @@ pub type BoxError = Box<dyn std::error::Error + Sync + Send>;
pub type Error = rmcp::ServiceError;
#[derive(Clone, Debug)]
pub struct McpMeta {
pub session_id: String,
}
impl McpMeta {
pub fn new(session_id: impl Into<String>) -> Self {
Self {
session_id: session_id.into(),
}
}
fn inject_into_extensions(&self, extensions: Extensions) -> Extensions {
inject_session_id_into_extensions(extensions, &self.session_id)
}
}
#[async_trait::async_trait]
pub trait McpClientTrait: Send + Sync {
async fn list_resources(
&self,
next_cursor: Option<String>,
cancel_token: CancellationToken,
) -> Result<ListResourcesResult, Error>;
async fn read_resource(
&self,
uri: &str,
cancel_token: CancellationToken,
) -> Result<ReadResourceResult, Error>;
async fn list_tools(
&self,
next_cursor: Option<String>,
@@ -61,27 +66,50 @@ pub trait McpClientTrait: Send + Sync {
&self,
name: &str,
arguments: Option<JsonObject>,
meta: McpMeta,
cancel_token: CancellationToken,
) -> Result<CallToolResult, Error>;
fn get_info(&self) -> Option<&InitializeResult>;
async fn list_resources(
&self,
_next_cursor: Option<String>,
_cancel_token: CancellationToken,
) -> Result<ListResourcesResult, Error> {
Err(Error::TransportClosed)
}
async fn read_resource(
&self,
_uri: &str,
_cancel_token: CancellationToken,
) -> Result<ReadResourceResult, Error> {
Err(Error::TransportClosed)
}
async fn list_prompts(
&self,
next_cursor: Option<String>,
cancel_token: CancellationToken,
) -> Result<ListPromptsResult, Error>;
_next_cursor: Option<String>,
_cancel_token: CancellationToken,
) -> Result<ListPromptsResult, Error> {
Err(Error::TransportClosed)
}
async fn get_prompt(
&self,
name: &str,
arguments: Value,
cancel_token: CancellationToken,
) -> Result<GetPromptResult, Error>;
_name: &str,
_arguments: Value,
_cancel_token: CancellationToken,
) -> Result<GetPromptResult, Error> {
Err(Error::TransportClosed)
}
async fn subscribe(&self) -> mpsc::Receiver<ServerNotification>;
async fn subscribe(&self) -> mpsc::Receiver<ServerNotification> {
mpsc::channel(1).1
}
fn get_info(&self) -> Option<&InitializeResult>;
async fn get_moim(&self) -> Option<String> {
async fn get_moim(&self, _session_id: &str) -> Option<String> {
None
}
}
@@ -379,7 +407,7 @@ impl McpClientTrait for McpClient {
ClientRequest::ListResourcesRequest(ListResourcesRequest {
params: Some(PaginatedRequestParam { cursor }),
method: Default::default(),
extensions: inject_session_into_extensions(Default::default()),
extensions: inject_current_session_id_into_extensions(Default::default()),
}),
cancel_token,
)
@@ -403,7 +431,7 @@ impl McpClientTrait for McpClient {
uri: uri.to_string(),
},
method: Default::default(),
extensions: inject_session_into_extensions(Default::default()),
extensions: inject_current_session_id_into_extensions(Default::default()),
}),
cancel_token,
)
@@ -425,7 +453,7 @@ impl McpClientTrait for McpClient {
ClientRequest::ListToolsRequest(ListToolsRequest {
params: Some(PaginatedRequestParam { cursor }),
method: Default::default(),
extensions: inject_session_into_extensions(Default::default()),
extensions: inject_current_session_id_into_extensions(Default::default()),
}),
cancel_token,
)
@@ -441,6 +469,7 @@ impl McpClientTrait for McpClient {
&self,
name: &str,
arguments: Option<JsonObject>,
meta: McpMeta,
cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let res = self
@@ -451,7 +480,7 @@ impl McpClientTrait for McpClient {
arguments,
},
method: Default::default(),
extensions: inject_session_into_extensions(Default::default()),
extensions: meta.inject_into_extensions(Default::default()),
}),
cancel_token,
)
@@ -473,7 +502,7 @@ impl McpClientTrait for McpClient {
ClientRequest::ListPromptsRequest(ListPromptsRequest {
params: Some(PaginatedRequestParam { cursor }),
method: Default::default(),
extensions: inject_session_into_extensions(Default::default()),
extensions: inject_current_session_id_into_extensions(Default::default()),
}),
cancel_token,
)
@@ -503,7 +532,7 @@ impl McpClientTrait for McpClient {
arguments,
},
method: Default::default(),
extensions: inject_session_into_extensions(Default::default()),
extensions: inject_current_session_id_into_extensions(Default::default()),
}),
cancel_token,
)
@@ -522,29 +551,34 @@ impl McpClientTrait for McpClient {
}
}
/// Replaces session ID, case-insensitively, in Extensions._meta.
fn inject_session_into_extensions(
mut extensions: rmcp::model::Extensions,
) -> rmcp::model::Extensions {
use rmcp::model::Meta;
/// Injects the given session_id into Extensions._meta.
fn inject_session_id_into_extensions(mut extensions: Extensions, session_id: &str) -> Extensions {
let mut meta_map = extensions
.get::<Meta>()
.map(|meta| meta.0.clone())
.unwrap_or_default();
if let Some(session_id) = crate::session_context::current_session_id() {
let mut meta_map = extensions
.get::<Meta>()
.map(|meta| meta.0.clone())
.unwrap_or_default();
// JsonObject is case-sensitive, so we use retain for case-insensitive removal
meta_map.retain(|k, _| !k.eq_ignore_ascii_case(SESSION_ID_HEADER));
// JsonObject is case-sensitive, so we use retain for case-insensitive removal
meta_map.retain(|k, _| !k.eq_ignore_ascii_case(SESSION_ID_HEADER));
meta_map.insert(SESSION_ID_HEADER.to_string(), Value::String(session_id));
extensions.insert(Meta(meta_map));
}
meta_map.insert(
SESSION_ID_HEADER.to_string(),
Value::String(session_id.to_string()),
);
extensions.insert(Meta(meta_map));
extensions
}
/// Injects session ID from task-local context into Extensions._meta.
fn inject_current_session_id_into_extensions(extensions: Extensions) -> Extensions {
if let Some(session_id) = crate::session_context::current_session_id() {
inject_session_id_into_extensions(extensions, &session_id)
} else {
extensions
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -556,7 +590,7 @@ mod tests {
let session_id = "test-session-789";
crate::session_context::with_session_id(Some(session_id.to_string()), async {
let extensions = inject_session_into_extensions(Default::default());
let extensions = inject_current_session_id_into_extensions(Default::default());
let meta = extensions.get::<Meta>().unwrap();
assert_eq!(
@@ -573,7 +607,7 @@ mod tests {
#[tokio::test]
async fn test_no_session_id_in_mcp_when_absent() {
let extensions = inject_session_into_extensions(Default::default());
let extensions = inject_current_session_id_into_extensions(Default::default());
let meta = extensions.get::<Meta>();
assert!(meta.is_none());
@@ -585,9 +619,9 @@ mod tests {
let session_id = "consistent-session-id";
crate::session_context::with_session_id(Some(session_id.to_string()), async {
let ext1 = inject_session_into_extensions(Default::default());
let ext2 = inject_session_into_extensions(Default::default());
let ext3 = inject_session_into_extensions(Default::default());
let ext1 = inject_current_session_id_into_extensions(Default::default());
let ext2 = inject_current_session_id_into_extensions(Default::default());
let ext3 = inject_current_session_id_into_extensions(Default::default());
for ext in [&ext1, &ext2, &ext3] {
assert_eq!(
@@ -620,7 +654,7 @@ mod tests {
.unwrap(),
);
let extensions = inject_session_into_extensions(extensions);
let extensions = inject_current_session_id_into_extensions(extensions);
let meta = extensions.get::<Meta>().unwrap();
assert_eq!(