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

fix(goose): propagate session_id across providers and MCP (#6584)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-01-22 09:28:56 +09:00
committed by GitHub
parent f3bae7ea7a
commit 67de49abbb
61 changed files with 1457 additions and 616 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/configure.rs
+10 -2
View File
@@ -25,6 +25,7 @@ use goose::providers::{create, providers, retry_operation, RetryConfig};
use goose::session::SessionType;
use serde_json::Value;
use std::collections::HashMap;
use uuid::Uuid;
// useful for light themes where there is no dicernible colour contrast between
// cursor-selected and cursor-unselected items.
@@ -682,8 +683,10 @@ pub async fn configure_provider_dialog() -> anyhow::Result<bool> {
let models_res = {
let temp_model_config = ModelConfig::new(&provider_meta.default_model)?;
let temp_provider = create(provider_name, temp_model_config).await?;
// Provider setup runs before any user session exists; use an ephemeral id.
let session_id = Uuid::new_v4().to_string();
retry_operation(&RetryConfig::default(), || async {
temp_provider.fetch_recommended_models().await
temp_provider.fetch_recommended_models(&session_id).await
})
.await
};
@@ -1655,9 +1658,11 @@ pub async fn handle_openrouter_auth() -> anyhow::Result<()> {
match create("openrouter", model_config).await {
Ok(provider) => {
// Simple test request
// Config verification runs before any user session exists; use an ephemeral id.
let session_id = Uuid::new_v4().to_string();
let test_result = provider
.complete(
&session_id,
"You are goose, an AI assistant.",
&[Message::user().with_text("Say 'Configuration test successful!'")],
&[],
@@ -1733,8 +1738,11 @@ pub async fn handle_tetrate_auth() -> anyhow::Result<()> {
match create("tetrate", model_config).await {
Ok(provider) => {
// Config verification runs before any user session exists; use an ephemeral id.
let session_id = Uuid::new_v4().to_string();
let test_result = provider
.complete(
&session_id,
"You are goose, an AI assistant.",
&[Message::user().with_text("Say 'Configuration test successful!'")],
&[],
crates/goose-cli/src/scenario_tests/mock_client.rs
+7 -2
View File
@@ -1,7 +1,7 @@
//! MockClient is a mock implementation of the McpClientTrait for testing purposes.
//! add a tool you want to have around and then add the client to the extension router
use goose::agents::mcp_client::{Error, McpClientTrait, McpMeta};
use goose::agents::mcp_client::{Error, McpClientTrait};
use rmcp::{
model::{
CallToolResult, Content, ErrorData, GetPromptResult, ListPromptsResult,
@@ -44,6 +44,7 @@ impl MockClient {
impl McpClientTrait for MockClient {
async fn list_resources(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancel_token: CancellationToken,
) -> Result<ListResourcesResult, Error> {
@@ -60,6 +61,7 @@ impl McpClientTrait for MockClient {
async fn read_resource(
&self,
_session_id: &str,
_uri: &str,
_cancel_token: CancellationToken,
) -> Result<ReadResourceResult, Error> {
@@ -68,6 +70,7 @@ impl McpClientTrait for MockClient {
async fn list_tools(
&self,
_session_id: &str,
_: Option<String>,
_cancel_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
@@ -92,9 +95,9 @@ impl McpClientTrait for MockClient {
async fn call_tool(
&self,
_session_id: &str,
name: &str,
arguments: Option<serde_json::Map<String, Value>>,
_meta: McpMeta,
_cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
if let Some(handler) = self.handlers.get(name) {
@@ -114,6 +117,7 @@ impl McpClientTrait for MockClient {
async fn list_prompts(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancel_token: CancellationToken,
) -> Result<ListPromptsResult, Error> {
@@ -126,6 +130,7 @@ impl McpClientTrait for MockClient {
async fn get_prompt(
&self,
_session_id: &str,
_name: &str,
_arguments: Value,
_cancel_token: CancellationToken,
crates/goose-cli/src/session/mod.rs
+27 -11
View File
@@ -193,15 +193,16 @@ pub enum PlannerResponseType {
/// to the user's message. The response is either a plan or a clarifying
/// question.
pub async fn classify_planner_response(
session_id: &str,
message_text: String,
provider: Arc<dyn Provider>,
) -> Result<PlannerResponseType> {
let prompt = format!("The text below is the output from an AI model which can either provide a plan or list of clarifying questions. Based on the text below, decide if the output is a \"plan\" or \"clarifying questions\".\n---\n{message_text}");
// Generate the description
let message = Message::user().with_text(&prompt);
let (result, _usage) = provider
.complete(
session_id,
"Reply only with the classification label: \"plan\" or \"clarifying questions\"",
&[message],
&[],
@@ -367,7 +368,7 @@ impl CliSession {
&mut self,
extension: Option<String>,
) -> Result<HashMap<String, Vec<String>>> {
let prompts = self.agent.list_extension_prompts().await;
let prompts = self.agent.list_extension_prompts(&self.session_id).await;
// Early validation if filtering by extension
if let Some(filter) = &extension {
@@ -388,7 +389,7 @@ impl CliSession {
}
pub async fn get_prompt_info(&mut self, name: &str) -> Result<Option<output::PromptInfo>> {
let prompts = self.agent.list_extension_prompts().await;
let prompts = self.agent.list_extension_prompts(&self.session_id).await;
// Find which extension has this prompt
for (extension, prompt_list) in prompts {
@@ -406,7 +407,11 @@ impl CliSession {
}
pub async fn get_prompt(&mut self, name: &str, arguments: Value) -> Result<Vec<PromptMessage>> {
Ok(self.agent.get_prompt(name, arguments).await?.messages)
Ok(self
.agent
.get_prompt(&self.session_id, name, arguments)
.await?
.messages)
}
/// Process a single message and get the response
@@ -728,7 +733,10 @@ impl CliSession {
println!("{}", console::style("Generating Recipe").green());
output::show_thinking();
let recipe = self.agent.create_recipe(self.messages.clone()).await;
let recipe = self
.agent
.create_recipe(&self.session_id, self.messages.clone())
.await;
output::hide_thinking();
match recipe {
@@ -782,16 +790,24 @@ impl CliSession {
plan_messages: Conversation,
reasoner: Arc<dyn Provider>,
) -> Result<(), anyhow::Error> {
let plan_prompt = self.agent.get_plan_prompt().await?;
let plan_prompt = self.agent.get_plan_prompt(&self.session_id).await?;
output::show_thinking();
let (plan_response, _usage) = reasoner
.complete(&plan_prompt, plan_messages.messages(), &[])
.complete(
&self.session_id,
&plan_prompt,
plan_messages.messages(),
&[],
)
.await?;
output::render_message(&plan_response, self.debug);
output::hide_thinking();
let planner_response_type =
classify_planner_response(plan_response.as_concat_text(), self.agent.provider().await?)
.await?;
let planner_response_type = classify_planner_response(
&self.session_id,
plan_response.as_concat_text(),
self.agent.provider().await?,
)
.await?;
match planner_response_type {
PlannerResponseType::Plan => {
@@ -1154,7 +1170,7 @@ impl CliSession {
/// This should be called before the interactive session starts
pub async fn update_completion_cache(&mut self) -> Result<()> {
// Get fresh data
let prompts = self.agent.list_extension_prompts().await;
let prompts = self.agent.list_extension_prompts(&self.session_id).await;
// Update the cache with write lock
let mut cache = self.completion_cache.write().unwrap();