diff --git a/app/src/ai/agent_sdk/ambient.rs b/app/src/ai/agent_sdk/ambient.rs index 549dd2c9976..efa173526e5 100644 --- a/app/src/ai/agent_sdk/ambient.rs +++ b/app/src/ai/agent_sdk/ambient.rs @@ -387,7 +387,7 @@ impl AmbientAgentRunner { vec![] }; - let team_scope = match super::common::resolve_object_scope(&args.scope, ctx) { + let team_scope = match super::common::resolve_environment_team_scope(&args.scope, ctx) { Ok(team_scope) => team_scope, Err(err) => { super::report_fatal_error(err, ctx); @@ -404,8 +404,11 @@ impl AmbientAgentRunner { environment_args.environment = Some(environment_id); } - let environment_id = match EnvironmentChoice::resolve_for_create(environment_args, ctx) - { + let environment_id = match EnvironmentChoice::resolve_for_create( + environment_args, + &team_scope, + ctx, + ) { Ok(EnvironmentChoice::None) => { eprintln!("Agent will run without an environment."); None diff --git a/app/src/ai/agent_sdk/common.rs b/app/src/ai/agent_sdk/common.rs index 0c47bbf37ef..be1c9286673 100644 --- a/app/src/ai/agent_sdk/common.rs +++ b/app/src/ai/agent_sdk/common.rs @@ -157,12 +157,12 @@ pub(super) fn resolve_team_scope( .map_err(|err| describe_team_resolution_error(err, ctx)) } -pub(super) fn resolve_object_scope( - object_scope: &ObjectScope, +pub(super) fn resolve_environment_team_scope( + scope: &ObjectScope, ctx: &AppContext, ) -> anyhow::Result { UserWorkspaces::as_ref(ctx) - .team_scope_for_cli_object(object_scope) + .team_scope_for_cli_object(scope) .map_err(|err| describe_team_resolution_error(err, ctx)) } @@ -234,6 +234,15 @@ pub(super) fn resolve_owner_for_team_scope( } } +pub(super) fn environment_is_visible_to_scope( + environment: &CloudAmbientAgentEnvironment, + team_scope: &(impl TeamScope + ?Sized), +) -> bool { + match environment.permissions().owner { + Owner::User { .. } => true, + Owner::Team { team_uid } => team_scope.team_uid() == Some(team_uid), + } +} /// Refresh workspace metadata before executing an operation. /// /// This ensures that team state is up-to-date before creating cloud objects or performing @@ -336,10 +345,11 @@ pub enum EnvironmentChoice { } impl EnvironmentChoice { - /// Resolve the environment to use when creating an agent integration. + /// Resolve the environment to use when creating an agent operation. /// Warp Drive *must* have been synced first. pub fn resolve_for_create( args: EnvironmentCreateArgs, + team_scope: &(impl TeamScope + ?Sized), ctx: &AppContext, ) -> Result { if args.no_environment { @@ -351,6 +361,7 @@ impl EnvironmentChoice { let mut synced_environments: Vec<(ServerId, &CloudAmbientAgentEnvironment)> = all_environments .iter() + .filter(|env| environment_is_visible_to_scope(env, team_scope)) .filter_map(|env| { if let SyncId::ServerId(server_id) = env.sync_id() { Some((server_id, env)) diff --git a/app/src/ai/agent_sdk/common_tests.rs b/app/src/ai/agent_sdk/common_tests.rs index 098ca3d9079..25f2ef2eb07 100644 --- a/app/src/ai/agent_sdk/common_tests.rs +++ b/app/src/ai/agent_sdk/common_tests.rs @@ -1,13 +1,19 @@ use std::collections::HashMap; use std::sync::Arc; +use warp_cli::environment::EnvironmentCreateArgs; +use warp_cli::scope::{ObjectScope, TeamSelection}; use warpui::App; use super::{ - classify_agent_mode_base_model_id, parse_ambient_task_id, validate_agent_mode_base_model_id, + EnvironmentChoice, classify_agent_mode_base_model_id, environment_is_visible_to_scope, + parse_ambient_task_id, resolve_environment_team_scope, validate_agent_mode_base_model_id, validate_agent_mode_base_model_id_for_scope, }; use crate::LaunchMode; +use crate::ai::cloud_environments::{ + AmbientAgentEnvironment, CloudAmbientAgentEnvironment, CloudAmbientAgentEnvironmentModel, +}; use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; use crate::ai::llms::{ AvailableLLMs, LLMContextWindow, LLMId, LLMInfo, LLMPreferences, LLMProvider, LLMUsageMetadata, @@ -17,20 +23,184 @@ use crate::ai::mcp::TemplatableMCPServerManager; use crate::auth::AuthStateProvider; use crate::auth::auth_manager::AuthManager; use crate::cloud_object::model::persistence::CloudModel; +use crate::cloud_object::{CloudObjectMetadata, CloudObjectPermissions, Owner}; use crate::network::NetworkStatus; use crate::server::cloud_objects::update_manager::UpdateManager; -use crate::server::ids::ServerId; +use crate::server::ids::{ServerId, SyncId}; use crate::server::server_api::ServerApiProvider; use crate::server::server_api::team::MockTeamClient; use crate::server::server_api::workspace::MockWorkspaceClient; use crate::server::sync_queue::SyncQueue; +use crate::settings::PrivacySettings; use crate::test_util::settings::initialize_settings_for_tests; use crate::workspaces::team::{Team, TeamVisibility}; use crate::workspaces::team_tester::TeamTesterStatus; use crate::workspaces::user_workspaces::{ - TeamContextForOperation, TeamlessScopeForTest, UserWorkspaces, + TeamContextForOperation, TeamScope, TeamlessScopeForTest, UserWorkspaces, }; use crate::workspaces::workspace::{Workspace, WorkspaceUid}; +fn environment_with_owner( + sync_id: SyncId, + name: &str, + owner: Owner, +) -> CloudAmbientAgentEnvironment { + let environment = AmbientAgentEnvironment::new( + name.to_string(), + None, + Vec::new(), + "ubuntu:latest".to_string(), + Vec::new(), + ); + let mut permissions = CloudObjectPermissions::mock_personal(); + permissions.owner = owner; + CloudAmbientAgentEnvironment::new( + sync_id, + CloudAmbientAgentEnvironmentModel::new(environment), + CloudObjectMetadata::mock(), + permissions, + ) +} + +#[test] +fn environment_scope_includes_personal_and_matching_team_environments() { + let selected_team_uid = ServerId::from(123); + let other_team_uid = ServerId::from(456); + let selected_scope = TeamContextForOperation::new_for_test(selected_team_uid); + let personal_environment = environment_with_owner( + SyncId::ServerId(ServerId::from(1)), + "Personal", + Owner::mock_current_user(), + ); + let selected_team_environment = environment_with_owner( + SyncId::ServerId(ServerId::from(2)), + "Selected team", + Owner::Team { + team_uid: selected_team_uid, + }, + ); + let other_team_environment = environment_with_owner( + SyncId::ServerId(ServerId::from(3)), + "Other team", + Owner::Team { + team_uid: other_team_uid, + }, + ); + + assert!(environment_is_visible_to_scope( + &personal_environment, + &selected_scope + )); + assert!(environment_is_visible_to_scope( + &selected_team_environment, + &selected_scope + )); + assert!(!environment_is_visible_to_scope( + &other_team_environment, + &selected_scope + )); +} + +#[test] +fn multi_team_personal_scope_includes_only_personal_environments() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + app.add_singleton_model(PrivacySettings::mock); + let user_workspaces = app.add_singleton_model(UserWorkspaces::default_mock); + user_workspaces.update(&mut app, |user_workspaces, ctx| { + user_workspaces.setup_test_workspace(ctx); + user_workspaces.update_current_workspace( + |workspace| { + let mut second_team = workspace.teams[0].clone(); + second_team.uid = ServerId::from(456); + second_team.name = "Second team".to_string(); + workspace.teams.push(second_team); + }, + ctx, + ); + }); + let implicit_scope = app.read(|ctx| { + resolve_environment_team_scope( + &ObjectScope { + team_selection: TeamSelection { team: None }, + personal: false, + }, + ctx, + ) + }); + let personal_scope = app + .read(|ctx| { + resolve_environment_team_scope( + &ObjectScope { + team_selection: TeamSelection { team: None }, + personal: true, + }, + ctx, + ) + }) + .expect("explicit personal scope should not require a sole team"); + let personal_environment = environment_with_owner( + SyncId::ServerId(ServerId::from(1)), + "Personal", + Owner::mock_current_user(), + ); + let team_environment = environment_with_owner( + SyncId::ServerId(ServerId::from(2)), + "Team", + Owner::Team { + team_uid: ServerId::from(123), + }, + ); + + assert!(implicit_scope.is_err()); + assert_eq!(personal_scope.team_uid(), None); + assert!(environment_is_visible_to_scope( + &personal_environment, + &personal_scope + )); + assert!(!environment_is_visible_to_scope( + &team_environment, + &personal_scope + )); + }); +} + +#[test] +fn explicit_environment_id_remains_resource_authoritative() { + App::test((), |mut app| async move { + let cloud_model = app.add_singleton_model(CloudModel::mock); + let server_id = ServerId::from(123); + let sync_id = SyncId::ServerId(server_id); + let team_environment = environment_with_owner( + sync_id, + "Other team", + Owner::Team { + team_uid: ServerId::from(456), + }, + ); + cloud_model.update(&mut app, |model, ctx| { + model.create_object(sync_id, team_environment, ctx); + }); + + let choice = app.update(|ctx| { + EnvironmentChoice::resolve_for_create( + EnvironmentCreateArgs { + environment: Some(server_id.to_string()), + no_environment: false, + }, + &TeamlessScopeForTest, + ctx, + ) + }); + + assert_eq!( + choice.unwrap(), + EnvironmentChoice::Environment { + id: server_id.to_string(), + name: "Other team".to_string(), + } + ); + }); +} #[test] fn parse_ambient_task_id_accepts_valid_ids() { diff --git a/app/src/ai/agent_sdk/environment.rs b/app/src/ai/agent_sdk/environment.rs index 2af6f353483..677cb511251 100644 --- a/app/src/ai/agent_sdk/environment.rs +++ b/app/src/ai/agent_sdk/environment.rs @@ -2,13 +2,14 @@ use std::collections::HashSet; use comfy_table::Cell; use cynic::QueryBuilder; +use futures::future; use inquire::error::InquireError; use inquire::{Confirm, Select}; use serde::Serialize; use warp_cli::GlobalOptions; use warp_cli::agent::OutputFormat; use warp_cli::environment::{EnvironmentCommand, ImageCommand}; -use warp_cli::scope::ObjectScope; +use warp_cli::scope::{ObjectScope, TeamSelection}; use warp_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus; use warp_graphql::queries::list_warp_dev_images::{ ListWarpDevImages, ListWarpDevImagesResult, ListWarpDevImagesVariables, @@ -63,8 +64,10 @@ pub fn run( ) -> anyhow::Result<()> { let runner = ctx.add_singleton_model(|_ctx| EnvironmentCommandRunner); match command { - EnvironmentCommand::List => { - runner.update(ctx, |runner, ctx| runner.list(global_options, ctx)); + EnvironmentCommand::List { team_selection } => { + runner.update(ctx, |runner, ctx| { + runner.list(global_options, team_selection, ctx) + }); Ok(()) } EnvironmentCommand::Create { @@ -184,24 +187,36 @@ impl EnvironmentCommandRunner { }); } - fn list(&self, global_options: GlobalOptions, ctx: &mut ModelContext) { - let initial_sync = UpdateManager::as_ref(ctx) - .initial_load_complete() - .with_timeout(WARP_DRIVE_SYNC_TIMEOUT); + fn list( + &self, + global_options: GlobalOptions, + team_selection: TeamSelection, + ctx: &mut ModelContext, + ) { + let refresh_future = super::common::refresh_workspace_metadata(ctx); + let warp_drive_sync_future = super::common::refresh_warp_drive(ctx); + let setup_future = future::try_join(refresh_future, warp_drive_sync_future); - ctx.spawn(initial_sync, move |_, result, ctx| { - if result.is_err() { - super::report_fatal_error( - anyhow::anyhow!("Timed out waiting for Warp Drive to sync"), - ctx, - ); + ctx.spawn(setup_future, move |_, result, ctx| { + if let Err(err) = result { + super::report_fatal_error(err, ctx); return; } + let team_scope = match super::common::resolve_team_scope(&team_selection, ctx) { + Ok(team_scope) => team_scope, + Err(err) => { + super::report_fatal_error(err, ctx); + return; + } + }; let environments = CloudAmbientAgentEnvironment::get_all(ctx); let environment_infos: Vec<_> = environments .iter() + .filter(|environment| { + super::common::environment_is_visible_to_scope(environment, &team_scope) + }) .map(|environment| { let name = environment.model().string_model.name.clone(); let description = environment.model().string_model.description.clone(); diff --git a/app/src/ai/agent_sdk/integration.rs b/app/src/ai/agent_sdk/integration.rs index 681527869c5..7c65cad9b6b 100644 --- a/app/src/ai/agent_sdk/integration.rs +++ b/app/src/ai/agent_sdk/integration.rs @@ -2,6 +2,7 @@ use futures::future; use warp_cli::GlobalOptions; use warp_cli::integration::{CreateIntegrationArgs, IntegrationCommand, UpdateIntegrationArgs}; use warp_cli::provider::ProviderType; +use warp_cli::scope::TeamSelection; use warp_graphql::mutations::create_simple_integration::CreateSimpleIntegrationOutput; use warp_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus; use warp_graphql::queries::get_simple_integrations::SimpleIntegrationsOutput; @@ -12,6 +13,7 @@ use super::common::{EnvironmentChoice, ResolveConfigurationError}; use super::integration_output; use super::oauth_flow::poll_oauth_until_terminal; use crate::server::server_api::ServerApiProvider; +use crate::server::team_scope::RequestTeamScope; pub fn run( ctx: &mut AppContext, @@ -26,41 +28,103 @@ pub fn run( IntegrationCommand::Update(args) => { runner.update(ctx, |runner, ctx| runner.update(args, ctx)); } - IntegrationCommand::List => { - runner.update(ctx, |runner, ctx| runner.list(global_options, ctx)); + IntegrationCommand::List { team_selection } => { + runner.update(ctx, |runner, ctx| { + runner.list(global_options, team_selection, ctx) + }); } } Ok(()) } struct IntegrationCommandRunner; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct IntegrationRetryState { + request_team_scope: RequestTeamScope, + attempt: u32, +} -impl IntegrationCommandRunner { - fn list(&self, global_options: GlobalOptions, ctx: &mut ModelContext) { - // Hardcoded set of providers that this client knows how to render. - let providers = vec![ProviderType::Linear, ProviderType::Slack]; - let provider_slugs: Vec = providers.into_iter().map(|p| p.slug()).collect(); +impl IntegrationRetryState { + fn new(request_team_scope: RequestTeamScope) -> Self { + Self { + request_team_scope, + attempt: 1, + } + } - let integrations_client = ServerApiProvider::as_ref(ctx).get_integrations_client(); + fn next(self) -> Self { + Self { + attempt: self.attempt + 1, + ..self + } + } - let list_future = async move { - integrations_client - .list_simple_integrations(provider_slugs) - .await - }; + fn continue_after_oauth( + self, + poll_result: anyhow::Result, + ) -> anyhow::Result { + match poll_result { + Ok(OauthConnectTxStatus::Completed) => Ok(self.next()), + Ok(OauthConnectTxStatus::Failed) => Err(anyhow::anyhow!("OAuth authorization failed.")), + Ok(OauthConnectTxStatus::Expired) => { + Err(anyhow::anyhow!("OAuth authorization expired.")) + } + Ok(OauthConnectTxStatus::Pending) | Ok(OauthConnectTxStatus::InProgress) => Err( + anyhow::anyhow!("Unexpected non-terminal OAuth status returned"), + ), + Err(err) => Err(anyhow::anyhow!("Error polling OAuth status: {err}")), + } + } +} - ctx.spawn( - list_future, - move |_, result: anyhow::Result, ctx| match result { - Ok(output) => { - integration_output::print_integrations(&output, global_options.output_format); - ctx.terminate_app(TerminationMode::ForceTerminate, None); - } +impl IntegrationCommandRunner { + fn list( + &self, + global_options: GlobalOptions, + team_selection: TeamSelection, + ctx: &mut ModelContext, + ) { + let refresh_future = super::common::refresh_workspace_metadata(ctx); + ctx.spawn(refresh_future, move |_, result, ctx| { + if let Err(err) = result { + super::report_fatal_error(err, ctx); + return; + } + let team_scope = match super::common::resolve_team_scope(&team_selection, ctx) { + Ok(team_scope) => team_scope, Err(err) => { - ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err))); + super::report_fatal_error(err, ctx); + return; } - }, - ); + }; + let request_team_scope = RequestTeamScope::from_scope(&team_scope); + let provider_slugs = [ProviderType::Linear, ProviderType::Slack] + .into_iter() + .map(|provider| provider.slug()) + .collect(); + let integrations_client = ServerApiProvider::as_ref(ctx).get_integrations_client(); + let list_future = async move { + integrations_client + .list_simple_integrations(request_team_scope, provider_slugs) + .await + }; + + ctx.spawn( + list_future, + move |_, result: anyhow::Result, ctx| match result { + Ok(output) => { + integration_output::print_integrations( + &output, + global_options.output_format, + ); + ctx.terminate_app(TerminationMode::ForceTerminate, None); + } + Err(err) => { + ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err))); + } + }, + ); + }); } fn create(&self, args: CreateIntegrationArgs, ctx: &mut ModelContext) { @@ -73,6 +137,14 @@ impl IntegrationCommandRunner { ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err))); return; } + let team_scope = match super::common::resolve_team_scope(&args.team_selection, ctx) { + Ok(team_scope) => team_scope, + Err(err) => { + ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err))); + return; + } + }; + let request_team_scope = RequestTeamScope::from_scope(&team_scope); let loaded_file = match args.config_file.file.as_deref() { Some(path) => match super::config_file::load_config_file(path) { @@ -154,29 +226,30 @@ impl IntegrationCommandRunner { environment_args.environment = merged_config.environment_id.take(); } - let environment_uid = match EnvironmentChoice::resolve_for_create(environment_args, ctx) - { - Ok(EnvironmentChoice::None) => { - eprintln!("Creating integration without an environment."); - None - } - Ok(EnvironmentChoice::Environment { id, .. }) => { - eprintln!("Creating integration with environment {id}."); - Some(id) - } - Err(ResolveConfigurationError::Canceled) => { - eprintln!("Integration creation canceled."); - ctx.terminate_app(TerminationMode::ForceTerminate, None); - return; - } - Err(err) => { - super::report_fatal_error(anyhow::anyhow!(err), ctx); - return; - } - }; + let environment_uid = + match EnvironmentChoice::resolve_for_create(environment_args, &team_scope, ctx) { + Ok(EnvironmentChoice::None) => { + eprintln!("Creating integration without an environment."); + None + } + Ok(EnvironmentChoice::Environment { id, .. }) => { + eprintln!("Creating integration with environment {id}."); + Some(id) + } + Err(ResolveConfigurationError::Canceled) => { + eprintln!("Integration creation canceled."); + ctx.terminate_app(TerminationMode::ForceTerminate, None); + return; + } + Err(err) => { + super::report_fatal_error(anyhow::anyhow!(err), ctx); + return; + } + }; runner.start_create_or_update_flow( ctx, + IntegrationRetryState::new(request_team_scope), integration_type, environment_uid, base_prompt, @@ -186,7 +259,6 @@ impl IntegrationCommandRunner { worker_host, enabled, is_update, - 1, ); }); } @@ -195,6 +267,7 @@ impl IntegrationCommandRunner { fn start_create_or_update_flow( &self, ctx: &mut ModelContext, + retry_state: IntegrationRetryState, integration_type: String, environment_uid: Option, base_prompt: Option, @@ -204,12 +277,10 @@ impl IntegrationCommandRunner { worker_host: Option, enabled: bool, is_update: bool, - attempt: u32, ) { const MAX_CREATE_ATTEMPTS: u32 = 8; let action = if is_update { "update" } else { "creation" }; - - if attempt > MAX_CREATE_ATTEMPTS { + if retry_state.attempt > MAX_CREATE_ATTEMPTS { ctx.terminate_app( TerminationMode::ForceTerminate, Some(Err(anyhow::anyhow!( @@ -234,6 +305,7 @@ impl IntegrationCommandRunner { let create_future = async move { integrations_client .create_or_update_simple_integration( + retry_state.request_team_scope, future_integration_type, future_is_update, future_environment_uid, @@ -263,8 +335,8 @@ impl IntegrationCommandRunner { println!("Authorize the provider here: {auth_url}\n"); ctx.open_url(&auth_url); - let integrations_client = ServerApiProvider::as_ref(ctx) - .get_integrations_client(); + let integrations_client = + ServerApiProvider::as_ref(ctx).get_integrations_client(); let tx_id = tx_id.into_inner(); let poll_future = @@ -279,17 +351,15 @@ impl IntegrationCommandRunner { let next_worker_host = worker_host.clone(); let next_enabled = enabled; let next_is_update = is_update; - let next_attempt = attempt + 1; ctx.spawn( poll_future, move |runner, poll_result, ctx| { - match poll_result { - Ok(OauthConnectTxStatus::Completed) => { - // Inner loop done; try create or update again (outer loop). - // This may happen multiple times if the user needs to authorize multiple services. + match retry_state.continue_after_oauth(poll_result) { + Ok(next_retry_state) => { runner.start_create_or_update_flow( ctx, + next_retry_state, next_integration_type, next_environment_uid, next_base_prompt, @@ -299,33 +369,12 @@ impl IntegrationCommandRunner { next_worker_host, next_enabled, next_is_update, - next_attempt, - ); - } - Ok(OauthConnectTxStatus::Failed) => { - ctx.terminate_app( - TerminationMode::ForceTerminate, - Some(Err(anyhow::anyhow!("OAuth authorization failed."))), - ); - } - Ok(OauthConnectTxStatus::Expired) => { - ctx.terminate_app( - TerminationMode::ForceTerminate, - Some(Err(anyhow::anyhow!("OAuth authorization expired."))), - ); - } - Ok(OauthConnectTxStatus::Pending) - | Ok(OauthConnectTxStatus::InProgress) => { - // Should not be returned by poll_oauth_until_terminal. - ctx.terminate_app( - TerminationMode::ForceTerminate, - Some(Err(anyhow::anyhow!("Unexpected non-terminal OAuth status returned"))), ); } Err(err) => { ctx.terminate_app( TerminationMode::ForceTerminate, - Some(Err(anyhow::anyhow!("Error polling OAuth status: {err}"))), + Some(Err(err)), ); } } @@ -386,6 +435,14 @@ impl IntegrationCommandRunner { ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err))); return; } + let team_scope = match super::common::resolve_team_scope(&args.team_selection, ctx) { + Ok(team_scope) => team_scope, + Err(err) => { + super::report_fatal_error(err, ctx); + return; + } + }; + let request_team_scope = RequestTeamScope::from_scope(&team_scope); let loaded_file = match args.config_file.file.as_deref() { Some(path) => match super::config_file::load_config_file(path) { @@ -490,6 +547,7 @@ impl IntegrationCommandRunner { // Explicitly requested to update without an environment. runner.start_create_or_update_flow( ctx, + IntegrationRetryState::new(request_team_scope), integration_type, Some(String::new()), base_prompt, @@ -499,7 +557,6 @@ impl IntegrationCommandRunner { worker_host, enabled, is_update, - 1, ); return; } @@ -508,6 +565,7 @@ impl IntegrationCommandRunner { runner.start_create_or_update_flow( ctx, + IntegrationRetryState::new(request_team_scope), integration_type, environment_uid, base_prompt, @@ -517,7 +575,6 @@ impl IntegrationCommandRunner { worker_host, enabled, is_update, - 1, ); }); } @@ -527,3 +584,7 @@ impl warpui::Entity for IntegrationCommandRunner { type Event = (); } impl SingletonEntity for IntegrationCommandRunner {} + +#[cfg(test)] +#[path = "integration_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/integration_tests.rs b/app/src/ai/agent_sdk/integration_tests.rs new file mode 100644 index 00000000000..d4d485c764e --- /dev/null +++ b/app/src/ai/agent_sdk/integration_tests.rs @@ -0,0 +1,26 @@ +use warp_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus; + +use super::IntegrationRetryState; +use crate::server::ids::ServerId; +use crate::server::team_scope::RequestTeamScope; +use crate::workspaces::user_workspaces::{TeamContextForOperation, TeamlessScopeForTest}; + +#[test] +fn completed_oauth_continuation_preserves_the_initiating_team_scope() { + let scope = TeamContextForOperation::new_for_test(ServerId::from(7)); + assert_completed_oauth_continuation(RequestTeamScope::from_scope(&scope)); +} + +#[test] +fn completed_oauth_continuation_preserves_the_initiating_teamless_scope() { + assert_completed_oauth_continuation(RequestTeamScope::from_scope(&TeamlessScopeForTest)); +} + +fn assert_completed_oauth_continuation(expected_scope: RequestTeamScope) { + let retry_state = IntegrationRetryState::new(expected_scope) + .continue_after_oauth(Ok(OauthConnectTxStatus::Completed)) + .unwrap(); + + assert_eq!(retry_state.attempt, 2); + assert_eq!(retry_state.request_team_scope, expected_scope); +} diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index ed7fe3eff3c..4b3fe0d59a6 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -1636,7 +1636,7 @@ fn command_requires_auth(command: &CliCommand) -> bool { AgentCommand::Skills(_) => true, }, CliCommand::Environment(environment_cmd) => match environment_cmd { - EnvironmentCommand::List => true, + EnvironmentCommand::List { .. } => true, EnvironmentCommand::Create { .. } => true, EnvironmentCommand::Delete { .. } => true, EnvironmentCommand::Update { .. } => true, @@ -1862,7 +1862,9 @@ fn command_to_telemetry_event(command: &CliCommand) -> CliTelemetryEvent { CliCommand::Agent(AgentCommand::Update(_)) => CliTelemetryEvent::AgentUpdate, CliCommand::Agent(AgentCommand::Delete(_)) => CliTelemetryEvent::AgentDelete, CliCommand::Agent(AgentCommand::Skills(_)) => CliTelemetryEvent::AgentSkills, - CliCommand::Environment(EnvironmentCommand::List) => CliTelemetryEvent::EnvironmentList, + CliCommand::Environment(EnvironmentCommand::List { .. }) => { + CliTelemetryEvent::EnvironmentList + } CliCommand::Environment(EnvironmentCommand::Create { .. }) => { CliTelemetryEvent::EnvironmentCreate } @@ -1927,7 +1929,7 @@ fn command_to_telemetry_event(command: &CliCommand) -> CliTelemetryEvent { CliCommand::Integration(integration_cmd) => match integration_cmd { IntegrationCommand::Create(_) => CliTelemetryEvent::IntegrationCreate, IntegrationCommand::Update(_) => CliTelemetryEvent::IntegrationUpdate, - IntegrationCommand::List => CliTelemetryEvent::IntegrationList, + IntegrationCommand::List { .. } => CliTelemetryEvent::IntegrationList, }, CliCommand::Schedule(c) => match c.subcommand() { None | Some(ScheduleSubcommand::Create(_)) => CliTelemetryEvent::ScheduleCreate, diff --git a/app/src/ai/agent_sdk/provider.rs b/app/src/ai/agent_sdk/provider.rs index 9c89d23f653..f1aabf31a99 100644 --- a/app/src/ai/agent_sdk/provider.rs +++ b/app/src/ai/agent_sdk/provider.rs @@ -1,15 +1,17 @@ //! Provider command for linking third-party services. + use comfy_table::Cell; use serde::Serialize; use warp_cli::GlobalOptions; use warp_cli::provider::{ProviderCommand, ProviderType}; +use warp_cli::scope::ObjectScope; use warp_core::channel::ChannelState; use warpui::platform::TerminationMode; use warpui::{AppContext, ModelContext, SingletonEntity}; use crate::ai::agent_sdk::common::describe_sole_team_error; use crate::ai::agent_sdk::output::{self, TableFormat}; -use crate::workspaces::user_workspaces::UserWorkspaces; +use crate::workspaces::user_workspaces::{SoleTeamError, TeamScope}; /// Handle provider-related CLI commands. pub fn run( @@ -19,9 +21,12 @@ pub fn run( ) -> anyhow::Result<()> { let runner = ctx.add_singleton_model(|_ctx| ProviderCommandRunner); match command { - ProviderCommand::Setup(args) => runner.update(ctx, |runner, ctx| { - runner.setup(args.provider_type, args.team, args.personal, ctx) - }), + ProviderCommand::Setup(args) => { + runner.update(ctx, |runner, ctx| { + runner.setup(args.provider_type, args.scope, ctx) + }); + Ok(()) + } ProviderCommand::List => runner.update(ctx, |runner, ctx| runner.list(global_options, ctx)), } } @@ -31,18 +36,30 @@ struct ProviderCommandRunner; impl ProviderCommandRunner { // This shouldn't need to be done, it's usually done as part of create - fn setup( + fn setup(&self, provider_type: ProviderType, scope: ObjectScope, ctx: &mut ModelContext) { + let refresh_future = super::common::refresh_workspace_metadata(ctx); + ctx.spawn(refresh_future, move |runner, result, ctx| { + if let Err(err) = result { + super::report_fatal_error(err, ctx); + return; + } + if let Err(err) = runner.finish_setup(provider_type, scope, ctx) { + super::report_fatal_error(err, ctx); + } + }); + } + + fn finish_setup( &self, provider_type: ProviderType, - team: bool, - personal: bool, + scope: ObjectScope, ctx: &mut ModelContext, ) -> anyhow::Result<()> { // Construct the OAuth connect URL let server_url = ChannelState::server_root_url(); - let mut use_team_auth = team; - if !team && !personal { + let mut use_team_auth = scope.is_team(); + if !scope.is_team() && !scope.personal { if provider_type.allowed_in_team_context() && provider_type.allowed_in_personal_context() { @@ -52,16 +69,17 @@ impl ProviderCommandRunner { )); } use_team_auth = provider_type.allowed_in_team_context(); - } else if personal { + } else if scope.personal { use_team_auth = false; } // TODO(bens): initiate the OAuth flow and use the login-less auth URL let slug = provider_type.slug(); let url = if use_team_auth { - let team_uid = UserWorkspaces::as_ref(ctx) - .sole_team_uid() - .map_err(|err| describe_sole_team_error(err, ctx))?; + let team_scope = super::common::resolve_team_scope(&scope.team_selection, ctx)?; + let team_uid = team_scope + .team_uid() + .ok_or_else(|| describe_sole_team_error(SoleTeamError::NoTeam, ctx))?; format!("{server_url}/oauth/connect/{slug}?principalType=team&principalId={team_uid}") } else { format!("{server_url}/oauth/connect/{slug}") @@ -153,3 +171,7 @@ impl TableFormat for ProviderInfo { ] } } + +#[cfg(test)] +#[path = "provider_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/provider_tests.rs b/app/src/ai/agent_sdk/provider_tests.rs new file mode 100644 index 00000000000..d96bb0a785d --- /dev/null +++ b/app/src/ai/agent_sdk/provider_tests.rs @@ -0,0 +1,93 @@ +use std::sync::Arc; + +use warp_cli::provider::ProviderType; +use warp_cli::scope::{ObjectScope, TeamSelection}; +use warpui::App; + +use super::ProviderCommandRunner; +use crate::auth::AuthStateProvider; +use crate::network::NetworkStatus; +use crate::server::ids::ServerId; +use crate::server::server_api::team::MockTeamClient; +use crate::settings::PrivacySettings; +use crate::workspaces::team::Team; +use crate::workspaces::team_tester::TeamTesterStatus; +use crate::workspaces::update_manager::TeamUpdateManager; +use crate::workspaces::user_workspaces::{ + UserWorkspaces, WorkspacesMetadataResponse, WorkspacesMetadataWithPricing, +}; +use crate::workspaces::workspace::{Workspace, WorkspaceUid}; + +#[test] +fn setup_resolves_team_scope_from_refreshed_workspace_metadata() { + App::test((), |mut app| async move { + let workspace_uid = WorkspaceUid::from(ServerId::from(1)); + let team_uid = ServerId::from(2); + let workspace = Workspace::from_local_cache( + workspace_uid, + "Test Workspace".to_string(), + Some(vec![Team::from_local_cache( + team_uid, + "Test Team".to_string(), + None, + None, + None, + None, + )]), + None, + ); + let mut team_client = MockTeamClient::new(); + team_client + .expect_workspaces_metadata() + .times(1) + .return_once(|| { + Ok(WorkspacesMetadataWithPricing { + metadata: WorkspacesMetadataResponse { + workspaces: vec![workspace], + joinable_teams: vec![], + experiments: None, + ai_credit_availability: None, + user_purchase_policy: None, + }, + pricing_info: None, + }) + }); + + app.add_singleton_model(|_| NetworkStatus::new()); + app.add_singleton_model(TeamTesterStatus::new); + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + app.add_singleton_model(PrivacySettings::mock); + app.add_singleton_model(UserWorkspaces::default_mock); + app.add_singleton_model(|ctx| TeamUpdateManager::new(Arc::new(team_client), None, ctx)); + + let (opened_url_sender, opened_url_receiver) = async_channel::bounded(1); + app.update(|ctx| { + ctx.set_before_open_url(move |url, _| { + opened_url_sender.try_send(url.to_string()).unwrap(); + url.to_string() + }); + }); + + let runner = app.add_model(|_| ProviderCommandRunner); + runner.update(&mut app, |runner, ctx| { + runner.setup( + ProviderType::Slack, + ObjectScope { + team_selection: TeamSelection { team: Some(None) }, + personal: false, + }, + ctx, + ); + }); + + assert!(matches!( + opened_url_receiver.try_recv(), + Err(async_channel::TryRecvError::Empty) + )); + let opened_url = opened_url_receiver.recv().await.unwrap(); + assert!(opened_url.ends_with(&format!( + "/oauth/connect/slack?principalType=team&principalId={team_uid}" + ))); + assert!(app.termination_result().is_none()); + }); +} diff --git a/app/src/ai/agent_sdk/schedule.rs b/app/src/ai/agent_sdk/schedule.rs index 2e1b2ba815d..6b16a9ddc38 100644 --- a/app/src/ai/agent_sdk/schedule.rs +++ b/app/src/ai/agent_sdk/schedule.rs @@ -51,6 +51,13 @@ fn create(ctx: &mut AppContext, args: CreateScheduleArgs) -> anyhow::Result<()> super::report_fatal_error(err, ctx); return; } + let team_scope = match super::common::resolve_environment_team_scope(&args.scope, ctx) { + Ok(team_scope) => team_scope, + Err(err) => { + super::report_fatal_error(err, ctx); + return; + } + }; let loaded_file = match args.config_file.file.as_deref() { Some(path) => match super::config_file::load_config_file(path) { @@ -73,30 +80,23 @@ fn create(ctx: &mut AppContext, args: CreateScheduleArgs) -> anyhow::Result<()> environment_args.environment = Some(environment_id); } - let environment_id = match EnvironmentChoice::resolve_for_create(environment_args, ctx) - { - Ok(EnvironmentChoice::None) => { - eprintln!("Scheduling agent to run without an environment."); - None - } - Ok(EnvironmentChoice::Environment { id, .. }) => Some(id), - Err(ResolveConfigurationError::Canceled) => { - ctx.terminate_app(TerminationMode::ForceTerminate, None); - return; - } - Err(err) => { - super::report_fatal_error(anyhow::anyhow!(err), ctx); - return; - } - }; + let environment_id = + match EnvironmentChoice::resolve_for_create(environment_args, &team_scope, ctx) { + Ok(EnvironmentChoice::None) => { + eprintln!("Scheduling agent to run without an environment."); + None + } + Ok(EnvironmentChoice::Environment { id, .. }) => Some(id), + Err(ResolveConfigurationError::Canceled) => { + ctx.terminate_app(TerminationMode::ForceTerminate, None); + return; + } + Err(err) => { + super::report_fatal_error(anyhow::anyhow!(err), ctx); + return; + } + }; - let team_scope = match super::common::resolve_object_scope(&args.scope, ctx) { - Ok(team_scope) => team_scope, - Err(err) => { - super::report_fatal_error(err, ctx); - return; - } - }; let owner = match super::common::resolve_owner_for_team_scope(&team_scope, ctx) { Ok(owner) => owner, Err(err) => { diff --git a/app/src/server/server_api.rs b/app/src/server/server_api.rs index ee8d40cc579..a1bfb578975 100644 --- a/app/src/server/server_api.rs +++ b/app/src/server/server_api.rs @@ -566,10 +566,14 @@ impl ServerApi { timeout, ) } - - fn send_graphql_request_for_team<'a, QF, O: warp_graphql::client::Operation + Send + 'a>( + pub fn send_team_scoped_graphql_request< + 'a, + QF, + O: warp_graphql::client::Operation + Send + 'a, + >( &'a self, operation: O, + timeout: Option, team_scope: RequestTeamScope, ) -> BoxFuture<'a, Result> where @@ -580,18 +584,29 @@ impl ServerApi { warp_server_client::graphql_helpers::send_team_scoped_graphql_request( &self.base_client, operation, - None, + timeout, team_uid, ) } None => warp_server_client::graphql_helpers::send_graphql_request( &self.base_client, operation, - None, + timeout, ), } } + fn send_graphql_request_for_team<'a, QF, O: warp_graphql::client::Operation + Send + 'a>( + &'a self, + operation: O, + team_scope: RequestTeamScope, + ) -> BoxFuture<'a, Result> + where + QF: 'a, + { + self.send_team_scoped_graphql_request(operation, None, team_scope) + } + fn team_uid_header_value(team_scope: RequestTeamScope) -> Option { team_scope .team_uid() diff --git a/app/src/server/server_api/integrations.rs b/app/src/server/server_api/integrations.rs index b85abac1bdf..e48f6c2094b 100644 --- a/app/src/server/server_api/integrations.rs +++ b/app/src/server/server_api/integrations.rs @@ -37,6 +37,7 @@ use super::ServerApi; use crate::channel::ChannelState; use crate::features::FeatureFlag; use crate::server::graphql::{get_request_context, get_user_facing_error_message}; +use crate::server::team_scope::RequestTeamScope; #[cfg(not(target_family = "wasm"))] pub trait IntegrationsClientBounds: Send + Sync {} @@ -79,6 +80,7 @@ pub trait IntegrationsClient: 'static + IntegrationsClientBounds { #[allow(clippy::too_many_arguments)] async fn create_or_update_simple_integration( &self, + team_scope: RequestTeamScope, integration_type: String, is_update: bool, environment_uid: Option, @@ -96,6 +98,7 @@ pub trait IntegrationsClient: 'static + IntegrationsClientBounds { /// regardless of whether the connection or integration currently exists. async fn list_simple_integrations( &self, + team_scope: RequestTeamScope, providers: Vec, ) -> Result; @@ -167,6 +170,7 @@ impl IntegrationsClient for ServerApi { #[allow(clippy::too_many_arguments)] async fn create_or_update_simple_integration( &self, + team_scope: RequestTeamScope, integration_type: String, is_update: bool, environment_uid: Option, @@ -193,7 +197,9 @@ impl IntegrationsClient for ServerApi { }; let operation = CreateSimpleIntegration::build(variables); - let response = self.send_graphql_request(operation, None).await?; + let response = self + .send_team_scoped_graphql_request(operation, None, team_scope) + .await?; match response.create_simple_integration { CreateSimpleIntegrationResult::CreateSimpleIntegrationOutput(output) => Ok(output), CreateSimpleIntegrationResult::UserFacingError(error) => { @@ -232,6 +238,7 @@ impl IntegrationsClient for ServerApi { async fn list_simple_integrations( &self, + team_scope: RequestTeamScope, providers: Vec, ) -> Result { let variables = SimpleIntegrationsVariables { @@ -240,7 +247,9 @@ impl IntegrationsClient for ServerApi { }; let operation = SimpleIntegrations::build(variables); - let response = self.send_graphql_request(operation, None).await?; + let response = self + .send_team_scoped_graphql_request(operation, None, team_scope) + .await?; match response.simple_integrations { SimpleIntegrationsResult::SimpleIntegrationsOutput(output) => Ok(output), diff --git a/crates/warp_cli/src/environment.rs b/crates/warp_cli/src/environment.rs index 31d8b67e9dd..ea70c83613f 100644 --- a/crates/warp_cli/src/environment.rs +++ b/crates/warp_cli/src/environment.rs @@ -1,6 +1,6 @@ use clap::{ArgAction, ArgGroup, Args, Subcommand}; -use crate::scope::ObjectScope; +use crate::scope::{ObjectScope, TeamSelection}; /// Maximum length for environment descriptions. const MAX_DESCRIPTION_LENGTH: usize = 240; @@ -24,7 +24,10 @@ fn validate_description(s: &str) -> Result { #[command(visible_alias = "e")] pub enum EnvironmentCommand { /// List cloud environments. - List, + List { + #[command(flatten)] + team_selection: TeamSelection, + }, /// Manage base images for cloud environments. #[command(subcommand)] Image(ImageCommand), @@ -104,7 +107,7 @@ pub enum EnvironmentCommand { impl EnvironmentCommand { pub(crate) fn as_str_for_tracing(&self) -> &'static str { match self { - EnvironmentCommand::List => "environment list", + EnvironmentCommand::List { .. } => "environment list", EnvironmentCommand::Image(_) => "environment image", EnvironmentCommand::Create { .. } => "environment create", EnvironmentCommand::Delete { .. } => "environment delete", diff --git a/crates/warp_cli/src/integration.rs b/crates/warp_cli/src/integration.rs index bcf95985d19..8075394f97a 100644 --- a/crates/warp_cli/src/integration.rs +++ b/crates/warp_cli/src/integration.rs @@ -5,6 +5,7 @@ use crate::environment::{EnvironmentCreateArgs, EnvironmentUpdateArgs}; use crate::mcp::MCPSpec; use crate::model::ModelArgs; use crate::provider::ProviderType; +use crate::scope::TeamSelection; /// Integration-related subcommands. #[derive(Debug, Clone, Subcommand)] @@ -15,7 +16,10 @@ pub enum IntegrationCommand { /// Update an integration. Update(UpdateIntegrationArgs), /// List simple integrations and their connection status. - List, + List { + #[command(flatten)] + team_selection: TeamSelection, + }, } impl IntegrationCommand { @@ -23,7 +27,7 @@ impl IntegrationCommand { match self { IntegrationCommand::Create(_) => "integration create", IntegrationCommand::Update(_) => "integration update", - IntegrationCommand::List => "integration list", + IntegrationCommand::List { .. } => "integration list", } } } @@ -33,6 +37,8 @@ pub struct CreateIntegrationArgs { /// Provider to create the integration for. #[arg(value_enum)] pub provider: ProviderType, + #[command(flatten)] + pub team_selection: TeamSelection, #[command(flatten)] pub model: ModelArgs, @@ -68,6 +74,8 @@ pub struct UpdateIntegrationArgs { /// Provider to update the integration for. #[arg(value_enum)] pub provider: ProviderType, + #[command(flatten)] + pub team_selection: TeamSelection, #[command(flatten)] pub model: ModelArgs, diff --git a/crates/warp_cli/src/lib_tests.rs b/crates/warp_cli/src/lib_tests.rs index 53bee4a1fce..6b8634120db 100644 --- a/crates/warp_cli/src/lib_tests.rs +++ b/crates/warp_cli/src/lib_tests.rs @@ -9,6 +9,7 @@ use crate::environment::{EnvironmentCommand, ImageCommand}; use crate::harness_support::{HarnessSupportCommand, TaskStatus}; use crate::integration::IntegrationCommand; use crate::memory_store::{MemoryCommand, MemoryStoreCommand}; +use crate::provider::ProviderCommand; use crate::schedule::ScheduleSubcommand; use crate::secret::{CodexMethod, CreateProvider, SecretCommand}; use crate::task::{MessageCommand, TaskCommand}; @@ -23,6 +24,106 @@ fn identifies_worker_subcommands() { assert!(!is_worker_invocation("--prompt")); } +#[test] +fn integration_commands_parse_team_selection() { + let cases = [ + ( + vec!["warp", "integration", "list", "--team=team-list"], + "team-list", + ), + ( + vec![ + "warp", + "integration", + "create", + "slack", + "--team=team-create", + ], + "team-create", + ), + ( + vec![ + "warp", + "integration", + "update", + "linear", + "--team=team-update", + ], + "team-update", + ), + ]; + + for (command, expected_team_uid) in cases { + let args = Args::try_parse_from(command).expect("integration team scope should parse"); + let Some(Command::CommandLine(boxed_cmd)) = args.command else { + panic!("Expected an integration command"); + }; + let team_selection = match boxed_cmd.as_ref() { + CliCommand::Integration(IntegrationCommand::List { team_selection }) => team_selection, + CliCommand::Integration(IntegrationCommand::Create(args)) => &args.team_selection, + CliCommand::Integration(IntegrationCommand::Update(args)) => &args.team_selection, + _ => panic!("Expected an integration command"), + }; + + assert_eq!(team_selection.requested_team_uid(), Some(expected_team_uid)); + } +} + +#[test] +fn integration_list_distinguishes_default_and_bare_team_selection() { + for (command, expected_team) in [ + (vec!["warp", "integration", "list"], None), + (vec!["warp", "integration", "list", "--team"], Some(None)), + ] { + let args = Args::try_parse_from(command).expect("integration list scope should parse"); + let Some(Command::CommandLine(boxed_cmd)) = args.command else { + panic!("Expected `warp integration list` command"); + }; + let CliCommand::Integration(IntegrationCommand::List { team_selection }) = + boxed_cmd.as_ref() + else { + panic!("Expected `warp integration list` command"); + }; + + assert_eq!(team_selection.team, expected_team); + } +} + +#[test] +fn provider_setup_parses_object_scope() { + let team = Args::try_parse_from(["warp", "provider", "setup", "slack", "--team=team-provider"]) + .expect("provider team scope should parse"); + let Some(Command::CommandLine(boxed_cmd)) = team.command else { + panic!("Expected `warp provider setup` command"); + }; + let CliCommand::Provider(ProviderCommand::Setup(args)) = boxed_cmd.as_ref() else { + panic!("Expected `warp provider setup` command"); + }; + assert_eq!(args.scope.requested_team_uid(), Some("team-provider")); + assert!(!args.scope.personal); + + let personal = Args::try_parse_from(["warp", "provider", "setup", "slack", "--personal"]) + .expect("provider personal scope should parse"); + let Some(Command::CommandLine(boxed_cmd)) = personal.command else { + panic!("Expected `warp provider setup` command"); + }; + let CliCommand::Provider(ProviderCommand::Setup(args)) = boxed_cmd.as_ref() else { + panic!("Expected `warp provider setup` command"); + }; + assert!(args.scope.personal); + assert!(!args.scope.is_team()); + + Args::try_parse_from([ + "warp", + "provider", + "setup", + "slack", + "--personal", + "--team=team-provider", + ]) + .expect_err("provider scopes should be mutually exclusive"); +} + /// Pins that each pair of constants names the same variable under both prefixes. A typo in /// either half would otherwise go unnoticed until a consumer read the wrong name. #[test] @@ -2305,6 +2406,43 @@ fn environment_image_list_parses() { assert!(matches!(image_cmd, ImageCommand::List)); } +fn parse_environment_list(args: &[&str]) -> crate::scope::TeamSelection { + let full_args = std::iter::once("warp") + .chain(["environment", "list"]) + .chain(args.iter().copied()); + let args = Args::try_parse_from(full_args).expect("environment list args should parse"); + + let Some(Command::CommandLine(boxed_cmd)) = args.command else { + panic!("Expected `warp environment list` command"); + }; + let CliCommand::Environment(EnvironmentCommand::List { team_selection }) = boxed_cmd.as_ref() + else { + panic!("Expected `warp environment list` command"); + }; + + team_selection.clone() +} + +#[test] +fn environment_list_defaults_to_implicit_team_selection() { + let team_selection = parse_environment_list(&[]); + + assert_eq!(team_selection.team, None); +} + +#[test] +fn environment_list_accepts_bare_team_selection() { + let team_selection = parse_environment_list(&["--team"]); + + assert_eq!(team_selection.team, Some(None)); +} + +#[test] +fn environment_list_accepts_explicit_team_selection() { + let team_selection = parse_environment_list(&["--team=123"]); + + assert_eq!(team_selection.team, Some(Some("123".to_string()))); +} #[test] fn environment_create_accepts_description() { diff --git a/crates/warp_cli/src/provider.rs b/crates/warp_cli/src/provider.rs index caebd29b7ef..bd3fe6b9762 100644 --- a/crates/warp_cli/src/provider.rs +++ b/crates/warp_cli/src/provider.rs @@ -1,4 +1,6 @@ -use clap::{ArgGroup, Args, Subcommand, ValueEnum}; +use clap::{Args, Subcommand, ValueEnum}; + +use crate::scope::ObjectScope; /// Provider-related subcommands. #[derive(Debug, Clone, Subcommand)] @@ -53,15 +55,10 @@ impl ProviderType { } #[derive(Debug, Clone, Args)] -#[command(group(ArgGroup::new("scope").required(false)))] pub struct SetupArgs { /// The type of provider to setup. pub provider_type: ProviderType, - /// Setup provider for a team - #[arg(long, group = "scope")] - pub team: bool, - /// Setup provider for a personal account - #[arg(long, conflicts_with = "team", group = "scope")] - pub personal: bool, + #[command(flatten)] + pub scope: ObjectScope, } diff --git a/crates/warp_server_client/src/base_client.rs b/crates/warp_server_client/src/base_client.rs index 99aa70e90a7..ac00b50ddc4 100644 --- a/crates/warp_server_client/src/base_client.rs +++ b/crates/warp_server_client/src/base_client.rs @@ -253,6 +253,14 @@ impl BaseClient { *self.ambient_agent_task_id.write() = task_id; } + #[cfg(test)] + pub(crate) fn set_ambient_workload_token_for_test(&self, token: &str) { + *self.ambient_workload_token.lock() = Some(warp_isolation_platform::WorkloadToken { + token: token.to_string(), + expires_at: None, + }); + } + /// Returns an ambient agent workload token when the current runtime can issue one. pub async fn get_or_create_ambient_workload_token(&self) -> Result> { if cfg!(target_family = "wasm") { diff --git a/crates/warp_server_client/src/graphql_helpers_tests.rs b/crates/warp_server_client/src/graphql_helpers_tests.rs index 002c65b78f8..00b59c4856b 100644 --- a/crates/warp_server_client/src/graphql_helpers_tests.rs +++ b/crates/warp_server_client/src/graphql_helpers_tests.rs @@ -18,18 +18,17 @@ use crate::base_client::{ fn base_client(auth_state: AuthState) -> (BaseClient, async_channel::Receiver) { let (event_sender, event_receiver) = async_channel::unbounded(); - ( - BaseClient::new( - Arc::new(http_client::Client::new()), - Arc::new(auth_state), - event_sender, - None, - GraphqlRoutingConfig::default(), - AuthenticatedGraphqlConfig::default(), - None, - ), - event_receiver, - ) + let base_client = BaseClient::new( + Arc::new(http_client::Client::new()), + Arc::new(auth_state), + event_sender, + None, + GraphqlRoutingConfig::default(), + AuthenticatedGraphqlConfig::default(), + None, + ); + base_client.set_ambient_workload_token_for_test("test-workload-token"); + (base_client, event_receiver) } #[test]