wip

Author: oz-agent

app/src/ai/agent_sdk/common.rs

@@ -17,14 +17,15 @@ use crate::ai::agent_sdk::driver::{AgentDriverError, WARP_DRIVE_SYNC_TIMEOUT};
 use crate::ai::ambient_agents::AmbientAgentTaskId;
 use crate::ai::cloud_environments::CloudAmbientAgentEnvironment;
 use crate::ai::llms::{LLMId, LLMPreferences};
+use crate::auth::UserUid;
 use crate::auth::auth_state::AuthStateProvider;
 use crate::cloud_object::{CloudObject, CloudObjectLookup as _, Owner};
 use crate::server::cloud_objects::update_manager::UpdateManager;
 use crate::server::ids::{ServerId, SyncId};
 use crate::server::server_api::ServerApiProvider;
 use crate::server::server_api::ai::AIClient;
 use crate::workspaces::update_manager::TeamUpdateManager;
-use crate::workspaces::user_workspaces::UserWorkspaces;
+use crate::workspaces::user_workspaces::{SoleTeamError, UserWorkspaces};
 
 /// How long to wait for workspace metadata to refresh.
 pub const WORKSPACE_METADATA_REFRESH_TIMEOUT: Duration = Duration::from_secs(10);
@@ -95,38 +96,56 @@ pub(super) fn set_ambient_task_context_from_run_id(
     Ok(task_id)
 }
 
-/// Resolve the owner of a new cloud object. This resolution is based on the CLI `--team` and `--personal` flags.
+pub(super) fn describe_sole_team_error(error: SoleTeamError) -> anyhow::Error {
+    match error {
+        SoleTeamError::NoTeam => anyhow::anyhow!("You are not on a team"),
+        SoleTeamError::MoreThanOneTeam { team_uids } => anyhow::anyhow!(
+            "You are on {} teams; specify one with --team <UID>: {}",
+            team_uids.len(),
+            team_uids
+                .iter()
+                .map(ServerId::to_string)
+                .collect::<Vec<_>>()
+                .join(", ")
+        ),
+    }
+}
+
+fn current_user_uid(ctx: &AppContext) -> anyhow::Result<UserUid> {
+    AuthStateProvider::as_ref(ctx)
+        .get()
+        .user_id()
+        .ok_or_else(|| anyhow::anyhow!("User should be logged in"))
+}
+
+/// Resolve the owner of a new cloud object, based on the CLI `--team` and `--personal` flags.
 ///
-/// If `team_flag` is true, attempts to get the current team UID (errors if not on a team).
-/// If `user_flag` is true, gets the current user's UID.
-/// Otherwise, defaults to team if available, falling back to user.
+/// With neither flag, a user on exactly one team gets a team object and a user on no team gets
+/// a personal one. A user on several teams is asked to choose rather than silently handed a
+/// personal object.
 pub fn resolve_owner(team_flag: bool, user_flag: bool, ctx: &AppContext) -> anyhow::Result<Owner> {
     if team_flag {
-        let team_id = UserWorkspaces::as_ref(ctx)
+        let team_uid = UserWorkspaces::as_ref(ctx)
             .sole_team_uid()
-            .ok_or_else(|| anyhow::anyhow!("User is not on a team"))?;
-        return Ok(Owner::Team { team_uid: team_id });
+            .map_err(describe_sole_team_error)?;
+        return Ok(Owner::Team { team_uid });
     }
 
     if user_flag {
-        let user_id = AuthStateProvider::as_ref(ctx)
-            .get()
-            .user_id()
-            .ok_or_else(|| anyhow::anyhow!("User should be logged in"))?;
-        return Ok(Owner::User { user_uid: user_id });
+        return Ok(Owner::User {
+            user_uid: current_user_uid(ctx)?,
+        });
     }
 
-    // Default: try team first, fall back to user
-    if let Some(team_uid) = UserWorkspaces::as_ref(ctx).sole_team_uid() {
-        return Ok(Owner::Team { team_uid });
+    match UserWorkspaces::as_ref(ctx).sole_team_uid() {
+        Ok(team_uid) => Ok(Owner::Team { team_uid }),
+        Err(SoleTeamError::NoTeam) => Ok(Owner::User {
+            user_uid: current_user_uid(ctx)?,
+        }),
+        Err(error @ SoleTeamError::MoreThanOneTeam { .. }) => {
+            Err(describe_sole_team_error(error))
+        }
     }
-
-    log::warn!("Tried to default to creating team object, team could not be found.");
-    let user_id = AuthStateProvider::as_ref(ctx)
-        .get()
-        .user_id()
-        .ok_or_else(|| anyhow::anyhow!("User should be logged in"))?;
-    Ok(Owner::User { user_uid: user_id })
 }
 
 /// Refresh workspace metadata before executing an operation.

app/src/ai/agent_sdk/provider.rs

@@ -7,6 +7,7 @@ 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;
 
@@ -58,12 +59,9 @@ impl ProviderCommandRunner {
         // 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 = match UserWorkspaces::as_ref(ctx).sole_team_uid() {
-                Some(uid) => uid,
-                None => {
-                    return Err(anyhow::anyhow!("User is not on a team"));
-                }
-            };
+            let team_uid = UserWorkspaces::as_ref(ctx)
+                .sole_team_uid()
+                .map_err(describe_sole_team_error)?;
             format!("{server_url}/oauth/connect/{slug}?principalType=team&principalId={team_uid}")
         } else {
             format!("{server_url}/oauth/connect/{slug}")

app/src/ai/blocklist/inline_action/orchestration_controls.rs

@@ -52,6 +52,7 @@ use crate::view_components::FilterableDropdown;
 use crate::view_components::dropdown::{
     Dropdown, DropdownAction, DropdownItemAction, DropdownStyle,
 };
+use crate::workspaces::user_workspaces::UserWorkspaces;
 
 // ── Shared constants ────────────────────────────────────────────────
 
@@ -279,13 +280,15 @@ fn oz_model_menu_items<A: OrchestrationControlAction, V: View>(
                 .find(|llm| llm.id.to_string() == row.id)
         })
         .collect();
+    let scope = UserWorkspaces::as_ref(ctx).team_context_for_view(ctx);
     available_model_menu_items(
         ordered_choices,
         move |llm| DropdownAction::select_action_and_close(A::model_changed(llm.id.to_string())),
         None,
         None,
         false,
         false,
+        &scope,
         ctx,
     )
 }

app/src/ai/blocklist/prompt/prompt_alert.rs

@@ -436,7 +436,7 @@ impl View for PromptAlertView {
         });
 
         let can_purchase_addon_credits = workspaces
-            .purchase_policy_for_team(current_team)
+            .purchase_policy()
             .is_some_and(|policy| policy.allows_purchases());
 
         let suggest_buy_credits = can_purchase_addon_credits

app/src/ai/custom_model_router_editor.rs

@@ -42,6 +42,7 @@ use crate::view_components::action_button::{
     ActionButton, ButtonSize, PrimaryTheme, SecondaryTheme,
 };
 use crate::view_components::dropdown::DropdownAction;
+use crate::workspaces::user_workspaces::UserWorkspaces;
 
 pub const HEADER_TEXT: &str = "Router Editor";
 
@@ -1048,6 +1049,7 @@ fn fill_filterable_dropdown<F>(
     // `set_filtered_items` keeps an empty selection blank rather than
     // auto-selecting the first model.
     dropdown.set_placeholder(MODEL_PLACEHOLDER, ctx);
+    let scope = UserWorkspaces::as_ref(ctx).team_context_for_view(ctx);
     let items = available_model_menu_items(
         LLMPreferences::as_ref(ctx)
             .get_base_llm_choices_for_agent_mode(ctx)
@@ -1058,6 +1060,7 @@ fn fill_filterable_dropdown<F>(
         None,
         false,
         false,
+        &scope,
         ctx,
     );
     dropdown.set_rich_items(items, ctx);

app/src/ai/document/ai_document_model.rs

@@ -47,7 +47,7 @@ use crate::terminal::TerminalView;
 use crate::terminal::model::session::Session;
 use crate::terminal::model::session::active_session::ActiveSession;
 use crate::throttle::throttle;
-use crate::workspaces::user_workspaces::UserWorkspaces;
+use crate::workspaces::user_workspaces::{SoleTeamError, UserWorkspaces};
 
 /// The frequency at which we check for modifications and save the AI document to the server.
 /// Uses the same 2-second period as notebooks for consistency.
@@ -1275,13 +1275,22 @@ impl AIDocumentModel {
     fn get_plan_owner(ctx: &AppContext) -> Option<Owner> {
         let is_service_account = AuthStateProvider::as_ref(ctx).get().is_service_account();
 
-        if is_service_account {
-            // If the SA doesn't have a team, we'll skip the plan sync in the caller
-            UserWorkspaces::as_ref(ctx)
-                .sole_team_uid()
-                .map(|team_uid| Owner::Team { team_uid })
-        } else {
-            UserWorkspaces::as_ref(ctx).personal_drive(ctx)
+        if !is_service_account {
+            return UserWorkspaces::as_ref(ctx).personal_drive(ctx);
+        }
+        match UserWorkspaces::as_ref(ctx).sole_team_uid() {
+            Ok(team_uid) => Some(Owner::Team { team_uid }),
+            // The caller skips plan sync without a team.
+            Err(SoleTeamError::NoTeam) => None,
+            // A service account is bound to exactly one team server-side, so several here means
+            // the client's view of its memberships disagrees with that.
+            Err(error @ SoleTeamError::MoreThanOneTeam { .. }) => {
+                report_error!(
+                    anyhow::Error::new(error)
+                        .context("Service account resolved to more than one team")
+                );
+                None
+            }
         }
     }
 

app/src/ai/execution_profiles/editor/mod.rs

@@ -1202,13 +1202,15 @@ impl ExecutionProfileEditorView {
                 .iter()
                 .any(|llm| matches!(llm.disable_reason, Some(DisableReason::RequiresUpgrade)));
 
+            let scope = UserWorkspaces::as_ref(ctx).team_context_for_view(ctx);
             let items = available_model_menu_items(
                 choices,
                 |llm| DropdownAction::select_action_and_close(create_action(llm.id.clone())),
                 None,
                 None,
                 false,
                 false,
+                &scope,
                 ctx,
             );
             dropdown.set_rich_items(items, ctx);
@@ -1250,6 +1252,7 @@ impl ExecutionProfileEditorView {
                 .get_coding_llm_choices(ctx)
                 .collect_vec();
 
+            let scope = UserWorkspaces::as_ref(ctx).team_context_for_view(ctx);
             let items = available_model_menu_items(
                 choices,
                 |llm| {
@@ -1261,6 +1264,7 @@ impl ExecutionProfileEditorView {
                 None,
                 false,
                 false,
+                &scope,
                 ctx,
             );
             dropdown.set_rich_items(items, ctx);

app/src/ai/execution_profiles/model_menu_items.rs

@@ -16,6 +16,7 @@ use crate::ai::llms::{
     should_show_gemini_enterprise_agent_platform_icon_for_model, should_show_key_icon_for_model,
 };
 use crate::menu::{MenuItem, MenuItemFields, MenuTooltipPosition};
+use crate::workspaces::user_workspaces::TeamScope;
 
 pub fn is_auto(llm: &LLMInfo) -> bool {
     llm.display_name.to_lowercase().contains("auto")
@@ -75,6 +76,7 @@ fn make_item_fields<A: Action + Clone>(
     model_id_to_add_profile_default_label_to: Option<&LLMId>,
     collapse_auto: bool,
     collapse_reasoning_variants: bool,
+    scope: &dyn TeamScope,
     app: &AppContext,
 ) -> MenuItem<A> {
     let is_auto_model = is_auto(llm);
@@ -88,7 +90,7 @@ fn make_item_fields<A: Action + Clone>(
     let is_using_bedrock = should_show_bedrock_icon_for_model(llm, app);
     let is_using_gemini_enterprise_agent_platform =
         should_show_gemini_enterprise_agent_platform_icon_for_model(llm, app);
-    let is_using_api_key = should_show_key_icon_for_model(llm, app);
+    let is_using_api_key = should_show_key_icon_for_model(llm, scope, app);
     let is_custom_router = is_custom_router_id(llm.id.as_str());
     let leading_icon = model_leading_icon(
         llm,
@@ -183,6 +185,7 @@ pub fn available_model_menu_items<A: Action + Clone>(
     position_id_fn: Option<&dyn Fn(&LLMId) -> String>,
     collapse_auto: bool,
     collapse_reasoning_variants: bool,
+    scope: &dyn TeamScope,
     app: &AppContext,
 ) -> Vec<MenuItem<A>> {
     choices
@@ -195,6 +198,7 @@ pub fn available_model_menu_items<A: Action + Clone>(
                 model_id_to_add_profile_default_label_to,
                 collapse_auto,
                 collapse_reasoning_variants,
+                scope,
                 app,
             )
         })