diff --git a/.docker/selfhost/schema.json b/.docker/selfhost/schema.json index 6f00d1a3e9..fea7344555 100644 --- a/.docker/selfhost/schema.json +++ b/.docker/selfhost/schema.json @@ -199,9 +199,9 @@ "description": "Whether require email verification before accessing restricted resources(not implemented).\n@default true", "default": true }, - "newAccountShareActionDelay": { + "newAccountActionDelay": { "type": "number", - "description": "Minimum account age in seconds before new accounts can invite members or create share links.\n@default 86400", + "description": "Minimum account age in seconds before new accounts can invite members, create invite links, or publish documents. Set to 0 to disable.\n@default 86400", "default": 86400 }, "trustedCloudflareHeaders": { diff --git a/Cargo.lock b/Cargo.lock index ed70708113..03d420c6f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4758,9 +4758,9 @@ dependencies = [ [[package]] name = "llm_adapter" -version = "0.2.20" +version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d13be366ea35d2a9966ad5770e3d070e495af4e1e6dd009638dee4b55ab45ed" +checksum = "02c6b4fa5178b8183331a7d51e8f7ece5421ecb966d3431f3d9ce3c724c6fcc0" dependencies = [ "base64", "jsonschema", diff --git a/packages/backend/native/index.d.ts b/packages/backend/native/index.d.ts index e080782cac..a9b036b896 100644 --- a/packages/backend/native/index.d.ts +++ b/packages/backend/native/index.d.ts @@ -34,6 +34,7 @@ export declare class BackendRuntime { claimInviteAbuseAction(actionId: string, workerId: string): Promise claimRetryableInviteAbuseActions(workerId: string, limit: number): Promise> markInviteAbuseAction(actionId: string, workerId: string, status: string, error?: string | undefined | null): Promise + evaluateWorkspaceActionV1(actorUserId: string, workspaceId: string): Promise assertWorkspaceInviteQuotaV1(input: RuntimeWorkspaceInviteQuotaInput): Promise commitWorkspaceInviteQuotaV1(reservationId: string, usage: RuntimeWorkspaceInviteQuotaUsage): Promise releaseWorkspaceInviteQuotaV1(reservationId: string): Promise @@ -1496,6 +1497,12 @@ export interface RuntimeVerificationTokenRecord { expiresAtMs: number } +export interface RuntimeWorkspaceActionDecision { + allowed: boolean + retryAfterSeconds?: number + reason?: string +} + export interface RuntimeWorkspaceArtifact { id: string workspaceId: string diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs b/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs index 83275e081d..e9ac386163 100644 --- a/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs +++ b/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs @@ -6,10 +6,10 @@ use llm_adapter::{ AttachmentKind, AttachmentSource, ModelFeature, ModelInput, ModelOutput, ModelRequirements, declared_model_matches, }, core::{ - CoreContent, CoreMessage, CoreRequest, CoreRole, CoreToolDefinition, EmbeddingRequest, ImageOptions, - ImageProviderOptions, ImageRequest, RerankCandidate, RerankRequest, StructuredRequest, + CoreContent, CoreMessage, CoreRequest, CoreRole, CoreToolChoice, CoreToolDefinition, EmbeddingRequest, + ImageOptions, ImageProviderOptions, ImageRequest, RerankCandidate, RerankRequest, StructuredRequest, }, - router::{ExecutablePreparedRoute, ExecutableRequest, dispatch_prepared_route}, + router::{ExecutablePreparedRoute, ExecutableRequest, ExecutableResponse, dispatch_prepared_route}, target::{ BackendCredential, BackendEndpoint, BackendOperation, BackendTargetInput, EgressPolicy, compile_backend_target, }, @@ -30,9 +30,11 @@ pub(super) async fn execute_probe( checks: Vec, ) -> RuntimeResult { let tested_at_ms = chrono::Utc::now().timestamp_millis(); - let mut requested = HashSet::new(); + let mut requested = Vec::new(); + let mut requested_set = HashSet::new(); for check in checks { - if !requested.insert((check.model_id.clone(), check.operation.clone())) { + let key = (check.model_id.clone(), check.operation.clone()); + if !requested_set.insert(key.clone()) { return Err(RuntimeError::invalid_input("duplicate BYOK probe check")); } if !matches!( @@ -41,6 +43,7 @@ pub(super) async fn execute_probe( ) { return Err(RuntimeError::invalid_input("unknown BYOK probe operation")); } + requested.push(key); } let mut models = Vec::new(); @@ -160,7 +163,49 @@ fn dispatch_check( Err(_) => return failed(checked_at, "invalid_probe_request"), }; match dispatch_prepared_route(&DefaultHttpClient::default(), &route) { - Ok(_) => verified(checked_at), + Ok(ExecutableResponse::Chat(response)) => { + let valid = if operation == "tool_calling" { + response + .message + .content + .iter() + .any(|content| matches!(content, CoreContent::ToolCall { name, .. } if name == "byok_probe")) + } else { + response + .message + .content + .iter() + .any(|content| matches!(content, CoreContent::Text { text } if !text.trim().is_empty())) + }; + if valid { + verified(checked_at) + } else { + failed(checked_at, "invalid_response") + } + } + Ok(ExecutableResponse::Structured(response)) => { + let valid = response.output_json.as_ref().is_some_and(|output| { + let ExecutableRequest::Structured(request) = &route.request else { + return false; + }; + llm_adapter::schema::validate_json_schema(&request.schema, output).is_ok() + }); + if valid { + verified(checked_at) + } else { + failed(checked_at, "invalid_response") + } + } + Ok(ExecutableResponse::Embedding(response)) if operation == "embedding" && !response.embeddings.is_empty() => { + verified(checked_at) + } + Ok(ExecutableResponse::Rerank(response)) if operation == "rerank" && !response.scores.is_empty() => { + verified(checked_at) + } + Ok(ExecutableResponse::Image(response)) if operation == "image" && !response.images.is_empty() => { + verified(checked_at) + } + Ok(_) => failed(checked_at, "invalid_response"), Err(error) => failed(checked_at, backend_error_kind(&error)), } } @@ -169,7 +214,11 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest { let message = CoreMessage { role: CoreRole::User, content: vec![CoreContent::Text { - text: "Reply with OK.".to_string(), + text: match operation { + "tool_calling" => "Call the byok_probe tool.".to_string(), + "structured" => "Return exactly {\"ok\":true}.".to_string(), + _ => "Reply with OK.".to_string(), + }, }], }; match operation { @@ -177,8 +226,8 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest { model: String::new(), messages: vec![message], stream: false, - max_tokens: Some(8), - temperature: Some(0.0), + max_tokens: Some(64), + temperature: None, tools: if operation == "tool_calling" { vec![CoreToolDefinition { name: "byok_probe".to_string(), @@ -188,7 +237,9 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest { } else { vec![] }, - tool_choice: None, + tool_choice: (operation == "tool_calling").then_some(CoreToolChoice::Specific { + name: "byok_probe".to_string(), + }), include: None, reasoning: None, response_schema: None, @@ -202,8 +253,8 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest { "required": ["ok"], "additionalProperties": false }), - max_tokens: Some(16), - temperature: Some(0.0), + max_tokens: Some(128), + temperature: None, reasoning: None, strict: Some(true), response_mime_type: Some("application/json".to_string()), @@ -426,7 +477,48 @@ mod tests { let mut stream = stream.unwrap(); let request = read_request(&mut stream); let responses = request.starts_with("POST /v1/responses "); - let body = if responses { + let embedding = request.starts_with("POST /v1/embeddings "); + let image = request.starts_with("POST /v1/images/generations "); + let rerank = request.contains("\"logprobs\":true"); + let tool_calling = request.contains("byok_probe"); + let body = if embedding { + json!({ + "model": "smoke-model", + "data": [{ "embedding": [0.1], "index": 0 }], + "usage": { "prompt_tokens": 1, "total_tokens": 1 } + }) + } else if image { + json!({ + "created": 0, + "data": [{ "url": "https://example.com/smoke.png" }] + }) + } else if rerank { + json!({ + "model": "smoke-model", + "choices": [{ + "logprobs": { "content": [{ + "top_logprobs": [ + { "token": "Yes", "logprob": 0.0 }, + { "token": "No", "logprob": -1.0 } + ] + }] } + }] + }) + } else if responses && tool_calling { + json!({ + "id": "resp_smoke", + "model": "smoke-model", + "status": "completed", + "output": [{ + "type": "function_call", + "id": "fc_smoke", + "call_id": "call_smoke", + "name": "byok_probe", + "arguments": "{}" + }], + "usage": { "input_tokens": 1, "output_tokens": 1, "total_tokens": 2 } + }) + } else if responses { json!({ "id": "resp_smoke", "model": "smoke-model", @@ -439,6 +531,25 @@ mod tests { }], "usage": { "input_tokens": 1, "output_tokens": 1, "total_tokens": 2 } }) + } else if tool_calling { + json!({ + "id": "chat_smoke", + "model": "smoke-model", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_smoke", + "type": "function", + "function": { "name": "byok_probe", "arguments": "{}" } + }] + }, + "finish_reason": "tool_calls" + }], + "usage": { "prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2 } + }) } else { json!({ "id": "chat_smoke", @@ -489,7 +600,7 @@ mod tests { #[test] fn openai_compatible_probe_smoke_uses_the_selected_dialect() { - let operations = ["chat", "structured", "tool_calling"]; + let operations = ["chat", "structured", "tool_calling", "embedding", "rerank", "image"]; let (endpoint, requests, server) = serve_openai_compatible(operations.len() * 2); for dialect in [OpenAiDialect::Responses, OpenAiDialect::ChatCompletions] { @@ -520,19 +631,39 @@ mod tests { .iter() .filter(|request| request.starts_with("POST /v1/responses ")) .count(), - operations.len() + 3 ); assert_eq!( requests .iter() .filter(|request| request.starts_with("POST /v1/chat/completions ")) .count(), - operations.len() + 5 + ); + assert_eq!( + requests + .iter() + .filter(|request| request.starts_with("POST /v1/embeddings ")) + .count(), + 2 + ); + assert_eq!( + requests + .iter() + .filter(|request| request.starts_with("POST /v1/images/generations ")) + .count(), + 2 ); assert!(requests.iter().all(|request| !request.contains("/models"))); assert_eq!( requests.iter().filter(|request| request.contains("byok_probe")).count(), 2 ); + assert!( + requests + .iter() + .filter(|request| !request.contains("\"logprobs\":true")) + .all(|request| !request.contains("\"temperature\"")) + ); } } diff --git a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/invite_abuse_actions.rs b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/invite_abuse_actions.rs index 175ca87369..0069d36d3c 100644 --- a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/invite_abuse_actions.rs +++ b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/invite_abuse_actions.rs @@ -5,7 +5,7 @@ use super::{ BackendRuntime, RuntimeError, RuntimeInviteAbuseClaimedAction, RuntimeResult, napi_error, workspace_subject_key, }; -async fn invite_abuse_user_quarantined_or_banned(pool: &PgPool, user_id: &str) -> RuntimeResult { +pub(super) async fn invite_abuse_user_quarantined_or_banned(pool: &PgPool, user_id: &str) -> RuntimeResult { let row: Option = sqlx::query_scalar( r#" SELECT 1 @@ -22,7 +22,7 @@ async fn invite_abuse_user_quarantined_or_banned(pool: &PgPool, user_id: &str) - Ok(row.is_some()) } -async fn invite_abuse_workspace_quarantined(pool: &PgPool, workspace_id: &str) -> RuntimeResult { +pub(super) async fn invite_abuse_workspace_quarantined(pool: &PgPool, workspace_id: &str) -> RuntimeResult { let row: Option = sqlx::query_scalar( r#" SELECT 1 diff --git a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs index 7eac481007..5113354df9 100644 --- a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs +++ b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs @@ -6,6 +6,7 @@ mod reservation; mod workspace_invite; mod workspace_invite_policy; +use invite_abuse_actions::{invite_abuse_user_quarantined_or_banned, invite_abuse_workspace_quarantined}; use mail_delivery::{build_mail_scopes, decision_from_violation as mail_decision_from_violation, mail_class}; use napi::Result; use reservation::{ @@ -15,8 +16,8 @@ use reservation::{ use sha2::{Digest, Sha256}; use workspace_invite_policy::{ ActorFacts, InviteAbuseDecision, InviteActivityFacts, QuotaFacts, WorkspaceFacts, build_invite_scopes, - evaluate_projection, high_confidence_invite_abuse, invite_commit_usage_for_scope, source_cohort_subject_key, - subject_hash, sum_domains, + evaluate_projection, high_confidence_invite_abuse, invite_commit_usage_for_scope, new_account_action_retry_after, + source_cohort_subject_key, subject_hash, sum_domains, }; #[cfg(test)] @@ -26,7 +27,8 @@ pub(super) use super::{ types::{ RuntimeInviteAbuseActionRequired, RuntimeInviteAbuseClaimedAction, RuntimeMailDeliveryQuotaDecision, RuntimeMailDeliveryQuotaInput, RuntimeQuotaSourceInput, RuntimeQuotaTargetDomainInput, - RuntimeWorkspaceInviteQuotaDecision, RuntimeWorkspaceInviteQuotaInput, RuntimeWorkspaceInviteQuotaUsage, + RuntimeWorkspaceActionDecision, RuntimeWorkspaceInviteQuotaDecision, RuntimeWorkspaceInviteQuotaInput, + RuntimeWorkspaceInviteQuotaUsage, }, }; diff --git a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs index 560006d032..d001c0e6dd 100644 --- a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs +++ b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs @@ -5,19 +5,32 @@ use sqlx::{PgPool, Row}; use super::{ ActorFacts, BackendRuntime, InviteAbuseDecision, InviteActivityFacts, InviteQuotaConfig, QuotaFacts, QuotaViolation, - RuntimeError, RuntimeInviteAbuseActionRequired, RuntimeResult, RuntimeWorkspaceInviteQuotaDecision, - RuntimeWorkspaceInviteQuotaInput, RuntimeWorkspaceInviteQuotaUsage, WorkspaceFacts, build_invite_scopes, - commit_reservation, evaluate_projection, high_confidence_invite_abuse, invite_commit_usage_for_scope, napi_error, - normalize_domain, release_reservation, reserve_scopes, short_hash, source_cohort_subject_key, source_prefix, - subject_hash, sum_domains, workspace_subject_key, + RuntimeError, RuntimeInviteAbuseActionRequired, RuntimeResult, RuntimeWorkspaceActionDecision, + RuntimeWorkspaceInviteQuotaDecision, RuntimeWorkspaceInviteQuotaInput, RuntimeWorkspaceInviteQuotaUsage, + WorkspaceFacts, build_invite_scopes, commit_reservation, evaluate_projection, high_confidence_invite_abuse, + invite_abuse_user_quarantined_or_banned, invite_abuse_workspace_quarantined, invite_commit_usage_for_scope, + napi_error, new_account_action_retry_after, normalize_domain, release_reservation, reserve_scopes, short_hash, + source_cohort_subject_key, source_prefix, subject_hash, sum_domains, workspace_subject_key, }; async fn load_actor(pool: &PgPool, user_id: &str) -> RuntimeResult { let row = sqlx::query( r#" - SELECT email, created_at, registered, email_verified IS NOT NULL AS email_verified, disabled + SELECT + users.email, + users.created_at, + users.registered, + users.email_verified IS NOT NULL AS email_verified, + users.disabled, + CASE + WHEN quota.known + AND NOT quota.stale + AND (quota.stale_after IS NULL OR quota.stale_after > clock_timestamp()) + THEN quota.plan + END AS quota_plan FROM users - WHERE id = $1 + LEFT JOIN effective_user_quota_states quota ON quota.user_id = users.id + WHERE users.id = $1 "#, ) .bind(user_id) @@ -32,6 +45,7 @@ async fn load_actor(pool: &PgPool, user_id: &str) -> RuntimeResult { registered: row.get("registered"), email_verified: row.get("email_verified"), disabled: row.get("disabled"), + quota_plan: row.get("quota_plan"), }) } @@ -285,6 +299,49 @@ fn decision_from_violation(violation: QuotaViolation, reason: &str) -> RuntimeWo #[napi_derive::napi] impl BackendRuntime { + #[napi] + pub async fn evaluate_workspace_action_v1( + &self, + actor_user_id: String, + workspace_id: String, + ) -> Result { + let runtime_config = self.config()?; + let pool = self.pool().await?; + if invite_abuse_user_quarantined_or_banned(&pool, &actor_user_id).await? { + return Ok(RuntimeWorkspaceActionDecision { + allowed: false, + retry_after_seconds: None, + reason: Some("abuse_subject".to_string()), + }); + } + if invite_abuse_workspace_quarantined(&pool, &workspace_id).await? { + return Ok(RuntimeWorkspaceActionDecision { + allowed: false, + retry_after_seconds: None, + reason: Some("abuse_workspace".to_string()), + }); + } + let now: DateTime = sqlx::query_scalar("SELECT clock_timestamp()") + .fetch_one(&pool) + .await + .map_err(|err| RuntimeError::database("failed to read database clock", err))?; + let actor = load_actor(&pool, &actor_user_id).await?; + let quota = load_quota(&pool, &workspace_id).await?; + let current_quota = quota.as_ref().filter(|quota| evaluate_projection(quota, now).is_none()); + let retry_after_seconds = new_account_action_retry_after( + runtime_config.deployment, + &runtime_config.invite_quota, + &actor, + current_quota, + now, + ); + Ok(RuntimeWorkspaceActionDecision { + allowed: retry_after_seconds.is_none(), + retry_after_seconds, + reason: retry_after_seconds.map(|_| "new_account_action_delay".to_string()), + }) + } + #[napi] pub async fn assert_workspace_invite_quota_v1( &self, @@ -388,6 +445,22 @@ impl BackendRuntime { action_required: None, }); } + if let Some(retry_after_seconds) = + new_account_action_retry_after(runtime_config.deployment, config, &actor, Some("a), now) + { + return Ok(RuntimeWorkspaceInviteQuotaDecision { + allowed: false, + reservation_id: None, + retry_after_seconds: Some(retry_after_seconds), + reason: Some("new_account_action_delay".to_string()), + scope_key: None, + window_seconds: None, + limit: None, + current: None, + requested: Some(input.target_count), + action_required: None, + }); + } if let Some(abuse_decision) = high_confidence_invite_abuse(&input, &actor, config) { let reason = abuse_decision.reason; let scope_key = match abuse_decision.subject_kind { diff --git a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite_policy.rs b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite_policy.rs index 56b3b539e3..e497050255 100644 --- a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite_policy.rs +++ b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite_policy.rs @@ -10,6 +10,7 @@ use super::{ InviteQuotaConfig, RuntimeQuotaTargetDomainInput, RuntimeWorkspaceInviteQuotaInput, ScopeLimit, bucket_seconds, high_risk_domain, napi_error, normalize_domain, scope, short_hash, source_prefix, workspace_subject_key, }; +use crate::llm::Deployment; #[derive(Clone, Debug)] pub(super) struct ActorFacts { @@ -18,6 +19,7 @@ pub(super) struct ActorFacts { pub(super) registered: bool, pub(super) email_verified: bool, pub(super) disabled: bool, + pub(super) quota_plan: Option, } #[derive(Clone, Debug)] @@ -108,9 +110,6 @@ fn base_invite_limits( let mut per_day = 15; let mut per_week = 30; - if account_age < Duration::hours(24) { - return (0, 0, 0, 0); - } if !actor.email_verified { single = 1; per_hour = 1; @@ -185,6 +184,31 @@ pub(super) fn evaluate_projection(quota: &QuotaFacts, now: DateTime) -> Opt None } +pub(super) fn new_account_action_retry_after( + deployment: Deployment, + config: &InviteQuotaConfig, + actor: &ActorFacts, + workspace_quota: Option<&QuotaFacts>, + now: DateTime, +) -> Option { + if deployment == Deployment::SelfHosted || config.new_account_action_delay_seconds <= 0 { + return None; + } + if actor + .quota_plan + .as_deref() + .is_some_and(|plan| matches!(plan, "pro" | "lifetime_pro" | "ai")) + || workspace_quota.map(|quota| quota.plan.as_str()).is_some_and(|plan| { + matches!(plan, "pro" | "lifetime_pro" | "ai") || plan.contains("team") && !plan.contains("trial") + }) + { + return None; + } + + let remaining = config.new_account_action_delay_seconds - (now - actor.created_at).num_seconds(); + (remaining > 0).then(|| remaining.min(i64::from(i32::MAX)) as i32) +} + pub(super) fn build_invite_scopes( input: &RuntimeWorkspaceInviteQuotaInput, actor: &ActorFacts, @@ -442,6 +466,7 @@ mod tests { registered: true, email_verified: true, disabled: false, + quota_plan: None, } } @@ -467,7 +492,7 @@ mod tests { } #[test] - fn seat_based_weekly_limit_binds_paid_team_and_high_risk_domain() { + fn invite_scopes_apply_plan_ceiling_domain_risk_and_graduated_limits() { let now = Utc.with_ymd_and_hms(2026, 7, 6, 0, 0, 0).single().unwrap(); let input = RuntimeWorkspaceInviteQuotaInput { actor_user_id: "u1".to_string(), @@ -500,6 +525,67 @@ mod tests { .find(|scope| scope.scope_key == "invite:quota_subject_domain:workspace:w1:qq.com") .unwrap(); assert_eq!(high_risk.limit, 5); + + let fresh_actor_scopes = build_invite_scopes( + &input, + &user(now - Duration::hours(1)), + &workspace(now - Duration::hours(1)), + "a("paid_team", 10), + &InviteActivityFacts::default(), + &invite_config(), + now, + ) + .unwrap(); + assert_eq!(fresh_actor_scopes[0].limit, 3); + + let mut fresh_actor = user(now - Duration::hours(1)); + assert_eq!( + new_account_action_retry_after( + Deployment::Cloud, + &invite_config(), + &fresh_actor, + Some("a("free", 3)), + now, + ), + Some(23 * 60 * 60) + ); + assert_eq!( + new_account_action_retry_after( + Deployment::SelfHosted, + &invite_config(), + &fresh_actor, + Some("a("free", 3)), + now, + ), + None + ); + assert_eq!( + new_account_action_retry_after( + Deployment::Cloud, + &invite_config(), + &fresh_actor, + Some("a("paid_team", 10)), + now, + ), + None + ); + let mut no_delay = invite_config(); + no_delay.new_account_action_delay_seconds = 0; + assert_eq!( + new_account_action_retry_after(Deployment::Cloud, &no_delay, &fresh_actor, Some("a("free", 3)), now,), + None + ); + fresh_actor.quota_plan = Some("pro".to_string()); + assert_eq!( + new_account_action_retry_after( + Deployment::Cloud, + &invite_config(), + &fresh_actor, + Some("a("free", 3)), + now, + ), + None + ); } #[test] diff --git a/packages/backend/native/src/runtime/config.rs b/packages/backend/native/src/runtime/config.rs index 59d42a89fb..a89fc78c41 100644 --- a/packages/backend/native/src/runtime/config.rs +++ b/packages/backend/native/src/runtime/config.rs @@ -319,6 +319,7 @@ impl TryFrom for CopilotManagedProfileConfig { #[derive(Clone, Debug)] pub(crate) struct InviteQuotaConfig { + pub(crate) new_account_action_delay_seconds: i64, pub(crate) high_risk_target_domains: Vec, pub(crate) subject_hash_salt: String, pub(crate) mail_class_mapping: BTreeMap, @@ -327,6 +328,7 @@ pub(crate) struct InviteQuotaConfig { impl Default for InviteQuotaConfig { fn default() -> Self { Self { + new_account_action_delay_seconds: 24 * 60 * 60, high_risk_target_domains: [ "qq.com", "proton.me", @@ -479,12 +481,19 @@ fn deployment_from_env() -> Deployment { #[derive(Default, Deserialize)] struct AppConfigFile { + auth: Option, db: Option, crypto: Option, copilot: Option, indexer: Option, } +#[derive(Default, Deserialize)] +#[serde(rename_all = "camelCase")] +struct AuthConfigFile { + new_account_action_delay: Option, +} + #[derive(Default, Deserialize)] #[serde(rename_all = "camelCase", default)] struct SearchRuntimeConfigFile { @@ -542,7 +551,11 @@ impl AppConfigFile { } fn invite_quota_config(&self) -> InviteQuotaConfig { - InviteQuotaConfig::default() + let mut config = InviteQuotaConfig::default(); + if let Some(delay) = self.auth.as_ref().and_then(|auth| auth.new_account_action_delay) { + config.new_account_action_delay_seconds = delay.max(0); + } + config } } @@ -985,14 +998,16 @@ mod tests { } #[test] - fn invite_quota_policy_is_internal_not_app_configurable() { + fn invite_abuse_policy_is_internal_while_action_delay_is_configurable() { let app_config = app_config_from_flat_overrides([ + ("auth.newAccountActionDelay", serde_json::json!(123)), ("auth.untrustedPolicyOverride", serde_json::json!("runtime-salt-v2")), ("auth.untrustedDomainList", serde_json::json!(["Example.COM."])), ]) .unwrap(); let config = app_config.invite_quota_config(); + assert_eq!(config.new_account_action_delay_seconds, 123); assert!(!config.high_risk_target_domains.contains(&"example.com".to_string())); assert_ne!(config.subject_hash_salt, "runtime-salt-v2"); assert_eq!( diff --git a/packages/backend/native/src/runtime/types.rs b/packages/backend/native/src/runtime/types.rs index fff6a6c664..0ccc090418 100644 --- a/packages/backend/native/src/runtime/types.rs +++ b/packages/backend/native/src/runtime/types.rs @@ -263,6 +263,13 @@ pub struct RuntimeWorkspaceInviteQuotaUsage { pub target_domains: Vec, } +#[napi_derive::napi(object)] +pub struct RuntimeWorkspaceActionDecision { + pub allowed: bool, + pub retry_after_seconds: Option, + pub reason: Option, +} + #[napi_derive::napi(object)] pub struct RuntimeInviteAbuseActionRequired { pub action: String, diff --git a/packages/backend/server/src/core/auth/config.ts b/packages/backend/server/src/core/auth/config.ts index 8bc07c6d00..de693945c9 100644 --- a/packages/backend/server/src/core/auth/config.ts +++ b/packages/backend/server/src/core/auth/config.ts @@ -18,7 +18,7 @@ export interface AuthConfig { allowSignupForOauth: boolean; requireEmailDomainVerification: boolean; requireEmailVerification: boolean; - newAccountShareActionDelay: number; + newAccountActionDelay: number; trustedCloudflareHeaders: boolean; signInRateLimit: ConfigItem<{ ttl: number; @@ -56,8 +56,8 @@ defineModuleConfig('auth', { desc: 'Whether require email verification before accessing restricted resources(not implemented).', default: true, }, - newAccountShareActionDelay: { - desc: 'Minimum account age in seconds before new accounts can invite members or create share links.', + newAccountActionDelay: { + desc: 'Minimum account age in seconds before new accounts can invite members, create invite links, or publish documents. Set to 0 to disable.', default: 24 * 60 * 60, shape: z.number().int().min(0), }, diff --git a/packages/backend/server/src/core/backend-runtime/provider.ts b/packages/backend/server/src/core/backend-runtime/provider.ts index 42e4431d2f..6ce331e8e2 100644 --- a/packages/backend/server/src/core/backend-runtime/provider.ts +++ b/packages/backend/server/src/core/backend-runtime/provider.ts @@ -121,6 +121,12 @@ export type RuntimeWorkspaceInviteQuotaUsage = { targetDomains: RuntimeQuotaTargetDomainInput[]; }; +export type RuntimeWorkspaceActionDecision = { + allowed: boolean; + retryAfterSeconds?: number; + reason?: string; +}; + export type RuntimeInviteAbuseAction = | 'ban_actor' | 'quarantine_actor' @@ -213,6 +219,10 @@ export type RuntimeMailDeliveryQuotaDecision = { }; type RuntimeQuotaMethods = RuntimeInstance & { + evaluateWorkspaceActionV1( + actorUserId: string, + workspaceId: string + ): Promise; assertWorkspaceInviteQuotaV1( input: RuntimeWorkspaceInviteQuotaInput ): Promise; @@ -328,6 +338,7 @@ export class BackendRuntimeProvider !updates.copilot && !updates.crypto && !updates.db && + !updates.auth && !updates.indexer && !updates.storages ) { @@ -537,6 +548,12 @@ export class BackendRuntimeProvider ); } + async evaluateWorkspaceActionV1(actorUserId: string, workspaceId: string) { + return await this.measured('evaluateWorkspaceActionV1', rt => + this.quotaRuntime(rt).evaluateWorkspaceActionV1(actorUserId, workspaceId) + ); + } + async commitWorkspaceInviteQuotaV1( reservationId: string, usage: RuntimeWorkspaceInviteQuotaUsage diff --git a/packages/backend/server/src/core/workspaces/__tests__/abuse.spec.ts b/packages/backend/server/src/core/workspaces/__tests__/abuse.spec.ts index fd90008374..73e7bcfbd3 100644 --- a/packages/backend/server/src/core/workspaces/__tests__/abuse.spec.ts +++ b/packages/backend/server/src/core/workspaces/__tests__/abuse.spec.ts @@ -15,6 +15,9 @@ import { Mockers } from '../../../__tests__/mocks'; import { Config } from '../../../base'; import { ActionForbidden, TooManyRequest } from '../../../base/error'; import { Models, WorkspaceRole } from '../../../models'; +import { BackendRuntimeProvider } from '../../backend-runtime'; +import { EntitlementService } from '../../entitlement'; +import { QuotaService } from '../../quota'; import { getAbuseRequestSource, InviteAbuseDispositionService, @@ -23,6 +26,7 @@ import { let app: TestingApp; const quota = { + assertWorkspaceActionAllowed: Sinon.stub(), assertWorkspaceInviteQuota: Sinon.stub(), commitWorkspaceInviteQuota: Sinon.stub(), releaseWorkspaceInviteQuota: Sinon.stub(), @@ -41,6 +45,7 @@ test.before(async () => { }); test.beforeEach(() => { + quota.assertWorkspaceActionAllowed.reset(); quota.assertWorkspaceInviteQuota.reset(); quota.commitWorkspaceInviteQuota.reset(); quota.releaseWorkspaceInviteQuota.reset(); @@ -345,8 +350,8 @@ test('workspace quarantine blocks invite link creation', async t => { updated_at = now() `; - const previousDelay = config.auth.newAccountShareActionDelay; - config.auth.newAccountShareActionDelay = 0; + const previousDelay = config.auth.newAccountActionDelay; + config.auth.newAccountActionDelay = 0; try { await app.login(owner); await t.throwsAsync( @@ -359,32 +364,58 @@ test('workspace quarantine blocks invite link creation', async t => { }) ); } finally { - config.auth.newAccountShareActionDelay = previousDelay; + config.auth.newAccountActionDelay = previousDelay; } }); -test('domain workspace name blocks invite link creation', async t => { - const config = app.get(Config); +test('workspace action admission applies exemption before content policy', async t => { + const db = app.get(PrismaClient); + const inviteQuota = new InviteQuotaAssertService( + app.get(Config), + app.get(QuotaService), + app.get(BackendRuntimeProvider), + app.get(InviteAbuseDispositionService) + ); const owner = await app.create(Mockers.User); + await db.user.update({ + where: { id: owner.id }, + data: { createdAt: new Date() }, + }); const workspace = await app.create(Mockers.Workspace, { owner, name: 'Join example.com', }); - const previousDelay = config.auth.newAccountShareActionDelay; - config.auth.newAccountShareActionDelay = 0; - try { - await app.login(owner); - await t.throwsAsync( - app.gql({ - query: createInviteLinkMutation, - variables: { - workspaceId: workspace.id, - expireTime: WorkspaceInviteLinkExpireTime.OneDay, - }, - }) - ); - } finally { - config.auth.newAccountShareActionDelay = previousDelay; - } + await t.throwsAsync( + inviteQuota.assertWorkspaceActionAllowed({ + actorUserId: owner.id, + workspaceId: workspace.id, + action: 'inviteMember', + }), + { instanceOf: ActionForbidden } + ); + + await app.get(EntitlementService).upsertAdminGrant({ + targetType: 'user', + targetId: owner.id, + plan: 'pro', + }); + await t.notThrowsAsync( + inviteQuota.assertWorkspaceActionAllowed({ + actorUserId: owner.id, + workspaceId: workspace.id, + action: 'inviteMember', + }) + ); + + await app.login(owner); + await t.throwsAsync( + app.gql({ + query: createInviteLinkMutation, + variables: { + workspaceId: workspace.id, + expireTime: WorkspaceInviteLinkExpireTime.OneDay, + }, + }) + ); }); diff --git a/packages/backend/server/src/core/workspaces/abuse.ts b/packages/backend/server/src/core/workspaces/abuse.ts index bcdde29b03..49cd8dd97f 100644 --- a/packages/backend/server/src/core/workspaces/abuse.ts +++ b/packages/backend/server/src/core/workspaces/abuse.ts @@ -39,14 +39,6 @@ declare global { } } -export function canUserExecuteLimitedActions( - user: { createdAt: Date }, - minimumAccountAgeMs: number -) { - if (minimumAccountAgeMs <= 0) return true; - return Date.now() - user.createdAt.getTime() >= minimumAccountAgeMs; -} - function parseAsn(value: string | undefined) { if (!value) { return; @@ -237,6 +229,28 @@ export class InviteQuotaAssertService { private readonly disposition: InviteAbuseDispositionService ) {} + async assertWorkspaceActionAllowed(input: { + actorUserId: string; + workspaceId: string; + action: 'inviteMember' | 'createInviteLink' | 'publishDoc'; + docId?: string; + }) { + const decision = await this.runtime.evaluateWorkspaceActionV1( + input.actorUserId, + input.workspaceId + ); + if (decision.allowed) return; + + this.logger.warn('Workspace action rejected', { + ...input, + reason: decision.reason, + retryAfter: decision.retryAfterSeconds, + }); + throw new ActionForbidden( + 'This feature is temporarily unavailable for you.' + ); + } + async assertWorkspaceInviteQuota(input: { actorUserId: string; workspaceId: string; @@ -384,7 +398,11 @@ export class InviteQuotaAssertService { private mapDecision( decision: RuntimeWorkspaceInviteQuotaDecision ): UserFriendlyError { - if (decision.reason === 'abuse_subject' || decision.actionRequired) { + if ( + decision.reason === 'abuse_subject' || + decision.reason === 'new_account_action_delay' || + decision.actionRequired + ) { return new ActionForbidden('This feature is temporarily unavailable.'); } return new TooManyRequest(); diff --git a/packages/backend/server/src/core/workspaces/resolvers/doc.ts b/packages/backend/server/src/core/workspaces/resolvers/doc.ts index 6668c05422..7123743def 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/doc.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/doc.ts @@ -15,9 +15,7 @@ import { Prisma, PrismaClient } from '@prisma/client'; import { SafeIntResolver } from 'graphql-scalars'; import { - ActionForbidden, Cache, - Config, DocActionDenied, DocDefaultRoleCanNotBeOwner, DocNotFound, @@ -46,7 +44,7 @@ import { PermissionAccess, } from '../../permission'; import { PublicUserType, WorkspaceUserType } from '../../user'; -import { canUserExecuteLimitedActions } from '../abuse'; +import { InviteQuotaAssertService } from '../abuse'; import { DocGrantsService } from '../doc-grants'; import { WorkspaceType } from '../types'; import { TimeBucket, TimeWindow } from './analytics-types'; @@ -302,51 +300,10 @@ export class WorkspaceDocResolver { private readonly models: Models, private readonly cache: Cache, private readonly event: EventBus, - private readonly config: Config, - private readonly runtime: BackendRuntimeProvider + private readonly runtime: BackendRuntimeProvider, + private readonly inviteQuota: InviteQuotaAssertService ) {} - private async assertCanShare( - userId: string, - context: { workspaceId: string; docId: string; action: 'publishDoc' } - ) { - if (await this.runtime.isInviteAbuseUserQuarantinedOrBanned(userId)) { - this.logger.warn('Share action blocked for quarantined actor', { - userId, - ...context, - }); - throw new ActionForbidden( - 'This feature is temporarily unavailable for you.' - ); - } - if ( - await this.runtime.isInviteAbuseWorkspaceQuarantined(context.workspaceId) - ) { - this.logger.warn('Share action blocked for quarantined workspace', { - userId, - ...context, - }); - throw new ActionForbidden( - 'This feature is temporarily unavailable for you.' - ); - } - const user = await this.models.user.get(userId); - const newAccountAgeMs = this.config.auth.newAccountShareActionDelay * 1000; - if (!user || !canUserExecuteLimitedActions(user, newAccountAgeMs)) { - this.logger.warn('Share action blocked for new account', { - userId, - email: user?.email, - createdAt: user?.createdAt, - accountAgeMs: user ? Date.now() - user.createdAt.getTime() : null, - minimumAccountAgeMs: newAccountAgeMs, - ...context, - }); - throw new ActionForbidden( - 'This feature is temporarily unavailable for you.' - ); - } - } - @ResolveField(() => WorkspaceDocMeta, { description: 'Cloud page metadata of workspace', complexity: 2, @@ -475,7 +432,8 @@ export class WorkspaceDocResolver { } await this.ac.user(user.id).doc(workspaceId, docId).assert('Doc.Publish'); - await this.assertCanShare(user.id, { + await this.inviteQuota.assertWorkspaceActionAllowed({ + actorUserId: user.id, workspaceId, docId, action: 'publishDoc', diff --git a/packages/backend/server/src/core/workspaces/resolvers/member.ts b/packages/backend/server/src/core/workspaces/resolvers/member.ts index 3ab1afb64d..fa6a61e8d2 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/member.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/member.ts @@ -1,4 +1,3 @@ -import { Logger } from '@nestjs/common'; import { Args, Context, @@ -39,7 +38,6 @@ import { import type { GraphqlContext } from '../../../base/graphql'; import { Models, type WorkspaceUserCompat } from '../../../models'; import { CurrentUser, Public } from '../../auth'; -import { BackendRuntimeProvider } from '../../backend-runtime'; import { containsUrlOrDomain } from '../../content-policy'; import { PermissionAccess, @@ -49,11 +47,7 @@ import { import { QuotaService } from '../../quota'; import { UserType } from '../../user'; import { validators } from '../../utils/validators'; -import { - canUserExecuteLimitedActions, - getAbuseRequestSource, - InviteQuotaAssertService, -} from '../abuse'; +import { getAbuseRequestSource, InviteQuotaAssertService } from '../abuse'; import { WorkspaceService } from '../service'; import { InvitationType, @@ -92,8 +86,6 @@ function aggregateTargetDomains(candidates: InviteCandidate[]) { */ @Resolver(() => WorkspaceType) export class WorkspaceMemberResolver { - private readonly logger = new Logger(WorkspaceMemberResolver.name); - constructor( private readonly cache: Cache, private readonly event: EventBus, @@ -105,55 +97,9 @@ export class WorkspaceMemberResolver { private readonly workspaceService: WorkspaceService, private readonly quota: QuotaService, private readonly config: Config, - private readonly inviteQuota: InviteQuotaAssertService, - private readonly runtime: BackendRuntimeProvider + private readonly inviteQuota: InviteQuotaAssertService ) {} - private async assertCanInviteOrShare( - userId: string, - context: { - workspaceId: string; - action: 'createInviteLink'; - } - ) { - if (await this.runtime.isInviteAbuseUserQuarantinedOrBanned(userId)) { - this.logger.warn('Share action blocked for quarantined actor', { - userId, - ...context, - }); - throw new ActionForbidden( - 'This feature is temporarily unavailable for you.' - ); - } - if ( - await this.runtime.isInviteAbuseWorkspaceQuarantined(context.workspaceId) - ) { - this.logger.warn('Share action blocked for quarantined workspace', { - userId, - ...context, - }); - throw new ActionForbidden( - 'This feature is temporarily unavailable for you.' - ); - } - // Member invites are owned by native quota; this guard stays for invite links until share/link actions migrate. - const user = await this.models.user.get(userId); - const newAccountAgeMs = this.config.auth.newAccountShareActionDelay * 1000; - if (!user || !canUserExecuteLimitedActions(user, newAccountAgeMs)) { - this.logger.warn('Share action blocked for new account', { - userId, - email: user?.email, - createdAt: user?.createdAt, - accountAgeMs: user ? Date.now() - user.createdAt.getTime() : null, - minimumAccountAgeMs: newAccountAgeMs, - ...context, - }); - throw new ActionForbidden( - 'This feature is temporarily unavailable for you.' - ); - } - } - private async assertWorkspaceNameCanInvite(workspaceId: string) { const workspace = await this.workspaceService.getWorkspaceInfo(workspaceId); if (containsUrlOrDomain(workspace.name)) { @@ -287,6 +233,12 @@ export class WorkspaceMemberResolver { return results; } + await this.inviteQuota.assertWorkspaceActionAllowed({ + actorUserId: me.id, + workspaceId, + action: 'inviteMember', + }); + // lock to prevent concurrent invite const lockFlag = `invite:${workspaceId}`; await using lock = await this.mutex.acquire(lockFlag); @@ -452,7 +404,8 @@ export class WorkspaceMemberResolver { .user(user.id) .workspace(workspaceId) .assert('Workspace.Users.Manage'); - await this.assertCanInviteOrShare(user.id, { + await this.inviteQuota.assertWorkspaceActionAllowed({ + actorUserId: user.id, workspaceId, action: 'createInviteLink', }); diff --git a/packages/frontend/admin/src/config.json b/packages/frontend/admin/src/config.json index f1f7c6cc96..a32b209aa7 100644 --- a/packages/frontend/admin/src/config.json +++ b/packages/frontend/admin/src/config.json @@ -87,9 +87,9 @@ "type": "Boolean", "desc": "Whether require email verification before accessing restricted resources(not implemented)." }, - "newAccountShareActionDelay": { + "newAccountActionDelay": { "type": "Number", - "desc": "Minimum account age in seconds before new accounts can invite members or create share links." + "desc": "Minimum account age in seconds before new accounts can invite members, create invite links, or publish documents. Set to 0 to disable." }, "trustedCloudflareHeaders": { "type": "Boolean", diff --git a/packages/frontend/admin/src/modules/settings/config.ts b/packages/frontend/admin/src/modules/settings/config.ts index b1197c15e4..d75037dded 100644 --- a/packages/frontend/admin/src/modules/settings/config.ts +++ b/packages/frontend/admin/src/modules/settings/config.ts @@ -58,9 +58,9 @@ export const KNOWN_CONFIG_GROUPS = [ 'allowSignup', 'allowSignupForOauth', { - key: 'newAccountShareActionDelay', + key: 'newAccountActionDelay', type: 'Number', - desc: 'Minimum account age in seconds before new accounts can invite members or create share links.', + desc: 'Minimum account age in seconds before accounts can invite members, create invite links, or publish documents. Set to 0 to disable.', }, // nested json object { diff --git a/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/add-key-modal.tsx b/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/add-key-modal.tsx index 7fb1269a36..64ed1ba12f 100644 --- a/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/add-key-modal.tsx +++ b/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/add-key-modal.tsx @@ -27,6 +27,7 @@ import { type ModelDeclaration, modelUseCases, probeChecks, + retainVerifiedCapabilities, } from './model-utils'; import type { ByokDefinition, ByokKey, ByokSettings, GqlFn } from './types'; import { ByokStorage } from './types'; @@ -138,7 +139,7 @@ export const AddKeyModal = ({ const invalidateTest = () => setTestStatus(null); const runProbe = useCallback(async () => { - if (!gql) return false; + if (!gql) return { passed: false, definition }; const canReuseServerCredential = editingKey?.storage === ByokStorage.server && !apiKey; const checks = probeChecks(models, includeImageProbe); @@ -159,21 +160,19 @@ export const AddKeyModal = ({ }, }); const probe = result.probeWorkspaceByokDraft; - const verifiedChecks = new Set( - probe.models.flatMap(model => - model.checks - .filter(check => check.status.kind === 'verified') - .map(check => `${model.modelId}\0${check.operation}`) - ) + const nextModels = retainVerifiedCapabilities(models, probe.models); + const nextDefinition = { ...definition, models: nextModels }; + const hasVerifiedCheck = probe.models.some(model => + model.checks.some(check => check.status.kind === 'verified') ); const passed = checks.length > 0 && probe.connection.kind === 'verified' && - checks.every(check => - verifiedChecks.has(`${check.modelId}\0${check.operation}`) - ); + hasVerifiedCheck && + nextModels.some(model => model.enabled && model.capabilities.length > 0); + if (passed) setModels(nextModels); setTestStatus(passed ? 'passed' : 'failed'); - return passed; + return { passed, definition: nextDefinition }; }, [ apiKey, definition, @@ -185,112 +184,118 @@ export const AddKeyModal = ({ workspaceId, ]); - const persist = useCallback(async () => { - if (!gql) return; - if (storage === ByokStorage.local) { - const saved = await upsertLocalKey(workspaceId, { - id: - editingKey?.storage === ByokStorage.local - ? editingKey.id - : crypto.randomUUID(), - provider, - name, - description, - credential: apiKey, - definition, - sortOrder: - editingKey?.storage === ByokStorage.local - ? editingKey.sortOrder - : localKeys.length, - enabled: profileEnabled, - }); - if (!saved) { - notify.error({ - title: byokT(t, 'notify.local-save-failed.title'), - message: byokT(t, 'notify.local-save-failed.message'), + const persist = useCallback( + async (persistedDefinition = definition) => { + if (!gql) return; + if (storage === ByokStorage.local) { + const saved = await upsertLocalKey(workspaceId, { + id: + editingKey?.storage === ByokStorage.local + ? editingKey.id + : crypto.randomUUID(), + provider, + name, + description, + credential: apiKey, + definition: persistedDefinition, + sortOrder: + editingKey?.storage === ByokStorage.local + ? editingKey.sortOrder + : localKeys.length, + enabled: profileEnabled, }); - return; - } - setLocalKeys(await readLocalKeys(workspaceId)); - } else if (editingKey?.storage === ByokStorage.server) { - if (editingKey.revision === undefined) { - notify.error({ - title: byokT(t, 'notify.reload-required.title'), - message: byokT(t, 'notify.reload-required.message'), + if (!saved) { + notify.error({ + title: byokT(t, 'notify.local-save-failed.title'), + message: byokT(t, 'notify.local-save-failed.message'), + }); + return; + } + setLocalKeys(await readLocalKeys(workspaceId)); + } else if (editingKey?.storage === ByokStorage.server) { + if (editingKey.revision === undefined) { + notify.error({ + title: byokT(t, 'notify.reload-required.title'), + message: byokT(t, 'notify.reload-required.message'), + }); + return; + } + await gql({ + query: replaceWorkspaceByokProfileMutation, + variables: { + input: { + workspaceId, + profileId: editingKey.id, + expectedRevision: editingKey.revision, + name, + description: description || null, + credential: apiKey || null, + definition: persistedDefinition, + enabled: profileEnabled, + }, + }, }); - return; + await onSaved(); + } else { + await gql({ + query: createWorkspaceByokProfileMutation, + variables: { + input: { + workspaceId, + provider, + name, + description: description || null, + credential: apiKey, + definition: persistedDefinition, + enabled: profileEnabled, + }, + }, + }); + await onSaved(); } - await gql({ - query: replaceWorkspaceByokProfileMutation, - variables: { - input: { - workspaceId, - profileId: editingKey.id, - expectedRevision: editingKey.revision, - name, - description: description || null, - credential: apiKey || null, - definition, - enabled: profileEnabled, - }, - }, - }); - await onSaved(); - } else { - await gql({ - query: createWorkspaceByokProfileMutation, - variables: { - input: { - workspaceId, - provider, - name, - description: description || null, - credential: apiKey, - definition, - enabled: profileEnabled, - }, - }, - }); - await onSaved(); - } - onOpenChange(false); - }, [ - apiKey, - definition, - description, - editingKey, - gql, - localKeys.length, - name, - onOpenChange, - onSaved, - provider, - profileEnabled, - setLocalKeys, - storage, - t, - workspaceId, - ]); + onOpenChange(false); + }, + [ + apiKey, + definition, + description, + editingKey, + gql, + localKeys.length, + name, + onOpenChange, + onSaved, + provider, + profileEnabled, + setLocalKeys, + storage, + t, + workspaceId, + ] + ); const connect = useCallback(async () => { if (busyRef.current) return; busyRef.current = true; setBusy(true); try { - const passed = testStatus === 'passed' || (await runProbe()); - if (!passed) { + const probe = + testStatus === 'passed' + ? { passed: true, definition } + : await runProbe(); + if (!probe.passed) { notify.error({ title: byokT(t, 'notify.test-failed.title'), message: byokT(t, 'notify.operation-failed.message'), }); return; } - await persist(); + await persist(probe.definition); } finally { busyRef.current = false; setBusy(false); } - }, [persist, runProbe, t, testStatus]); + }, [definition, persist, runProbe, t, testStatus]); const testConnection = useCallback(async () => { if (busyRef.current) return; @@ -309,7 +314,10 @@ export const AddKeyModal = ({ !!name.trim() && hasCredential && models.length > 0 && - models.every(model => model.modelId.trim() && model.capabilities.length) && + models.every( + model => + model.modelId.trim() && (!model.enabled || model.capabilities.length) + ) && new Set(models.map(model => model.modelId.trim())).size === models.length && (!customEndpoint || (!!endpoint.trim() && dialect !== null)); diff --git a/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.spec.ts b/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.spec.ts index c5ea485c86..83b2db8e9a 100644 --- a/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.spec.ts +++ b/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.spec.ts @@ -11,6 +11,7 @@ import { capabilitiesForUseCases, type ModelDeclaration, modelUseCases, + retainVerifiedCapabilities, } from './model-utils'; describe('BYOK model capabilities', () => { @@ -60,4 +61,53 @@ describe('BYOK model capabilities', () => { capabilitiesForUseCases(model, ['chat', 'actions', 'vision']) ).toEqual([capability]); }); + + test('drops capabilities that imply failed uses while keeping independent uses', () => { + const embeddingCapability = { + input: [ByokModelInput.text], + output: [ByokModelOutput.embedding], + features: [], + attachmentKinds: [], + attachmentSources: [], + }; + const model: ModelDeclaration = { + modelId: 'multimodal-tools', + enabled: true, + capabilities: [ + { + input: [ByokModelInput.text, ByokModelInput.image], + output: [ByokModelOutput.text], + features: [ByokModelFeature.tool_calling], + attachmentKinds: [ByokAttachmentKind.image], + attachmentSources: [ + ByokAttachmentSource.url, + ByokAttachmentSource.data, + ByokAttachmentSource.bytes, + ByokAttachmentSource.file_handle, + ], + }, + embeddingCapability, + ], + }; + + const retained = retainVerifiedCapabilities( + [model], + [ + { + modelId: model.modelId, + checks: [ + { operation: 'chat', status: { kind: 'failed' } }, + { operation: 'tool_calling', status: { kind: 'verified' } }, + { operation: 'embedding', status: { kind: 'verified' } }, + ], + }, + ] + )[0]; + + expect(retained).toEqual({ + ...model, + capabilities: [embeddingCapability], + }); + expect(modelUseCases(retained)).toEqual(['embedding']); + }); }); diff --git a/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.ts b/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.ts index bf3fd054c9..183eee8cee 100644 --- a/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.ts +++ b/packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.ts @@ -163,6 +163,55 @@ export function probeChecks(models: ModelDeclaration[], includeImage: boolean) { ); } +export function retainVerifiedCapabilities( + models: ModelDeclaration[], + probeModels: Array<{ + modelId: string; + checks: Array<{ operation: string; status: { kind: string } }>; + }> +) { + return models.map(model => { + const probe = probeModels.find(item => item.modelId === model.modelId); + if (!probe) return model; + const failedOperations = new Set( + probe.checks + .filter(check => check.status.kind !== 'verified') + .map(check => check.operation) + ); + const selectedUseCases = modelUseCases(model); + const verifiedUseCases = selectedUseCases.filter(useCase => { + const operation = useCase === 'actions' ? 'tool_calling' : useCase; + if (failedOperations.has(operation)) return false; + + const represented = modelUseCases({ + ...model, + capabilities: [capabilityForUseCase(useCase)], + }); + return represented.every( + representedUseCase => + !failedOperations.has( + representedUseCase === 'actions' + ? 'tool_calling' + : representedUseCase + ) + ); + }); + const rebuiltCapabilities = + selectedUseCases.length > 0 && + verifiedUseCases.length === selectedUseCases.length + ? model.capabilities + : verifiedUseCases.map(capabilityForUseCase); + const capabilities = rebuiltCapabilities.length + ? rebuiltCapabilities + : model.capabilities; + return { + ...model, + enabled: rebuiltCapabilities.length > 0 && model.enabled, + capabilities, + }; + }); +} + export function catalogModels(settings: ByokSettings, provider: ByokProvider) { return ( settings.catalog.providers.find(item => item.provider === provider) diff --git a/scripts/set-version.sh b/scripts/set-version.sh index 38cde249f3..6c2e1dfc8c 100755 --- a/scripts/set-version.sh +++ b/scripts/set-version.sh @@ -147,7 +147,6 @@ echo "iOS MARKETING_VERSION: $ios_new_version (app version: $new_version)" update_app_version_in_helm_charts ".github/helm/affine/Chart.yaml" "$new_version" update_app_version_in_helm_charts ".github/helm/affine/charts/graphql/Chart.yaml" "$new_version" update_app_version_in_helm_charts ".github/helm/affine/charts/front/Chart.yaml" "$new_version" -update_app_version_in_helm_charts ".github/helm/affine/charts/doc/Chart.yaml" "$new_version" update_app_stream_version "packages/frontend/apps/electron/resources/affine.metainfo.xml" "$new_version"