mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-10 11:19:50 +08:00
Merge pull request #688 from AAEE86/feat/agent-identity-support
feat(codex): support agent identity accounts
This commit is contained in:
ZheFox
@@ -0,0 +1,784 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::{
|
||||
engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
|
||||
Engine as _,
|
||||
};
|
||||
use chrono::{DateTime, SecondsFormat, Utc};
|
||||
use crypto_box::SecretKey as Curve25519SecretKey;
|
||||
use ed25519_dalek::{pkcs8::DecodePrivateKey, Signer, SigningKey};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Map, Value};
|
||||
use sha2::{Digest, Sha512};
|
||||
use url::Url;
|
||||
|
||||
use super::oauth_refresh::{
|
||||
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthRefreshAdapter,
|
||||
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use super::snapshot::GatewayProviderTransportSnapshot;
|
||||
|
||||
pub const CODEX_AGENT_IDENTITY_AUTH_MODE: &str = "agentIdentity";
|
||||
pub const CODEX_AGENT_IDENTITY_PROVIDER_TYPE: &str = "codex";
|
||||
pub const CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE: &str = "codex_agent_identity";
|
||||
pub const CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID: &str =
|
||||
"codex:agent-identity-task-register";
|
||||
const CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL: &str = "https://auth.openai.com/api/accounts";
|
||||
const AUTHORIZATION_HEADER: &str = "authorization";
|
||||
const ASSERTION_PREFIX: &str = "AgentAssertion ";
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AgentIdentityCredentials {
|
||||
runtime_id: String,
|
||||
signing_key: SigningKey,
|
||||
task_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AgentTaskRegistrationResponse {
|
||||
#[serde(default)]
|
||||
task_id: Option<String>,
|
||||
#[serde(default, rename = "taskId")]
|
||||
task_id_camel: Option<String>,
|
||||
#[serde(default)]
|
||||
encrypted_task_id: Option<String>,
|
||||
#[serde(default, rename = "encryptedTaskId")]
|
||||
encrypted_task_id_camel: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CodexAgentIdentityRefreshAdapter {
|
||||
auth_api_base_url: String,
|
||||
}
|
||||
|
||||
impl Default for CodexAgentIdentityRefreshAdapter {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
auth_api_base_url: CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexAgentIdentityRefreshAdapter {
|
||||
pub fn with_auth_api_base_url_for_tests(mut self, base_url: impl Into<String>) -> Self {
|
||||
self.auth_api_base_url = base_url.into();
|
||||
self
|
||||
}
|
||||
|
||||
fn config_from_transport(transport: &GatewayProviderTransportSnapshot) -> Option<Value> {
|
||||
transport
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(|value| serde_json::from_str::<Value>(value).ok())
|
||||
}
|
||||
|
||||
fn config_from_entry(entry: &CachedOAuthEntry) -> Option<Value> {
|
||||
entry
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE)
|
||||
.then(|| entry.metadata.clone())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
fn resolve_from_config(config: &Value) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
let credentials = agent_identity_credentials(config).ok()?;
|
||||
let task_id = credentials.task_id.as_deref()?;
|
||||
let value = build_agent_assertion(&credentials, task_id, Utc::now()).ok()?;
|
||||
Some(LocalResolvedOAuthRequestAuth::Header {
|
||||
name: AUTHORIZATION_HEADER.to_string(),
|
||||
value,
|
||||
})
|
||||
}
|
||||
|
||||
fn preferred_config(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<Value> {
|
||||
Self::config_from_transport(transport).or_else(|| entry.and_then(Self::config_from_entry))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_codex_agent_identity_auth_config_value(config: &Value) -> bool {
|
||||
let Some(root) = config.as_object() else {
|
||||
return false;
|
||||
};
|
||||
let nested = agent_identity_nested_object(root);
|
||||
let mode = string_from_maps(root, nested, &["auth_mode", "authMode"]);
|
||||
mode.as_deref()
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case(CODEX_AGENT_IDENTITY_AUTH_MODE))
|
||||
|| (nested.is_some()
|
||||
&& string_from_maps(root, nested, &["agent_runtime_id", "agentRuntimeId"]).is_some()
|
||||
&& string_from_maps(root, nested, &["agent_private_key", "agentPrivateKey"]).is_some())
|
||||
}
|
||||
|
||||
pub fn is_codex_agent_identity_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||
transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(CODEX_AGENT_IDENTITY_PROVIDER_TYPE)
|
||||
&& transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
|
||||
&& CodexAgentIdentityRefreshAdapter::config_from_transport(transport)
|
||||
.as_ref()
|
||||
.is_some_and(is_codex_agent_identity_auth_config_value)
|
||||
}
|
||||
|
||||
pub fn is_codex_agent_identity_cached_entry(entry: &CachedOAuthEntry) -> bool {
|
||||
entry
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE)
|
||||
}
|
||||
|
||||
pub fn validate_codex_agent_identity_auth_config(config: &Value) -> Result<(), String> {
|
||||
agent_identity_credentials(config).map(|_| ())
|
||||
}
|
||||
|
||||
/// Returns whether an upstream response proves that the registered Agent Identity task is no
|
||||
/// longer usable. Only this condition should trigger task registration again; an arbitrary 401
|
||||
/// can instead mean that the account itself has lost access.
|
||||
pub fn is_codex_agent_identity_invalid_task_response(
|
||||
status_code: u16,
|
||||
response_text: Option<&str>,
|
||||
) -> bool {
|
||||
if status_code != 401 {
|
||||
return false;
|
||||
}
|
||||
let Some(response_text) = response_text else {
|
||||
return false;
|
||||
};
|
||||
let lower = response_text.to_ascii_lowercase();
|
||||
let compact = lower
|
||||
.chars()
|
||||
.filter(|character| !character.is_ascii_whitespace())
|
||||
.collect::<String>();
|
||||
if [
|
||||
r#""code":"invalid_task_id""#,
|
||||
r#""code":"task_not_found""#,
|
||||
r#""code":"task_expired""#,
|
||||
r#""error":"invalid_task_id""#,
|
||||
]
|
||||
.iter()
|
||||
.any(|marker| compact.contains(marker))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
[
|
||||
"invalid task_id",
|
||||
"invalid task id",
|
||||
"task_id is invalid",
|
||||
"task id is invalid",
|
||||
"task not found",
|
||||
"task expired",
|
||||
"unknown task_id",
|
||||
"unknown task id",
|
||||
]
|
||||
.iter()
|
||||
.any(|marker| lower.contains(marker))
|
||||
}
|
||||
|
||||
fn agent_identity_nested_object(root: &Map<String, Value>) -> Option<&Map<String, Value>> {
|
||||
root.get("agent_identity")
|
||||
.or_else(|| root.get("agentIdentity"))
|
||||
.and_then(Value::as_object)
|
||||
}
|
||||
|
||||
fn string_from_map(map: &Map<String, Value>, keys: &[&str]) -> Option<String> {
|
||||
keys.iter().find_map(|key| {
|
||||
map.get(*key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn string_from_maps(
|
||||
root: &Map<String, Value>,
|
||||
nested: Option<&Map<String, Value>>,
|
||||
keys: &[&str],
|
||||
) -> Option<String> {
|
||||
nested
|
||||
.and_then(|nested| string_from_map(nested, keys))
|
||||
.or_else(|| string_from_map(root, keys))
|
||||
}
|
||||
|
||||
fn agent_identity_credentials(config: &Value) -> Result<AgentIdentityCredentials, String> {
|
||||
let root = config
|
||||
.as_object()
|
||||
.ok_or_else(|| "Agent Identity auth_config must be a JSON object".to_string())?;
|
||||
if !is_codex_agent_identity_auth_config_value(config) {
|
||||
return Err("Codex Agent Identity auth_mode must be agentIdentity".to_string());
|
||||
}
|
||||
let nested = agent_identity_nested_object(root);
|
||||
let runtime_id = string_from_maps(root, nested, &["agent_runtime_id", "agentRuntimeId"])
|
||||
.ok_or_else(|| "Agent Identity agent_runtime_id is required".to_string())?;
|
||||
let encoded_private_key =
|
||||
string_from_maps(root, nested, &["agent_private_key", "agentPrivateKey"])
|
||||
.ok_or_else(|| "Agent Identity agent_private_key is required".to_string())?;
|
||||
let private_key_der = STANDARD
|
||||
.decode(encoded_private_key)
|
||||
.map_err(|_| "Agent Identity agent_private_key must be base64 PKCS#8".to_string())?;
|
||||
let signing_key = SigningKey::from_pkcs8_der(&private_key_der).map_err(|_| {
|
||||
"Agent Identity agent_private_key must be an Ed25519 PKCS#8 key".to_string()
|
||||
})?;
|
||||
|
||||
Ok(AgentIdentityCredentials {
|
||||
runtime_id,
|
||||
signing_key,
|
||||
task_id: string_from_maps(root, nested, &["task_id", "taskId"]),
|
||||
})
|
||||
}
|
||||
|
||||
fn agent_identity_timestamp(now: DateTime<Utc>) -> String {
|
||||
now.to_rfc3339_opts(SecondsFormat::Secs, true)
|
||||
}
|
||||
|
||||
fn build_agent_assertion(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
task_id: &str,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<String, String> {
|
||||
let task_id = task_id.trim();
|
||||
if task_id.is_empty() {
|
||||
return Err("Agent Identity task_id is required".to_string());
|
||||
}
|
||||
let timestamp = agent_identity_timestamp(now);
|
||||
let payload = format!("{}:{task_id}:{timestamp}", credentials.runtime_id);
|
||||
let signature = credentials.signing_key.sign(payload.as_bytes());
|
||||
let envelope = json!({
|
||||
"agent_runtime_id": credentials.runtime_id,
|
||||
"task_id": task_id,
|
||||
"timestamp": timestamp,
|
||||
"signature": STANDARD.encode(signature.to_bytes()),
|
||||
});
|
||||
let encoded = serde_json::to_vec(&envelope)
|
||||
.map_err(|_| "failed to serialize Agent Identity assertion".to_string())?;
|
||||
Ok(format!(
|
||||
"{ASSERTION_PREFIX}{}",
|
||||
URL_SAFE_NO_PAD.encode(encoded)
|
||||
))
|
||||
}
|
||||
|
||||
fn build_task_registration_signature(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
now: DateTime<Utc>,
|
||||
) -> (String, String) {
|
||||
let timestamp = agent_identity_timestamp(now);
|
||||
let payload = format!("{}:{timestamp}", credentials.runtime_id);
|
||||
let signature = credentials.signing_key.sign(payload.as_bytes());
|
||||
(timestamp, STANDARD.encode(signature.to_bytes()))
|
||||
}
|
||||
|
||||
fn task_registration_url(base_url: &str, runtime_id: &str) -> Result<String, String> {
|
||||
let mut url = Url::parse(base_url.trim())
|
||||
.map_err(|_| "Agent Identity auth API base URL is invalid".to_string())?;
|
||||
let mut segments = url
|
||||
.path_segments_mut()
|
||||
.map_err(|_| "Agent Identity auth API base URL cannot be a base URL".to_string())?;
|
||||
segments.pop_if_empty();
|
||||
for segment in ["v1", "agent", runtime_id, "task", "register"] {
|
||||
segments.push(segment);
|
||||
}
|
||||
drop(segments);
|
||||
Ok(url.into())
|
||||
}
|
||||
|
||||
fn task_id_from_registration_response(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
body: &str,
|
||||
) -> Result<String, String> {
|
||||
let response = serde_json::from_str::<AgentTaskRegistrationResponse>(body)
|
||||
.map_err(|_| "Agent Identity task registration returned invalid JSON".to_string())?;
|
||||
for task_id in [response.task_id, response.task_id_camel]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
let task_id = task_id.trim();
|
||||
if !task_id.is_empty() {
|
||||
return Ok(task_id.to_string());
|
||||
}
|
||||
}
|
||||
let encrypted_task_id = [response.encrypted_task_id, response.encrypted_task_id_camel]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find(|value| !value.trim().is_empty())
|
||||
.ok_or_else(|| "Agent Identity task registration omitted task_id".to_string())?;
|
||||
decrypt_agent_task_id(credentials, encrypted_task_id.as_str())
|
||||
}
|
||||
|
||||
fn decrypt_agent_task_id(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
encrypted_task_id: &str,
|
||||
) -> Result<String, String> {
|
||||
let ciphertext = STANDARD
|
||||
.decode(encrypted_task_id.trim())
|
||||
.map_err(|_| "Agent Identity encrypted_task_id must be base64".to_string())?;
|
||||
let seed = credentials.signing_key.to_bytes();
|
||||
let digest = Sha512::digest(seed);
|
||||
let mut curve_private_key = [0u8; 32];
|
||||
curve_private_key.copy_from_slice(&digest[..32]);
|
||||
let secret_key = Curve25519SecretKey::from_bytes(curve_private_key);
|
||||
let plaintext = secret_key
|
||||
.unseal(&ciphertext)
|
||||
.map_err(|_| "Agent Identity encrypted_task_id could not be decrypted".to_string())?;
|
||||
let task_id = String::from_utf8(plaintext)
|
||||
.map_err(|_| "Agent Identity decrypted task_id is invalid".to_string())?;
|
||||
let task_id = task_id.trim();
|
||||
if task_id.is_empty() {
|
||||
return Err("Agent Identity decrypted task_id is empty".to_string());
|
||||
}
|
||||
Ok(task_id.to_string())
|
||||
}
|
||||
|
||||
fn with_agent_identity_task_id(config: &Value, task_id: String) -> Result<Value, String> {
|
||||
let mut root = config
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or_else(|| "Agent Identity auth_config must be a JSON object".to_string())?;
|
||||
let nested_key = ["agent_identity", "agentIdentity"]
|
||||
.into_iter()
|
||||
.find(|key| root.get(*key).and_then(Value::as_object).is_some());
|
||||
if let Some(nested_key) = nested_key {
|
||||
let nested = root
|
||||
.get_mut(nested_key)
|
||||
.and_then(Value::as_object_mut)
|
||||
.expect("Agent Identity nested config was checked as an object");
|
||||
nested.insert("task_id".to_string(), Value::String(task_id.clone()));
|
||||
nested.remove("taskId");
|
||||
}
|
||||
root.insert("task_id".to_string(), Value::String(task_id));
|
||||
root.remove("taskId");
|
||||
Ok(Value::Object(root))
|
||||
}
|
||||
|
||||
fn agent_identity_refresh_error(message: impl Into<String>) -> LocalOAuthRefreshError {
|
||||
LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type: CODEX_AGENT_IDENTITY_PROVIDER_TYPE,
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LocalOAuthRefreshAdapter for CodexAgentIdentityRefreshAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
CODEX_AGENT_IDENTITY_PROVIDER_TYPE
|
||||
}
|
||||
|
||||
fn supports(&self, transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||
is_codex_agent_identity_transport(transport)
|
||||
}
|
||||
|
||||
fn resolve_cached(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
Self::config_from_entry(entry).and_then(|config| Self::resolve_from_config(&config))
|
||||
}
|
||||
|
||||
fn resolve_without_refresh(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
Self::config_from_transport(transport).and_then(|config| Self::resolve_from_config(&config))
|
||||
}
|
||||
|
||||
fn should_refresh(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> bool {
|
||||
if !self.supports(transport) {
|
||||
return false;
|
||||
}
|
||||
Self::preferred_config(transport, entry)
|
||||
.and_then(|config| agent_identity_credentials(&config).ok())
|
||||
.is_some_and(|credentials| credentials.task_id.is_none())
|
||||
}
|
||||
|
||||
async fn refresh(
|
||||
&self,
|
||||
executor: &dyn LocalOAuthHttpExecutor,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
|
||||
let Some(config) = Self::preferred_config(transport, entry) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let credentials =
|
||||
agent_identity_credentials(&config).map_err(agent_identity_refresh_error)?;
|
||||
let (timestamp, signature) = build_task_registration_signature(&credentials, Utc::now());
|
||||
let url = task_registration_url(&self.auth_api_base_url, credentials.runtime_id.as_str())
|
||||
.map_err(agent_identity_refresh_error)?;
|
||||
let request = LocalOAuthHttpRequest {
|
||||
request_id: CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID,
|
||||
method: reqwest::Method::POST,
|
||||
url,
|
||||
headers: BTreeMap::from([
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
]),
|
||||
json_body: Some(json!({
|
||||
"timestamp": timestamp,
|
||||
"signature": signature,
|
||||
})),
|
||||
body_bytes: None,
|
||||
};
|
||||
let response = executor
|
||||
.execute(CODEX_AGENT_IDENTITY_PROVIDER_TYPE, transport, &request)
|
||||
.await?;
|
||||
if !(200..300).contains(&response.status_code) {
|
||||
return Err(LocalOAuthRefreshError::HttpStatus {
|
||||
provider_type: CODEX_AGENT_IDENTITY_PROVIDER_TYPE,
|
||||
status_code: response.status_code,
|
||||
// Registration responses can echo credential material. Keep the stored/logged
|
||||
// error deliberately generic instead of forwarding an upstream body excerpt.
|
||||
body_excerpt: format!(
|
||||
"Agent Identity task registration returned HTTP {}",
|
||||
response.status_code
|
||||
),
|
||||
});
|
||||
}
|
||||
let task_id = task_id_from_registration_response(&credentials, response.body_text.as_str())
|
||||
.map_err(agent_identity_refresh_error)?;
|
||||
let config = with_agent_identity_task_id(&config, task_id.clone())
|
||||
.map_err(agent_identity_refresh_error)?;
|
||||
let updated_credentials =
|
||||
agent_identity_credentials(&config).map_err(agent_identity_refresh_error)?;
|
||||
let auth_header_value = build_agent_assertion(&updated_credentials, &task_id, Utc::now())
|
||||
.map_err(agent_identity_refresh_error)?;
|
||||
|
||||
Ok(Some(CachedOAuthEntry {
|
||||
provider_type: CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE.to_string(),
|
||||
auth_header_name: AUTHORIZATION_HEADER.to_string(),
|
||||
auth_header_value,
|
||||
expires_at_unix_secs: None,
|
||||
metadata: Some(config),
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crypto_box::{aead::rand_core::OsRng, PublicKey};
|
||||
use ed25519_dalek::{pkcs8::EncodePrivateKey, Signature, Verifier};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
agent_identity_credentials, build_agent_assertion, decrypt_agent_task_id,
|
||||
is_codex_agent_identity_auth_config_value, is_codex_agent_identity_invalid_task_response,
|
||||
task_id_from_registration_response, with_agent_identity_task_id,
|
||||
CodexAgentIdentityRefreshAdapter, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
|
||||
};
|
||||
use crate::oauth_refresh::{
|
||||
LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthHttpResponse,
|
||||
LocalOAuthRefreshAdapter, LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use base64::{
|
||||
engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
|
||||
Engine as _,
|
||||
};
|
||||
use chrono::{TimeZone, Utc};
|
||||
use ed25519_dalek::SigningKey;
|
||||
use sha2::Digest;
|
||||
|
||||
fn test_auth_config(task_id: Option<&str>) -> serde_json::Value {
|
||||
let signing_key = SigningKey::from_bytes(&[7u8; 32]);
|
||||
let private_key_der = signing_key
|
||||
.to_pkcs8_der()
|
||||
.expect("test key should encode as PKCS#8");
|
||||
let mut config = serde_json::Map::from_iter([
|
||||
("provider_type".to_string(), json!("codex")),
|
||||
("auth_mode".to_string(), json!("agentIdentity")),
|
||||
("agent_runtime_id".to_string(), json!("runtime-test")),
|
||||
(
|
||||
"agent_private_key".to_string(),
|
||||
json!(STANDARD.encode(private_key_der.as_bytes())),
|
||||
),
|
||||
]);
|
||||
if let Some(task_id) = task_id {
|
||||
config.insert("task_id".to_string(), json!(task_id));
|
||||
}
|
||||
serde_json::Value::Object(config)
|
||||
}
|
||||
|
||||
fn sample_transport(config: serde_json::Value) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "Codex".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:responses".to_string(),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://chatgpt.com/backend-api/codex".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "Agent Identity".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
allow_auth_channel_mismatch_formats: None,
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
upstream_metadata: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: Some(config.to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_verifiable_agent_assertion() {
|
||||
let config = test_auth_config(Some("task-test"));
|
||||
let credentials = agent_identity_credentials(&config).expect("credentials should parse");
|
||||
let now = Utc.with_ymd_and_hms(2030, 1, 2, 3, 4, 5).unwrap();
|
||||
|
||||
let assertion =
|
||||
build_agent_assertion(&credentials, "task-test", now).expect("assertion should build");
|
||||
let encoded = assertion
|
||||
.strip_prefix("AgentAssertion ")
|
||||
.expect("assertion should have its scheme");
|
||||
let envelope: serde_json::Value = serde_json::from_slice(
|
||||
&URL_SAFE_NO_PAD
|
||||
.decode(encoded)
|
||||
.expect("assertion should be URL-safe base64"),
|
||||
)
|
||||
.expect("assertion envelope should be JSON");
|
||||
|
||||
assert_eq!(envelope["agent_runtime_id"], "runtime-test");
|
||||
assert_eq!(envelope["task_id"], "task-test");
|
||||
assert_eq!(envelope["timestamp"], "2030-01-02T03:04:05Z");
|
||||
let signature = Signature::from_slice(
|
||||
&STANDARD
|
||||
.decode(envelope["signature"].as_str().unwrap())
|
||||
.expect("signature should be base64"),
|
||||
)
|
||||
.expect("signature should be valid length");
|
||||
credentials
|
||||
.signing_key
|
||||
.verifying_key()
|
||||
.verify(b"runtime-test:task-test:2030-01-02T03:04:05Z", &signature)
|
||||
.expect("assertion signature should verify");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decrypts_sealed_task_registration_response() {
|
||||
let config = test_auth_config(None);
|
||||
let credentials = agent_identity_credentials(&config).expect("credentials should parse");
|
||||
let digest = sha2::Sha512::digest(credentials.signing_key.to_bytes());
|
||||
let mut curve_private_key = [0u8; 32];
|
||||
curve_private_key.copy_from_slice(&digest[..32]);
|
||||
let secret_key = crypto_box::SecretKey::from_bytes(curve_private_key);
|
||||
let public_key = PublicKey::from(&secret_key);
|
||||
let encrypted = public_key
|
||||
.seal(&mut OsRng, b"task-encrypted")
|
||||
.expect("sealed task should encrypt");
|
||||
|
||||
assert_eq!(
|
||||
decrypt_agent_task_id(&credentials, &STANDARD.encode(encrypted))
|
||||
.expect("sealed task should decrypt"),
|
||||
"task-encrypted"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_nested_agent_identity_export_shape() {
|
||||
let nested = test_auth_config(Some("task-nested"));
|
||||
let config = json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_identity": nested,
|
||||
});
|
||||
|
||||
assert!(is_codex_agent_identity_auth_config_value(&config));
|
||||
assert_eq!(
|
||||
agent_identity_credentials(&config)
|
||||
.expect("nested credentials should parse")
|
||||
.task_id
|
||||
.as_deref(),
|
||||
Some("task-nested")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn synchronizes_replaced_task_id_across_nested_and_flat_import_fields() {
|
||||
let config = json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"taskId": "old-root-task",
|
||||
"agent_identity": {
|
||||
"agent_runtime_id": "runtime-test",
|
||||
"agent_private_key": "placeholder",
|
||||
"taskId": "old-nested-task"
|
||||
}
|
||||
});
|
||||
|
||||
let updated =
|
||||
with_agent_identity_task_id(&config, "new-task".to_string()).expect("task updates");
|
||||
|
||||
assert_eq!(updated["task_id"], "new-task");
|
||||
assert!(updated.get("taskId").is_none());
|
||||
assert_eq!(updated["agent_identity"]["task_id"], "new-task");
|
||||
assert!(updated["agent_identity"].get("taskId").is_none());
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RecordingExecutor {
|
||||
requests: Arc<Mutex<Vec<LocalOAuthHttpRequest>>>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl LocalOAuthHttpExecutor for RecordingExecutor {
|
||||
async fn execute(
|
||||
&self,
|
||||
_provider_type: &'static str,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
request: &LocalOAuthHttpRequest,
|
||||
) -> Result<LocalOAuthHttpResponse, LocalOAuthRefreshError> {
|
||||
self.requests
|
||||
.lock()
|
||||
.expect("recording lock should hold")
|
||||
.push(request.clone());
|
||||
Ok(LocalOAuthHttpResponse {
|
||||
status_code: 200,
|
||||
body_text: r#"{"task_id":"task-registered"}"#.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn registers_missing_task_then_builds_a_fresh_assertion_from_cache() {
|
||||
let config = test_auth_config(None);
|
||||
let transport = sample_transport(config);
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let executor = RecordingExecutor {
|
||||
requests: Arc::clone(&requests),
|
||||
};
|
||||
let adapter = CodexAgentIdentityRefreshAdapter::default()
|
||||
.with_auth_api_base_url_for_tests("https://auth.test/api/accounts");
|
||||
|
||||
assert!(adapter.should_refresh(&transport, None));
|
||||
let entry = adapter
|
||||
.refresh(&executor, &transport, None)
|
||||
.await
|
||||
.expect("registration should succeed")
|
||||
.expect("registration should return a cache entry");
|
||||
assert_eq!(
|
||||
entry.provider_type,
|
||||
CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE
|
||||
);
|
||||
assert!(entry.auth_header_value.starts_with("AgentAssertion "));
|
||||
assert_eq!(
|
||||
entry
|
||||
.metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("task_id")),
|
||||
Some(&json!("task-registered"))
|
||||
);
|
||||
let cached_auth = adapter
|
||||
.resolve_cached(&transport, &entry)
|
||||
.expect("cached task should create a new assertion");
|
||||
assert!(matches!(
|
||||
cached_auth,
|
||||
LocalResolvedOAuthRequestAuth::Header { ref name, ref value }
|
||||
if name == "authorization" && value.starts_with("AgentAssertion ")
|
||||
));
|
||||
|
||||
let requests = requests.lock().expect("recording lock should hold");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].request_id,
|
||||
super::CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].url,
|
||||
"https://auth.test/api/accounts/v1/agent/runtime-test/task/register"
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].json_body.as_ref().unwrap()["timestamp"]
|
||||
.as_str()
|
||||
.is_some(),
|
||||
true
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].json_body.as_ref().unwrap()["signature"]
|
||||
.as_str()
|
||||
.is_some(),
|
||||
true
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registration_response_accepts_encrypted_task_aliases() {
|
||||
let config = test_auth_config(Some("task-original"));
|
||||
let credentials = agent_identity_credentials(&config).expect("credentials should parse");
|
||||
|
||||
assert_eq!(
|
||||
task_id_from_registration_response(&credentials, r#"{"taskId":"task-camel"}"#)
|
||||
.expect("camel task id should parse"),
|
||||
"task-camel"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_only_agent_task_specific_unauthorized_responses() {
|
||||
assert!(is_codex_agent_identity_invalid_task_response(
|
||||
401,
|
||||
Some(r#"{"error": {"code": "invalid_task_id"}}"#),
|
||||
));
|
||||
assert!(is_codex_agent_identity_invalid_task_response(
|
||||
401,
|
||||
Some("Agent task expired; register a new task"),
|
||||
));
|
||||
assert!(!is_codex_agent_identity_invalid_task_response(
|
||||
401,
|
||||
Some(r#"{"error": {"code": "invalid_api_key"}}"#),
|
||||
));
|
||||
assert!(!is_codex_agent_identity_invalid_task_response(
|
||||
403,
|
||||
Some(r#"{"error": {"code": "invalid_task_id"}}"#),
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -22,6 +22,7 @@ const UNSAFE_AUTH_CONFIG_QUERY_NAMES: &[&str] = &[
|
||||
];
|
||||
const SENSITIVE_AUTH_CONFIG_KEYS: &[&str] = &[
|
||||
"access_token",
|
||||
"agent_private_key",
|
||||
"api_key",
|
||||
"apikey",
|
||||
"authorization",
|
||||
@@ -106,6 +107,18 @@ pub fn apply_local_auth_config_header_overrides(
|
||||
let mut overrides = BTreeMap::new();
|
||||
collect_auth_config_header_overrides(object, &mut overrides);
|
||||
for (key, value) in overrides {
|
||||
// Agent Identity assertions are generated for the current request. A stale
|
||||
// imported Authorization header must never replace one later in planning.
|
||||
if key.eq_ignore_ascii_case("authorization")
|
||||
&& headers.get("authorization").is_some_and(|current| {
|
||||
current
|
||||
.trim_start()
|
||||
.to_ascii_lowercase()
|
||||
.starts_with("agentassertion ")
|
||||
})
|
||||
{
|
||||
continue;
|
||||
}
|
||||
headers.insert(key, value);
|
||||
}
|
||||
}
|
||||
@@ -606,6 +619,35 @@ mod tests {
|
||||
assert!(!headers.contains_key("host"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_dynamic_agent_assertion_over_static_imported_authorization() {
|
||||
let mut headers = std::collections::BTreeMap::from([(
|
||||
"authorization".to_string(),
|
||||
"AgentAssertion signed-for-this-request".to_string(),
|
||||
)]);
|
||||
|
||||
apply_local_auth_config_header_overrides(
|
||||
&mut headers,
|
||||
Some(
|
||||
r#"{
|
||||
"headers": {
|
||||
"authorization": "Bearer stale-imported-session",
|
||||
"chatgpt-account-id": "acct-1"
|
||||
}
|
||||
}"#,
|
||||
),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("authorization"),
|
||||
Some(&"AgentAssertion signed-for-this-request".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id"),
|
||||
Some(&"acct-1".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_invalid_auth_config_header_values_when_applying_overrides() {
|
||||
let mut headers = std::collections::BTreeMap::new();
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod agent_identity;
|
||||
pub mod antigravity;
|
||||
pub mod auth;
|
||||
mod auth_config;
|
||||
@@ -28,6 +29,13 @@ mod video;
|
||||
pub mod windsurf;
|
||||
|
||||
pub use aether_oauth as oauth;
|
||||
pub use agent_identity::{
|
||||
is_codex_agent_identity_auth_config_value, is_codex_agent_identity_cached_entry,
|
||||
is_codex_agent_identity_invalid_task_response, is_codex_agent_identity_transport,
|
||||
validate_codex_agent_identity_auth_config, CodexAgentIdentityRefreshAdapter,
|
||||
CODEX_AGENT_IDENTITY_AUTH_MODE, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
|
||||
CODEX_AGENT_IDENTITY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID,
|
||||
};
|
||||
pub use auth::{build_passthrough_headers, ensure_upstream_auth_header};
|
||||
pub use auth_config::apply_local_auth_config_header_overrides;
|
||||
pub use cache::{provider_transport_snapshot_looks_refreshed, ProviderTransportSnapshotCacheKey};
|
||||
|
||||
@@ -13,6 +13,7 @@ use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::agent_identity::{is_codex_agent_identity_transport, CodexAgentIdentityRefreshAdapter};
|
||||
use super::generic_oauth::supports_local_generic_oauth_request_auth_resolution;
|
||||
pub use super::generic_oauth::GenericOAuthRefreshAdapter;
|
||||
use super::kiro::{
|
||||
@@ -341,6 +342,7 @@ impl LocalOAuthRefreshCoordinator {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
adapters: vec![
|
||||
Arc::new(CodexAgentIdentityRefreshAdapter::default()),
|
||||
Arc::new(KiroOAuthRefreshAdapter::default()),
|
||||
Arc::new(VertexServiceAccountRefreshAdapter),
|
||||
Arc::new(GenericOAuthRefreshAdapter::default()),
|
||||
@@ -550,7 +552,8 @@ impl LocalOAuthResolution {
|
||||
pub fn supports_local_oauth_request_auth_resolution(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
supports_local_kiro_request_auth_resolution(transport)
|
||||
is_codex_agent_identity_transport(transport)
|
||||
|| supports_local_kiro_request_auth_resolution(transport)
|
||||
|| supports_local_vertex_service_account_auth_resolution(transport)
|
||||
|| supports_local_generic_oauth_request_auth_resolution(transport)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user