admin/CC-Switch
1
0
Fork 0
mirror of https://github.com/farion1231/cc-switch.git synced 2026-08-04 03:32:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(pi): harden gateway publication and retry responses

Browse Source
This commit is contained in:
SaladDay
2026-08-03 03:23:48 +00:00
parent 1e9e59051a
commit 40e087c072
5 changed files with 1083 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/pi_config/document.rs
+204 -31
View File
@@ -357,10 +357,19 @@ fn serialize_models_mutation(
/// Unknown root fields, unowned provider entries, comments, and formatting are
/// preserved by the CST patch. An optimistic fingerprint check prevents a
/// Pi/user write observed before replacement from being silently overwritten.
#[cfg(test)]
pub(crate) fn apply_pi_provider_patch(
path: &Path,
patch: &IndexMap<String, Option<Value>>,
) -> Result<(), AppError> {
apply_pi_provider_patch_checked(path, None, patch).map(|_| ())
}
fn apply_pi_provider_patch_checked(
path: &Path,
expected: Option<&IndexMap<String, Option<Value>>>,
patch: &IndexMap<String, Option<Value>>,
) -> Result<bool, AppError> {
let lock = path_lock(path)?;
let _guard = lock_path(&lock)?;
if let Some(parent) = path.parent() {
@@ -369,6 +378,28 @@ pub(crate) fn apply_pi_provider_patch(
for _ in 0..MAX_MUTATION_ATTEMPTS {
let before = read_models_bytes(path)?;
if let Some(expected) = expected {
let observed = provider_values_from_bytes(path, before.as_deref(), expected.keys())?;
if let Some((provider_key, expected_value)) =
expected.iter().find(|(provider_key, expected_value)| {
observed.get(*provider_key) != Some(*expected_value)
})
{
return Err(AppError::Conflict(format!(
"Pi provider key '{provider_key}' changed since directory/catalog preflight \
(expected {}, observed {})",
provider_value_label(expected_value),
provider_value_label(
observed
.get(provider_key)
.expect("every requested provider key is observed")
)
)));
}
}
if before.is_none() && patch.values().all(Option::is_none) {
return Ok(false);
}
let serialized = serialize_models_mutation(path, before.as_deref(), &|document| {
let providers = document
.as_object_mut()
@@ -387,6 +418,9 @@ pub(crate) fn apply_pi_provider_patch(
}
Ok(())
})?;
if before.as_deref() == Some(serialized.as_slice()) {
return Ok(false);
}
match compare_exchange_shared_file_bytes(
path,
@@ -396,7 +430,7 @@ pub(crate) fn apply_pi_provider_patch(
None,
"Pi models file",
) {
Ok(_) => return Ok(()),
Ok(_) => return Ok(true),
Err(AppError::Conflict(_)) => continue,
Err(error) => return Err(error),
}
@@ -408,53 +442,130 @@ pub(crate) fn apply_pi_provider_patch(
)))
}
#[derive(Debug)]
fn provider_values_from_bytes<'a>(
path: &Path,
bytes: Option<&[u8]>,
provider_keys: impl IntoIterator<Item = &'a String>,
) -> Result<IndexMap<String, Option<Value>>, AppError> {
let document = parse_pi_models_document(path, bytes)?;
Ok(provider_keys
.into_iter()
.map(|key| {
let value = document
.providers()
.get(key)
.map(|entry| entry.value.clone());
(key.clone(), value)
})
.collect())
}
fn provider_value_label(value: &Option<Value>) -> &'static str {
if value.is_some() {
"present"
} else {
"absent"
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct PiProviderValuesSnapshot {
pub file_existed: bool,
pub values: IndexMap<String, Option<Value>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct PiProviderPatchReceipt {
path: PathBuf,
before: PiProviderValuesSnapshot,
attempted: IndexMap<String, Option<Value>>,
attempted_file_existed: bool,
}
impl PiProviderPatchReceipt {
pub(crate) fn attempted_values(&self) -> &IndexMap<String, Option<Value>> {
&self.attempted
}
pub(crate) fn attempted_snapshot(&self) -> PiProviderValuesSnapshot {
PiProviderValuesSnapshot {
file_existed: self.attempted_file_existed,
values: self.attempted.clone(),
}
}
/// Restore only provider keys which still contain this operation's exact
/// attempted values. Unrelated root/provider edits are preserved, while a
/// concurrent edit of an owned key turns compensation into an explicit
/// conflict instead of being overwritten.
pub(crate) fn rollback(&self) -> Result<(), AppError> {
apply_pi_provider_patch_checked(&self.path, Some(&self.attempted), &self.before.values)?;
remove_new_empty_document(&self.path, &self.before)
}
}
pub(crate) fn snapshot_pi_provider_values(
path: &Path,
provider_keys: impl IntoIterator<Item = String>,
) -> Result<PiProviderValuesSnapshot, AppError> {
let bytes = read_models_bytes(path)?;
let file_existed = bytes.is_some();
let document = parse_pi_models_document(path, bytes.as_deref())?;
Ok(PiProviderValuesSnapshot {
file_existed,
values: provider_keys
.into_iter()
.map(|key| {
let value = document
.providers()
.get(&key)
.map(|entry| entry.value.clone());
(key, value)
})
.collect(),
values: provider_values_from_bytes(
path,
bytes.as_deref(),
provider_keys.into_iter().collect::<Vec<_>>().iter(),
)?,
})
}
pub(crate) fn current_pi_provider_values(
/// Publish an exact-key patch only if every preflighted provider value still
/// matches. Whole-file CAS retries preserve unrelated edits, but re-check the
/// provider precondition before every retry.
pub(crate) fn apply_pi_provider_patch_with_receipt(
path: &Path,
provider_keys: impl IntoIterator<Item = String>,
) -> Result<IndexMap<String, Option<Value>>, AppError> {
Ok(snapshot_pi_provider_values(path, provider_keys)?.values)
before: &PiProviderValuesSnapshot,
patch: &IndexMap<String, Option<Value>>,
) -> Result<PiProviderPatchReceipt, AppError> {
if patch
.keys()
.any(|provider_key| !before.values.contains_key(provider_key))
{
return Err(AppError::InvalidInput(
"Pi provider patch contains a key which was not preflighted".to_string(),
));
}
let changed = apply_pi_provider_patch_checked(path, Some(&before.values), patch)?;
let attempted = before
.values
.iter()
.map(|(provider_key, previous)| {
(
provider_key.clone(),
patch
.get(provider_key)
.cloned()
.unwrap_or_else(|| previous.clone()),
)
})
.collect();
Ok(PiProviderPatchReceipt {
path: path.to_path_buf(),
before: before.clone(),
attempted,
attempted_file_existed: before.file_existed || changed,
})
}
/// Restore an exact-key snapshot and remove a file which this operation
/// created only when its bytes are still the canonical empty document.
///
/// If Pi or the user added any other content, the file is retained. The
/// revision-checked delete also preserves a writer which races after the
/// emptiness check.
pub(crate) fn restore_pi_provider_values(
/// Remove a file which this operation created only when its bytes are still
/// the canonical empty document. If Pi or the user added any other content,
/// the file is retained.
fn remove_new_empty_document(
path: &Path,
snapshot: &PiProviderValuesSnapshot,
before: &PiProviderValuesSnapshot,
) -> Result<(), AppError> {
apply_pi_provider_patch(path, &snapshot.values)?;
if snapshot.file_existed {
if before.file_existed {
return Ok(());
}
@@ -548,12 +659,13 @@ mod tests {
let path = temp.path().join("models.json");
let before =
snapshot_pi_provider_values(&path, ["managed".to_string()]).expect("snapshot absence");
apply_pi_provider_patch(
let receipt = apply_pi_provider_patch_with_receipt(
&path,
&before,
&IndexMap::from([("managed".to_string(), Some(serde_json::json!({"api": "x"})))]),
)
.expect("publish");
restore_pi_provider_values(&path, &before).expect("restore absence");
receipt.rollback().expect("restore absence");
assert!(
!path.exists(),
"rollback must not leave an empty shadow file"
@@ -561,8 +673,9 @@ mod tests {
let before =
snapshot_pi_provider_values(&path, ["managed".to_string()]).expect("snapshot absence");
apply_pi_provider_patch(
let receipt = apply_pi_provider_patch_with_receipt(
&path,
&before,
&IndexMap::from([("managed".to_string(), Some(serde_json::json!({"api": "x"})))]),
)
.expect("publish");
@@ -571,7 +684,7 @@ mod tests {
"{\n \"providers\": {\"managed\": {\"api\": \"x\"}},\n \"external\": true\n}\n",
)
.expect("external root update");
restore_pi_provider_values(&path, &before).expect("restore managed key");
receipt.rollback().expect("restore managed key");
let restored: Value =
serde_json::from_slice(&fs::read(&path).expect("external file retained"))
.expect("parse restored");
@@ -582,6 +695,66 @@ mod tests {
.is_some_and(serde_json::Map::is_empty));
}
#[test]
fn checked_provider_publish_rejects_a_key_changed_after_preflight() {
let temp = tempfile::tempdir().expect("tempdir");
let path = temp.path().join("models.json");
fs::write(&path, r#"{"providers":{}}"#).expect("seed");
let before =
snapshot_pi_provider_values(&path, ["managed".to_string()]).expect("preflight");
crate::pi_config::shared_file::replace_before_next_compare_exchange(
&path,
br#"{"providers":{"managed":{"api":"external"}}}"#,
);
let error = apply_pi_provider_patch_with_receipt(
&path,
&before,
&IndexMap::from([(
"managed".to_string(),
Some(serde_json::json!({"api": "attempted"})),
)]),
)
.expect_err("external key must win");
assert!(matches!(error, AppError::Conflict(_)));
let document = read_pi_models_document(&path).expect("external document");
assert_eq!(document.providers()["managed"].value["api"], "external");
}
#[test]
fn provider_receipt_rollback_rejects_a_newer_exact_key_but_preserves_it() {
let temp = tempfile::tempdir().expect("tempdir");
let path = temp.path().join("models.json");
fs::write(
&path,
r#"{"root":"preserved","providers":{"managed":{"api":"before"}}}"#,
)
.expect("seed");
let before =
snapshot_pi_provider_values(&path, ["managed".to_string()]).expect("preflight");
let receipt = apply_pi_provider_patch_with_receipt(
&path,
&before,
&IndexMap::from([(
"managed".to_string(),
Some(serde_json::json!({"api": "attempted"})),
)]),
)
.expect("publish");
crate::pi_config::shared_file::replace_before_next_compare_exchange(
&path,
br#"{"root":"external","providers":{"managed":{"api":"external"}}}"#,
);
let error = receipt
.rollback()
.expect_err("rollback must not overwrite the newer key");
assert!(matches!(error, AppError::Conflict(_)));
let document: Value =
serde_json::from_slice(&fs::read(&path).expect("read external")).expect("parse");
assert_eq!(document["root"], "external");
assert_eq!(document["providers"]["managed"]["api"], "external");
}
#[test]
fn external_rename_during_patch_is_reparsed_before_owned_fields_change() {
let temp = tempfile::tempdir().expect("tempdir");
src-tauri/src/proxy/pi_handler.rs
+116 -10
View File
@@ -91,6 +91,7 @@ pub(crate) async fn handle_pi_native(
crate::proxy::extract_session_id(&incoming_headers, &request_json, "pi").session_id;
let mut protocol_anchor = ProtocolAnchor::for_primary(attempts.first());
let mut last_error = None;
let mut pending_retryable: Option<PendingRetryableResponse> = None;
record_request_start(&state).await;
let mut index = 0;
@@ -152,6 +153,30 @@ pub(crate) async fn handle_pi_native(
if !network_budget.begin_send() {
break;
}
if let Some(pending) = pending_retryable.take() {
if pending.provider_id != provider_id {
record_provider_result(
&state,
route.catalog_epoch,
&pending.provider_id,
pending.used_half_open_permit,
false,
Some(format!(
"Pi upstream returned retryable status {}",
pending.response.status()
)),
)
.await;
} else {
debug_assert_eq!(
pending.used_half_open_permit, permit.used_half_open_permit,
"one provider group must retain one circuit-breaker permit"
);
}
// A real later send has now begun, so the earlier fallback
// response is no longer client-visible.
drop(pending);
}
let send = crate::proxy::http_client::get()
.request(method.clone(), materialized.url.clone())
.headers(outgoing_headers)
@@ -186,6 +211,14 @@ pub(crate) async fn handle_pi_native(
if retryable_status(status) && network_budget.has_remaining() {
saw_health_failure = true;
last_error = Some(format!("Pi upstream returned retryable status {status}"));
let selected_is_failover = materialized.is_failover;
pending_retryable = Some(PendingRetryableResponse {
response,
materialized,
provider_id: provider_id.clone(),
used_half_open_permit: permit.used_half_open_permit,
selected_is_failover,
});
continue;
}
let selected_is_failover = materialized.is_failover;
@@ -251,25 +284,98 @@ pub(crate) async fn handle_pi_native(
}
}
release_or_record_provider(
&state,
route.catalog_epoch,
&provider_id,
permit.used_half_open_permit,
reached_network,
saw_health_failure,
last_error.clone(),
)
.await;
if pending_retryable
.as_ref()
.is_none_or(|pending| pending.provider_id != provider_id)
{
release_or_record_provider(
&state,
route.catalog_epoch,
&provider_id,
permit.used_half_open_permit,
reached_network,
saw_health_failure,
last_error.clone(),
)
.await;
}
index = provider_end;
}
if let Some(pending) = pending_retryable {
let status = pending.response.status();
let provider_id = pending.provider_id.clone();
let used_half_open_permit = pending.used_half_open_permit;
let selected_is_failover = pending.selected_is_failover;
match prepare_response(
state.clone(),
pending.response,
pending.materialized,
route.catalog_epoch,
model_id.clone(),
session_id.clone(),
started,
is_streaming,
route.app_config.streaming_first_byte_timeout,
route.app_config.streaming_idle_timeout,
route.app_config.non_streaming_timeout,
used_half_open_permit,
selected_is_failover,
)
.await
{
Ok(prepared) => {
if !prepared.finalization_deferred {
record_provider_result(
&state,
route.catalog_epoch,
&provider_id,
used_half_open_permit,
false,
Some(format!("Pi upstream returned {status}")),
)
.await;
record_request_finish(
&state,
status.is_success(),
selected_is_failover,
(!status.is_success()).then(|| format!("Pi upstream returned {status}")),
)
.await;
}
return Ok(prepared.response);
}
Err(error) => {
record_provider_result(
&state,
route.catalog_epoch,
&provider_id,
used_half_open_permit,
false,
Some(error.to_string()),
)
.await;
record_request_finish(&state, false, selected_is_failover, Some(error.to_string()))
.await;
return Err(error);
}
}
}
let error =
last_error.unwrap_or_else(|| "no wire-compatible Pi candidate was available".to_string());
record_request_finish(&state, false, false, Some(error.clone())).await;
Err(ProxyError::ForwardFailed(error))
}
struct PendingRetryableResponse {
response: reqwest::Response,
materialized: PiMaterializedAttempt,
provider_id: String,
used_half_open_permit: bool,
selected_is_failover: bool,
}
#[derive(Debug)]
struct NetworkAttemptBudget {
remaining: usize,
src-tauri/src/proxy/pi_runtime.rs
+17
View File
@@ -698,6 +698,23 @@ impl PiRuntimeStore {
self.lease(server_generation).is_some()
}
/// Retain the credential of the last published generation for exact
/// native-projection compensation even while an odd epoch fences new
/// admission. This does not mint or rotate credentials and never crosses
/// IPC; it is only an ownership witness for restoring `models.json`.
pub(crate) fn retained_gateway_token(
&self,
server_generation: u64,
) -> Option<crate::settings::GatewayToken> {
self.publication
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.current
.as_ref()
.filter(|snapshot| snapshot.server_generation == server_generation)
.map(|snapshot| snapshot.gateway_token.clone())
}
pub(crate) async fn admission_guard(
self: &Arc<Self>,
server_generation: u64,
src-tauri/src/services/pi_catalog.rs
+150 -98
View File
@@ -11,8 +11,8 @@ use crate::database::{
};
use crate::error::AppError;
use crate::pi_config::document::{
apply_pi_provider_patch, current_pi_provider_values, restore_pi_provider_values,
snapshot_pi_provider_values, PiProviderValuesSnapshot,
apply_pi_provider_patch_with_receipt, snapshot_pi_provider_values, PiProviderPatchReceipt,
PiProviderValuesSnapshot,
};
use crate::pi_config::model::{
effective_pi_model, validate_pi_managed_provider, PiManagedProviderConfig, PiManagementStatus,
@@ -31,7 +31,6 @@ use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::BTreeMap;
use std::path::PathBuf;
const PI_APP: &str = "pi";
@@ -94,6 +93,8 @@ pub(crate) struct PiCatalogMutationResult {
pub provider_id: Option<String>,
#[serde(skip)]
native_defaults_receipt: Option<PiNativeDefaultsReceipt>,
#[serde(skip)]
native_patch_receipt: Option<PiProviderPatchReceipt>,
}
pub(crate) struct PiCatalogCoordinator;
@@ -102,7 +103,8 @@ struct PiCatalogSnapshot {
aggregates: IndexMap<String, ProviderAggregate>,
projections: Vec<PiProviderProjection>,
db_current: Option<String>,
models_path: PathBuf,
// Capture validates the complete shared projection before any DB write.
// Exact compensation itself uses per-operation before/attempted receipts.
native: PiProviderValuesSnapshot,
}
@@ -195,16 +197,21 @@ impl PiCatalogCoordinator {
pub(crate) fn reconcile_portable_import(state: &AppState) -> Result<(), AppError> {
Self::run_with_runtime_reconcile(state, None, || {
let models_path = get_pi_models_path()?;
let native_defaults_receipt = Self::reconcile_portable_catalog_at(
state,
&models_path,
|provider_id, provider_key, config| {
state
.proxy_service
.project_pi_provider_value(provider_id, provider_key, config)
},
)?;
Ok(success(None).with_native_defaults_receipt(native_defaults_receipt))
let (native_defaults_receipt, native_patch_receipt) =
Self::reconcile_portable_catalog_at(
state,
&models_path,
|provider_id, provider_key, config| {
state.proxy_service.project_pi_provider_value(
provider_id,
provider_key,
config,
)
},
)?;
Ok(success(None)
.with_native_defaults_receipt(native_defaults_receipt)
.with_native_patch_receipt(native_patch_receipt))
})
.map(|_| ())
}
@@ -229,17 +236,35 @@ impl PiCatalogCoordinator {
.ok()
.and_then(|result| result.native_defaults_receipt.as_ref())
.cloned();
let native_patch_receipt = result
.as_ref()
.ok()
.and_then(|result| result.native_patch_receipt.as_ref())
.cloned();
let mut expected_native = snapshot.native.clone();
if let Some(receipt) = native_patch_receipt.as_ref() {
for (provider_key, attempted) in receipt.attempted_values() {
expected_native
.values
.insert(provider_key.clone(), attempted.clone());
}
}
let reconcile = futures::executor::block_on(
state
.proxy_service
.reconcile_pi_runtime_at_epoch(catalog_epoch),
.reconcile_pi_runtime_at_epoch_with_native_precondition(
catalog_epoch,
Some(&expected_native),
),
);
match (result, reconcile) {
(Ok(result), Ok(_)) => Ok(result),
(Ok(_), Err(error)) => {
if let Err(rollback_error) =
snapshot.restore(state, native_defaults_receipt.as_ref())
{
if let Err(rollback_error) = snapshot.restore(
state,
native_defaults_receipt.as_ref(),
native_patch_receipt.as_ref(),
) {
let _ = futures::executor::block_on(
state.proxy_service.close_pi_runtime_at_epoch(catalog_epoch),
);
@@ -253,7 +278,10 @@ impl PiCatalogCoordinator {
match futures::executor::block_on(
state
.proxy_service
.reconcile_pi_runtime_at_epoch(catalog_epoch),
.reconcile_pi_runtime_at_epoch_with_native_precondition(
catalog_epoch,
Some(&snapshot.native),
),
) {
Ok(_) => Err(authority_error(
PiCatalogAuthority::PreviousRestored,
@@ -275,12 +303,17 @@ impl PiCatalogCoordinator {
}
}
(Err(error), Ok(_)) => Err(error),
(Err(error), Err(reconcile_error)) => Err(authority_error(
PiCatalogAuthority::ProjectionPending,
format!(
"{error}; additionally failed to reconcile Pi admission: {reconcile_error}"
),
)),
(Err(error), Err(reconcile_error)) => {
let _ = futures::executor::block_on(
state.proxy_service.close_pi_runtime_at_epoch(catalog_epoch),
);
Err(authority_error(
PiCatalogAuthority::ProjectionPending,
format!(
"{error}; additionally failed to reconcile Pi admission: {reconcile_error}"
),
))
}
}
}
@@ -288,7 +321,13 @@ impl PiCatalogCoordinator {
state: &AppState,
models_path: &std::path::Path,
mut project: impl FnMut(&str, &str, &PiManagedProviderConfig) -> Result<Value, AppError>,
) -> Result<Option<PiNativeDefaultsReceipt>, AppError> {
) -> Result<
(
Option<PiNativeDefaultsReceipt>,
Option<PiProviderPatchReceipt>,
),
AppError,
> {
let providers = state.db.get_all_providers(PI_APP)?;
let manifest = state.db.get_pi_projection_manifest()?;
let claimed_keys = manifest
@@ -357,12 +396,13 @@ impl PiCatalogCoordinator {
});
}
let before_file = current_pi_provider_values(
let before_file = snapshot_pi_provider_values(
models_path,
plans.iter().map(|plan| plan.provider_key.clone()),
)?;
for plan in plans.iter().filter(|plan| plan.needs_claim) {
if before_file
.values
.get(&plan.provider_key)
.and_then(Option::as_ref)
.is_some()
@@ -442,25 +482,21 @@ impl PiCatalogCoordinator {
.iter()
.map(|plan| (plan.provider_key.clone(), Some(plan.projected.clone())))
.collect::<IndexMap<_, _>>();
if !patch.is_empty() {
if let Err(error) = apply_pi_provider_patch(models_path, &patch) {
let file_restored = apply_pi_provider_patch(models_path, &before_file);
let claim_error = compensate_portable_projection_ledger(
state,
&newly_claimed,
&released_claims,
error,
);
return Err(if file_restored.is_ok() {
claim_error
} else {
authority_error(
PiCatalogAuthority::ProjectionPending,
format!("{claim_error}; native catalog compensation failed"),
)
});
let native_patch_receipt = if patch.is_empty() {
None
} else {
match apply_pi_provider_patch_with_receipt(models_path, &before_file, &patch) {
Ok(receipt) => Some(receipt),
Err(error) => {
return Err(compensate_portable_projection_ledger(
state,
&newly_claimed,
&released_claims,
error,
));
}
}
}
};
let native_selection = match (
previous_defaults.default_provider.as_deref(),
@@ -505,7 +541,9 @@ impl PiCatalogCoordinator {
match Self::set_default(state, &current_provider, &model_id) {
Ok(result) => native_defaults_receipt = result.native_defaults_receipt,
Err(error) => {
let file_restored = apply_pi_provider_patch(models_path, &before_file);
let file_restored = native_patch_receipt
.as_ref()
.map_or(Ok(()), PiProviderPatchReceipt::rollback);
let claims_restored = compensate_portable_projection_ledger(
state,
&newly_claimed,
@@ -524,7 +562,7 @@ impl PiCatalogCoordinator {
}
}
}
Ok(native_defaults_receipt)
Ok((native_defaults_receipt, native_patch_receipt))
}
pub(crate) fn inspect_native(
@@ -584,8 +622,9 @@ impl PiCatalogCoordinator {
}
let models_path = get_pi_models_path()?;
let before_file = current_pi_provider_values(&models_path, [provider_key.to_string()])?;
let before_file = snapshot_pi_provider_values(&models_path, [provider_key.to_string()])?;
if before_file
.values
.get(provider_key)
.and_then(Option::as_ref)
.is_some()
@@ -608,15 +647,18 @@ impl PiCatalogCoordinator {
)?;
let projection = IndexMap::from([(provider_key.to_string(), Some(projected))]);
if let Err(projection_error) = apply_pi_provider_patch(&models_path, &projection) {
return Err(compensate_created_provider(
state,
&models_path,
&provider_id,
&before_file,
projection_error,
));
}
let native_patch_receipt =
match apply_pi_provider_patch_with_receipt(&models_path, &before_file, &projection) {
Ok(receipt) => receipt,
Err(projection_error) => {
return Err(compensate_created_provider(
state,
&provider_id,
None,
projection_error,
));
}
};
let no_selected_provider = previous_local.is_none() && previous_db.is_none();
let native_defaults_empty = previous_defaults.default_provider.is_none()
@@ -636,16 +678,15 @@ impl PiCatalogCoordinator {
Err(error) => {
return Err(compensate_created_provider(
state,
&models_path,
&provider_id,
&before_file,
Some(&native_patch_receipt),
error,
));
}
}
}
Ok(result)
Ok(result.with_native_patch_receipt(Some(native_patch_receipt)))
}
fn update(
@@ -668,7 +709,7 @@ impl PiCatalogCoordinator {
let was_current = previous_db.as_deref() == Some(&provider_id);
let models_path = get_pi_models_path()?;
let before_file =
current_pi_provider_values(&models_path, [projection.provider_key.clone()])?;
snapshot_pi_provider_values(&models_path, [projection.provider_key.clone()])?;
let projected = state.proxy_service.project_pi_provider_value(
&provider_id,
&projection.provider_key,
@@ -680,17 +721,20 @@ impl PiCatalogCoordinator {
.db
.update_pi_catalog_provider(&key, &ProviderRowUpdate::from_input(&input)?)?;
let patch = IndexMap::from([(projection.provider_key.clone(), Some(projected))]);
if let Err(error) = apply_pi_provider_patch(&models_path, &patch) {
return Err(compensate_existing_provider(
state,
&models_path,
&previous,
was_current,
&projection,
&before_file,
error,
));
}
let native_patch_receipt =
match apply_pi_provider_patch_with_receipt(&models_path, &before_file, &patch) {
Ok(receipt) => receipt,
Err(error) => {
return Err(compensate_existing_provider(
state,
&previous,
was_current,
&projection,
None,
error,
));
}
};
let mut result = success(Some(provider_id.clone()));
if previous_defaults.default_provider.as_deref() == Some(&projection.provider_key) {
if let Some(default_model) = previous_defaults.default_model.as_deref() {
@@ -708,11 +752,10 @@ impl PiCatalogCoordinator {
Err(error) => {
return Err(compensate_existing_provider(
state,
&models_path,
&previous,
was_current,
&projection,
&before_file,
Some(&native_patch_receipt),
error,
));
}
@@ -720,7 +763,7 @@ impl PiCatalogCoordinator {
}
}
}
Ok(result)
Ok(result.with_native_patch_receipt(Some(native_patch_receipt)))
}
fn delete(state: &AppState, provider_id: &str) -> Result<PiCatalogMutationResult, AppError> {
@@ -742,22 +785,26 @@ impl PiCatalogCoordinator {
let models_path = get_pi_models_path()?;
let before_file =
current_pi_provider_values(&models_path, [projection.provider_key.clone()])?;
snapshot_pi_provider_values(&models_path, [projection.provider_key.clone()])?;
let was_current = db_current.as_deref() == Some(provider_id);
state.db.delete_pi_catalog_provider(provider_id)?;
let patch = IndexMap::from([(projection.provider_key.clone(), None)]);
if let Err(error) = apply_pi_provider_patch(&models_path, &patch) {
return Err(compensate_existing_provider(
state,
&models_path,
&previous,
was_current,
&projection,
&before_file,
error,
));
}
Ok(success(Some(provider_id.to_string())))
let native_patch_receipt =
match apply_pi_provider_patch_with_receipt(&models_path, &before_file, &patch) {
Ok(receipt) => receipt,
Err(error) => {
return Err(compensate_existing_provider(
state,
&previous,
was_current,
&projection,
None,
error,
));
}
};
Ok(success(Some(provider_id.to_string()))
.with_native_patch_receipt(Some(native_patch_receipt)))
}
fn add_endpoint(
@@ -1002,13 +1049,12 @@ fn normalize_endpoint(value: &str) -> Result<String, AppError> {
fn compensate_created_provider(
state: &AppState,
models_path: &std::path::Path,
provider_id: &str,
before_file: &IndexMap<String, Option<Value>>,
native_patch_receipt: Option<&PiProviderPatchReceipt>,
cause: AppError,
) -> AppError {
let db_restored = state.db.delete_pi_catalog_provider(provider_id);
let file_restored = apply_pi_provider_patch(models_path, before_file);
let file_restored = native_patch_receipt.map_or(Ok(()), PiProviderPatchReceipt::rollback);
if db_restored.is_ok() && file_restored.is_ok() {
authority_error(PiCatalogAuthority::PreviousRestored, cause.to_string())
} else {
@@ -1060,20 +1106,18 @@ fn compensate_portable_projection_ledger(
}
}
#[allow(clippy::too_many_arguments)]
fn compensate_existing_provider(
state: &AppState,
models_path: &std::path::Path,
previous: &ProviderAggregate,
was_current: bool,
projection: &crate::database::PiProviderProjection,
before_file: &IndexMap<String, Option<Value>>,
native_patch_receipt: Option<&PiProviderPatchReceipt>,
cause: AppError,
) -> AppError {
let db_restored = state
.db
.restore_pi_catalog_provider(previous, was_current, Some(projection));
let file_restored = apply_pi_provider_patch(models_path, before_file);
let file_restored = native_patch_receipt.map_or(Ok(()), PiProviderPatchReceipt::rollback);
if db_restored.is_ok() && file_restored.is_ok() {
authority_error(PiCatalogAuthority::PreviousRestored, cause.to_string())
} else {
@@ -1119,7 +1163,6 @@ impl PiCatalogSnapshot {
aggregates,
projections,
db_current: state.db.get_current_provider(PI_APP)?,
models_path,
native,
})
}
@@ -1128,6 +1171,7 @@ impl PiCatalogSnapshot {
&self,
state: &AppState,
native_defaults_receipt: Option<&PiNativeDefaultsReceipt>,
native_patch_receipt: Option<&PiProviderPatchReceipt>,
) -> Result<(), AppError> {
let mut failures = Vec::new();
if let Err(error) = state.db.restore_pi_catalog_snapshot(
@@ -1137,8 +1181,10 @@ impl PiCatalogSnapshot {
) {
failures.push(format!("database={error}"));
}
if let Err(error) = restore_pi_provider_values(&self.models_path, &self.native) {
failures.push(format!("models={error}"));
if let Some(receipt) = native_patch_receipt {
if let Err(error) = receipt.rollback() {
failures.push(format!("models={error}"));
}
}
if let Some(receipt) = native_defaults_receipt {
if let Err(error) = receipt.rollback() {
@@ -1180,6 +1226,7 @@ fn success(provider_id: Option<String>) -> PiCatalogMutationResult {
authority: PiCatalogAuthority::Published,
provider_id,
native_defaults_receipt: None,
native_patch_receipt: None,
}
}
@@ -1188,6 +1235,11 @@ impl PiCatalogMutationResult {
self.native_defaults_receipt = receipt;
self
}
fn with_native_patch_receipt(mut self, receipt: Option<PiProviderPatchReceipt>) -> Self {
self.native_patch_receipt = receipt;
self
}
}
fn authority_error(authority: PiCatalogAuthority, message: impl std::fmt::Display) -> AppError {
src-tauri/src/services/proxy.rs
+596 -78
View File
@@ -611,6 +611,15 @@ impl ProxyService {
pub(crate) async fn reconcile_pi_runtime_at_epoch(
&self,
catalog_epoch: u64,
) -> Result<Vec<String>, AppError> {
self.reconcile_pi_runtime_at_epoch_with_native_precondition(catalog_epoch, None)
.await
}
pub(crate) async fn reconcile_pi_runtime_at_epoch_with_native_precondition(
&self,
catalog_epoch: u64,
expected_native: Option<&crate::pi_config::document::PiProviderValuesSnapshot>,
) -> Result<Vec<String>, AppError> {
#[cfg(test)]
if self.fail_next_pi_reconcile.swap(false, Ordering::AcqRel) {
@@ -642,11 +651,54 @@ impl ProxyService {
token,
app_config,
)?;
crate::pi_config::document::apply_pi_provider_patch(
&crate::pi_config::native::get_pi_models_path()?,
&build.projection_patch,
)?;
self.pi_runtime.publish(build.snapshot).await?;
let models_path = crate::pi_config::native::get_pi_models_path()?;
let native_receipt = if build.projection_patch.is_empty() {
None
} else {
let expected = match expected_native {
Some(expected) => {
crate::pi_config::document::PiProviderValuesSnapshot {
file_existed: expected.file_existed,
values: build
.projection_patch
.keys()
.map(|provider_key| {
expected
.values
.get(provider_key)
.cloned()
.map(|value| (provider_key.clone(), value))
.ok_or_else(|| {
AppError::Config(format!(
"Pi runtime provider key '{provider_key}' was not included in the catalog preflight"
))
})
})
.collect::<Result<_, _>>()?,
}
}
None => self.preflight_pi_owned_projection_at(
&models_path,
&build.projection_patch,
)?,
};
Some(
crate::pi_config::document::apply_pi_provider_patch_with_receipt(
&models_path,
&expected,
&build.projection_patch,
)?,
)
};
if let Err(error) = self.pi_runtime.publish(build.snapshot).await {
let rollback = native_receipt.as_ref().map_or(
Ok(()),
crate::pi_config::document::PiProviderPatchReceipt::rollback,
);
return Err(AppError::Config(format!(
"failed to publish Pi runtime after native projection: {error}; native rollback={rollback:?}"
)));
}
Ok(build.direct_only_provider_ids)
}
@@ -692,12 +744,126 @@ impl ProxyService {
fn restore_pi_direct_projection_at(
&self,
models_path: &std::path::Path,
) -> Result<(), AppError> {
let patch = direct_pi_projection_patch(self.db.as_ref())?;
crate::pi_config::document::apply_pi_provider_patch(models_path, &patch)
) -> Result<crate::pi_config::document::PiProviderPatchReceipt, AppError> {
if self
.pi_listener
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.is_none()
{
return self.confirm_pi_projection_is_already_direct_at(models_path);
}
let gateway = self.current_pi_gateway_projection()?;
self.restore_pi_direct_projection_at_with_expected_gateway(models_path, &gateway)
}
fn restore_pi_direct_projection(&self) -> Result<(), AppError> {
fn confirm_pi_projection_is_already_direct_at(
&self,
models_path: &std::path::Path,
) -> Result<crate::pi_config::document::PiProviderPatchReceipt, AppError> {
let direct = direct_pi_projection_patch(self.db.as_ref())?;
let before = crate::pi_config::document::snapshot_pi_provider_values(
models_path,
direct.keys().cloned(),
)?;
if let Some(provider_key) = direct.iter().find_map(|(provider_key, expected)| {
(before.values.get(provider_key) != Some(expected)).then_some(provider_key)
}) {
return Err(AppError::Conflict(format!(
"Pi listener is unavailable and provider key '{provider_key}' is not already in its direct projection"
)));
}
crate::pi_config::document::apply_pi_provider_patch_with_receipt(
models_path,
&before,
&direct,
)
}
fn current_pi_gateway_projection(
&self,
) -> Result<indexmap::IndexMap<String, Option<Value>>, AppError> {
let listener = self
.pi_listener
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
.ok_or_else(|| {
AppError::Conflict(
"cannot restore Pi direct projection without its active listener identity"
.to_string(),
)
})?;
let gateway_token = crate::settings::get_pi_gateway_token().or_else(|settings_error| {
self.pi_runtime
.retained_gateway_token(listener.server_generation)
.ok_or(settings_error)
})?;
let gateway = build_pi_runtime(
self.db.as_ref(),
listener.server_generation,
0,
&listener.gateway_origin,
gateway_token,
crate::settings::get_pi_app_proxy_config(),
)?
.projection_patch;
Ok(gateway)
}
fn preflight_pi_owned_projection_at(
&self,
models_path: &std::path::Path,
expected_gateway: &indexmap::IndexMap<String, Option<Value>>,
) -> Result<crate::pi_config::document::PiProviderValuesSnapshot, AppError> {
let direct = direct_pi_projection_patch(self.db.as_ref())?;
let direct_keys = direct.keys().collect::<std::collections::BTreeSet<_>>();
let gateway_keys = expected_gateway
.keys()
.collect::<std::collections::BTreeSet<_>>();
if direct_keys != gateway_keys {
return Err(AppError::Conflict(
"Pi direct and gateway projections cover different exact-key ownership".to_string(),
));
}
let before = crate::pi_config::document::snapshot_pi_provider_values(
models_path,
direct.keys().cloned(),
)?;
for (provider_key, direct_value) in &direct {
let observed = before
.values
.get(provider_key)
.expect("every owned key was preflighted");
let gateway_value = expected_gateway
.get(provider_key)
.expect("gateway ownership keys were checked");
if observed != direct_value && observed != gateway_value {
return Err(AppError::Conflict(format!(
"Pi provider key '{provider_key}' changed outside CC Switch"
)));
}
}
Ok(before)
}
fn restore_pi_direct_projection_at_with_expected_gateway(
&self,
models_path: &std::path::Path,
expected_gateway: &indexmap::IndexMap<String, Option<Value>>,
) -> Result<crate::pi_config::document::PiProviderPatchReceipt, AppError> {
let direct = direct_pi_projection_patch(self.db.as_ref())?;
let before = self.preflight_pi_owned_projection_at(models_path, expected_gateway)?;
crate::pi_config::document::apply_pi_provider_patch_with_receipt(
models_path,
&before,
&direct,
)
}
fn restore_pi_direct_projection(
&self,
) -> Result<crate::pi_config::document::PiProviderPatchReceipt, AppError> {
let models_path = crate::pi_config::native::get_pi_models_path()?;
self.restore_pi_direct_projection_at(&models_path)
}
@@ -751,19 +917,21 @@ impl ProxyService {
if !existing.pi_takeover_enabled {
crate::settings::update_settings(next)?;
if let Err(error) =
crate::pi_config::document::apply_pi_provider_patch(&new_models_path, &direct_patch)
{
let settings_restored = crate::settings::update_settings(existing.clone()).is_ok();
let new_native_restored = crate::pi_config::document::restore_pi_provider_values(
let native_receipt =
match crate::pi_config::document::apply_pi_provider_patch_with_receipt(
&new_models_path,
&new_native_before,
)
.is_ok();
return Err(AppError::Config(format!(
"failed to publish managed Pi providers in the new directory: {error}; rollback: settings={settings_restored}, native={new_native_restored}"
&direct_patch,
) {
Ok(receipt) => receipt,
Err(error) => {
let settings_restored =
crate::settings::update_settings(existing.clone()).is_ok();
return Err(AppError::Config(format!(
"failed to publish managed Pi providers in the new directory: {error}; rollback: settings={settings_restored}"
)));
}
}
};
if let Err(error) =
PromptService::reconcile_pi_native_under_guard(self.db.as_ref(), &prompt_guard)
{
@@ -772,11 +940,7 @@ impl ProxyService {
.db
.save_prompt_selection(AppType::Pi.as_str(), &previous_prompts)
.is_ok();
let new_native_restored = crate::pi_config::document::restore_pi_provider_values(
&new_models_path,
&new_native_before,
)
.is_ok();
let new_native_restored = native_receipt.rollback().is_ok();
return Err(AppError::Config(format!(
"failed to reconcile Pi prompts in the new directory: {error}; rollback: settings={settings_restored}, prompts={prompts_restored}, native={new_native_restored}"
)));
@@ -789,11 +953,7 @@ impl ProxyService {
.db
.save_prompt_selection(AppType::Pi.as_str(), &previous_prompts)
.is_ok();
let new_native_restored = crate::pi_config::document::restore_pi_provider_values(
&new_models_path,
&new_native_before,
)
.is_ok();
let new_native_restored = native_receipt.rollback().is_ok();
let skills_restored = settings_restored
&& crate::services::skill_deployment::PiSkillDeploymentService::reconcile_all(
&self.db,
@@ -810,13 +970,22 @@ impl ProxyService {
// process that already loaded the old gateway projection therefore
// cannot enter a catalog whose directory ownership is in flight.
let epoch = self.pi_runtime.begin_mutation().await;
if let Err(error) = self.restore_pi_direct_projection_at(&old_models_path) {
let _ = self.pi_runtime.republish_current(epoch).await;
return Err(error);
}
let old_direct_receipt = match self.restore_pi_direct_projection_at(&old_models_path) {
Ok(receipt) => receipt,
Err(error) => {
let _ = self.pi_runtime.republish_current(epoch).await;
return Err(error);
}
};
if let Err(error) = crate::settings::update_settings(next) {
let runtime_restored = self.reconcile_pi_runtime_at_epoch(epoch).await;
let expected_old_direct = old_direct_receipt.attempted_snapshot();
let runtime_restored = self
.reconcile_pi_runtime_at_epoch_with_native_precondition(
epoch,
Some(&expected_old_direct),
)
.await;
return Err(AppError::Config(if runtime_restored.is_ok() {
format!(
"failed to save the new Pi directory; the previous gateway projection was restored: {error}"
@@ -836,19 +1005,24 @@ impl ProxyService {
.db
.save_prompt_selection(AppType::Pi.as_str(), &previous_prompts)
.is_ok();
let new_native_restored = crate::pi_config::document::restore_pi_provider_values(
&new_models_path,
&new_native_before,
)
.is_ok();
let skills_restored = settings_restored
&& crate::services::skill_deployment::PiSkillDeploymentService::reconcile_all(
&self.db,
)
.is_ok();
let runtime_restored = self.reconcile_pi_runtime_at_epoch(epoch).await.is_ok();
let runtime_restored = if settings_restored {
let expected_old_direct = old_direct_receipt.attempted_snapshot();
self.reconcile_pi_runtime_at_epoch_with_native_precondition(
epoch,
Some(&expected_old_direct),
)
.await
.is_ok()
} else {
self.pi_runtime.republish_current(epoch).await.is_ok()
};
return Err(AppError::Config(format!(
"failed to reconcile Pi prompts in the new directory: {error}; rollback: settings={settings_restored}, prompts={prompts_restored}, native={new_native_restored}, skills={skills_restored}, gateway={runtime_restored}"
"failed to reconcile Pi prompts in the new directory: {error}; rollback: settings={settings_restored}, prompts={prompts_restored}, skills={skills_restored}, gateway={runtime_restored}"
)));
}
@@ -860,45 +1034,54 @@ impl ProxyService {
.db
.save_prompt_selection(AppType::Pi.as_str(), &previous_prompts)
.is_ok();
let new_native_restored = crate::pi_config::document::restore_pi_provider_values(
&new_models_path,
&new_native_before,
)
.is_ok();
let skills_restored = settings_restored
&& crate::services::skill_deployment::PiSkillDeploymentService::reconcile_all(
&self.db,
)
.is_ok();
let runtime_restored = self.reconcile_pi_runtime_at_epoch(epoch).await.is_ok();
let runtime_restored = if settings_restored {
let expected_old_direct = old_direct_receipt.attempted_snapshot();
self.reconcile_pi_runtime_at_epoch_with_native_precondition(
epoch,
Some(&expected_old_direct),
)
.await
.is_ok()
} else {
self.pi_runtime.republish_current(epoch).await.is_ok()
};
return Err(AppError::Config(format!(
"failed to reconcile Pi Skills in the new directory: {error}; rollback: settings={settings_restored}, prompts={prompts_restored}, native={new_native_restored}, skills={skills_restored}, gateway={runtime_restored}"
"failed to reconcile Pi Skills in the new directory: {error}; rollback: settings={settings_restored}, prompts={prompts_restored}, skills={skills_restored}, gateway={runtime_restored}"
)));
}
if let Err(error) = self.reconcile_pi_runtime_at_epoch(epoch).await {
if let Err(error) = self
.reconcile_pi_runtime_at_epoch_with_native_precondition(epoch, Some(&new_native_before))
.await
{
let settings_restored = crate::settings::update_settings(existing.clone()).is_ok();
let prompts_restored = self
.db
.save_prompt_selection(AppType::Pi.as_str(), &previous_prompts)
.is_ok();
let new_native_restored = crate::pi_config::document::restore_pi_provider_values(
&new_models_path,
&new_native_before,
)
.is_ok();
let skills_restored = settings_restored
&& crate::services::skill_deployment::PiSkillDeploymentService::reconcile_all(
&self.db,
)
.is_ok();
let old_gateway_restored = if settings_restored {
self.reconcile_pi_runtime_at_epoch(epoch).await.is_ok()
let expected_old_direct = old_direct_receipt.attempted_snapshot();
self.reconcile_pi_runtime_at_epoch_with_native_precondition(
epoch,
Some(&expected_old_direct),
)
.await
.is_ok()
} else {
self.pi_runtime.republish_current(epoch).await.is_ok()
};
return Err(AppError::Config(format!(
"failed to publish Pi in the new native directory: {error}; rollback: native={new_native_restored}, settings={settings_restored}, prompts={prompts_restored}, skills={skills_restored}, old_gateway={old_gateway_restored}"
"failed to publish Pi in the new native directory: {error}; rollback: settings={settings_restored}, prompts={prompts_restored}, skills={skills_restored}, old_gateway={old_gateway_restored}"
)));
}
@@ -910,11 +1093,26 @@ impl ProxyService {
return Ok(());
}
let epoch = self.pi_runtime.begin_mutation().await;
if let Err(error) = self.restore_pi_direct_projection() {
let _ = self.pi_runtime.republish_current(epoch).await;
return Err(error);
let direct_receipt = match self.restore_pi_direct_projection() {
Ok(receipt) => receipt,
Err(error) => {
let _ = self.pi_runtime.republish_current(epoch).await;
return Err(error);
}
};
if let Err(error) = self.pi_runtime.close(epoch).await {
let expected_direct = direct_receipt.attempted_snapshot();
let rollback = self
.reconcile_pi_runtime_at_epoch_with_native_precondition(
epoch,
Some(&expected_direct),
)
.await;
return Err(AppError::Config(format!(
"failed to close Pi admission after direct projection: {error}; gateway rollback={rollback:?}"
)));
}
self.pi_runtime.close(epoch).await
Ok(())
}
/// Enter the portable-import boundary while the caller holds Pi's switch
@@ -958,16 +1156,32 @@ impl ProxyService {
pub(crate) async fn rotate_pi_gateway_token(&self) -> Result<(), AppError> {
let _guard = self.switch_locks.lock_for_app(AppType::Pi.as_str()).await;
let previous = crate::settings::get_or_create_pi_gateway_token()?;
let takeover_enabled = crate::settings::pi_takeover_enabled();
let native_before = if takeover_enabled {
let models_path = crate::pi_config::native::get_pi_models_path()?;
let gateway = self.current_pi_gateway_projection()?;
Some(self.preflight_pi_owned_projection_at(&models_path, &gateway)?)
} else {
None
};
crate::settings::reset_pi_gateway_token()?;
if !crate::settings::pi_takeover_enabled() {
if !takeover_enabled {
return Ok(());
}
let native_before = native_before.expect("enabled takeover has a native preflight");
let epoch = self.pi_runtime.begin_mutation().await;
if let Err(error) = self.reconcile_pi_runtime_at_epoch(epoch).await {
if let Err(error) = self
.reconcile_pi_runtime_at_epoch_with_native_precondition(epoch, Some(&native_before))
.await
{
let token_restored = crate::settings::replace_pi_gateway_token(previous);
let runtime_restored = if token_restored.is_ok() {
let rollback_epoch = self.pi_runtime.begin_mutation().await;
self.reconcile_pi_runtime_at_epoch(rollback_epoch).await
self.reconcile_pi_runtime_at_epoch_with_native_precondition(
rollback_epoch,
Some(&native_before),
)
.await
} else {
Err(AppError::Config(
"failed to restore the previous Pi gateway credential".to_string(),
@@ -1561,7 +1775,7 @@ impl ProxyService {
let direct_projection = self.restore_pi_direct_projection();
match direct_projection {
Ok(()) => {
Ok(_) => {
let epoch = self.pi_runtime.begin_mutation().await;
let admission_closed = self.pi_runtime.close(epoch).await;
return Err(if admission_closed.is_ok() {
@@ -1597,16 +1811,25 @@ impl ProxyService {
return Ok(());
}
let epoch = self.pi_runtime.begin_mutation().await;
if let Err(error) = self.restore_pi_direct_projection() {
let _ = self.pi_runtime.republish_current(epoch).await;
return Err(format!(
"failed to restore Pi's direct native projection: {error}"
));
}
let direct_receipt = match self.restore_pi_direct_projection() {
Ok(receipt) => receipt,
Err(error) => {
let _ = self.pi_runtime.republish_current(epoch).await;
return Err(format!(
"failed to restore Pi's direct native projection: {error}"
));
}
};
if let Err(error) = crate::settings::set_pi_takeover_enabled(false) {
// Desired state did not change; rebuild the gateway projection and
// restore admission so the native file cannot be left lying.
let _ = self.reconcile_pi_runtime_at_epoch(epoch).await;
let expected_direct = direct_receipt.attempted_snapshot();
let _ = self
.reconcile_pi_runtime_at_epoch_with_native_precondition(
epoch,
Some(&expected_direct),
)
.await;
return Err(format!("failed to persist Pi takeover disable: {error}"));
}
@@ -4164,16 +4387,17 @@ mod tests {
"apiKey": "native-key",
"models": [{"id": "native-model"}]
});
let expected_gateway = serde_json::json!({
"api": "openai-responses",
"baseUrl": "http://127.0.0.1:15721/pi/route",
"apiKey": "gateway-token",
"models": [{"id": "model-a"}]
});
std::fs::write(
&path,
serde_json::to_vec_pretty(&serde_json::json!({
"providers": {
"managed-pi": {
"api": "openai-responses",
"baseUrl": "http://127.0.0.1:15721/pi/route",
"apiKey": "gateway-token",
"models": [{"id": "model-a"}]
},
"managed-pi": expected_gateway.clone(),
"native": unowned.clone()
}
}))
@@ -4183,7 +4407,10 @@ mod tests {
state
.proxy_service
.restore_pi_direct_projection_at(&path)
.restore_pi_direct_projection_at_with_expected_gateway(
&path,
&indexmap::IndexMap::from([("managed-pi".to_string(), Some(expected_gateway))]),
)
.expect("restore direct Pi projection");
let restored: Value =
@@ -4452,6 +4679,91 @@ mod tests {
);
}
#[tokio::test]
#[serial]
async fn changing_pi_directory_does_not_overwrite_an_exact_key_changed_after_preflight() {
let home = TempHome::new();
crate::settings::reload_settings().expect("reload isolated settings");
let old_dir = home.dir.path().join("old-pi");
let new_dir = home.dir.path().join("new-pi");
std::fs::create_dir_all(&old_dir).expect("old Pi directory");
std::fs::create_dir_all(&new_dir).expect("new Pi directory");
let direct = json!({
"name": "Managed Pi",
"api": "openai-responses",
"baseUrl": "https://managed.example/v1",
"apiKey": "managed-key",
"models": [{"id": "model-a", "name": "Model A"}]
});
let external = json!({
"name": "External Pi",
"api": "openai-responses",
"baseUrl": "https://external.example/v1",
"apiKey": "external-key",
"models": [{"id": "external-model", "name": "External"}]
});
let new_models_path = new_dir.join("models.json");
std::fs::write(
&new_models_path,
serde_json::to_vec(&json!({"providers": {}})).expect("serialize direct document"),
)
.expect("seed direct document");
let mut settings = crate::settings::get_settings();
settings.pi_config_dir = Some(old_dir.to_string_lossy().into_owned());
settings.pi_takeover_enabled = false;
crate::settings::update_settings(settings).expect("old Pi settings");
let db = Arc::new(Database::memory().expect("database"));
db.create_pi_catalog_provider(
NewProviderAggregate::from_input(
AppType::Pi.as_str(),
ProviderMutationInput {
id: "managed-pi".to_string(),
name: "Managed Pi".to_string(),
settings_config: direct,
website_url: None,
category: None,
created_at: None,
sort_index: Some(0),
notes: None,
meta: None,
icon: Some("pi".to_string()),
icon_color: None,
in_failover_queue: false,
},
)
.expect("aggregate"),
"managed-pi",
)
.expect("managed provider");
crate::pi_config::shared_file::replace_before_next_compare_exchange(
&new_models_path,
&serde_json::to_vec(&json!({"providers": {"managed-pi": external.clone()}}))
.expect("serialize external document"),
);
let service = ProxyService::new(db);
let switch_guard = service.lock_switch_for_app(AppType::Pi.as_str()).await;
let existing = crate::settings::get_settings();
let mut next = existing.clone();
next.pi_config_dir = Some(new_dir.to_string_lossy().into_owned());
let error = service
.replace_settings_with_pi_directory_boundary_under_lock(&switch_guard, &existing, next)
.await
.expect_err("the post-preflight native edit must win");
assert!(error.to_string().contains("changed since"));
assert_eq!(
crate::settings::get_settings().pi_config_dir,
existing.pi_config_dir
);
let live: Value = serde_json::from_slice(
&std::fs::read(&new_models_path).expect("read external native document"),
)
.expect("parse external native document");
assert_eq!(live.pointer("/providers/managed-pi"), Some(&external));
}
#[tokio::test]
#[serial]
async fn changing_pi_directory_without_takeover_reconciles_missing_agents_truth() {
@@ -4544,6 +4856,143 @@ mod tests {
);
}
#[tokio::test]
#[serial]
async fn pi_gateway_returns_the_last_real_retryable_upstream_response_when_no_later_send_occurs(
) {
let home = TempHome::new();
crate::settings::reload_settings().expect("reload isolated settings");
let pi_dir = home.dir.path().join("pi");
std::fs::create_dir_all(&pi_dir).expect("Pi directory");
let mut settings = crate::settings::get_settings();
settings.pi_config_dir = Some(pi_dir.to_string_lossy().into_owned());
settings.pi_takeover_enabled = false;
crate::settings::update_settings(settings).expect("Pi settings");
let upstream_hits = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hits = upstream_hits.clone();
let upstream = axum::Router::new().fallback(axum::routing::any(move || {
let hits = hits.clone();
async move {
hits.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let mut response = axum::response::Response::builder()
.status(http::StatusCode::TOO_MANY_REQUESTS)
.header("x-upstream-marker", "real-429")
.header(http::header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(
r#"{"error":{"message":"upstream rate limit"}}"#,
))
.expect("upstream response");
response
.headers_mut()
.insert("retry-after", http::HeaderValue::from_static("17"));
response
}
}));
let upstream_listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
.await
.expect("upstream listener");
let upstream_address = upstream_listener.local_addr().expect("upstream address");
let upstream_task = tokio::spawn(async move {
axum::serve(upstream_listener, upstream)
.await
.expect("upstream server");
});
let db = Arc::new(Database::memory().expect("database"));
use_ephemeral_proxy_port(&db).await;
let direct = json!({
"name": "Retryable upstream",
"api": "openai-responses",
"baseUrl": format!("http://{upstream_address}/v1"),
"apiKey": "upstream-key",
"models": [{"id": "model-a", "name": "Model A"}]
});
db.create_pi_catalog_provider(
NewProviderAggregate::from_input(
AppType::Pi.as_str(),
ProviderMutationInput {
id: "retryable-upstream".to_string(),
name: "Retryable upstream".to_string(),
settings_config: direct.clone(),
website_url: None,
category: None,
created_at: None,
sort_index: Some(0),
notes: None,
meta: None,
icon: Some("pi".to_string()),
icon_color: None,
in_failover_queue: false,
},
)
.expect("aggregate"),
"retryable-upstream",
)
.expect("managed provider");
std::fs::write(
pi_dir.join("models.json"),
serde_json::to_vec(&json!({
"providers": {"retryable-upstream": direct}
}))
.expect("serialize native catalog"),
)
.expect("native catalog");
let service = ProxyService::new(db);
service
.set_takeover_for_app(AppType::Pi.as_str(), true)
.await
.expect("enable Pi gateway");
let projected: Value = serde_json::from_slice(
&std::fs::read(pi_dir.join("models.json")).expect("projected catalog"),
)
.expect("parse projected catalog");
let base_url = projected
.pointer("/providers/retryable-upstream/baseUrl")
.and_then(Value::as_str)
.expect("gateway base URL");
let gateway_token = projected
.pointer("/providers/retryable-upstream/apiKey")
.and_then(Value::as_str)
.expect("gateway token");
let response = reqwest::Client::new()
.post(format!("{base_url}/responses"))
.bearer_auth(gateway_token)
.json(&json!({"model": "model-a", "input": "hello"}))
.send()
.await
.expect("gateway response");
let status = response.status();
let marker = response
.headers()
.get("x-upstream-marker")
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let retry_after = response
.headers()
.get("retry-after")
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let body = response.text().await.expect("upstream body");
upstream_task.abort();
assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS);
assert_eq!(marker.as_deref(), Some("real-429"));
assert_eq!(retry_after.as_deref(), Some("17"));
assert_eq!(body, r#"{"error":{"message":"upstream rate limit"}}"#);
assert_eq!(
upstream_hits.load(std::sync::atomic::Ordering::SeqCst),
1,
"local candidate exhaustion must not invent a phantom retry"
);
service
.set_takeover_for_app(AppType::Pi.as_str(), false)
.await
.expect("disable Pi gateway");
}
#[tokio::test]
#[serial]
async fn failed_listener_rebind_recovers_pi_on_the_previous_listener() {
@@ -4716,6 +5165,75 @@ mod tests {
.expect("explicit disable clears desired state");
}
#[tokio::test]
#[serial]
async fn listener_absence_does_not_authorize_overwriting_a_non_direct_native_key() {
let home = TempHome::new();
crate::settings::reload_settings().expect("reload isolated settings");
let pi_dir = home.dir.path().join("pi");
std::fs::create_dir_all(&pi_dir).expect("Pi directory");
let mut settings = crate::settings::get_settings();
settings.pi_config_dir = Some(pi_dir.to_string_lossy().into_owned());
settings.pi_takeover_enabled = true;
crate::settings::update_settings(settings).expect("Pi settings");
let db = Arc::new(Database::memory().expect("database"));
let direct = json!({
"name": "Managed Pi",
"api": "openai-responses",
"baseUrl": "https://managed.example/v1",
"apiKey": "managed-key",
"models": [{"id": "model-a", "name": "Model A"}]
});
db.create_pi_catalog_provider(
NewProviderAggregate::from_input(
AppType::Pi.as_str(),
ProviderMutationInput {
id: "managed-pi".to_string(),
name: "Managed Pi".to_string(),
settings_config: direct,
website_url: None,
category: None,
created_at: None,
sort_index: Some(0),
notes: None,
meta: None,
icon: Some("pi".to_string()),
icon_color: None,
in_failover_queue: false,
},
)
.expect("aggregate"),
"managed-pi",
)
.expect("managed provider");
let external = json!({
"name": "External edit",
"api": "openai-responses",
"baseUrl": "https://external.example/v1",
"apiKey": "external-key",
"models": [{"id": "external-model"}]
});
let native_bytes = serde_json::to_vec(&json!({
"providers": {"managed-pi": external}
}))
.expect("serialize external catalog");
let models_path = pi_dir.join("models.json");
std::fs::write(&models_path, &native_bytes).expect("external native catalog");
let service = ProxyService::new(db);
let error = service
.set_takeover_for_app(AppType::Pi.as_str(), false)
.await
.expect_err("no listener is not ownership evidence");
assert!(error.contains("not already in its direct projection"));
assert!(crate::settings::pi_takeover_enabled());
assert_eq!(
std::fs::read(&models_path).expect("external catalog remains"),
native_bytes
);
}
fn seed_codex_model_template() {
let codex_dir = crate::codex_config::get_codex_config_dir();
std::fs::create_dir_all(&codex_dir).expect("create codex dir");