Send account id when available (#1767)

For users with multiple accounts we need to specify the account to use.
This commit is contained in:
pakrym-oai
2025-07-31 15:40:19 -07:00
committed by GitHub
parent 6ce0a5875b
commit 0935e6a875
3 changed files with 21 additions and 2 deletions

View File

@@ -192,7 +192,7 @@ impl ModelClient {
loop {
attempt += 1;
let req_builder = self
let mut req_builder = self
.client
.post(format!("{base_url}/responses"))
.header("OpenAI-Beta", "responses=experimental")
@@ -201,6 +201,12 @@ impl ModelClient {
.header(reqwest::header::ACCEPT, "text/event-stream")
.json(&payload);
if auth.mode == AuthMode::ChatGPT {
if let Some(account_id) = auth.get_account_id().await {
req_builder = req_builder.header("chatgpt-account-id", account_id);
}
}
let req_builder = self.provider.apply_http_headers(req_builder);
let res = req_builder.send().await;

View File

@@ -237,6 +237,7 @@ async fn chatgpt_auth_sends_correct_request() {
let request_session_id = request.headers.get("session_id").unwrap();
let request_originator = request.headers.get("originator").unwrap();
let request_authorization = request.headers.get("authorization").unwrap();
let request_chatgpt_account_id = request.headers.get("chatgpt-account-id").unwrap();
let request_body = request.body_json::<serde_json::Value>().unwrap();
assert!(current_session_id.is_some());
@@ -249,6 +250,7 @@ async fn chatgpt_auth_sends_correct_request() {
request_authorization.to_str().unwrap(),
"Bearer Access Token"
);
assert_eq!(request_chatgpt_account_id.to_str().unwrap(), "account_id");
assert!(!request_body["store"].as_bool().unwrap());
assert!(request_body["stream"].as_bool().unwrap());
assert_eq!(
@@ -332,7 +334,7 @@ fn auth_from_token(id_token: String) -> CodexAuth {
id_token,
access_token: "Access Token".to_string(),
refresh_token: "test".to_string(),
account_id: None,
account_id: Some("account_id".to_string()),
}),
last_refresh: Some(Utc::now()),
}),

View File

@@ -123,6 +123,17 @@ impl CodexAuth {
}
}
}
pub async fn get_account_id(&self) -> Option<String> {
match self.mode {
AuthMode::ApiKey => None,
AuthMode::ChatGPT => {
let token_data = self.get_token_data().await.ok()?;
token_data.account_id.clone()
}
}
}
}
// Loads the available auth information from the auth.json or OPENAI_API_KEY environment variable.