Skip to main content

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