fix(gateway): add provider key concurrency and cache affinity modes

This commit is contained in:
ZheFox
2026-09-01 08:06:58 +08:00
parent 36daba7a34
commit 9631b229b3
26 changed files with 1176 additions and 154 deletions
@@ -38,7 +38,17 @@ impl ProviderOAuthAdapter for AntigravityProviderOAuthAdapter {
state: &str,
code_challenge: Option<&str>,
) -> Result<crate::core::OAuthAuthorizeResponse, crate::core::OAuthError> {
self.inner.build_authorize_url(ctx, state, code_challenge)
let mut response = self.inner.build_authorize_url(ctx, state, code_challenge)?;
let mut url = url::Url::parse(&response.authorize_url).map_err(|_| {
crate::core::OAuthError::invalid_request("authorize_url must be absolute")
})?;
{
let mut query = url.query_pairs_mut();
query.append_pair("access_type", "offline");
query.append_pair("prompt", "consent");
}
response.authorize_url = url.to_string();
Ok(response)
}
async fn exchange_code(
@@ -112,20 +122,8 @@ mod tests {
struct UnusedExecutor;
#[async_trait]
impl OAuthHttpExecutor for UnusedExecutor {
async fn execute(
&self,
_request: OAuthHttpRequest,
) -> Result<OAuthHttpResponse, crate::core::OAuthError> {
unreachable!("metadata probe should not execute network requests")
}
}
#[tokio::test]
async fn antigravity_probe_marks_forbidden_metadata_invalid() {
let adapter = AntigravityProviderOAuthAdapter::default();
let ctx = ProviderOAuthTransportContext {
fn transport_context() -> ProviderOAuthTransportContext {
ProviderOAuthTransportContext {
provider_id: String::new(),
provider_type: "antigravity".to_string(),
endpoint_id: None,
@@ -137,7 +135,46 @@ mod tests {
endpoint_config: None,
key_config: None,
network: crate::network::OAuthNetworkContext::provider_operation(None),
};
}
}
#[async_trait]
impl OAuthHttpExecutor for UnusedExecutor {
async fn execute(
&self,
_request: OAuthHttpRequest,
) -> Result<OAuthHttpResponse, crate::core::OAuthError> {
unreachable!("metadata probe should not execute network requests")
}
}
#[test]
fn antigravity_authorize_requests_offline_refresh_token() {
let adapter = AntigravityProviderOAuthAdapter::default();
let response = adapter
.build_authorize_url(&transport_context(), "state-1", Some("challenge-1"))
.expect("authorize url should build");
let url = url::Url::parse(&response.authorize_url).expect("authorize url should parse");
let query = url.query_pairs().collect::<BTreeMap<_, _>>();
assert_eq!(
query.get("access_type").map(|value| value.as_ref()),
Some("offline")
);
assert_eq!(
query.get("prompt").map(|value| value.as_ref()),
Some("consent")
);
assert_eq!(
query.get("code_challenge").map(|value| value.as_ref()),
Some("challenge-1")
);
}
#[tokio::test]
async fn antigravity_probe_marks_forbidden_metadata_invalid() {
let adapter = AntigravityProviderOAuthAdapter::default();
let ctx = transport_context();
let account = ProviderOAuthAccount {
provider_type: "antigravity".to_string(),
access_token: "access-token".to_string(),
@@ -374,10 +374,9 @@ impl ProviderOAuthAdapter for GenericProviderOAuthAdapter {
ctx: &ProviderOAuthTransportContext,
account: &ProviderOAuthAccount,
) -> Result<ProviderOAuthTokenSet, OAuthError> {
let refresh_token = account
.auth_config
.get("refresh_token")
.and_then(Value::as_str)
let refresh_token = ["refresh_token", "refreshToken"]
.iter()
.find_map(|field| account.auth_config.get(*field).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| OAuthError::invalid_request("auth_config missing refresh_token"))?;