mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-08-04 19:45:34 +08:00
fix(pi): close final gateway ownership regressions
This commit is contained in:
@@ -274,75 +274,13 @@ async fn set_pi_auto_failover_enabled(
|
||||
state: &AppState,
|
||||
enabled: bool,
|
||||
) -> Result<(), String> {
|
||||
let _guard = state
|
||||
.proxy_service
|
||||
.lock_switch_for_app(AppType::Pi.as_str())
|
||||
.await;
|
||||
let previous_config = crate::settings::get_pi_proxy_settings();
|
||||
if enabled && !crate::settings::pi_takeover_enabled() {
|
||||
return Err("Pi gateway takeover must be enabled before failover".to_string());
|
||||
}
|
||||
|
||||
let mut auto_added = None;
|
||||
if enabled
|
||||
&& state
|
||||
.db
|
||||
.get_failover_queue("pi")
|
||||
.map_err(|error| error.to_string())?
|
||||
.is_empty()
|
||||
{
|
||||
let current =
|
||||
crate::services::pi_catalog::PiCatalogCoordinator::current_native_provider(state)
|
||||
.map_err(|error| error.to_string())?
|
||||
.ok_or_else(|| {
|
||||
"Pi failover queue is empty and no current provider is selected".to_string()
|
||||
})?;
|
||||
state
|
||||
.db
|
||||
.add_to_failover_queue("pi", ¤t)
|
||||
.map_err(|error| error.to_string())?;
|
||||
auto_added = Some(current);
|
||||
}
|
||||
|
||||
let mut next = previous_config.clone();
|
||||
next.auto_failover_enabled = enabled;
|
||||
let epoch = state.proxy_service.begin_pi_catalog_mutation().await;
|
||||
if let Err(error) = crate::settings::update_pi_proxy_settings(next) {
|
||||
if let Some(provider_id) = auto_added {
|
||||
let _ = state.db.remove_from_failover_queue("pi", &provider_id);
|
||||
}
|
||||
let _ = state
|
||||
.proxy_service
|
||||
.reconcile_pi_runtime_at_epoch(epoch)
|
||||
.await;
|
||||
return Err(error.to_string());
|
||||
}
|
||||
if let Err(error) = state
|
||||
.proxy_service
|
||||
.reconcile_pi_runtime_at_epoch(epoch)
|
||||
.await
|
||||
{
|
||||
let _ = crate::settings::update_pi_proxy_settings(previous_config);
|
||||
if let Some(provider_id) = auto_added {
|
||||
let _ = state.db.remove_from_failover_queue("pi", &provider_id);
|
||||
}
|
||||
let rollback_epoch = state.proxy_service.begin_pi_catalog_mutation().await;
|
||||
let _ = state
|
||||
.proxy_service
|
||||
.reconcile_pi_runtime_at_epoch(rollback_epoch)
|
||||
.await;
|
||||
return Err(format!(
|
||||
"Pi failover preference changed but runtime publication failed: {error}"
|
||||
));
|
||||
}
|
||||
let selected_provider = set_pi_auto_failover_enabled_inner(state, enabled).await?;
|
||||
|
||||
let _ = app.emit(
|
||||
"provider-switched",
|
||||
serde_json::json!({
|
||||
"appType": "pi",
|
||||
"providerId":
|
||||
crate::services::pi_catalog::PiCatalogCoordinator::current_native_provider(state)
|
||||
.map_err(|error| error.to_string())?,
|
||||
"providerId": selected_provider,
|
||||
"source": "failoverPreferenceChanged"
|
||||
}),
|
||||
);
|
||||
@@ -353,3 +291,271 @@ async fn set_pi_auto_failover_enabled(
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn set_pi_auto_failover_enabled_inner(
|
||||
state: &AppState,
|
||||
enabled: bool,
|
||||
) -> Result<Option<String>, String> {
|
||||
let guard = state
|
||||
.proxy_service
|
||||
.lock_switch_for_app(AppType::Pi.as_str())
|
||||
.await;
|
||||
let previous_config = crate::settings::get_pi_proxy_settings();
|
||||
if enabled && !crate::settings::pi_takeover_enabled() {
|
||||
return Err("Pi gateway takeover must be enabled before failover".to_string());
|
||||
}
|
||||
|
||||
let previous_provider =
|
||||
crate::services::pi_catalog::PiCatalogCoordinator::current_native_provider(state)
|
||||
.map_err(|error| error.to_string())?;
|
||||
let mut auto_added = None;
|
||||
let mut switched_primary = false;
|
||||
let selected_provider = if enabled {
|
||||
let previous_provider = previous_provider.clone().ok_or_else(|| {
|
||||
"Pi has no current provider, so failover cannot select queue P1".to_string()
|
||||
})?;
|
||||
let mut queue = state
|
||||
.db
|
||||
.get_failover_queue("pi")
|
||||
.map_err(|error| error.to_string())?;
|
||||
if queue.is_empty() {
|
||||
state
|
||||
.db
|
||||
.add_to_failover_queue("pi", &previous_provider)
|
||||
.map_err(|error| error.to_string())?;
|
||||
auto_added = Some(previous_provider.clone());
|
||||
queue = state
|
||||
.db
|
||||
.get_failover_queue("pi")
|
||||
.map_err(|error| error.to_string())?;
|
||||
}
|
||||
let primary = queue
|
||||
.first()
|
||||
.map(|item| item.provider_id.clone())
|
||||
.ok_or_else(|| "Pi failover queue is empty".to_string())?;
|
||||
if primary != previous_provider {
|
||||
if let Err(error) = set_pi_default_under_switch_guard(state, &guard, &primary) {
|
||||
if let Some(provider_id) = auto_added.take() {
|
||||
let _ = state.db.remove_from_failover_queue("pi", &provider_id);
|
||||
}
|
||||
return Err(error);
|
||||
}
|
||||
switched_primary = true;
|
||||
}
|
||||
Some(primary)
|
||||
} else {
|
||||
previous_provider.clone()
|
||||
};
|
||||
|
||||
let mut next = previous_config.clone();
|
||||
next.auto_failover_enabled = enabled;
|
||||
let epoch = state.proxy_service.begin_pi_catalog_mutation().await;
|
||||
if let Err(error) = crate::settings::update_pi_proxy_settings(next) {
|
||||
if let Some(provider_id) = auto_added.take() {
|
||||
let _ = state.db.remove_from_failover_queue("pi", &provider_id);
|
||||
}
|
||||
let rollback = rollback_pi_failover_primary(
|
||||
state,
|
||||
&guard,
|
||||
previous_provider.as_deref(),
|
||||
switched_primary,
|
||||
Some(epoch),
|
||||
)
|
||||
.await;
|
||||
return Err(with_pi_failover_rollback(error.to_string(), rollback));
|
||||
}
|
||||
if let Err(error) = state
|
||||
.proxy_service
|
||||
.reconcile_pi_runtime_at_epoch(epoch)
|
||||
.await
|
||||
{
|
||||
let settings_rollback = crate::settings::update_pi_proxy_settings(previous_config)
|
||||
.err()
|
||||
.map(|rollback_error| rollback_error.to_string());
|
||||
if let Some(provider_id) = auto_added.take() {
|
||||
let _ = state.db.remove_from_failover_queue("pi", &provider_id);
|
||||
}
|
||||
let primary_rollback = rollback_pi_failover_primary(
|
||||
state,
|
||||
&guard,
|
||||
previous_provider.as_deref(),
|
||||
switched_primary,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let rollback = settings_rollback.or(primary_rollback);
|
||||
return Err(with_pi_failover_rollback(
|
||||
format!("Pi failover preference changed but runtime publication failed: {error}"),
|
||||
rollback,
|
||||
));
|
||||
}
|
||||
|
||||
Ok(selected_provider)
|
||||
}
|
||||
|
||||
fn set_pi_default_under_switch_guard(
|
||||
state: &AppState,
|
||||
guard: &tokio::sync::OwnedMutexGuard<()>,
|
||||
provider_id: &str,
|
||||
) -> Result<(), String> {
|
||||
let aggregate = state
|
||||
.db
|
||||
.get_provider_aggregate(AppType::Pi.as_str(), provider_id)
|
||||
.map_err(|error| error.to_string())?
|
||||
.ok_or_else(|| format!("Pi provider does not exist: {provider_id}"))?;
|
||||
let config: crate::pi_config::model::PiManagedProviderConfig =
|
||||
serde_json::from_value(aggregate.provider.settings_config)
|
||||
.map_err(|error| format!("managed Pi provider '{provider_id}' is invalid: {error}"))?;
|
||||
let model_id = config
|
||||
.models
|
||||
.first()
|
||||
.map(|model| model.id.clone())
|
||||
.ok_or_else(|| format!("Pi provider '{provider_id}' has no selectable models"))?;
|
||||
crate::services::pi_catalog::PiCatalogCoordinator::apply_under_switch_guard(
|
||||
state,
|
||||
guard,
|
||||
crate::services::pi_catalog::PiCatalogMutation::SetDefault {
|
||||
provider_id: provider_id.to_string(),
|
||||
model_id,
|
||||
},
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
async fn rollback_pi_failover_primary(
|
||||
state: &AppState,
|
||||
guard: &tokio::sync::OwnedMutexGuard<()>,
|
||||
previous_provider: Option<&str>,
|
||||
switched_primary: bool,
|
||||
pending_epoch: Option<u64>,
|
||||
) -> Option<String> {
|
||||
if switched_primary {
|
||||
let previous_provider =
|
||||
previous_provider.expect("switching P1 requires a previous Pi provider");
|
||||
return set_pi_default_under_switch_guard(state, guard, previous_provider)
|
||||
.err()
|
||||
.map(|error| format!("primary rollback failed: {error}"));
|
||||
}
|
||||
|
||||
let epoch = match pending_epoch {
|
||||
Some(epoch) => epoch,
|
||||
None => state.proxy_service.begin_pi_catalog_mutation().await,
|
||||
};
|
||||
state
|
||||
.proxy_service
|
||||
.reconcile_pi_runtime_at_epoch(epoch)
|
||||
.await
|
||||
.err()
|
||||
.map(|error| format!("runtime rollback failed: {error}"))
|
||||
}
|
||||
|
||||
fn with_pi_failover_rollback(error: String, rollback: Option<String>) -> String {
|
||||
match rollback {
|
||||
Some(rollback) => format!("{error}; {rollback}"),
|
||||
None => error,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::database::Database;
|
||||
use crate::provider::ProviderMutationInput;
|
||||
use crate::services::pi_catalog::{PiCatalogCoordinator, PiCatalogMutation};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
|
||||
struct TestHome(Option<std::ffi::OsString>);
|
||||
|
||||
impl TestHome {
|
||||
fn install(path: &std::path::Path) -> Result<Self, crate::error::AppError> {
|
||||
let previous = std::env::var_os("CC_SWITCH_TEST_HOME");
|
||||
std::env::set_var("CC_SWITCH_TEST_HOME", path);
|
||||
crate::settings::reload_settings()?;
|
||||
Ok(Self(previous))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestHome {
|
||||
fn drop(&mut self) {
|
||||
match self.0.take() {
|
||||
Some(value) => std::env::set_var("CC_SWITCH_TEST_HOME", value),
|
||||
None => std::env::remove_var("CC_SWITCH_TEST_HOME"),
|
||||
}
|
||||
let _ = crate::settings::reload_settings();
|
||||
}
|
||||
}
|
||||
|
||||
fn managed_input(id: &str) -> ProviderMutationInput {
|
||||
ProviderMutationInput {
|
||||
id: id.to_string(),
|
||||
name: id.to_string(),
|
||||
settings_config: json!({
|
||||
"name": id,
|
||||
"api": "openai-responses",
|
||||
"baseUrl": format!("https://{id}.example/v1"),
|
||||
"apiKey": "literal-key",
|
||||
"models": [{"id": format!("{id}-model"), "name": id}]
|
||||
}),
|
||||
website_url: None,
|
||||
category: None,
|
||||
created_at: Some(1),
|
||||
sort_index: Some(0),
|
||||
notes: None,
|
||||
meta: None,
|
||||
icon: Some("pi".to_string()),
|
||||
icon_color: None,
|
||||
in_failover_queue: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn enabling_pi_failover_selects_queue_p1() -> Result<(), crate::error::AppError> {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let _home = TestHome::install(temp.path())?;
|
||||
let mut settings = crate::settings::get_settings();
|
||||
settings.pi_config_dir = Some(temp.path().join("pi-agent").to_string_lossy().into_owned());
|
||||
settings.pi_takeover_enabled = false;
|
||||
crate::settings::update_settings(settings)?;
|
||||
|
||||
let state = AppState::new(Arc::new(Database::memory()?));
|
||||
for (provider_id, activate_if_first) in [("provider-a", true), ("provider-b", false)] {
|
||||
PiCatalogCoordinator::apply(
|
||||
&state,
|
||||
PiCatalogMutation::CreateProvider {
|
||||
input: managed_input(provider_id),
|
||||
provider_key: provider_id.to_string(),
|
||||
activate_if_first,
|
||||
},
|
||||
)?;
|
||||
}
|
||||
state.db.add_to_failover_queue("pi", "provider-b")?;
|
||||
let mut proxy_config = state.db.get_global_proxy_config().await?;
|
||||
proxy_config.listen_port = 0;
|
||||
state.db.update_global_proxy_config(proxy_config).await?;
|
||||
state
|
||||
.proxy_service
|
||||
.set_takeover_for_app("pi", true)
|
||||
.await
|
||||
.map_err(crate::error::AppError::Message)?;
|
||||
|
||||
let selected = set_pi_auto_failover_enabled_inner(&state, true)
|
||||
.await
|
||||
.map_err(crate::error::AppError::Message)?;
|
||||
|
||||
assert_eq!(selected.as_deref(), Some("provider-b"));
|
||||
assert_eq!(
|
||||
PiCatalogCoordinator::current_native_provider(&state)?.as_deref(),
|
||||
Some("provider-b")
|
||||
);
|
||||
assert!(crate::settings::get_pi_proxy_settings().auto_failover_enabled);
|
||||
state
|
||||
.proxy_service
|
||||
.set_takeover_for_app("pi", false)
|
||||
.await
|
||||
.map_err(crate::error::AppError::Message)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -165,6 +165,7 @@ pub async fn update_proxy_config_for_app(
|
||||
circuit_timeout_seconds: config.circuit_timeout_seconds,
|
||||
circuit_error_rate_threshold: config.circuit_error_rate_threshold,
|
||||
circuit_min_requests: config.circuit_min_requests,
|
||||
..previous.clone()
|
||||
};
|
||||
let epoch = state.proxy_service.begin_pi_catalog_mutation().await;
|
||||
if let Err(error) = crate::settings::update_pi_proxy_settings(next) {
|
||||
@@ -210,6 +211,9 @@ async fn get_default_cost_multiplier_internal(
|
||||
state: &AppState,
|
||||
app_type: &str,
|
||||
) -> Result<String, AppError> {
|
||||
if app_type == "pi" {
|
||||
return Ok(crate::settings::get_pi_default_cost_multiplier());
|
||||
}
|
||||
let db = &state.db;
|
||||
db.get_default_cost_multiplier(app_type).await
|
||||
}
|
||||
@@ -238,6 +242,9 @@ async fn set_default_cost_multiplier_internal(
|
||||
app_type: &str,
|
||||
value: &str,
|
||||
) -> Result<(), AppError> {
|
||||
if app_type == "pi" {
|
||||
return crate::settings::set_pi_default_cost_multiplier(value);
|
||||
}
|
||||
let db = &state.db;
|
||||
db.set_default_cost_multiplier(app_type, value).await
|
||||
}
|
||||
@@ -267,6 +274,9 @@ async fn get_pricing_model_source_internal(
|
||||
state: &AppState,
|
||||
app_type: &str,
|
||||
) -> Result<String, AppError> {
|
||||
if app_type == "pi" {
|
||||
return Ok(crate::settings::get_pi_pricing_model_source());
|
||||
}
|
||||
let db = &state.db;
|
||||
db.get_pricing_model_source(app_type).await
|
||||
}
|
||||
@@ -295,6 +305,9 @@ async fn set_pricing_model_source_internal(
|
||||
app_type: &str,
|
||||
value: &str,
|
||||
) -> Result<(), AppError> {
|
||||
if app_type == "pi" {
|
||||
return crate::settings::set_pi_pricing_model_source(value);
|
||||
}
|
||||
let db = &state.db;
|
||||
db.set_pricing_model_source(app_type, value).await
|
||||
}
|
||||
@@ -511,3 +524,61 @@ pub async fn get_circuit_breaker_stats(
|
||||
let _ = (state, provider_id, app_type);
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::database::{lock_conn, Database};
|
||||
use std::sync::Arc;
|
||||
|
||||
struct TestHome(Option<std::ffi::OsString>);
|
||||
|
||||
impl TestHome {
|
||||
fn install(path: &std::path::Path) -> Result<Self, AppError> {
|
||||
let previous = std::env::var_os("CC_SWITCH_TEST_HOME");
|
||||
std::env::set_var("CC_SWITCH_TEST_HOME", path);
|
||||
crate::settings::reload_settings()?;
|
||||
Ok(Self(previous))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestHome {
|
||||
fn drop(&mut self) {
|
||||
match self.0.take() {
|
||||
Some(value) => std::env::set_var("CC_SWITCH_TEST_HOME", value),
|
||||
None => std::env::remove_var("CC_SWITCH_TEST_HOME"),
|
||||
}
|
||||
let _ = crate::settings::reload_settings();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn pi_pricing_round_trips_without_out_of_schema_proxy_row() -> Result<(), AppError> {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let _home = TestHome::install(temp.path())?;
|
||||
let state = AppState::new(Arc::new(Database::memory()?));
|
||||
|
||||
set_default_cost_multiplier_test_hook(&state, "pi", "1.25").await?;
|
||||
set_pricing_model_source_test_hook(&state, "pi", "request").await?;
|
||||
|
||||
assert_eq!(
|
||||
get_default_cost_multiplier_test_hook(&state, "pi").await?,
|
||||
"1.25"
|
||||
);
|
||||
assert_eq!(
|
||||
get_pricing_model_source_test_hook(&state, "pi").await?,
|
||||
"request"
|
||||
);
|
||||
let conn = lock_conn!(state.db.conn);
|
||||
let pi_rows: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM proxy_config WHERE app_type = 'pi'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
assert_eq!(pi_rows, 0);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,8 @@
|
||||
use crate::database::SkillDeploymentMethod;
|
||||
use crate::database::{lock_conn, Database, PiProviderProjection, SkillDeployment};
|
||||
use crate::error::AppError;
|
||||
use rusqlite::{params, Connection, OpenFlags, OptionalExtension};
|
||||
use rusqlite::{Connection, OpenFlags, OptionalExtension};
|
||||
use std::fmt::Write as _;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
@@ -47,9 +48,9 @@ impl Database {
|
||||
/// Import a user-portable SQL backup while retaining this device's evidence.
|
||||
pub(crate) fn import_portable_sql(&self, source_path: &Path) -> Result<String, AppError> {
|
||||
let local = self.capture_pi_device_local_state()?;
|
||||
let backup_id = self.import_sql(source_path)?;
|
||||
self.replace_pi_device_local_state(&local)?;
|
||||
Ok(backup_id)
|
||||
let sql =
|
||||
fs::read_to_string(source_path).map_err(|error| AppError::io(source_path, error))?;
|
||||
self.import_sql_string(&append_pi_device_local_state(&sql, &local)?)
|
||||
}
|
||||
|
||||
/// Import a cloud-sync snapshot while retaining this device's evidence.
|
||||
@@ -58,9 +59,7 @@ impl Database {
|
||||
sql: &str,
|
||||
) -> Result<String, AppError> {
|
||||
let local = self.capture_pi_device_local_state()?;
|
||||
let backup_id = self.import_sql_string_for_sync(sql)?;
|
||||
self.replace_pi_device_local_state(&local)?;
|
||||
Ok(backup_id)
|
||||
self.import_sql_string_for_sync(&append_pi_device_local_state(sql, &local)?)
|
||||
}
|
||||
|
||||
/// Fail closed before the legacy whole-database restore can import
|
||||
@@ -181,64 +180,13 @@ impl Database {
|
||||
skill_deployments,
|
||||
})
|
||||
}
|
||||
|
||||
fn replace_pi_device_local_state(&self, local: &PiDeviceLocalState) -> Result<(), AppError> {
|
||||
let mut conn = lock_conn!(self.conn);
|
||||
let transaction = conn
|
||||
.transaction()
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
transaction
|
||||
.execute("DELETE FROM pi_provider_projections", [])
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
transaction
|
||||
.execute("DELETE FROM skill_deployments WHERE app_type = 'pi'", [])
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
|
||||
for projection in &local.projections {
|
||||
transaction
|
||||
.execute(
|
||||
"INSERT INTO pi_provider_projections
|
||||
(provider_id, provider_key, created_at, updated_at)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
params![
|
||||
projection.provider_id,
|
||||
projection.provider_key,
|
||||
projection.created_at,
|
||||
projection.updated_at
|
||||
],
|
||||
)
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
}
|
||||
for deployment in &local.skill_deployments {
|
||||
transaction
|
||||
.execute(
|
||||
"INSERT INTO skill_deployments (
|
||||
app_type, skill_id, destination, destination_key, method,
|
||||
source_identity, deployed_digest, created_at, updated_at
|
||||
) VALUES ('pi', ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
|
||||
params![
|
||||
deployment.skill_id,
|
||||
deployment.destination,
|
||||
deployment.destination_key,
|
||||
deployment.method.as_str(),
|
||||
deployment.source_identity,
|
||||
deployment.deployed_digest,
|
||||
deployment.created_at,
|
||||
deployment.updated_at
|
||||
],
|
||||
)
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.map_err(|error| AppError::Database(error.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
fn pi_device_local_rows_exist(conn: &Connection) -> Result<bool, AppError> {
|
||||
let table_exists = |name: &str| -> Result<bool, AppError> {
|
||||
let schema_object_exists = |name: &str| -> Result<bool, AppError> {
|
||||
conn.query_row(
|
||||
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?1",
|
||||
"SELECT 1 FROM sqlite_master
|
||||
WHERE name = ?1 COLLATE NOCASE",
|
||||
[name],
|
||||
|_| Ok(()),
|
||||
)
|
||||
@@ -247,7 +195,7 @@ fn pi_device_local_rows_exist(conn: &Connection) -> Result<bool, AppError> {
|
||||
.map_err(|error| AppError::Database(error.to_string()))
|
||||
};
|
||||
|
||||
let projections_exist = if table_exists("pi_provider_projections")? {
|
||||
let projections_exist = if schema_object_exists("pi_provider_projections")? {
|
||||
conn.query_row(
|
||||
"SELECT EXISTS(SELECT 1 FROM pi_provider_projections LIMIT 1)",
|
||||
[],
|
||||
@@ -261,7 +209,7 @@ fn pi_device_local_rows_exist(conn: &Connection) -> Result<bool, AppError> {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
if table_exists("skill_deployments")? {
|
||||
if schema_object_exists("skill_deployments")? {
|
||||
return conn
|
||||
.query_row(
|
||||
"SELECT EXISTS(
|
||||
@@ -277,6 +225,165 @@ fn pi_device_local_rows_exist(conn: &Connection) -> Result<bool, AppError> {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
/// Merge receiving-device evidence into the same temporary database that the
|
||||
/// generic importer validates and publishes. No live database replacement can
|
||||
/// occur before these statements have succeeded.
|
||||
fn append_pi_device_local_state(sql: &str, local: &PiDeviceLocalState) -> Result<String, AppError> {
|
||||
ensure_sql_append_boundary(sql)?;
|
||||
let mut merged = String::with_capacity(sql.len() + 1024);
|
||||
merged.push_str(sql);
|
||||
// The leading newline closes a trailing `--` comment; the standalone
|
||||
// semicolon terminates a final statement that omitted its delimiter.
|
||||
merged.push_str(
|
||||
"\n;\n\
|
||||
DROP TABLE IF EXISTS pi_provider_projections;\n\
|
||||
CREATE TABLE pi_provider_projections (\n\
|
||||
provider_id TEXT PRIMARY KEY,\n\
|
||||
provider_key TEXT NOT NULL UNIQUE,\n\
|
||||
created_at INTEGER NOT NULL,\n\
|
||||
updated_at INTEGER NOT NULL\n\
|
||||
);\n\
|
||||
DROP TABLE IF EXISTS skill_deployments;\n\
|
||||
CREATE TABLE skill_deployments (\n\
|
||||
app_type TEXT NOT NULL CHECK (app_type = 'pi'),\n\
|
||||
skill_id TEXT NOT NULL,\n\
|
||||
destination TEXT NOT NULL,\n\
|
||||
destination_key TEXT NOT NULL,\n\
|
||||
method TEXT NOT NULL CHECK (method IN ('symlink', 'copy')),\n\
|
||||
source_identity TEXT NOT NULL,\n\
|
||||
deployed_digest TEXT,\n\
|
||||
created_at INTEGER NOT NULL,\n\
|
||||
updated_at INTEGER NOT NULL,\n\
|
||||
PRIMARY KEY (app_type, skill_id, destination_key),\n\
|
||||
UNIQUE (app_type, destination_key)\n\
|
||||
);\n",
|
||||
);
|
||||
|
||||
for projection in &local.projections {
|
||||
writeln!(
|
||||
merged,
|
||||
"INSERT INTO pi_provider_projections \
|
||||
(provider_id, provider_key, created_at, updated_at) \
|
||||
VALUES ({}, {}, {}, {});",
|
||||
sql_text(&projection.provider_id)?,
|
||||
sql_text(&projection.provider_key)?,
|
||||
projection.created_at,
|
||||
projection.updated_at,
|
||||
)
|
||||
.expect("writing to String cannot fail");
|
||||
}
|
||||
for deployment in &local.skill_deployments {
|
||||
writeln!(
|
||||
merged,
|
||||
"INSERT INTO skill_deployments \
|
||||
(app_type, skill_id, destination, destination_key, method, \
|
||||
source_identity, deployed_digest, created_at, updated_at) \
|
||||
VALUES ('pi', {}, {}, {}, {}, {}, {}, {}, {});",
|
||||
sql_text(&deployment.skill_id)?,
|
||||
sql_text(&deployment.destination)?,
|
||||
sql_text(&deployment.destination_key)?,
|
||||
sql_text(deployment.method.as_str())?,
|
||||
sql_text(&deployment.source_identity)?,
|
||||
sql_optional_text(deployment.deployed_digest.as_deref())?,
|
||||
deployment.created_at,
|
||||
deployment.updated_at,
|
||||
)
|
||||
.expect("writing to String cannot fail");
|
||||
}
|
||||
Ok(merged)
|
||||
}
|
||||
|
||||
fn sql_optional_text(value: Option<&str>) -> Result<String, AppError> {
|
||||
value.map_or_else(|| Ok("NULL".to_string()), sql_text)
|
||||
}
|
||||
|
||||
fn sql_text(value: &str) -> Result<String, AppError> {
|
||||
if value.contains('\0') {
|
||||
return Err(AppError::InvalidInput(
|
||||
"device-local Pi ownership text cannot contain NUL".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(format!("'{}'", value.replace('\'', "''")))
|
||||
}
|
||||
|
||||
/// SQLite accepts an unterminated block comment at EOF. Reject such input so
|
||||
/// it cannot swallow the receiving-device statements appended above. Other
|
||||
/// unterminated quoted forms are rejected here as a clearer pre-publish error.
|
||||
fn ensure_sql_append_boundary(sql: &str) -> Result<(), AppError> {
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
enum State {
|
||||
Normal,
|
||||
SingleQuote,
|
||||
DoubleQuote,
|
||||
Backtick,
|
||||
Bracket,
|
||||
LineComment,
|
||||
BlockComment,
|
||||
}
|
||||
|
||||
let bytes = sql.as_bytes();
|
||||
let mut state = State::Normal;
|
||||
let mut cursor = 0;
|
||||
while cursor < bytes.len() {
|
||||
let current = bytes[cursor];
|
||||
let next = bytes.get(cursor + 1).copied();
|
||||
match state {
|
||||
State::Normal => match (current, next) {
|
||||
(b'\'', _) => state = State::SingleQuote,
|
||||
(b'"', _) => state = State::DoubleQuote,
|
||||
(b'`', _) => state = State::Backtick,
|
||||
(b'[', _) => state = State::Bracket,
|
||||
(b'-', Some(b'-')) => {
|
||||
state = State::LineComment;
|
||||
cursor += 1;
|
||||
}
|
||||
(b'/', Some(b'*')) => {
|
||||
state = State::BlockComment;
|
||||
cursor += 1;
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
State::SingleQuote if current == b'\'' => {
|
||||
if next == Some(b'\'') {
|
||||
cursor += 1;
|
||||
} else {
|
||||
state = State::Normal;
|
||||
}
|
||||
}
|
||||
State::DoubleQuote if current == b'"' => {
|
||||
if next == Some(b'"') {
|
||||
cursor += 1;
|
||||
} else {
|
||||
state = State::Normal;
|
||||
}
|
||||
}
|
||||
State::Backtick if current == b'`' => {
|
||||
if next == Some(b'`') {
|
||||
cursor += 1;
|
||||
} else {
|
||||
state = State::Normal;
|
||||
}
|
||||
}
|
||||
State::Bracket if current == b']' => state = State::Normal,
|
||||
State::LineComment if current == b'\n' => state = State::Normal,
|
||||
State::BlockComment if current == b'*' && next == Some(b'/') => {
|
||||
state = State::Normal;
|
||||
cursor += 1;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
cursor += 1;
|
||||
}
|
||||
|
||||
if matches!(state, State::Normal | State::LineComment) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(AppError::InvalidInput(
|
||||
"portable SQL ends inside an unterminated quoted value or comment".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn binary_restore_ownership_error(source_zh: &str, source_en: &str) -> AppError {
|
||||
AppError::localized(
|
||||
"pi.binary_restore_device_ownership_unsupported",
|
||||
@@ -369,6 +476,20 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn add_restore_poison_trigger(sql: &str) -> String {
|
||||
let insertion = sql
|
||||
.rfind("COMMIT;")
|
||||
.expect("CC Switch dump must contain a final COMMIT");
|
||||
let mut poisoned = sql.to_string();
|
||||
poisoned.insert_str(
|
||||
insertion,
|
||||
"CREATE TRIGGER poison_pi_restore \
|
||||
BEFORE INSERT ON pi_provider_projections \
|
||||
BEGIN SELECT RAISE(ABORT, 'poisoned local restore'); END;\n",
|
||||
);
|
||||
poisoned
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn portable_sql_scrubber_handles_multiline_quoted_values() -> Result<(), AppError> {
|
||||
let sql = concat!(
|
||||
@@ -414,7 +535,7 @@ mod tests {
|
||||
remote.claim_pi_projection_key("remote-provider", "shared-key")?;
|
||||
remote.save_pi_skill_deployment(&deployment("remote-skill", "remote-destination"))?;
|
||||
// Model an older remote snapshot created before portable row scrubbing.
|
||||
let remote_sql = remote.export_sql_string()?;
|
||||
let remote_sql = add_restore_poison_trigger(&remote.export_sql_string()?);
|
||||
|
||||
let local = Database::memory()?;
|
||||
local.claim_pi_projection_key("local-provider", "shared-key")?;
|
||||
@@ -443,7 +564,11 @@ mod tests {
|
||||
remote.claim_pi_projection_key("remote-provider", "remote-key")?;
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let path = temp.path().join("portable.sql");
|
||||
fs::write(&path, remote.export_sql_string()?).expect("write SQL backup");
|
||||
fs::write(
|
||||
&path,
|
||||
add_restore_poison_trigger(&remote.export_sql_string()?),
|
||||
)
|
||||
.expect("write SQL backup");
|
||||
|
||||
let local = Database::memory()?;
|
||||
local.claim_pi_projection_key("local-provider", "local-key")?;
|
||||
@@ -482,11 +607,11 @@ mod tests {
|
||||
let conn =
|
||||
Connection::open(path).map_err(|error| AppError::Database(error.to_string()))?;
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE pi_provider_projections (
|
||||
"CREATE TABLE PI_PROVIDER_PROJECTIONS (
|
||||
provider_id TEXT PRIMARY KEY,
|
||||
provider_key TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE skill_deployments (
|
||||
CREATE TABLE SKILL_DEPLOYMENTS (
|
||||
app_type TEXT NOT NULL
|
||||
);",
|
||||
)
|
||||
@@ -531,6 +656,47 @@ mod tests {
|
||||
)
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
assert!(pi_device_local_rows_exist(&skill_only)?);
|
||||
|
||||
let view_backed =
|
||||
Connection::open_in_memory().map_err(|error| AppError::Database(error.to_string()))?;
|
||||
view_backed
|
||||
.execute_batch(
|
||||
"CREATE VIEW PI_PROVIDER_PROJECTIONS AS
|
||||
SELECT 'view-provider' AS provider_id, 'view-key' AS provider_key;",
|
||||
)
|
||||
.map_err(|error| AppError::Database(error.to_string()))?;
|
||||
assert!(
|
||||
pi_device_local_rows_exist(&view_backed)?,
|
||||
"a reserved-name view must not bypass binary ownership rejection"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn portable_import_rejects_an_unclosed_comment_before_live_replacement() -> Result<(), AppError>
|
||||
{
|
||||
let remote = Database::memory()?;
|
||||
seed_portable_provider(&remote)?;
|
||||
let malicious = format!("{}/*", remote.export_sql_string()?);
|
||||
|
||||
let local = Database::memory()?;
|
||||
local.claim_pi_projection_key("local-provider", "local-key")?;
|
||||
let error = local
|
||||
.import_portable_sql_string_for_sync(&malicious)
|
||||
.expect_err("the appended ownership program must not be swallowed");
|
||||
assert!(error.to_string().contains("unterminated"));
|
||||
assert_eq!(
|
||||
local
|
||||
.get_pi_projection("local-provider")?
|
||||
.map(|projection| projection.provider_key),
|
||||
Some("local-key".to_string())
|
||||
);
|
||||
assert!(
|
||||
local
|
||||
.get_provider_by_id("portable-sentinel", "codex")?
|
||||
.is_none(),
|
||||
"the remote database must not be published"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,8 +10,15 @@ use std::net::IpAddr;
|
||||
use std::sync::RwLock;
|
||||
use std::time::Duration;
|
||||
|
||||
/// 全局 HTTP 客户端实例
|
||||
static GLOBAL_CLIENT: OnceCell<RwLock<Client>> = OnceCell::new();
|
||||
#[derive(Clone)]
|
||||
struct GlobalClients {
|
||||
standard: Client,
|
||||
no_redirect: Client,
|
||||
}
|
||||
|
||||
/// 全局 HTTP 客户端实例。Pi 网关使用同一代理配置下的 no-redirect 客户端,
|
||||
/// 防止 307/308 把凭证、自定义头和请求体重放到另一个 origin。
|
||||
static GLOBAL_CLIENTS: OnceCell<RwLock<GlobalClients>> = OnceCell::new();
|
||||
|
||||
/// 当前代理 URL(用于日志和状态查询)
|
||||
static CURRENT_PROXY_URL: OnceCell<RwLock<Option<String>>> = OnceCell::new();
|
||||
@@ -52,10 +59,10 @@ fn get_proxy_port() -> u16 {
|
||||
/// 传入 None 或空字符串表示直连
|
||||
pub fn init(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
let effective_url = proxy_url.filter(|s| !s.trim().is_empty());
|
||||
let client = build_client(effective_url)?;
|
||||
let clients = build_clients(effective_url)?;
|
||||
|
||||
// 尝试初始化全局客户端,如果已存在则记录警告并使用 apply_proxy 更新
|
||||
if GLOBAL_CLIENT.set(RwLock::new(client.clone())).is_err() {
|
||||
if GLOBAL_CLIENTS.set(RwLock::new(clients)).is_err() {
|
||||
log::warn!(
|
||||
"[GlobalProxy] [GP-003] Already initialized, updating instead: {}",
|
||||
effective_url
|
||||
@@ -91,8 +98,8 @@ pub fn init(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
/// 验证成功返回 Ok(()),失败返回错误信息
|
||||
pub fn validate_proxy(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
let effective_url = proxy_url.filter(|s| !s.trim().is_empty());
|
||||
// 只调用 build_client 来验证,但不应用
|
||||
build_client(effective_url)?;
|
||||
// 同时验证标准与 no-redirect 客户端,保证应用配置时不会只更新一半。
|
||||
build_clients(effective_url)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -105,15 +112,15 @@ pub fn validate_proxy(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
/// * `proxy_url` - 代理 URL,None 或空字符串表示直连
|
||||
pub fn apply_proxy(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
let effective_url = proxy_url.filter(|s| !s.trim().is_empty());
|
||||
let new_client = build_client(effective_url)?;
|
||||
let new_clients = build_clients(effective_url)?;
|
||||
|
||||
// 更新客户端
|
||||
if let Some(lock) = GLOBAL_CLIENT.get() {
|
||||
let mut client = lock.write().map_err(|e| {
|
||||
if let Some(lock) = GLOBAL_CLIENTS.get() {
|
||||
let mut clients = lock.write().map_err(|e| {
|
||||
log::error!("[GlobalProxy] [GP-001] Failed to acquire write lock: {e}");
|
||||
"Failed to update proxy: lock poisoned".to_string()
|
||||
})?;
|
||||
*client = new_client;
|
||||
*clients = new_clients;
|
||||
} else {
|
||||
// 如果还没初始化,则初始化
|
||||
return init(proxy_url);
|
||||
@@ -148,54 +155,42 @@ pub fn apply_proxy(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
/// * `proxy_url` - 新的代理 URL,None 或空字符串表示直连
|
||||
#[allow(dead_code)]
|
||||
pub fn update_proxy(proxy_url: Option<&str>) -> Result<(), String> {
|
||||
let effective_url = proxy_url.filter(|s| !s.trim().is_empty());
|
||||
let new_client = build_client(effective_url)?;
|
||||
|
||||
// 更新客户端
|
||||
if let Some(lock) = GLOBAL_CLIENT.get() {
|
||||
let mut client = lock.write().map_err(|e| {
|
||||
log::error!("[GlobalProxy] [GP-001] Failed to acquire write lock: {e}");
|
||||
"Failed to update proxy: lock poisoned".to_string()
|
||||
})?;
|
||||
*client = new_client;
|
||||
} else {
|
||||
// 如果还没初始化,则初始化
|
||||
return init(proxy_url);
|
||||
}
|
||||
|
||||
// 更新代理 URL 记录
|
||||
if let Some(lock) = CURRENT_PROXY_URL.get() {
|
||||
let mut url = lock.write().map_err(|e| {
|
||||
log::error!("[GlobalProxy] [GP-002] Failed to acquire URL write lock: {e}");
|
||||
"Failed to update proxy URL record: lock poisoned".to_string()
|
||||
})?;
|
||||
*url = effective_url.map(|s| s.to_string());
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"[GlobalProxy] Updated: {}",
|
||||
effective_url
|
||||
.map(mask_url)
|
||||
.unwrap_or_else(|| "direct connection".to_string())
|
||||
);
|
||||
|
||||
Ok(())
|
||||
apply_proxy(proxy_url)
|
||||
}
|
||||
|
||||
/// 获取全局 HTTP 客户端
|
||||
///
|
||||
/// 返回配置了代理的客户端(如果已配置代理),否则返回跟随系统代理的客户端。
|
||||
pub fn get() -> Client {
|
||||
GLOBAL_CLIENT
|
||||
GLOBAL_CLIENTS
|
||||
.get()
|
||||
.and_then(|lock| lock.read().ok())
|
||||
.map(|c| c.clone())
|
||||
.map(|clients| clients.standard.clone())
|
||||
.unwrap_or_else(|| {
|
||||
log::warn!("[GlobalProxy] [GP-004] Client not initialized, using fallback");
|
||||
build_client(None).unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
/// 获取禁用自动重定向的全局客户端。
|
||||
///
|
||||
/// 与 `get` 不同,这个安全边界在锁毒化或构建失败时 fail closed;调用方不得
|
||||
/// 回退到会自动跟随跨 origin 重定向的默认客户端。
|
||||
pub(crate) fn get_no_redirect() -> Result<Client, String> {
|
||||
if let Some(lock) = GLOBAL_CLIENTS.get() {
|
||||
return lock
|
||||
.read()
|
||||
.map(|clients| clients.no_redirect.clone())
|
||||
.map_err(|error| {
|
||||
log::error!("[GlobalProxy] [GP-005] Failed to acquire read lock: {error}");
|
||||
"Failed to read no-redirect HTTP client: lock poisoned".to_string()
|
||||
});
|
||||
}
|
||||
|
||||
log::warn!("[GlobalProxy] [GP-006] Client not initialized, using no-redirect fallback");
|
||||
build_no_redirect_client(None)
|
||||
}
|
||||
|
||||
/// 获取当前代理 URL
|
||||
///
|
||||
/// 返回当前配置的代理 URL,None 表示直连。
|
||||
@@ -214,6 +209,24 @@ pub fn is_proxy_enabled() -> bool {
|
||||
|
||||
/// 构建 HTTP 客户端
|
||||
fn build_client(proxy_url: Option<&str>) -> Result<Client, String> {
|
||||
build_client_with_redirect_policy(proxy_url, true)
|
||||
}
|
||||
|
||||
fn build_no_redirect_client(proxy_url: Option<&str>) -> Result<Client, String> {
|
||||
build_client_with_redirect_policy(proxy_url, false)
|
||||
}
|
||||
|
||||
fn build_clients(proxy_url: Option<&str>) -> Result<GlobalClients, String> {
|
||||
Ok(GlobalClients {
|
||||
standard: build_client(proxy_url)?,
|
||||
no_redirect: build_no_redirect_client(proxy_url)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_client_with_redirect_policy(
|
||||
proxy_url: Option<&str>,
|
||||
follow_redirects: bool,
|
||||
) -> Result<Client, String> {
|
||||
let mut builder = Client::builder()
|
||||
.timeout(Duration::from_secs(600))
|
||||
.connect_timeout(Duration::from_secs(30))
|
||||
@@ -225,6 +238,9 @@ fn build_client(proxy_url: Option<&str>) -> Result<Client, String> {
|
||||
.no_brotli()
|
||||
.no_deflate()
|
||||
.no_zstd();
|
||||
if !follow_redirects {
|
||||
builder = builder.redirect(reqwest::redirect::Policy::none());
|
||||
}
|
||||
|
||||
// 有代理地址则使用代理,否则跟随系统代理
|
||||
if let Some(url) = proxy_url {
|
||||
@@ -337,7 +353,9 @@ pub fn mask_url(url: &str) -> String {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
fn env_lock() -> &'static Mutex<()> {
|
||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
@@ -392,6 +410,52 @@ mod tests {
|
||||
assert!(result.is_err(), "Should reject invalid proxy scheme");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_redirect_client_exposes_redirect_without_replaying_request() {
|
||||
let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
|
||||
.await
|
||||
.expect("bind redirect server");
|
||||
let address = listener.local_addr().expect("server address");
|
||||
let hits = Arc::new(AtomicUsize::new(0));
|
||||
let server_hits = Arc::clone(&hits);
|
||||
let server = tokio::spawn(async move {
|
||||
loop {
|
||||
let (mut stream, _) = listener.accept().await.expect("accept request");
|
||||
let hit = server_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let mut request = [0_u8; 2048];
|
||||
let _ = stream.read(&mut request).await.expect("read request");
|
||||
let response = if hit == 0 {
|
||||
format!(
|
||||
"HTTP/1.1 307 Temporary Redirect\r\n\
|
||||
Location: http://{address}/redirected\r\n\
|
||||
Content-Length: 0\r\nConnection: close\r\n\r\n"
|
||||
)
|
||||
} else {
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string()
|
||||
};
|
||||
stream
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("write response");
|
||||
}
|
||||
});
|
||||
|
||||
let client = build_no_redirect_client(None).expect("build no-redirect client");
|
||||
let response = client
|
||||
.get(format!("http://{address}/initial"))
|
||||
.send()
|
||||
.await
|
||||
.expect("send request");
|
||||
assert_eq!(response.status(), reqwest::StatusCode::TEMPORARY_REDIRECT);
|
||||
tokio::time::sleep(Duration::from_millis(25)).await;
|
||||
assert_eq!(
|
||||
hits.load(Ordering::SeqCst),
|
||||
1,
|
||||
"the redirected endpoint must not receive a replay"
|
||||
);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_proxy_points_to_loopback() {
|
||||
// 设置 CC Switch 代理端口为 15721(默认值)
|
||||
|
||||
@@ -172,7 +172,10 @@ pub(crate) async fn handle_pi_native(
|
||||
// response is no longer client-visible.
|
||||
drop(pending);
|
||||
}
|
||||
let send = crate::proxy::http_client::get()
|
||||
let client = crate::proxy::http_client::get_no_redirect().map_err(|error| {
|
||||
ProxyError::ForwardFailed(format!("Pi gateway HTTP client is unavailable: {error}"))
|
||||
})?;
|
||||
let send = client
|
||||
.request(method.clone(), materialized.url.clone())
|
||||
.headers(outgoing_headers)
|
||||
.body(body.clone())
|
||||
|
||||
@@ -362,14 +362,17 @@ impl<'a> UsageLogger<'a> {
|
||||
} else {
|
||||
app_type
|
||||
};
|
||||
let default_multiplier_raw =
|
||||
let default_multiplier_raw = if default_app_type == "pi" {
|
||||
crate::settings::get_pi_default_cost_multiplier()
|
||||
} else {
|
||||
match self.db.get_default_cost_multiplier(default_app_type).await {
|
||||
Ok(value) => value,
|
||||
Err(e) => {
|
||||
log::warn!("[USG-003] 获取默认倍率失败 (app_type={app_type}): {e}");
|
||||
"1".to_string()
|
||||
}
|
||||
};
|
||||
}
|
||||
};
|
||||
let default_multiplier = match Decimal::from_str(&default_multiplier_raw) {
|
||||
Ok(value) => value,
|
||||
Err(e) => {
|
||||
@@ -380,14 +383,17 @@ impl<'a> UsageLogger<'a> {
|
||||
}
|
||||
};
|
||||
|
||||
let default_pricing_source_raw =
|
||||
let default_pricing_source_raw = if default_app_type == "pi" {
|
||||
crate::settings::get_pi_pricing_model_source()
|
||||
} else {
|
||||
match self.db.get_pricing_model_source(default_app_type).await {
|
||||
Ok(value) => value,
|
||||
Err(e) => {
|
||||
log::warn!("[USG-003] 获取默认计费模式失败 (app_type={app_type}): {e}");
|
||||
PRICING_SOURCE_RESPONSE.to_string()
|
||||
}
|
||||
};
|
||||
}
|
||||
};
|
||||
let default_pricing_source = if default_pricing_source_raw == PRICING_SOURCE_RESPONSE
|
||||
|| default_pricing_source_raw == PRICING_SOURCE_REQUEST
|
||||
{
|
||||
|
||||
@@ -159,36 +159,60 @@ impl PiCatalogCoordinator {
|
||||
pub(crate) fn apply(
|
||||
state: &AppState,
|
||||
mutation: PiCatalogMutation,
|
||||
) -> Result<PiCatalogMutationResult, AppError> {
|
||||
let switch_guard = futures::executor::block_on(
|
||||
state
|
||||
.proxy_service
|
||||
.lock_switch_for_app(AppType::Pi.as_str()),
|
||||
);
|
||||
Self::apply_under_switch_guard(state, &switch_guard, mutation)
|
||||
}
|
||||
|
||||
/// Compose a Pi catalog mutation with a caller-owned switch boundary.
|
||||
///
|
||||
/// Multi-step control-plane operations (for example selecting failover P1
|
||||
/// and enabling its runtime policy) use this entry point so they do not
|
||||
/// release the shared ownership boundary or deadlock by acquiring it twice.
|
||||
pub(crate) fn apply_under_switch_guard(
|
||||
state: &AppState,
|
||||
switch_guard: &tokio::sync::OwnedMutexGuard<()>,
|
||||
mutation: PiCatalogMutation,
|
||||
) -> Result<PiCatalogMutationResult, AppError> {
|
||||
let additional_native_key = match &mutation {
|
||||
PiCatalogMutation::CreateProvider { provider_key, .. }
|
||||
| PiCatalogMutation::ImportNative { provider_key, .. } => Some(provider_key.clone()),
|
||||
_ => None,
|
||||
};
|
||||
Self::run_with_runtime_reconcile(state, additional_native_key.as_deref(), || match mutation
|
||||
{
|
||||
PiCatalogMutation::CreateProvider {
|
||||
input,
|
||||
provider_key,
|
||||
activate_if_first,
|
||||
} => Self::create(state, input, provider_key, activate_if_first),
|
||||
PiCatalogMutation::UpdateProvider { input } => Self::update(state, input),
|
||||
PiCatalogMutation::DeleteProvider { provider_id } => Self::delete(state, &provider_id),
|
||||
PiCatalogMutation::AddEndpoint { provider_id, url } => {
|
||||
Self::add_endpoint(state, &provider_id, &url)
|
||||
}
|
||||
PiCatalogMutation::RemoveEndpoint { provider_id, url } => {
|
||||
Self::remove_endpoint(state, &provider_id, &url)
|
||||
}
|
||||
PiCatalogMutation::ImportNative {
|
||||
provider_key,
|
||||
expected_fingerprint,
|
||||
} => Self::import_native(state, &provider_key, &expected_fingerprint),
|
||||
PiCatalogMutation::SetDefault {
|
||||
provider_id,
|
||||
model_id,
|
||||
} => Self::set_default(state, &provider_id, &model_id),
|
||||
})
|
||||
Self::run_with_runtime_reconcile_locked(
|
||||
state,
|
||||
switch_guard,
|
||||
additional_native_key.as_deref(),
|
||||
|| match mutation {
|
||||
PiCatalogMutation::CreateProvider {
|
||||
input,
|
||||
provider_key,
|
||||
activate_if_first,
|
||||
} => Self::create(state, input, provider_key, activate_if_first),
|
||||
PiCatalogMutation::UpdateProvider { input } => Self::update(state, input),
|
||||
PiCatalogMutation::DeleteProvider { provider_id } => {
|
||||
Self::delete(state, &provider_id)
|
||||
}
|
||||
PiCatalogMutation::AddEndpoint { provider_id, url } => {
|
||||
Self::add_endpoint(state, &provider_id, &url)
|
||||
}
|
||||
PiCatalogMutation::RemoveEndpoint { provider_id, url } => {
|
||||
Self::remove_endpoint(state, &provider_id, &url)
|
||||
}
|
||||
PiCatalogMutation::ImportNative {
|
||||
provider_key,
|
||||
expected_fingerprint,
|
||||
} => Self::import_native(state, &provider_key, &expected_fingerprint),
|
||||
PiCatalogMutation::SetDefault {
|
||||
provider_id,
|
||||
model_id,
|
||||
} => Self::set_default(state, &provider_id, &model_id),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Reconcile portable provider rows with this device's exact-key ledger.
|
||||
@@ -224,11 +248,25 @@ impl PiCatalogCoordinator {
|
||||
additional_native_key: Option<&str>,
|
||||
operation: impl FnOnce() -> Result<PiCatalogMutationResult, AppError>,
|
||||
) -> Result<PiCatalogMutationResult, AppError> {
|
||||
let _switch_guard = futures::executor::block_on(
|
||||
let switch_guard = futures::executor::block_on(
|
||||
state
|
||||
.proxy_service
|
||||
.lock_switch_for_app(AppType::Pi.as_str()),
|
||||
);
|
||||
Self::run_with_runtime_reconcile_locked(
|
||||
state,
|
||||
&switch_guard,
|
||||
additional_native_key,
|
||||
operation,
|
||||
)
|
||||
}
|
||||
|
||||
fn run_with_runtime_reconcile_locked(
|
||||
state: &AppState,
|
||||
_switch_guard: &tokio::sync::OwnedMutexGuard<()>,
|
||||
additional_native_key: Option<&str>,
|
||||
operation: impl FnOnce() -> Result<PiCatalogMutationResult, AppError>,
|
||||
) -> Result<PiCatalogMutationResult, AppError> {
|
||||
Self::reconcile_current_indexes_from_native(state)?;
|
||||
let snapshot = PiCatalogSnapshot::capture(state, additional_native_key)?;
|
||||
let catalog_epoch =
|
||||
|
||||
@@ -181,6 +181,7 @@ pub struct PiPromptTemplateService;
|
||||
|
||||
impl PiPromptTemplateService {
|
||||
pub fn list() -> Result<Vec<PiPromptTemplate>, AppError> {
|
||||
let _guard = lock_instruction_files()?;
|
||||
Self::list_at(&get_pi_agent_dir()?.join("prompts"))
|
||||
}
|
||||
|
||||
@@ -189,6 +190,7 @@ impl PiPromptTemplateService {
|
||||
expected_revision: &str,
|
||||
content: &str,
|
||||
) -> Result<PiPromptTemplate, AppError> {
|
||||
let _guard = lock_instruction_files()?;
|
||||
Self::upsert_at(
|
||||
&get_pi_agent_dir()?.join("prompts"),
|
||||
slug,
|
||||
@@ -199,6 +201,7 @@ impl PiPromptTemplateService {
|
||||
|
||||
pub fn delete(slug: &str, expected_revision: &str) -> Result<bool, AppError> {
|
||||
validate_template_slug(slug)?;
|
||||
let _guard = lock_instruction_files()?;
|
||||
delete_shared_file(
|
||||
&template_path(&get_pi_agent_dir()?.join("prompts"), slug),
|
||||
expected_revision,
|
||||
|
||||
@@ -459,7 +459,7 @@ impl ProfileService {
|
||||
.set_current_profile_id(scope.as_str(), Some(profile_id))?;
|
||||
|
||||
// 当前分组内所有接管已关闭;若其它应用也无接管,可停止代理服务。
|
||||
let should_stop_proxy = !state.db.is_live_takeover_active_sync();
|
||||
let should_stop_proxy = !state.proxy_service.shared_listener_is_desired_sync();
|
||||
|
||||
Ok((warnings, should_stop_proxy))
|
||||
}
|
||||
|
||||
@@ -576,6 +576,25 @@ impl ProxyService {
|
||||
self.switch_locks.lock_for_app(app_type).await
|
||||
}
|
||||
|
||||
async fn shared_listener_is_desired(&self) -> Result<bool, String> {
|
||||
let legacy_takeover = self
|
||||
.db
|
||||
.is_live_takeover_active()
|
||||
.await
|
||||
.map_err(|error| format!("检查接管状态失败: {error}"))?;
|
||||
Ok(shared_listener_is_desired(
|
||||
legacy_takeover,
|
||||
crate::settings::pi_takeover_enabled(),
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn shared_listener_is_desired_sync(&self) -> bool {
|
||||
shared_listener_is_desired(
|
||||
self.db.is_live_takeover_active_sync(),
|
||||
crate::settings::pi_takeover_enabled(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn begin_pi_catalog_mutation(&self) -> u64 {
|
||||
self.pi_runtime.begin_mutation().await
|
||||
}
|
||||
@@ -1808,7 +1827,7 @@ impl ProxyService {
|
||||
.await
|
||||
.map_err(|e| format!("检查接管状态失败: {e}"))?;
|
||||
|
||||
if !any_enabled {
|
||||
if !shared_listener_is_desired(any_enabled, crate::settings::pi_takeover_enabled()) {
|
||||
let _ = self.db.set_live_takeover_active(false).await;
|
||||
|
||||
if self.is_running().await {
|
||||
@@ -1938,13 +1957,7 @@ impl ProxyService {
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
|
||||
let status = self.get_takeover_status().await?;
|
||||
if !status.claude
|
||||
&& !status.codex
|
||||
&& !status.gemini
|
||||
&& !status.grokbuild
|
||||
&& self.is_running().await
|
||||
{
|
||||
if !self.shared_listener_is_desired().await? && self.is_running().await {
|
||||
let _ = self.stop().await;
|
||||
}
|
||||
Ok(())
|
||||
@@ -4371,6 +4384,10 @@ impl ProxyService {
|
||||
}
|
||||
}
|
||||
|
||||
fn shared_listener_is_desired(legacy_takeover: bool, pi_takeover: bool) -> bool {
|
||||
legacy_takeover || pi_takeover
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -4381,6 +4398,14 @@ mod tests {
|
||||
use std::env;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn shared_listener_stays_live_until_both_takeover_domains_are_disabled() {
|
||||
assert!(!shared_listener_is_desired(false, false));
|
||||
assert!(shared_listener_is_desired(true, false));
|
||||
assert!(shared_listener_is_desired(false, true));
|
||||
assert!(shared_listener_is_desired(true, true));
|
||||
}
|
||||
|
||||
struct TempHome {
|
||||
#[allow(dead_code)]
|
||||
dir: TempDir,
|
||||
|
||||
@@ -24,6 +24,10 @@ pub struct CustomEndpoint {
|
||||
pub struct PiProxySettings {
|
||||
#[serde(default)]
|
||||
pub auto_failover_enabled: bool,
|
||||
#[serde(default = "default_pi_cost_multiplier")]
|
||||
pub default_cost_multiplier: String,
|
||||
#[serde(default = "default_pi_pricing_model_source")]
|
||||
pub pricing_model_source: String,
|
||||
#[serde(default = "default_pi_max_retries")]
|
||||
pub max_retries: u32,
|
||||
#[serde(default = "default_pi_first_byte_timeout")]
|
||||
@@ -44,6 +48,12 @@ pub struct PiProxySettings {
|
||||
pub circuit_min_requests: u32,
|
||||
}
|
||||
|
||||
fn default_pi_cost_multiplier() -> String {
|
||||
"1".to_string()
|
||||
}
|
||||
fn default_pi_pricing_model_source() -> String {
|
||||
crate::database::PRICING_SOURCE_RESPONSE.to_string()
|
||||
}
|
||||
const fn default_pi_max_retries() -> u32 {
|
||||
3
|
||||
}
|
||||
@@ -76,6 +86,8 @@ impl Default for PiProxySettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
auto_failover_enabled: false,
|
||||
default_cost_multiplier: default_pi_cost_multiplier(),
|
||||
pricing_model_source: default_pi_pricing_model_source(),
|
||||
max_retries: default_pi_max_retries(),
|
||||
streaming_first_byte_timeout: default_pi_first_byte_timeout(),
|
||||
streaming_idle_timeout: default_pi_idle_timeout(),
|
||||
@@ -91,6 +103,8 @@ impl Default for PiProxySettings {
|
||||
|
||||
impl PiProxySettings {
|
||||
pub(crate) fn validate(&self) -> Result<(), AppError> {
|
||||
crate::database::validate_cost_multiplier(&self.default_cost_multiplier)?;
|
||||
crate::database::validate_pricing_source(&self.pricing_model_source)?;
|
||||
if !self.circuit_error_rate_threshold.is_finite()
|
||||
|| !(0.0..=1.0).contains(&self.circuit_error_rate_threshold)
|
||||
{
|
||||
@@ -1152,6 +1166,25 @@ pub(crate) fn update_pi_proxy_settings(settings: PiProxySettings) -> Result<(),
|
||||
mutate_settings(|current| current.pi_proxy = settings)
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_default_cost_multiplier() -> String {
|
||||
get_pi_proxy_settings().default_cost_multiplier
|
||||
}
|
||||
|
||||
pub(crate) fn set_pi_default_cost_multiplier(value: &str) -> Result<(), AppError> {
|
||||
crate::database::validate_cost_multiplier(value)?;
|
||||
let value = value.trim().to_string();
|
||||
mutate_settings(move |settings| settings.pi_proxy.default_cost_multiplier = value)
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_pricing_model_source() -> String {
|
||||
get_pi_proxy_settings().pricing_model_source
|
||||
}
|
||||
|
||||
pub(crate) fn set_pi_pricing_model_source(value: &str) -> Result<(), AppError> {
|
||||
let value = crate::database::validate_pricing_source(value)?.to_string();
|
||||
mutate_settings(move |settings| settings.pi_proxy.pricing_model_source = value)
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_app_proxy_config() -> crate::proxy::types::AppProxyConfig {
|
||||
get_pi_proxy_settings().app_config(pi_takeover_enabled())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user