fix(pi): close final gateway ownership regressions

This commit is contained in:
SaladDay
2026-08-03 14:22:17 +00:00
parent f35b47c1cf
commit 63b1a20eef
12 changed files with 832 additions and 215 deletions
+270 -64
View File
@@ -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", &current)
.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(())
}
}
+71
View File
@@ -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(())
}
}
+233 -67
View File
@@ -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(())
}
}
+109 -45
View File
@@ -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` - 代理 URLNone 或空字符串表示直连
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` - 新的代理 URLNone 或空字符串表示直连
#[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(默认值)
+4 -1
View File
@@ -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())
+10 -4
View File
@@ -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
{
+63 -25
View File
@@ -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,
+1 -1
View File
@@ -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))
}
+33 -8
View File
@@ -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,
+33
View File
@@ -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())
}