mirror of
https://github.com/openai/codex.git
synced 2026-11-13 06:45:07 +00:00
Support turn cost telemetry for custom model providers (#39785)
## What changed - Route turn-cost queries for non-OpenAI providers through the configured provider endpoint and authentication, while retaining the existing OpenAI API-key path and excluding Amazon Bedrock. - Observe turns only when their model provider matches the worker's provider. - Retry custom-provider authentication failures during periodic availability probes and ensure client authentication takes precedence over provider headers. ## Testing - Add coverage for provider matching, custom-provider authentication retries, ChatGPT-auth rejection, and header precedence. GitOrigin-RevId: 04a7b28e8e3e18a510ae6fface5193af40d114c0
This commit is contained in:
1 parent
90c67e6f33
commit
0cc80b8db5
4 files changed
+390
-114
No files matched your search
@@ -62,21 +62,37 @@ impl Client {
|
||||
url.set_path("/v1/analytics/codex/turn-costs");
|
||||
url.set_query(None);
|
||||
url.set_fragment(None);
|
||||
let url = url.to_string();
|
||||
let mut headers = self.headers();
|
||||
for header_name in ["openai-organization", "openai-project"] {
|
||||
if let Some(value) = provider_headers.get(header_name) {
|
||||
headers.insert(header_name, value.clone());
|
||||
}
|
||||
}
|
||||
let provider_scope_headers = provider_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
["openai-organization", "openai-project"]
|
||||
.iter()
|
||||
.any(|allowed| name.as_str().eq_ignore_ascii_case(allowed))
|
||||
})
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect();
|
||||
self.query_api_key_turn_costs_at(url.as_ref(), turn_ids, &provider_scope_headers)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Queries an API-key turn-cost endpoint chosen by the caller, attaching
|
||||
/// the supplied provider headers in addition to this client's auth.
|
||||
pub async fn query_api_key_turn_costs_at(
|
||||
&self,
|
||||
url: &str,
|
||||
turn_ids: &[String],
|
||||
provider_headers: &HeaderMap,
|
||||
) -> Result<Vec<ApiKeyTurnCost>, RequestError> {
|
||||
let mut headers = provider_headers.clone();
|
||||
headers.extend(self.headers());
|
||||
let request = self
|
||||
.request(Method::POST, &url)
|
||||
.request(Method::POST, url)
|
||||
.headers(headers)
|
||||
.header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
|
||||
.json(&ApiKeyTurnCostsRequest { turn_ids });
|
||||
let (body, content_type) = self.exec_request_detailed(request, "POST", &url).await?;
|
||||
let (body, content_type) = self.exec_request_detailed(request, "POST", url).await?;
|
||||
let response: ApiKeyTurnCostsResponse = self
|
||||
.decode_json(&url, &content_type, &body)
|
||||
.decode_json(url, &content_type, &body)
|
||||
.map_err(RequestError::Other)?;
|
||||
Ok(response.turns)
|
||||
}
|
||||
@@ -185,4 +201,46 @@ mod tests {
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn custom_turn_cost_queries_apply_client_auth_after_provider_headers() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/analytics/codex/turn-costs"))
|
||||
.and(header("authorization", "Bearer sk-test"))
|
||||
.and(header("chatgpt-account-id", "account-test"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"turns": []
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let auth = CodexAuth::from_api_key("sk-test");
|
||||
let client = Client::from_auth(
|
||||
server.uri(),
|
||||
&auth,
|
||||
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
|
||||
)
|
||||
.with_chatgpt_account_id("account-test");
|
||||
let mut provider_headers = HeaderMap::new();
|
||||
provider_headers.insert(
|
||||
"authorization",
|
||||
HeaderValue::from_static("Bearer provider-override"),
|
||||
);
|
||||
provider_headers.insert(
|
||||
"chatgpt-account-id",
|
||||
HeaderValue::from_static("provider-account-override"),
|
||||
);
|
||||
let costs = client
|
||||
.query_api_key_turn_costs_at(
|
||||
&format!("{}/analytics/codex/turn-costs", server.uri()),
|
||||
&["turn-one".to_string()],
|
||||
&provider_headers,
|
||||
)
|
||||
.await
|
||||
.expect("query custom-provider turn costs");
|
||||
|
||||
assert_eq!(costs, Vec::new());
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user