rainy_sdk/endpoints/
cowork.rs

1//! Cowork endpoint for retrieving subscription capabilities
2//!
3//! This endpoint validates the API key and returns the user's
4//! Cowork plan, available models, and feature access.
5
6use crate::{
7    cowork::{CoworkCapabilities, CoworkFeatures, CoworkPlan, CoworkProfile, CoworkUsage},
8    error::Result,
9    RainyClient,
10};
11
12impl RainyClient {
13    /// Retrieve Cowork capabilities for the current API key.
14    ///
15    /// This method validates the API key and returns information about:
16    /// - Subscription plan (Free, GoPlus, Plus, Pro, ProPlus)
17    /// - Available AI models
18    /// - Feature access (web research, document export, etc.)
19    /// - Usage tracking
20    ///
21    /// # Returns
22    ///
23    /// A `Result` containing `CoworkCapabilities` on success, or a `RainyError` on failure.
24    pub async fn get_cowork_capabilities(&self) -> Result<CoworkCapabilities> {
25        match self.get_cowork_profile().await {
26            Ok(profile) => {
27                // In a real app, features/models might come from the API too,
28                // or be derived from the plan ID.
29                // For now, we simulate them based on plan ID or use defaults.
30                // Assuming the API returns features/models in the profile response is better,
31                // but if not, logic goes here.
32                // The new CoworkProfile struct does NOT have models/features, so we map them here.
33
34                let models = match profile.plan.id.as_str() {
35                    "free" => vec!["gemini-2.0-flash".to_string()],
36                    "go" | "go_plus" => {
37                        vec![
38                            "gemini-2.0-flash".to_string(),
39                            "gemini-2.5-flash".to_string(),
40                        ]
41                    }
42                    "plus" => vec![
43                        "gemini-2.0-flash".to_string(),
44                        "gemini-2.5-flash".to_string(),
45                        "gemini-2.5-pro".to_string(),
46                    ],
47                    "pro" | "pro_plus" => vec![
48                        "gemini-2.0-flash".to_string(),
49                        "gemini-2.5-flash".to_string(),
50                        "gemini-2.5-pro".to_string(),
51                        "claude-sonnet-4".to_string(),
52                    ],
53                    _ => vec![],
54                };
55
56                let features = CoworkFeatures {
57                    web_research: profile.plan.id != "free",
58                    document_export: profile.plan.id != "free",
59                    image_analysis: true,
60                    priority_support: profile.plan.id.contains("pro"),
61                };
62
63                Ok(CoworkCapabilities {
64                    profile,
65                    features,
66                    is_valid: true,
67                    models,
68                    upgrade_message: None,
69                })
70            }
71            Err(_) => Ok(CoworkCapabilities::free()),
72        }
73    }
74
75    /// Check if the current API key grants paid access.
76    pub async fn has_paid_plan(&self) -> bool {
77        match self.get_cowork_capabilities().await {
78            Ok(caps) => caps.profile.plan.is_paid(),
79            Err(_) => false,
80        }
81    }
82
83    /// Get available models for Cowork based on subscription plan.
84    pub async fn get_cowork_models(&self) -> Result<Vec<String>> {
85        let caps = self.get_cowork_capabilities().await?;
86        Ok(caps.models)
87    }
88
89    /// Check if a specific feature is available.
90    pub async fn can_use_feature(&self, feature: &str) -> bool {
91        match self.get_cowork_capabilities().await {
92            Ok(caps) => caps.can_use_feature(feature),
93            Err(_) => false,
94        }
95    }
96
97    /// Check if a specific model is available for the current plan.
98    pub async fn can_use_model(&self, model: &str) -> bool {
99        match self.get_cowork_capabilities().await {
100            Ok(caps) => caps.can_use_model(model),
101            Err(_) => false,
102        }
103    }
104
105    /// Check if user can make another request based on usage limits.
106    pub async fn can_make_request(&self) -> bool {
107        match self.get_cowork_capabilities().await {
108            Ok(caps) => caps.can_make_request(),
109            Err(_) => false,
110        }
111    }
112}
113
114/// Offline capabilities based on cached plan data.
115pub fn get_offline_capabilities(cached_plan: Option<CoworkPlan>) -> CoworkCapabilities {
116    match cached_plan {
117        Some(plan) if plan.is_paid() => {
118            // Reconstruct minimal capabilities
119            CoworkCapabilities {
120                profile: CoworkProfile {
121                    name: "Offline User".to_string(),
122                    email: "".to_string(),
123                    plan,
124                    usage: CoworkUsage::default(),
125                },
126                features: CoworkFeatures::default(), // Pessimistic offline features
127                is_valid: true,
128                models: vec!["gemini-2.5-flash".to_string()],
129                upgrade_message: None,
130            }
131        }
132        _ => CoworkCapabilities::free(),
133    }
134}