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

feat: AgentManager - foundation for unified execution (#4389) (#4684)

Browse Source
This commit is contained in:
tlongwell-block
2025-09-24 16:15:26 -04:00
committed by GitHub
parent ad045bc94c
commit 9d40422e2d
34 changed files with 896 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/commands/agent.rs
+1 -49
View File
@@ -1,13 +1,7 @@
use std::sync::Arc;
use crate::configuration;
use crate::state;
use anyhow::Result;
use axum::middleware;
use etcetera::{choose_app_strategy, AppStrategy};
use goose::agents::Agent;
use goose::config::APP_STRATEGY;
use goose::scheduler_factory::SchedulerFactory;
use goose_server::auth::check_token;
use tower_http::cors::{Any, CorsLayer};
use tracing::info;
@@ -32,49 +26,7 @@ pub async fn run() -> Result<()> {
let secret_key =
std::env::var("GOOSE_SERVER__SECRET_KEY").unwrap_or_else(|_| "test".to_string());
let new_agent = Agent::new();
// Only initialize provider and extensions when running in standalone goosed mode
// This prevents breaking the Electron app which manages its own provider setup
if std::env::var("GOOSE_STANDALONE_MODE").unwrap_or_else(|_| "false".to_string()) == "true" {
tracing::info!("Running in standalone mode - initializing provider and extensions");
// Initialize provider like the CLI does
let config = goose::config::Config::global();
let provider_name: String = config
.get_param("GOOSE_PROVIDER")
.expect("No provider configured. Run 'goose configure' first");
let model_name: String = config
.get_param("GOOSE_MODEL")
.expect("No model configured. Run 'goose configure' first");
let model_config = goose::model::ModelConfig::new(&model_name)
.expect("Failed to create model configuration");
let provider = goose::providers::create(&provider_name, model_config)
.expect("Failed to create provider");
new_agent
.update_provider(provider)
.await
.expect("Failed to update agent provider");
}
let agent_ref = Arc::new(new_agent);
let app_state = state::AppState::new(agent_ref.clone());
let schedule_file_path = choose_app_strategy(APP_STRATEGY.clone())?
.data_dir()
.join("schedules.json");
let scheduler_instance = SchedulerFactory::create(schedule_file_path).await?;
app_state.set_scheduler(scheduler_instance.clone()).await;
// NEW: Provide scheduler access to the agent
agent_ref.set_scheduler(scheduler_instance).await;
let app_state = state::AppState::new().await?;
let cors = CorsLayer::new()
.allow_origin(Any)
crates/goose-server/src/routes/agent.rs
+26 -35
View File
@@ -1,5 +1,4 @@
use crate::state::AppState;
use axum::response::IntoResponse;
use axum::{
extract::{Query, State},
http::StatusCode,
@@ -28,7 +27,6 @@ use tracing::error;
#[derive(Deserialize, utoipa::ToSchema)]
pub struct ExtendPromptRequest {
extension: String,
#[allow(dead_code)]
session_id: String,
}
@@ -40,7 +38,6 @@ pub struct ExtendPromptResponse {
#[derive(Deserialize, utoipa::ToSchema)]
pub struct AddSubRecipesRequest {
sub_recipes: Vec<SubRecipe>,
#[allow(dead_code)]
session_id: String,
}
@@ -53,27 +50,23 @@ pub struct AddSubRecipesResponse {
pub struct UpdateProviderRequest {
provider: String,
model: Option<String>,
#[allow(dead_code)]
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct SessionConfigRequest {
response: Option<Response>,
#[allow(dead_code)]
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct GetToolsQuery {
extension_name: Option<String>,
#[allow(dead_code)]
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct UpdateRouterToolSelectorRequest {
#[allow(dead_code)]
session_id: String,
}
@@ -116,8 +109,6 @@ async fn start_agent(
State(state): State<Arc<AppState>>,
Json(payload): Json<StartAgentRequest>,
) -> Result<Json<StartAgentResponse>, StatusCode> {
state.reset().await;
let session_id = session::generate_session_id();
let counter = state.session_counter.fetch_add(1, Ordering::SeqCst) + 1;
@@ -203,7 +194,7 @@ async fn add_sub_recipes(
State(state): State<Arc<AppState>>,
Json(payload): Json<AddSubRecipesRequest>,
) -> Result<Json<AddSubRecipesResponse>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(payload.session_id).await?;
agent.add_sub_recipes(payload.sub_recipes.clone()).await;
Ok(Json(AddSubRecipesResponse { success: true }))
}
@@ -222,7 +213,7 @@ async fn extend_prompt(
State(state): State<Arc<AppState>>,
Json(payload): Json<ExtendPromptRequest>,
) -> Result<Json<ExtendPromptResponse>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(payload.session_id).await?;
agent.extend_system_prompt(payload.extension.clone()).await;
Ok(Json(ExtendPromptResponse { success: true }))
}
@@ -247,7 +238,7 @@ async fn get_tools(
) -> Result<Json<Vec<ToolInfo>>, StatusCode> {
let config = Config::global();
let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string());
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(query.session_id).await?;
let permission_manager = PermissionManager::default();
let mut tools: Vec<ToolInfo> = agent
@@ -298,35 +289,37 @@ async fn get_tools(
async fn update_agent_provider(
State(state): State<Arc<AppState>>,
Json(payload): Json<UpdateProviderRequest>,
) -> Result<StatusCode, impl IntoResponse> {
let agent = state.get_agent().await;
) -> Result<StatusCode, StatusCode> {
let agent = state
.get_agent_for_route(payload.session_id.clone())
.await?;
let config = Config::global();
let model = match payload
.model
.or_else(|| config.get_param("GOOSE_MODEL").ok())
{
Some(m) => m,
None => return Err((StatusCode::BAD_REQUEST, "No model specified".to_string())),
None => {
tracing::error!("No model specified");
return Err(StatusCode::BAD_REQUEST);
}
};
let model_config = ModelConfig::new(&model).map_err(|e| {
(
StatusCode::BAD_REQUEST,
format!("Invalid model config: {}", e),
)
tracing::error!("Invalid model config: {}", e);
StatusCode::BAD_REQUEST
})?;
let new_provider = create(&payload.provider, model_config).map_err(|e| {
(
StatusCode::BAD_REQUEST,
format!("Failed to create provider: {}", e),
)
tracing::error!("Failed to create provider: {}", e);
StatusCode::BAD_REQUEST
})?;
agent
.update_provider(new_provider)
.await
.map_err(|_e| (StatusCode::INTERNAL_SERVER_ERROR, String::new()))?;
agent.update_provider(new_provider).await.map_err(|e| {
tracing::error!("Failed to update provider: {}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(StatusCode::OK)
}
@@ -344,17 +337,15 @@ async fn update_agent_provider(
)]
async fn update_router_tool_selector(
State(state): State<Arc<AppState>>,
Json(_payload): Json<UpdateRouterToolSelectorRequest>,
) -> Result<Json<String>, Json<ErrorResponse>> {
let agent = state.get_agent().await;
Json(payload): Json<UpdateRouterToolSelectorRequest>,
) -> Result<Json<String>, StatusCode> {
let agent = state.get_agent_for_route(payload.session_id).await?;
agent
.update_router_tool_selector(None, Some(true))
.await
.map_err(|e| {
tracing::error!("Failed to update tool selection strategy: {}", e);
Json(ErrorResponse {
error: format!("Failed to update tool selection strategy: {}", e),
})
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(Json(
@@ -376,8 +367,8 @@ async fn update_router_tool_selector(
async fn update_session_config(
State(state): State<Arc<AppState>>,
Json(payload): Json<SessionConfigRequest>,
) -> Result<Json<String>, Json<ErrorResponse>> {
let agent = state.get_agent().await;
) -> Result<Json<String>, StatusCode> {
let agent = state.get_agent_for_route(payload.session_id).await?;
if let Some(response) = payload.response {
agent.add_final_output_tool(response).await;
crates/goose-server/src/routes/audio.rs
+13 -35
View File
@@ -391,13 +391,13 @@ pub fn routes(state: Arc<AppState>) -> Router {
mod tests {
use super::*;
use axum::{body::Body, http::Request};
use serde_json::json;
use tower::ServiceExt;
#[tokio::test]
#[tokio::test(flavor = "multi_thread")]
async fn test_transcribe_endpoint_requires_auth() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
let state = AppState::new().await.unwrap();
let app = routes(state);
// Test without auth header
let request = Request::builder()
.uri("/audio/transcribe")
@@ -413,40 +413,18 @@ mod tests {
.unwrap();
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert!(
response.status() == StatusCode::PRECONDITION_FAILED
|| response.status() == StatusCode::UNAUTHORIZED
);
}
#[tokio::test]
#[tokio::test(flavor = "multi_thread")]
async fn test_transcribe_endpoint_validates_size() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
let app = routes(state);
// Create a large base64 string (simulating > 25MB audio)
let large_audio = BASE64.encode(vec![0u8; MAX_AUDIO_SIZE_BYTES + 1]);
let request = Request::builder()
.uri("/audio/transcribe")
.method("POST")
.header("content-type", "application/json")
.header("x-secret-key", "test-secret")
.body(Body::from(
serde_json::to_string(&serde_json::json!({
"audio": large_audio,
"mime_type": "audio/webm"
}))
.unwrap(),
))
.unwrap();
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
}
#[tokio::test]
async fn test_transcribe_endpoint_validates_mime_type() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
let state = AppState::new().await.unwrap();
let app = routes(state);
let large_data = "a".repeat(30 * 1024 * 1024); // 30MB
let request = Request::builder()
.uri("/audio/transcribe")
.method("POST")
@@ -468,9 +446,9 @@ mod tests {
);
}
#[tokio::test]
async fn test_transcribe_endpoint_handles_invalid_base64() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
#[tokio::test(flavor = "multi_thread")]
async fn test_transcribe_endpoint_validates_mime_type() {
let state = AppState::new().await.unwrap();
let app = routes(state);
let request = Request::builder()
crates/goose-server/src/routes/context.rs
+3 -1
View File
@@ -13,6 +13,8 @@ pub struct ContextManageRequest {
pub messages: Vec<Message>,
/// Operation to perform: "truncation" or "summarize"
pub manage_action: String,
/// Optional session ID for session-specific agent
pub session_id: String,
}
/// Response from context management operations
@@ -44,7 +46,7 @@ async fn manage_context(
State(state): State<Arc<AppState>>,
Json(request): Json<ContextManageRequest>,
) -> Result<Json<ContextManageResponse>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(request.session_id).await?;
let mut processed_messages = Conversation::new_unvalidated(vec![]);
let mut token_counts: Vec<usize> = vec![];
crates/goose-server/src/routes/extension.rs
+27 -22
View File
@@ -96,33 +96,31 @@ struct ExtensionResponse {
message: Option<String>,
}
/// Request structure for adding an extension, combining session_id with the extension config
#[derive(Deserialize)]
struct AddExtensionRequest {
session_id: String,
#[serde(flatten)]
config: ExtensionConfigRequest,
}
/// Handler for adding a new extension configuration.
async fn add_extension(
State(state): State<Arc<AppState>>,
raw: axum::extract::Json<serde_json::Value>,
Json(request): Json<AddExtensionRequest>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
// Log the raw request for debugging
// Log the request for debugging
tracing::info!(
"Received extension request: {}",
serde_json::to_string_pretty(&raw.0).unwrap()
"Received extension request for session: {}",
request.session_id
);
// Try to parse into our enum
let request: ExtensionConfigRequest = match serde_json::from_value(raw.0.clone()) {
Ok(req) => req,
Err(e) => {
tracing::error!("Failed to parse extension request: {}", e);
tracing::error!(
"Raw request was: {}",
serde_json::to_string_pretty(&raw.0).unwrap()
);
return Err(StatusCode::UNPROCESSABLE_ENTITY);
}
};
let session_id = request.session_id.clone();
let extension_request = request.config;
// If this is a Stdio extension that uses npx, check for Node.js installation
#[cfg(target_os = "windows")]
if let ExtensionConfigRequest::Stdio { cmd, .. } = &request {
if let ExtensionConfigRequest::Stdio { cmd, .. } = &extension_request {
if cmd.ends_with("npx.cmd") || cmd.ends_with("npx") {
// Check if Node.js is installed in standard locations
let node_exists = std::path::Path::new(r"C:\Program Files\nodejs\node.exe").exists()
@@ -175,7 +173,7 @@ async fn add_extension(
}
// Construct ExtensionConfig with Envs populated from keyring based on provided env_keys.
let extension_config: ExtensionConfig = match request {
let extension_config: ExtensionConfig = match extension_request {
ExtensionConfigRequest::Sse {
name,
uri,
@@ -267,7 +265,7 @@ async fn add_extension(
},
};
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(session_id).await?;
let response = agent.add_extension(extension_config).await;
// Respond with the result.
@@ -289,13 +287,20 @@ async fn add_extension(
}
}
#[derive(Deserialize)]
struct RemoveExtensionRequest {
name: String,
session_id: String,
}
/// Handler for removing an extension by name
async fn remove_extension(
State(state): State<Arc<AppState>>,
Json(name): Json<String>,
Json(request): Json<RemoveExtensionRequest>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
let agent = state.get_agent().await;
match agent.remove_extension(&name).await {
let agent = state.get_agent_for_route(request.session_id).await?;
match agent.remove_extension(&request.name).await {
Ok(_) => Ok(Json(ExtensionResponse {
error: false,
message: None,
crates/goose-server/src/routes/recipe.rs
+4 -8
View File
@@ -25,6 +25,7 @@ pub struct CreateRecipeRequest {
activities: Option<Vec<String>>,
#[serde(default)]
author: Option<AuthorRequest>,
session_id: String,
}
#[derive(Debug, Deserialize, ToSchema)]
@@ -108,13 +109,13 @@ pub struct ListRecipeResponse {
async fn create_recipe(
State(state): State<Arc<AppState>>,
Json(request): Json<CreateRecipeRequest>,
) -> Result<Json<CreateRecipeResponse>, (StatusCode, Json<CreateRecipeResponse>)> {
) -> Result<Json<CreateRecipeResponse>, StatusCode> {
tracing::info!(
"Recipe creation request received with {} messages",
request.messages.len()
);
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(request.session_id).await?;
// Create base recipe from agent state and messages
let recipe_result = agent
@@ -143,12 +144,7 @@ async fn create_recipe(
}
Err(e) => {
tracing::error!("Error details: {:?}", e);
let error_message = format!("Recipe creation failed: {}", e);
let error_response = CreateRecipeResponse {
recipe: None,
error: Some(error_message),
};
Err((StatusCode::BAD_REQUEST, Json(error_response)))
Err(StatusCode::BAD_REQUEST)
}
}
}
crates/goose-server/src/routes/reply.rs
+32 -23
View File
@@ -10,6 +10,7 @@ use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::conversation::message::{Message, MessageContent};
use goose::conversation::Conversation;
use goose::execution::SessionExecutionMode;
use goose::{
agents::{AgentEvent, SessionConfig},
permission::permission_confirmation::PrincipalType,
@@ -86,7 +87,7 @@ fn track_tool_telemetry(content: &MessageContent, all_messages: &[Message]) {
#[derive(Debug, Deserialize, Serialize)]
struct ChatRequest {
messages: Vec<Message>,
session_id: Option<String>,
session_id: String,
recipe_name: Option<String>,
recipe_version: Option<String>,
}
@@ -178,9 +179,9 @@ async fn reply_handler(
"Session started"
);
if let (Some(recipe_name), Some(session_id)) =
(request.recipe_name.clone(), request.session_id.clone())
{
let session_id = request.session_id.clone();
if let Some(recipe_name) = request.recipe_name.clone() {
if state.mark_recipe_run_if_absent(&session_id).await {
let recipe_version = request
.recipe_version
@@ -204,16 +205,28 @@ async fn reply_handler(
let messages = Conversation::new_unvalidated(request.messages);
let session_id = request.session_id.ok_or_else(|| {
tracing::error!("session_id is required but was not provided");
StatusCode::BAD_REQUEST
})?;
let task_cancel = cancel_token.clone();
let task_tx = tx.clone();
drop(tokio::spawn(async move {
let agent = state.get_agent().await;
let agent = match state
.get_agent(session_id.clone(), SessionExecutionMode::Interactive)
.await
{
Ok(agent) => agent,
Err(e) => {
tracing::error!("Failed to get session agent: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: format!("Failed to get session agent: {}", e),
},
&task_tx,
&task_cancel,
)
.await;
return;
}
};
// Load session metadata to get the working directory and other config
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
@@ -453,7 +466,6 @@ pub struct PermissionConfirmationRequest {
#[serde(default = "default_principal_type")]
principal_type: PrincipalType,
action: String,
#[allow(dead_code)]
session_id: String,
}
@@ -475,7 +487,7 @@ pub async fn confirm_permission(
State(state): State<Arc<AppState>>,
Json(request): Json<PermissionConfirmationRequest>,
) -> Result<Json<Value>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(request.session_id).await?;
let permission = match request.action.as_str() {
"always_allow" => Permission::AlwaysAllow,
"allow_once" => Permission::AllowOnce,
@@ -499,7 +511,6 @@ pub async fn confirm_permission(
struct ToolResultRequest {
id: String,
result: ToolResult<Vec<Content>>,
#[allow(dead_code)]
session_id: String,
}
@@ -524,7 +535,7 @@ async fn submit_tool_result(
}
};
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(payload.session_id).await?;
agent.handle_tool_result(payload.id, payload.result).await;
Ok(Json(json!({"status": "ok"})))
}
@@ -548,7 +559,6 @@ mod tests {
use super::*;
use goose::conversation::message::Message;
use goose::{
agents::Agent,
model::ModelConfig,
providers::{
base::{Provider, ProviderUsage, Usage},
@@ -589,18 +599,17 @@ mod tests {
use super::*;
use axum::{body::Body, http::Request};
use goose::conversation::message::Message;
use std::sync::Arc;
use serde_json::json;
use tower::ServiceExt;
#[tokio::test]
#[tokio::test(flavor = "multi_thread")]
async fn test_reply_endpoint() {
let mock_model_config = ModelConfig::new("test-model").unwrap();
let mock_provider = Arc::new(MockProvider {
let mock_provider = MockProvider {
model_config: mock_model_config,
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let state = AppState::new(Arc::new(agent));
};
let state = AppState::new().await.unwrap();
let app = routes(state);
@@ -612,7 +621,7 @@ mod tests {
.body(Body::from(
serde_json::to_string(&ChatRequest {
messages: vec![Message::user().with_text("test message")],
session_id: Some("test-session".to_string()),
session_id: "test-session".to_string(),
recipe_name: None,
recipe_version: None,
})
crates/goose-server/src/routes/session.rs
+1
View File
@@ -60,6 +60,7 @@ pub struct SessionInsights {
}
#[derive(Serialize, ToSchema, Debug)]
#[allow(dead_code)]
#[serde(rename_all = "camelCase")]
pub struct ActivityHeatmapCell {
pub week: usize,
crates/goose-server/src/state.rs
+33 -61
View File
@@ -1,18 +1,15 @@
use goose::agents::Agent;
use axum::http::StatusCode;
use goose::execution::manager::AgentManager;
use goose::execution::SessionExecutionMode;
use goose::scheduler_trait::SchedulerTrait;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::RwLock;
type AgentRef = Arc<Agent>;
#[derive(Clone)]
pub struct AppState {
agent: Arc<RwLock<AgentRef>>,
pub scheduler: Arc<RwLock<Option<Arc<dyn SchedulerTrait>>>>,
pub(crate) agent_manager: Arc<AgentManager>,
pub recipe_file_hash_map: Arc<Mutex<HashMap<String, PathBuf>>>,
pub session_counter: Arc<AtomicUsize>,
/// Tracks sessions that have already emitted recipe telemetry to prevent double counting.
@@ -20,31 +17,18 @@ pub struct AppState {
}
impl AppState {
pub fn new(agent: AgentRef) -> Arc<AppState> {
Arc::new(Self {
agent: Arc::new(RwLock::new(agent)),
scheduler: Arc::new(RwLock::new(None)),
pub async fn new() -> anyhow::Result<Arc<AppState>> {
let agent_manager = Arc::new(AgentManager::new(None).await?);
Ok(Arc::new(Self {
agent_manager,
recipe_file_hash_map: Arc::new(Mutex::new(HashMap::new())),
session_counter: Arc::new(AtomicUsize::new(0)),
recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())),
})
}
pub async fn get_agent(&self) -> AgentRef {
self.agent.read().await.clone()
}
pub async fn set_scheduler(&self, sched: Arc<dyn SchedulerTrait>) {
let mut guard = self.scheduler.write().await;
*guard = Some(sched);
}))
}
pub async fn scheduler(&self) -> Result<Arc<dyn SchedulerTrait>, anyhow::Error> {
self.scheduler
.read()
.await
.clone()
.ok_or_else(|| anyhow::anyhow!("Scheduler not initialized"))
self.agent_manager.scheduler().await
}
pub async fn set_recipe_file_hash_map(&self, hash_map: HashMap<String, PathBuf>) {
@@ -52,41 +36,6 @@ impl AppState {
*map = hash_map;
}
pub async fn reset(&self) {
let mut agent = self.agent.write().await;
let new_agent = Agent::new();
// Only initialize provider when running in standalone goosed mode
// This prevents breaking the Electron app which manages its own provider setup
if std::env::var("GOOSE_STANDALONE_MODE").unwrap_or_else(|_| "false".to_string()) == "true"
{
tracing::info!("Running in standalone mode - initializing provider");
let config = goose::config::Config::global();
let provider_name: String = config
.get_param("GOOSE_PROVIDER")
.expect("No provider configured. Run 'goose configure' first");
let model_name: String = config
.get_param("GOOSE_MODEL")
.expect("No model configured. Run 'goose configure' first");
let model_config = goose::model::ModelConfig::new(&model_name)
.expect("Failed to create model configuration");
let provider = goose::providers::create(&provider_name, model_config)
.expect("Failed to create provider");
new_agent
.update_provider(provider)
.await
.expect("Failed to update agent provider");
}
*agent = Arc::new(new_agent);
}
pub async fn mark_recipe_run_if_absent(&self, session_id: &str) -> bool {
let mut sessions = self.recipe_session_tracker.lock().await;
if sessions.contains(session_id) {
@@ -96,4 +45,27 @@ impl AppState {
true
}
}
pub async fn get_agent(
&self,
session_id: String,
mode: SessionExecutionMode,
) -> anyhow::Result<Arc<goose::agents::Agent>> {
self.agent_manager
.get_or_create_agent(session_id, mode)
.await
}
/// Get agent for route handlers - always uses Interactive mode and converts any error to 500
pub async fn get_agent_for_route(
&self,
session_id: String,
) -> Result<Arc<goose::agents::Agent>, StatusCode> {
self.get_agent(session_id, SessionExecutionMode::Interactive)
.await
.map_err(|e| {
tracing::error!("Failed to get agent: {}", e);
StatusCode::INTERNAL_SERVER_ERROR
})
}
}
crates/goose-server/tests/pricing_api_test.rs
-41
View File
@@ -1,41 +0,0 @@
use axum::http::StatusCode;
use axum::Router;
use axum::{body::Body, http::Request};
use etcetera::AppStrategy;
use serde_json::json;
use std::sync::Arc;
use tower::ServiceExt;
async fn create_test_app() -> Router {
let agent = Arc::new(goose::agents::Agent::default());
let state = goose_server::AppState::new(agent);
// Add scheduler setup like in the existing tests
let sched_storage_path = etcetera::choose_app_strategy(goose::config::APP_STRATEGY.clone())
.unwrap()
.data_dir()
.join("schedules.json");
let sched = goose::scheduler_factory::SchedulerFactory::create_legacy(sched_storage_path)
.await
.unwrap();
state.set_scheduler(sched).await;
goose_server::routes::config_management::routes(state)
}
#[tokio::test]
async fn test_pricing_endpoint_basic() {
// Basic test to ensure pricing endpoint responds correctly
let app = create_test_app().await;
let request = Request::builder()
.uri("/config/pricing")
.method("POST")
.header("content-type", "application/json")
.header("x-secret-key", "test")
.body(Body::from(json!({"configured_only": true}).to_string()))
.unwrap();
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}