diff --git a/src/client/access_token.rs b/src/client/access_token.rs index ccbfd06..a1d6939 100644 --- a/src/client/access_token.rs +++ b/src/client/access_token.rs @@ -4,29 +4,44 @@ use indexmap::IndexMap; use parking_lot::RwLock; use std::sync::LazyLock; -static ACCESS_TOKENS: LazyLock>> = +type AccessTokenEntry = (String, i64, Option); + +static ACCESS_TOKENS: LazyLock>> = LazyLock::new(|| RwLock::new(IndexMap::new())); pub fn get_access_token(client_name: &str) -> Result { ACCESS_TOKENS .read() .get(client_name) - .map(|(token, _)| token.clone()) + .map(|(token, _, _)| token.clone()) .ok_or_else(|| anyhow!("Invalid access token")) } +pub fn get_access_token_account_id(client_name: &str) -> Option { + ACCESS_TOKENS + .read() + .get(client_name) + .and_then(|(_, _, account_id)| account_id.clone()) +} + pub fn is_valid_access_token(client_name: &str) -> bool { let access_tokens = ACCESS_TOKENS.read(); - let (token, expires_at) = match access_tokens.get(client_name) { + let (token, expires_at, _) = match access_tokens.get(client_name) { Some(v) => v, None => return false, }; !token.is_empty() && Utc::now().timestamp() < *expires_at } -pub fn set_access_token(client_name: &str, token: String, expires_at: i64) { +pub fn set_access_token( + client_name: &str, + token: String, + expires_at: i64, + account_id: Option, +) { let mut access_tokens = ACCESS_TOKENS.write(); let entry = access_tokens.entry(client_name.to_string()).or_default(); entry.0 = token; entry.1 = expires_at; + entry.2 = account_id; } diff --git a/src/client/mod.rs b/src/client/mod.rs index 0cd9365..1b3361c 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -4,6 +4,7 @@ mod common; mod gemini_oauth; mod message; pub mod oauth; +mod openai_oauth; #[macro_use] mod macros; mod model; diff --git a/src/client/oauth.rs b/src/client/oauth.rs index e2e9e4d..8845d8c 100644 --- a/src/client/oauth.rs +++ b/src/client/oauth.rs @@ -57,6 +57,14 @@ pub trait OAuthProvider: Send + Sync { fn fixed_redirect_uri(&self) -> Option { None } + + fn extract_account_id(&self, _response: &Value) -> Option { + None + } + + fn include_state_in_token_exchange(&self) -> bool { + true + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -64,6 +72,8 @@ pub struct OAuthTokens { pub access_token: String, pub refresh_token: String, pub expires_at: i64, + #[serde(default)] + pub account_id: Option, } pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> { @@ -137,18 +147,17 @@ pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> } let client = ReqwestClient::new(); - let request = build_token_request( - &client, - provider, - &[ - ("grant_type", "authorization_code"), - ("client_id", provider.client_id()), - ("code", &code), - ("code_verifier", &code_verifier), - ("redirect_uri", &redirect_uri), - ("state", &state), - ], - ); + let mut token_params = vec![ + ("grant_type", "authorization_code"), + ("client_id", provider.client_id()), + ("code", code.as_str()), + ("code_verifier", code_verifier.as_str()), + ("redirect_uri", redirect_uri.as_str()), + ]; + if provider.include_state_in_token_exchange() { + token_params.push(("state", state.as_str())); + } + let request = build_token_request(&client, provider, &token_params); let response: Value = request.send().await?.json().await?; @@ -166,10 +175,13 @@ pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> let expires_at = Utc::now().timestamp() + expires_in; + let account_id = provider.extract_account_id(&response); + let tokens = OAuthTokens { access_token, refresh_token, expires_at, + account_id, }; save_oauth_tokens(client_name, &tokens)?; @@ -231,10 +243,15 @@ pub async fn refresh_oauth_token( let expires_at = Utc::now().timestamp() + expires_in; + let account_id = provider + .extract_account_id(&response) + .or_else(|| tokens.account_id.clone()); + let new_tokens = OAuthTokens { access_token, refresh_token, expires_at, + account_id, }; save_oauth_tokens(client_name, &new_tokens)?; @@ -262,7 +279,12 @@ pub async fn prepare_oauth_access_token( tokens }; - set_access_token(client_name, tokens.access_token.clone(), tokens.expires_at); + set_access_token( + client_name, + tokens.access_token.clone(), + tokens.expires_at, + tokens.account_id.clone(), + ); Ok(true) } @@ -371,6 +393,7 @@ pub fn get_oauth_provider(provider_type: &str) -> Option> match provider_type { "claude" => Some(Box::new(super::claude_oauth::ClaudeOAuthProvider)), "gemini" => Some(Box::new(super::gemini_oauth::GeminiOAuthProvider)), + "openai" => Some(Box::new(super::openai_oauth::OpenAIOAuthProvider)), _ => None, } } @@ -409,7 +432,11 @@ fn client_config_info(client_config: &ClientConfig) -> (&str, &'static str, Opti "claude", c.auth.as_deref(), ), - ClientConfig::OpenAIConfig(c) => (c.name.as_deref().unwrap_or("openai"), "openai", None), + ClientConfig::OpenAIConfig(c) => ( + c.name.as_deref().unwrap_or("openai"), + "openai", + c.auth.as_deref(), + ), ClientConfig::OpenAICompatibleConfig(c) => ( c.name.as_deref().unwrap_or("openai-compatible"), "openai-compatible", diff --git a/src/client/openai.rs b/src/client/openai.rs index f5a4e2c..54ddae6 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -1,13 +1,17 @@ +use super::access_token::{get_access_token, get_access_token_account_id}; +use super::oauth::{self, OAuthProvider}; +use super::openai_oauth::OpenAIOAuthProvider; use super::*; use crate::utils::strip_think_tag; use anyhow::{Context, Result, bail}; -use reqwest::RequestBuilder; +use reqwest::{Client as ReqwestClient, RequestBuilder}; use serde::Deserialize; use serde_json::{Value, json}; const API_BASE: &str = "https://api.openai.com/v1"; +const CODEX_API_ENDPOINT: &str = "https://chatgpt.com/backend-api/codex/responses"; #[derive(Debug, Clone, Deserialize, Default)] pub struct OpenAIConfig { @@ -15,6 +19,7 @@ pub struct OpenAIConfig { pub api_key: Option, pub api_base: Option, pub organization_id: Option, + pub auth: Option, #[serde(default)] pub models: Vec, pub patch: Option, @@ -25,36 +30,131 @@ impl OpenAIClient { config_get_fn!(api_key, get_api_key); config_get_fn!(api_base, get_api_base); - create_client_config!([("api_key", "API Key", None, true)]); + create_oauth_supported_client_config!(); } -impl_client_trait!( - OpenAIClient, - ( - prepare_chat_completions, - openai_chat_completions, - openai_chat_completions_streaming - ), - (prepare_embeddings, openai_embeddings), - (noop_prepare_rerank, noop_rerank), -); +#[async_trait::async_trait] +impl Client for OpenAIClient { + client_common_fns!(); -fn prepare_chat_completions( + fn supports_oauth(&self) -> bool { + self.config.auth.as_deref() == Some("oauth") + } + + async fn chat_completions_inner( + &self, + client: &ReqwestClient, + data: ChatCompletionsData, + ) -> Result { + let uses_codex = + self.config.auth.as_deref() == Some("oauth") && self.get_api_base().is_err(); + let request_data = prepare_chat_completions(self, client, data).await?; + let builder = self.request_builder(client, request_data); + if uses_codex { + openai_responses_chat_completions(builder, self.model()).await + } else { + openai_chat_completions(builder, self.model()).await + } + } + + async fn chat_completions_streaming_inner( + &self, + client: &ReqwestClient, + handler: &mut SseHandler, + data: ChatCompletionsData, + ) -> Result<()> { + let uses_codex = + self.config.auth.as_deref() == Some("oauth") && self.get_api_base().is_err(); + let request_data = prepare_chat_completions(self, client, data).await?; + let builder = self.request_builder(client, request_data); + + if uses_codex { + openai_responses_streaming(builder, handler).await + } else { + openai_chat_completions_streaming(builder, handler, self.model()).await + } + } + + async fn embeddings_inner( + &self, + client: &ReqwestClient, + data: &EmbeddingsData, + ) -> Result { + let request_data = prepare_embeddings(self, client, data).await?; + let builder = self.request_builder(client, request_data); + openai_embeddings(builder, self.model()).await + } + + async fn rerank_inner( + &self, + client: &ReqwestClient, + data: &RerankData, + ) -> Result { + let request_data = noop_prepare_rerank(self, data)?; + let builder = self.request_builder(client, request_data); + noop_rerank(builder, self.model()).await + } +} + +async fn prepare_chat_completions( self_: &OpenAIClient, + client: &ReqwestClient, data: ChatCompletionsData, ) -> Result { - let api_key = self_.get_api_key()?; - let api_base = self_ - .get_api_base() - .unwrap_or_else(|_| API_BASE.to_string()); + let uses_oauth = self_.config.auth.as_deref() == Some("oauth"); + let has_custom_base = self_.get_api_base().is_ok(); - let url = format!("{}/chat/completions", api_base.trim_end_matches('/')); + let uses_codex = uses_oauth && !has_custom_base; - let body = openai_build_chat_completions_body(data, &self_.model); + let url = if uses_codex { + CODEX_API_ENDPOINT.to_string() + } else { + let api_base = self_ + .get_api_base() + .unwrap_or_else(|_| API_BASE.to_string()); + format!("{}/chat/completions", api_base.trim_end_matches('/')) + }; + + let body = if uses_codex { + openai_build_responses_body(data, &self_.model) + } else { + openai_build_chat_completions_body(data, &self_.model) + }; let mut request_data = RequestData::new(url, body); - request_data.bearer_auth(api_key); + if uses_oauth { + let provider = OpenAIOAuthProvider; + let ready = oauth::prepare_oauth_access_token(client, &provider, self_.name()).await?; + + if !ready { + bail!( + "OAuth configured but no tokens found for '{}'. Run: 'coyote --authenticate {}' or '.authenticate' in the REPL", + self_.name(), + self_.name() + ); + } + + let token = get_access_token(self_.name())?; + request_data.bearer_auth(token); + + if let Some(account_id) = get_access_token_account_id(self_.name()) { + request_data.header("ChatGPT-Account-Id", account_id); + } + + for (key, value) in provider.extra_request_headers() { + request_data.header(key, value); + } + } else if let Ok(api_key) = self_.get_api_key() { + request_data.bearer_auth(api_key); + } else { + bail!( + "No authentication configured for '{}'. Set `api_key` or use `auth: oauth` with `coyote --authenticate {}`.", + self_.name(), + self_.name() + ); + } + if let Some(organization_id) = &self_.config.organization_id { request_data.header("OpenAI-Organization", organization_id); } @@ -62,8 +162,11 @@ fn prepare_chat_completions( Ok(request_data) } -fn prepare_embeddings(self_: &OpenAIClient, data: &EmbeddingsData) -> Result { - let api_key = self_.get_api_key()?; +async fn prepare_embeddings( + self_: &OpenAIClient, + client: &ReqwestClient, + data: &EmbeddingsData, +) -> Result { let api_base = self_ .get_api_base() .unwrap_or_else(|_| API_BASE.to_string()); @@ -74,7 +177,30 @@ fn prepare_embeddings(self_: &OpenAIClient, data: &EmbeddingsData) -> Result Option { Some(value.to_string()) } } + +pub fn openai_build_responses_body(data: ChatCompletionsData, model: &Model) -> Value { + let ChatCompletionsData { + messages, + temperature, + top_p, + functions, + stream, + } = data; + + let messages_len = messages.len(); + let input: Vec = messages + .into_iter() + .enumerate() + .flat_map(|(i, message)| { + let Message { role, content } = message; + match content { + MessageContent::ToolCalls(MessageContentToolCalls { + tool_results, + text: _, + sequence: _, + }) => tool_results + .into_iter() + .flat_map(|tool_result| { + vec![ + json!({ + "type": "function_call", + "call_id": tool_result.call.id, + "name": tool_result.call.name, + "arguments": tool_result.call.arguments.to_string(), + }), + json!({ + "type": "function_call_output", + "call_id": tool_result.call.id, + "output": tool_result.output.to_string(), + }), + ] + }) + .collect(), + MessageContent::Text(text) if role.is_assistant() && i != messages_len - 1 => { + vec![json!({ "role": role, "content": strip_think_tag(&text) })] + } + _ => vec![json!({ "role": role, "content": content })], + } + }) + .collect(); + + let mut body = json!({ + "model": &model.real_name(), + "input": input, + "store": false, + }); + + if let Some(v) = model.max_tokens_param() { + body["max_output_tokens"] = v.into(); + } + if let Some(v) = temperature { + body["temperature"] = v.into(); + } + if let Some(v) = top_p { + body["top_p"] = v.into(); + } + if stream { + body["stream"] = true.into(); + } + if let Some(functions) = functions { + body["tools"] = functions + .iter() + .map(|v| { + let mut tool = serde_json::to_value(v).unwrap_or_default(); + tool["type"] = "function".into(); + tool + }) + .collect(); + } + body +} + +pub async fn openai_responses_chat_completions( + builder: RequestBuilder, + _model: &Model, +) -> Result { + let res = builder.send().await?; + let status = res.status(); + let data: Value = res.json().await?; + + if !status.is_success() { + catch_error(&data, status.as_u16())?; + } + + debug!("non-stream-data: {data}"); + openai_extract_responses(&data) +} + +pub fn openai_extract_responses(data: &Value) -> Result { + let mut text = String::new(); + let mut tool_calls = vec![]; + + if let Some(output) = data["output"].as_array() { + for item in output { + match item["type"].as_str() { + Some("message") => { + if let Some(content) = item["content"].as_array() { + for part in content { + if part["type"].as_str() == Some("output_text") + && let Some(t) = part["text"].as_str() + { + text.push_str(t); + } + } + } + } + Some("function_call") => { + if let (Some(name), Some(arguments_str), Some(call_id)) = ( + item["name"].as_str(), + item["arguments"].as_str(), + item["call_id"].as_str(), + ) { + let arguments: Value = arguments_str.parse().with_context(|| { + format!("Tool call '{name}' has non-JSON arguments '{arguments_str}'") + })?; + tool_calls.push(ToolCall::new( + name.to_string(), + arguments, + Some(call_id.to_string()), + )); + } + } + _ => {} + } + } + } + + if text.is_empty() && tool_calls.is_empty() { + bail!("Invalid response data: {data}"); + } + Ok(ChatCompletionsOutput { text, tool_calls }) +} + +pub async fn openai_responses_streaming( + builder: RequestBuilder, + handler: &mut SseHandler, +) -> Result<()> { + let handle = |message: SseMessage| -> Result { + if message.data == "[DONE]" { + return Ok(true); + } + let data: Value = serde_json::from_str(&message.data)?; + debug!("stream-data: {data}"); + + match data["type"].as_str() { + Some("response.output_text.delta") => { + if let Some(delta) = data["delta"].as_str().filter(|v| !v.is_empty()) { + handler.text(delta)?; + } + } + Some("response.output_item.done") => { + let item = &data["item"]; + if item["type"].as_str() == Some("function_call") + && let (Some(name), Some(arguments_str), Some(call_id)) = ( + item["name"].as_str(), + item["arguments"].as_str(), + item["call_id"].as_str(), + ) + { + let arguments: Value = arguments_str.parse().with_context(|| { + format!("Tool call '{name}' has non-JSON arguments '{arguments_str}'") + })?; + handler.tool_call(ToolCall::new( + name.to_string(), + arguments, + Some(call_id.to_string()), + ))?; + } + } + Some("response.completed") => { + return Ok(true); + } + _ => {} + } + Ok(false) + }; + + sse_stream(builder, handle).await +} diff --git a/src/client/openai_oauth.rs b/src/client/openai_oauth.rs new file mode 100644 index 0000000..45126ce --- /dev/null +++ b/src/client/openai_oauth.rs @@ -0,0 +1,79 @@ +use super::oauth::{OAuthProvider, TokenRequestFormat}; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use serde_json::Value; + +pub struct OpenAIOAuthProvider; + +impl OAuthProvider for OpenAIOAuthProvider { + fn provider_name(&self) -> &str { + "openai" + } + + fn client_id(&self) -> &str { + "app_EMoamEEZ73f0CkXaXp7hrann" + } + + fn authorize_url(&self) -> &str { + "https://auth.openai.com/oauth/authorize" + } + + fn token_url(&self) -> &str { + "https://auth.openai.com/oauth/token" + } + + fn redirect_uri(&self) -> &str { + "http://localhost:1455/auth/callback" + } + + fn scopes(&self) -> &str { + "openid profile email offline_access" + } + + fn token_request_format(&self) -> TokenRequestFormat { + TokenRequestFormat::FormUrlEncoded + } + + fn extra_authorize_params(&self) -> Vec<(&str, &str)> { + vec![ + ("id_token_add_organizations", "true"), + ("codex_cli_simplified_flow", "true"), + ] + } + + fn fixed_redirect_uri(&self) -> Option { + Some("http://localhost:1455/auth/callback".to_string()) + } + + fn include_state_in_token_exchange(&self) -> bool { + false + } + + fn extract_account_id(&self, response: &Value) -> Option { + let id_token = response["id_token"].as_str().unwrap_or_default(); + let access_token = response["access_token"].as_str().unwrap_or_default(); + extract_account_id_from_jwt(id_token).or_else(|| extract_account_id_from_jwt(access_token)) + } +} + +fn extract_account_id_from_jwt(token: &str) -> Option { + let parts: Vec<&str> = token.splitn(3, '.').collect(); + if parts.len() < 2 { + return None; + } + let decoded = URL_SAFE_NO_PAD.decode(parts[1]).ok()?; + let claims: Value = serde_json::from_slice(&decoded).ok()?; + claims["chatgpt_account_id"] + .as_str() + .map(|s| s.to_string()) + .or_else(|| { + claims["https://api.openai.com/auth"]["chatgpt_account_id"] + .as_str() + .map(|s| s.to_string()) + }) + .or_else(|| { + claims["organizations"][0]["id"] + .as_str() + .map(|s| s.to_string()) + }) +} diff --git a/src/client/stream.rs b/src/client/stream.rs index daf308e..c4db9f5 100644 --- a/src/client/stream.rs +++ b/src/client/stream.rs @@ -232,8 +232,8 @@ where .map(|value| value.to_string()); let is_event_stream = content_type .as_deref() - .map(|ct| ct.starts_with("text/event-stream")) - .unwrap_or(false); + .map(|ct| ct.is_empty() || ct.starts_with("text/event-stream")) + .unwrap_or(true); if !is_event_stream { let header_value = content_type.unwrap_or_default(); let text = res.text().await?; diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs index 06fc43b..e7bf287 100644 --- a/src/client/vertexai.rs +++ b/src/client/vertexai.rs @@ -483,7 +483,7 @@ pub async fn prepare_gcloud_access_token( let expires_at = Utc::now() + Duration::try_seconds(expires_in) .ok_or_else(|| anyhow!("Failed to parse expires_in of access_token"))?; - set_access_token(client_name, token, expires_at.timestamp()) + set_access_token(client_name, token, expires_at.timestamp(), None) } Ok(()) }