mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-06 13:44:21 +08:00
508 lines
16 KiB
Rust
508 lines
16 KiB
Rust
//! Implements Cursor account information services.
|
|
use axum::{
|
|
body::{to_bytes, Body},
|
|
extract::Extension,
|
|
http::{header, HeaderValue, Request, Response},
|
|
};
|
|
use prost::Message;
|
|
use serde_json::Value;
|
|
|
|
use crate::{api::cursor::proxy, local_app, Result};
|
|
|
|
use super::entitlement::FreeEntitlementCache;
|
|
|
|
const LOCAL_AUTH_ID: &str = "cursor-local-user";
|
|
const LOCAL_EMAIL: &str = "cursor@ai.com";
|
|
const LOCAL_ULTRA_PLAN_INCLUDED_CENTS: i32 = 20_000;
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct GetEmailResponse {
|
|
#[prost(string, tag = "1")]
|
|
email: String,
|
|
#[prost(int32, tag = "2")]
|
|
sign_up_type: i32,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct GetUserMetaResponse {
|
|
#[prost(string, tag = "1")]
|
|
email: String,
|
|
#[prost(int32, tag = "2")]
|
|
sign_up_type: i32,
|
|
#[prost(int64, tag = "3")]
|
|
user_id: i64,
|
|
#[prost(string, optional, tag = "4")]
|
|
workos_id: Option<String>,
|
|
#[prost(string, optional, tag = "5")]
|
|
profile_picture_url: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct GetMeResponse {
|
|
#[prost(string, tag = "1")]
|
|
auth_id: String,
|
|
#[prost(int32, tag = "2")]
|
|
user_id: i32,
|
|
#[prost(string, optional, tag = "3")]
|
|
email: Option<String>,
|
|
#[prost(string, optional, tag = "4")]
|
|
first_name: Option<String>,
|
|
#[prost(string, optional, tag = "5")]
|
|
last_name: Option<String>,
|
|
#[prost(string, optional, tag = "8")]
|
|
created_at: Option<String>,
|
|
#[prost(bool, optional, tag = "9")]
|
|
is_enterprise_user: Option<bool>,
|
|
#[prost(string, optional, tag = "11")]
|
|
email_domain_type: Option<String>,
|
|
#[prost(string, optional, tag = "12")]
|
|
country: Option<String>,
|
|
#[prost(string, optional, tag = "13")]
|
|
profile_picture_url: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct GetUserProfileResponse {
|
|
#[prost(bool, optional, tag = "4")]
|
|
public_visibility_allowed: Option<bool>,
|
|
#[prost(string, optional, tag = "5")]
|
|
max_visibility: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct GetCurrentPeriodUsageResponse {
|
|
#[prost(int64, tag = "1")]
|
|
billing_cycle_start: i64,
|
|
#[prost(int64, tag = "2")]
|
|
billing_cycle_end: i64,
|
|
#[prost(message, optional, tag = "3")]
|
|
plan_usage: Option<PlanUsage>,
|
|
#[prost(message, optional, tag = "4")]
|
|
spend_limit_usage: Option<SpendLimitUsage>,
|
|
#[prost(int32, optional, tag = "5")]
|
|
display_threshold: Option<i32>,
|
|
#[prost(bool, tag = "6")]
|
|
enabled: bool,
|
|
#[prost(string, tag = "7")]
|
|
display_message: String,
|
|
#[prost(string, optional, tag = "11")]
|
|
auto_model_selected_display_message: Option<String>,
|
|
#[prost(string, optional, tag = "12")]
|
|
named_model_selected_display_message: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct PlanUsage {
|
|
#[prost(int32, tag = "1")]
|
|
total_spend: i32,
|
|
#[prost(int32, tag = "2")]
|
|
included_spend: i32,
|
|
#[prost(int32, tag = "4")]
|
|
remaining: i32,
|
|
#[prost(int32, tag = "5")]
|
|
limit: i32,
|
|
#[prost(bool, optional, tag = "6")]
|
|
remaining_bonus: Option<bool>,
|
|
#[prost(string, optional, tag = "7")]
|
|
bonus_tooltip: Option<String>,
|
|
#[prost(int32, optional, tag = "8")]
|
|
auto_spend: Option<i32>,
|
|
#[prost(int32, optional, tag = "9")]
|
|
api_spend: Option<i32>,
|
|
#[prost(double, optional, tag = "12")]
|
|
auto_percent_used: Option<f64>,
|
|
#[prost(double, optional, tag = "13")]
|
|
api_percent_used: Option<f64>,
|
|
#[prost(double, optional, tag = "14")]
|
|
total_percent_used: Option<f64>,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct SpendLimitUsage {
|
|
#[prost(string, tag = "8")]
|
|
limit_type: String,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct GetUsageLimitStatusAndActiveGrantsResponse {
|
|
#[prost(message, optional, tag = "1")]
|
|
usage_limit_policy_status: Option<UsageLimitPolicyStatus>,
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Message)]
|
|
struct UsageLimitPolicyStatus {
|
|
#[prost(bool, tag = "1")]
|
|
is_in_slow_pool: bool,
|
|
#[prost(map = "string, string", tag = "5")]
|
|
features: std::collections::HashMap<String, String>,
|
|
#[prost(bool, tag = "6")]
|
|
can_configure_spend_limit: bool,
|
|
#[prost(bool, tag = "8")]
|
|
has_pending_request: bool,
|
|
#[prost(string, repeated, tag = "9")]
|
|
allowed_model_ids: Vec<String>,
|
|
#[prost(string, repeated, tag = "10")]
|
|
allowed_model_tags: Vec<String>,
|
|
}
|
|
|
|
#[derive(Clone, Copy, PartialEq, Message)]
|
|
struct Empty {}
|
|
|
|
pub async fn get_email(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_forward(upstream, request, || {
|
|
proto(GetEmailResponse {
|
|
email: LOCAL_EMAIL.into(),
|
|
sign_up_type: 3,
|
|
})
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn get_user_meta(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_forward(upstream, request, || {
|
|
proto(GetUserMetaResponse {
|
|
email: LOCAL_EMAIL.into(),
|
|
sign_up_type: 3,
|
|
user_id: 1,
|
|
workos_id: None,
|
|
profile_picture_url: None,
|
|
})
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn get_me(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_forward(upstream, request, || {
|
|
proto(GetMeResponse {
|
|
auth_id: LOCAL_AUTH_ID.into(),
|
|
user_id: 1,
|
|
email: Some(LOCAL_EMAIL.into()),
|
|
first_name: Some("Cursor".into()),
|
|
last_name: Some("Local".into()),
|
|
created_at: Some(chrono::Utc::now().to_rfc3339()),
|
|
is_enterprise_user: Some(false),
|
|
email_domain_type: Some("personal".into()),
|
|
country: Some("US".into()),
|
|
profile_picture_url: None,
|
|
})
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn get_teams(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_forward(upstream, request, || proto(Empty {})).await
|
|
}
|
|
|
|
pub async fn get_user_profile(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_forward(upstream, request, || {
|
|
proto(GetUserProfileResponse {
|
|
public_visibility_allowed: Some(true),
|
|
max_visibility: Some("PUBLIC".into()),
|
|
})
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn current_period_usage(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
Extension(free_entitlements): Extension<FreeEntitlementCache>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_confirmed_free_or_forward(upstream, &free_entitlements, request, || {
|
|
let now = chrono::Utc::now();
|
|
proto(GetCurrentPeriodUsageResponse {
|
|
billing_cycle_start: (now - chrono::Duration::days(30)).timestamp_millis(),
|
|
billing_cycle_end: (now + chrono::Duration::days(10 * 365)).timestamp_millis(),
|
|
plan_usage: Some(PlanUsage {
|
|
total_spend: 0,
|
|
included_spend: LOCAL_ULTRA_PLAN_INCLUDED_CENTS,
|
|
remaining: LOCAL_ULTRA_PLAN_INCLUDED_CENTS,
|
|
limit: LOCAL_ULTRA_PLAN_INCLUDED_CENTS,
|
|
remaining_bonus: Some(false),
|
|
bonus_tooltip: Some("Ultra local account mock is active.".into()),
|
|
auto_spend: Some(0),
|
|
api_spend: Some(0),
|
|
auto_percent_used: Some(0.0),
|
|
api_percent_used: Some(0.0),
|
|
total_percent_used: Some(0.0),
|
|
}),
|
|
spend_limit_usage: Some(SpendLimitUsage {
|
|
limit_type: "user".into(),
|
|
}),
|
|
display_threshold: Some(99_999_999),
|
|
enabled: true,
|
|
display_message: "Ultra plan active".into(),
|
|
auto_model_selected_display_message: Some("Ultra plan active".into()),
|
|
named_model_selected_display_message: Some("Ultra plan active".into()),
|
|
})
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn usage_limit_status(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
Extension(free_entitlements): Extension<FreeEntitlementCache>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
local_or_confirmed_free_or_forward(upstream, &free_entitlements, request, || {
|
|
proto(GetUsageLimitStatusAndActiveGrantsResponse {
|
|
usage_limit_policy_status: Some(UsageLimitPolicyStatus {
|
|
is_in_slow_pool: false,
|
|
features: Default::default(),
|
|
can_configure_spend_limit: true,
|
|
has_pending_request: false,
|
|
allowed_model_ids: Vec::new(),
|
|
allowed_model_tags: Vec::new(),
|
|
}),
|
|
})
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn stripe_profile(
|
|
Extension(upstream): Extension<proxy::CursorProxy>,
|
|
Extension(free_entitlements): Extension<FreeEntitlementCache>,
|
|
request: Request<Body>,
|
|
) -> Result<Response<Body>> {
|
|
if local_app::request_uses_local_cursor_token(request.headers()) {
|
|
return local_stripe_profile(request.headers().get(header::ORIGIN).cloned());
|
|
}
|
|
|
|
let request_headers = request.headers().clone();
|
|
let origin = request.headers().get(header::ORIGIN).cloned();
|
|
let upstream_response = match proxy::forward_buffered(&upstream, request).await {
|
|
Ok(response) => response,
|
|
Err(error) if free_entitlements.is_confirmed_free(&request_headers) => {
|
|
tracing::warn!(%error, "using cached Free entitlement after Stripe upstream failure");
|
|
return local_stripe_profile(origin);
|
|
}
|
|
Err(error) => return Err(error),
|
|
};
|
|
if !upstream_response.status.is_success() {
|
|
if should_fallback_to_cached_free(
|
|
upstream_response.status,
|
|
&free_entitlements,
|
|
&request_headers,
|
|
) {
|
|
tracing::warn!(
|
|
status = %upstream_response.status,
|
|
"using cached Free entitlement after Stripe upstream failure"
|
|
);
|
|
return local_stripe_profile(origin);
|
|
}
|
|
return Ok(upstream_response.into_response());
|
|
}
|
|
|
|
let Some(membership_type) = membership_type(&upstream_response.body) else {
|
|
return Ok(upstream_response.into_response());
|
|
};
|
|
let observed = free_entitlements.observe_membership(&request_headers, &membership_type);
|
|
if observed && membership_type.eq_ignore_ascii_case("free") {
|
|
return local_stripe_profile(origin);
|
|
}
|
|
|
|
Ok(upstream_response.into_response())
|
|
}
|
|
|
|
fn should_fallback_to_cached_free(
|
|
status: axum::http::StatusCode,
|
|
free_entitlements: &FreeEntitlementCache,
|
|
headers: &axum::http::HeaderMap,
|
|
) -> bool {
|
|
status.is_server_error() && free_entitlements.is_confirmed_free(headers)
|
|
}
|
|
|
|
fn membership_type(body: &[u8]) -> Option<String> {
|
|
let profile: Value = serde_json::from_slice(body).ok()?;
|
|
let membership_type = profile.get("membershipType")?.as_str()?.trim();
|
|
(!membership_type.is_empty()).then(|| membership_type.to_owned())
|
|
}
|
|
|
|
fn local_stripe_profile(origin: Option<HeaderValue>) -> Result<Response<Body>> {
|
|
let mut response = json(ultra_profile())?;
|
|
if let Some(origin) = origin {
|
|
response
|
|
.headers_mut()
|
|
.insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, origin);
|
|
response.headers_mut().insert(
|
|
header::ACCESS_CONTROL_ALLOW_CREDENTIALS,
|
|
HeaderValue::from_static("true"),
|
|
);
|
|
response
|
|
.headers_mut()
|
|
.insert(header::VARY, HeaderValue::from_static("Origin"));
|
|
}
|
|
Ok(response)
|
|
}
|
|
|
|
async fn local_or_confirmed_free_or_forward(
|
|
upstream: proxy::CursorProxy,
|
|
free_entitlements: &FreeEntitlementCache,
|
|
request: Request<Body>,
|
|
local: impl FnOnce() -> Result<Response<Body>>,
|
|
) -> Result<Response<Body>> {
|
|
if local_app::request_uses_local_cursor_token(request.headers())
|
|
|| free_entitlements.is_confirmed_free(request.headers())
|
|
{
|
|
consume_body(request).await?;
|
|
return local();
|
|
}
|
|
proxy::forward(Extension(upstream), request).await
|
|
}
|
|
|
|
async fn local_or_forward(
|
|
upstream: proxy::CursorProxy,
|
|
request: Request<Body>,
|
|
local: impl FnOnce() -> Result<Response<Body>>,
|
|
) -> Result<Response<Body>> {
|
|
if local_app::request_uses_local_cursor_token(request.headers()) {
|
|
consume_body(request).await?;
|
|
return local();
|
|
}
|
|
proxy::forward(Extension(upstream), request).await
|
|
}
|
|
|
|
async fn consume_body(request: Request<Body>) -> Result<()> {
|
|
to_bytes(request.into_body(), usize::MAX)
|
|
.await
|
|
.map_err(|error| crate::Error::Protocol(format!("cannot read request body: {error}")))?;
|
|
Ok(())
|
|
}
|
|
|
|
fn proto(message: impl Message) -> Result<Response<Body>> {
|
|
response("application/proto", message.encode_to_vec())
|
|
}
|
|
|
|
fn json(value: Value) -> Result<Response<Body>> {
|
|
response("application/json", serde_json::to_vec(&value)?)
|
|
}
|
|
|
|
fn response(content_type: &'static str, body: Vec<u8>) -> Result<Response<Body>> {
|
|
let length = body.len();
|
|
let mut response = Response::new(Body::from(body));
|
|
response.headers_mut().insert(
|
|
header::CONTENT_TYPE,
|
|
axum::http::HeaderValue::from_static(content_type),
|
|
);
|
|
response.headers_mut().insert(
|
|
header::CONTENT_LENGTH,
|
|
length
|
|
.to_string()
|
|
.parse()
|
|
.expect("body length is always a valid header value"),
|
|
);
|
|
Ok(response)
|
|
}
|
|
|
|
fn ultra_profile() -> Value {
|
|
serde_json::json!({
|
|
"membershipType": "ultra",
|
|
"individualMembershipType": "ultra",
|
|
"subscriptionStatus": "active",
|
|
"lastPaymentFailed": false,
|
|
"pendingCancellationDate": null,
|
|
"daysRemainingOnTrial": 0,
|
|
"paymentId": LOCAL_AUTH_ID,
|
|
"isTeamMember": false
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::{
|
|
atomic::{AtomicBool, Ordering},
|
|
Arc,
|
|
};
|
|
|
|
use axum::body::Bytes;
|
|
use futures_util::stream;
|
|
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn local_account_response_consumes_request_body_before_replying() {
|
|
let polled = Arc::new(AtomicBool::new(false));
|
|
let observed = polled.clone();
|
|
let body = Body::from_stream(stream::once(async move {
|
|
observed.store(true, Ordering::SeqCst);
|
|
Ok::<_, std::convert::Infallible>(Bytes::from_static(b"request"))
|
|
}));
|
|
|
|
consume_body(Request::new(body)).await.unwrap();
|
|
|
|
assert!(polled.load(Ordering::SeqCst));
|
|
}
|
|
|
|
#[test]
|
|
fn cached_free_fallback_accepts_server_failures_but_not_auth_failures() {
|
|
let cache = FreeEntitlementCache::default();
|
|
let mut headers = axum::http::HeaderMap::new();
|
|
headers.insert(
|
|
header::AUTHORIZATION,
|
|
HeaderValue::from_static("Bearer official-free-token"),
|
|
);
|
|
assert!(cache.observe_membership(&headers, "free"));
|
|
|
|
assert!(should_fallback_to_cached_free(
|
|
axum::http::StatusCode::BAD_GATEWAY,
|
|
&cache,
|
|
&headers
|
|
));
|
|
assert!(!should_fallback_to_cached_free(
|
|
axum::http::StatusCode::UNAUTHORIZED,
|
|
&cache,
|
|
&headers
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn reads_only_a_non_empty_membership_type() {
|
|
assert_eq!(
|
|
membership_type(br#"{"membershipType":"free"}"#).as_deref(),
|
|
Some("free")
|
|
);
|
|
assert_eq!(membership_type(br#"{"membershipType":""}"#), None);
|
|
assert_eq!(membership_type(br#"{"subscriptionStatus":"active"}"#), None);
|
|
assert_eq!(membership_type(b"not-json"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn local_stripe_profile_allows_the_cursor_app_origin() {
|
|
let origin = HeaderValue::from_static("vscode-file://vscode-app");
|
|
let response = local_stripe_profile(Some(origin.clone())).unwrap();
|
|
assert_eq!(
|
|
response.headers().get(header::ACCESS_CONTROL_ALLOW_ORIGIN),
|
|
Some(&origin)
|
|
);
|
|
assert_eq!(
|
|
response
|
|
.headers()
|
|
.get(header::ACCESS_CONTROL_ALLOW_CREDENTIALS),
|
|
Some(&HeaderValue::from_static("true"))
|
|
);
|
|
assert_eq!(
|
|
response.headers().get(header::VARY),
|
|
Some(&HeaderValue::from_static("Origin"))
|
|
);
|
|
assert_eq!(
|
|
response.headers().get(header::CONTENT_TYPE),
|
|
Some(&HeaderValue::from_static("application/json"))
|
|
);
|
|
}
|
|
}
|