Fix frontend extension session state and tool routes (#8464)
Signed-off-by: sunilkumarvalmiki <g.sunilkumarvalmiki@gmail.com> Signed-off-by: jh-block <jhugo@block.xyz> Signed-off-by: Douwe Osinga <douwe@squareup.com> Signed-off-by: Michael Neale <michael.neale@gmail.com> Signed-off-by: Angie Jones <jones.angie@gmail.com> Signed-off-by: dependabot[bot] <support@github.com> Signed-off-by: Vincenzo Palazzo <vincenzopalazzodev@gmail.com> Signed-off-by: Adam Miller <admiller@redhat.com> Signed-off-by: morgmart <98432065+morgmart@users.noreply.github.com> Signed-off-by: Andrew Harvard <aharvard@squareup.com> Signed-off-by: Matt Toohey <contact@matttoohey.com> Signed-off-by: Wes <wesb@block.xyz> Signed-off-by: Clay Delk <clay.delk@gmail.com> Signed-off-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: Taylor Ho <taylorkmho@gmail.com> Signed-off-by: Alex Hancock <alexhancock@block.xyz> Signed-off-by: Treebird <treebird@treebird.dev> Signed-off-by: tulsi <tulsi@block.xyz> Signed-off-by: vonbai <nswanqi@gmail.com> Signed-off-by: olaservo <olahungerford@gmail.com> Signed-off-by: Abhijay Jain <Abhijay007j@gmail.com> Signed-off-by: Rodolfo Olivieri <rolivier@redhat.com> Signed-off-by: fresh3nough <anonwurcod@proton.me> Signed-off-by: Jheison Martinez Bolivar <jheison.mb@gmail.com> Co-authored-by: Matt Toohey <contact@matttoohey.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Lifei Zhou <lifei@squareup.com> Co-authored-by: jh-block <jhugo@block.xyz> Co-authored-by: Alex Hancock <alexhancock@block.xyz> Co-authored-by: Jack Amadeo <jackamadeo@block.xyz> Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Michael Neale <michael.neale@gmail.com> Co-authored-by: Angie Jones <jones.angie@gmail.com> Co-authored-by: SomeSolutionsArchitect <139817767+SomeSolutionsArchitect@users.noreply.github.com> Co-authored-by: Christian <chvargas@wfscorp.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Vincenzo Palazzo <vincenzopalazzodev@gmail.com> Co-authored-by: Adam Miller <admiller@redhat.com> Co-authored-by: Jack Amadeo <jackamadeo@squareup.com> Co-authored-by: morgmart <98432065+morgmart@users.noreply.github.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Andrew Harvard <aharvard@squareup.com> Co-authored-by: Kalvin C <kalvinnchau@users.noreply.github.com> Co-authored-by: Adewale Abati <acekyd01@gmail.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Kalvin Chau <kalvin@block.xyz> Co-authored-by: Clay Delk <clay.delk@gmail.com> Co-authored-by: dorien-koelemeijer <62866702+dorien-koelemeijer@users.noreply.github.com> Co-authored-by: Wes <wesbillman@users.noreply.github.com> Co-authored-by: Monroe Williams <monroe@pobox.com> Co-authored-by: Spike Wang <spike@spikewang.me> Co-authored-by: Taylor Ho <taylorkmho@gmail.com> Co-authored-by: Treebird <treebird@treebird.dev> Co-authored-by: Sarthak Bhardwaj <100398847+SarthakB11@users.noreply.github.com> Co-authored-by: Douwe Osinga <douwe@block.xyz> Co-authored-by: tulsi <tulsi@block.xyz> Co-authored-by: Vonbai <107612985+vonbai@users.noreply.github.com> Co-authored-by: Ola Hungerford <olahungerford@gmail.com> Co-authored-by: Lucas Kim <45152028+Raptors65@users.noreply.github.com> Co-authored-by: Abhijay Jain <abhijay007j@gmail.com> Co-authored-by: Rodolfo Olivieri <rolivier@redhat.com> Co-authored-by: fre$h <anonwurcod@proton.me> Co-authored-by: Jheison Martinez Bolivar <78370841+JheisonMB@users.noreply.github.com>
This commit is contained in:
g.sunilkumar
@@ -1118,4 +1118,192 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
mod frontend_extension_tests {
|
||||
use super::*;
|
||||
use goose::agents::{AgentConfig, ExtensionConfig};
|
||||
use goose::config::permission::PermissionManager;
|
||||
use goose::config::GooseMode;
|
||||
use goose::session::session_manager::SessionType;
|
||||
use goose::session::{
|
||||
EnabledExtensionsState, ExtensionData, ExtensionState, SessionManager,
|
||||
};
|
||||
use rmcp::model::Tool;
|
||||
use rmcp::object;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn frontend_extension_with_tool(name: &str, tool_name: &str) -> ExtensionConfig {
|
||||
ExtensionConfig::Frontend {
|
||||
name: name.to_string(),
|
||||
description: format!("Frontend test extension {name}"),
|
||||
tools: vec![Tool::new(
|
||||
tool_name.to_string(),
|
||||
format!("Run {tool_name} from the frontend"),
|
||||
object!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message": { "type": "string" }
|
||||
},
|
||||
"required": ["message"]
|
||||
}),
|
||||
)],
|
||||
instructions: Some(format!("Use the {tool_name} tool.")),
|
||||
bundled: None,
|
||||
available_tools: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
fn frontend_extension() -> ExtensionConfig {
|
||||
frontend_extension_with_tool("frontend-e2e", "frontend__echo")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_frontend_extensions_are_persisted_listed_and_removed() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let data_dir = temp_dir.path().to_path_buf();
|
||||
let session_manager = Arc::new(SessionManager::new(data_dir.clone()));
|
||||
let permission_manager = Arc::new(PermissionManager::new(data_dir));
|
||||
let agent = Agent::with_config(AgentConfig::new(
|
||||
session_manager.clone(),
|
||||
permission_manager,
|
||||
None,
|
||||
GooseMode::default(),
|
||||
false,
|
||||
GoosePlatform::GooseDesktop,
|
||||
));
|
||||
|
||||
let session = session_manager
|
||||
.create_session(
|
||||
std::env::current_dir().unwrap(),
|
||||
"frontend-extension-test".to_string(),
|
||||
SessionType::Hidden,
|
||||
GooseMode::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
agent
|
||||
.add_extension(frontend_extension(), &session.id)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let listed_tools = agent.list_tools(&session.id, None).await;
|
||||
assert!(listed_tools
|
||||
.iter()
|
||||
.any(|tool| tool.name == "frontend__echo"));
|
||||
|
||||
let filtered_tools = agent
|
||||
.list_tools(&session.id, Some("frontend-e2e".to_string()))
|
||||
.await;
|
||||
assert_eq!(filtered_tools.len(), 1);
|
||||
assert_eq!(filtered_tools[0].name, "frontend__echo");
|
||||
|
||||
let extension_names = agent.list_extensions().await;
|
||||
assert!(extension_names.iter().any(|name| name == "frontend-e2e"));
|
||||
|
||||
let persisted_session = session_manager
|
||||
.get_session(&session.id, false)
|
||||
.await
|
||||
.unwrap();
|
||||
let persisted_extensions =
|
||||
EnabledExtensionsState::from_extension_data(&persisted_session.extension_data)
|
||||
.unwrap()
|
||||
.extensions;
|
||||
assert!(persisted_extensions
|
||||
.iter()
|
||||
.any(|extension| extension.name() == "frontend-e2e"));
|
||||
|
||||
agent
|
||||
.remove_extension("frontend-e2e", &session.id)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let listed_tools = agent.list_tools(&session.id, None).await;
|
||||
assert!(!listed_tools
|
||||
.iter()
|
||||
.any(|tool| tool.name == "frontend__echo"));
|
||||
|
||||
let persisted_session = session_manager
|
||||
.get_session(&session.id, false)
|
||||
.await
|
||||
.unwrap();
|
||||
let persisted_extensions =
|
||||
EnabledExtensionsState::from_extension_data(&persisted_session.extension_data)
|
||||
.unwrap()
|
||||
.extensions;
|
||||
assert!(persisted_extensions
|
||||
.iter()
|
||||
.all(|extension| extension.name() != "frontend-e2e"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_frontend_session_load_keeps_all_tools() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let data_dir = temp_dir.path().to_path_buf();
|
||||
let session_manager = Arc::new(SessionManager::new(data_dir.clone()));
|
||||
let permission_manager = Arc::new(PermissionManager::new(data_dir));
|
||||
let agent = Arc::new(Agent::with_config(AgentConfig::new(
|
||||
session_manager.clone(),
|
||||
permission_manager,
|
||||
None,
|
||||
GooseMode::default(),
|
||||
false,
|
||||
GoosePlatform::GooseDesktop,
|
||||
)));
|
||||
|
||||
let session = session_manager
|
||||
.create_session(
|
||||
std::env::current_dir().unwrap(),
|
||||
"frontend-extension-load-test".to_string(),
|
||||
SessionType::Hidden,
|
||||
GooseMode::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let expected_tools = (0..12)
|
||||
.map(|index| format!("frontend__tool_{index}"))
|
||||
.collect::<Vec<_>>();
|
||||
let extensions = expected_tools
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, tool_name)| {
|
||||
frontend_extension_with_tool(&format!("frontend-{index}"), tool_name)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let mut extension_data = ExtensionData::new();
|
||||
EnabledExtensionsState::new(extensions)
|
||||
.to_extension_data(&mut extension_data)
|
||||
.unwrap();
|
||||
session_manager
|
||||
.update(&session.id)
|
||||
.extension_data(extension_data)
|
||||
.apply()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let session = session_manager
|
||||
.get_session(&session.id, false)
|
||||
.await
|
||||
.unwrap();
|
||||
let load_results = agent.load_extensions_from_session(&session).await;
|
||||
assert!(
|
||||
load_results.iter().all(|result| result.success),
|
||||
"failed to load frontend extensions: {load_results:?}",
|
||||
);
|
||||
|
||||
let listed_tools = agent.list_tools(&session.id, None).await;
|
||||
for tool_name in expected_tools {
|
||||
assert!(
|
||||
listed_tools.iter().any(|tool| tool.name == tool_name),
|
||||
"expected listed frontend tool {tool_name}",
|
||||
);
|
||||
assert!(
|
||||
agent.is_frontend_tool(&tool_name).await,
|
||||
"expected frontend dispatch state for {tool_name}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user