i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
64ad0f694b682952b600e44d2e48ea4f7057846b
Aether/crates/aether-provider-transport/src/grok.rs

1321 lines
46 KiB
Rust
Raw Blame History

use std::collections::BTreeMap;
use aether_contracts::{
ResolvedTransportProfile, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_HTTP_MODE_AUTO,
TRANSPORT_POOL_SCOPE_KEY,
};
use serde_json::Value;
use serde_json::{json, Map};
use uuid::Uuid;
use crate::rules::apply_local_header_rules_with_request_headers;
use crate::snapshot::GatewayProviderTransportSnapshot;
pub const GROK_INTERNAL_HEADER: &str = "x-aether-grok-runtime";
pub const GROK_DEFAULT_BASE_URL: &str = "https://grok.com";
pub const GROK_CHAT_PATH: &str = "/rest/app-chat/conversations/new";
pub const GROK_RATE_LIMITS_PATH: &str = "/rest/rate-limits";
pub const GROK_IMAGE_EDIT_MODEL_NAME: &str = "imagine-image-edit";
pub const GROK_IMAGE_EDIT_MODEL_KIND: &str = "imagine";
pub const GROK_DEFAULT_BROWSER_PROFILE: &str = "chrome136";
pub const GROK_DEFAULT_USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36";
const GROK_IMAGE_GENERATION_MAX_COUNT: u64 = 4;
const GROK_SEC_CH_UA_PLATFORM: &str = r#""macOS""#;
const GROK_STATSIG_ID: &str = "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GrokBrowserProfileMetadata {
pub profile_id: String,
pub browser_profile: String,
pub user_agent: String,
pub sec_ch_ua: String,
pub sec_ch_ua_platform: String,
}
#[derive(Debug, Clone)]
pub struct GrokHeaderInput<'a> {
pub transport: &'a GatewayProviderTransportSnapshot,
pub transport_profile: Option<&'a ResolvedTransportProfile>,
pub request_headers: Option<&'a http::HeaderMap>,
pub content_type: &'a str,
pub accept: &'a str,
pub header_rules: Option<&'a Value>,
pub provider_request_body: &'a Value,
pub original_request_body: &'a Value,
}
pub fn is_grok_provider_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok")
}
pub fn grok_base_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
if base_url.is_empty() {
GROK_DEFAULT_BASE_URL.to_string()
} else {
base_url.to_string()
}
}
pub fn build_grok_upstream_url(transport: &GatewayProviderTransportSnapshot, path: &str) -> String {
let base_url = grok_base_url(&transport.endpoint.base_url);
let path = path.trim();
if path.starts_with('/') {
format!("{base_url}{path}")
} else {
format!("{base_url}/{path}")
}
}
pub fn resolve_grok_session_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let cookie = grok_cookie_from_transport(transport)?;
Some(("cookie".to_string(), cookie))
}
pub fn grok_browser_profile_metadata(
profile_id: Option<&str>,
) -> Option<GrokBrowserProfileMetadata> {
let raw = profile_id
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(GROK_DEFAULT_BROWSER_PROFILE);
let normalized = raw.trim().to_ascii_lowercase().replace(['_', '-', ' '], "");
let version = supported_grok_chrome_profile_version(&normalized)?;
let profile_id = format!("chrome{version}");
Some(GrokBrowserProfileMetadata {
profile_id: profile_id.clone(),
browser_profile: profile_id,
user_agent: format!(
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/{version}.0.0.0 Safari/537.36"
),
sec_ch_ua: format!(
r#""Google Chrome";v="{version}", "Chromium";v="{version}", "Not(A:Brand";v="24""#
),
sec_ch_ua_platform: GROK_SEC_CH_UA_PLATFORM.to_string(),
})
}
pub fn grok_browser_profile_id_from_user_agent(user_agent: &str) -> Option<String> {
chrome_major_version_from_user_agent(user_agent, "Chrome/")
.or_else(|| chrome_major_version_from_user_agent(user_agent, "Chromium/"))
.and_then(|version| supported_grok_chrome_profile_version(&format!("chrome{version}")))
.map(|version| format!("chrome{version}"))
}
pub fn grok_browser_profile_metadata_from_resolved_transport_profile(
profile: &ResolvedTransportProfile,
) -> Option<GrokBrowserProfileMetadata> {
if !profile
.backend
.trim()
.eq_ignore_ascii_case(TRANSPORT_BACKEND_BROWSER_WREQ)
{
return None;
}
let profile_id = profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(Value::as_str)
.or(Some(profile.profile_id.as_str()));
grok_browser_profile_metadata(profile_id)
}
pub fn grok_browser_resolved_transport_profile_from_auth_config(
auth_config: &Map<String, Value>,
source: &str,
) -> Option<ResolvedTransportProfile> {
let browser_profile = grok_browser_profile_id_from_auth_config(auth_config)?;
grok_browser_resolved_transport_profile(browser_profile.as_deref(), source)
}
pub fn grok_browser_resolved_transport_profile(
profile_id: Option<&str>,
source: &str,
) -> Option<ResolvedTransportProfile> {
let metadata = grok_browser_profile_metadata(profile_id)?;
Some(ResolvedTransportProfile {
profile_id: metadata.profile_id.clone(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: Some(json!({
"browser_profile": metadata.browser_profile,
"source": source,
})),
})
}
pub fn grok_browser_transport_fingerprint_from_auth_config(
auth_config: &Map<String, Value>,
) -> Option<Value> {
let transport_profile =
grok_browser_resolved_transport_profile_from_auth_config(auth_config, "grok_import")?;
Some(json!({ "transport_profile": transport_profile }))
}
fn grok_browser_profile_id_from_auth_config(
auth_config: &Map<String, Value>,
) -> Option<Option<String>> {
if let Some(browser_profile) = grok_config_string_from_object(
Some(auth_config),
&[
"browser_profile",
"browserProfile",
"browser",
"impersonate",
],
) {
return Some(Some(browser_profile));
}
let Some(user_agent) =
grok_config_string_from_object(Some(auth_config), &["user_agent", "userAgent"])
else {
return Some(None);
};
grok_browser_profile_id_from_user_agent(&user_agent).map(Some)
}
fn chrome_major_version_from_user_agent(user_agent: &str, marker: &str) -> Option<u16> {
let user_agent_lower = user_agent.to_ascii_lowercase();
let marker_lower = marker.to_ascii_lowercase();
let index = user_agent_lower.find(&marker_lower)?;
let version_start = index + marker.len();
let digits = user_agent[version_start..]
.chars()
.take_while(|ch| ch.is_ascii_digit())
.collect::<String>();
(!digits.is_empty())
.then_some(digits)
.and_then(|digits| digits.parse::<u16>().ok())
}
fn supported_grok_chrome_profile_version(normalized: &str) -> Option<u16> {
let version = normalized.strip_prefix("chrome")?.parse::<u16>().ok()?;
match version {
100 | 101 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 114 | 116 | 117 | 118 | 119 | 120
| 123 | 124 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137
| 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 => Some(version),
_ => None,
}
}
pub fn build_grok_browser_headers(input: GrokHeaderInput<'_>) -> Option<BTreeMap<String, String>> {
let cookie = grok_cookie_from_transport(input.transport)?;
let base_url = grok_base_url(&input.transport.endpoint.base_url);
let browser_profile =
grok_browser_profile_metadata_from_resolved_transport_profile(input.transport_profile?)?;
let mut headers = BTreeMap::from([
("accept".to_string(), input.accept.to_string()),
(
"accept-language".to_string(),
"zh-CN,zh;q=0.9,en;q=0.8,en-US;q=0.7".to_string(),
),
(
"baggage".to_string(),
"sentry-environment=production,sentry-release=d6add6fb0460641fd482d767a335ef72b9b6abb8,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c".to_string(),
),
("content-type".to_string(), input.content_type.to_string()),
("cookie".to_string(), cookie),
(GROK_INTERNAL_HEADER.to_string(), "1".to_string()),
("origin".to_string(), base_url.clone()),
("priority".to_string(), "u=1, i".to_string()),
("referer".to_string(), format!("{base_url}/")),
("sec-ch-ua".to_string(), browser_profile.sec_ch_ua),
("sec-ch-ua-mobile".to_string(), "?0".to_string()),
(
"sec-ch-ua-platform".to_string(),
browser_profile.sec_ch_ua_platform,
),
("sec-fetch-dest".to_string(), "empty".to_string()),
("sec-fetch-mode".to_string(), "cors".to_string()),
("sec-fetch-site".to_string(), "same-origin".to_string()),
("user-agent".to_string(), browser_profile.user_agent),
("x-statsig-id".to_string(), GROK_STATSIG_ID.to_string()),
("x-xai-request-id".to_string(), Uuid::new_v4().to_string()),
]);
if !apply_local_header_rules_with_request_headers(
&mut headers,
input.header_rules,
&["cookie", "content-type", "accept", GROK_INTERNAL_HEADER],
input.provider_request_body,
Some(input.original_request_body),
input.request_headers,
) {
return None;
}
Some(headers)
}
pub fn build_grok_app_chat_body(
client_api_format: &str,
model: Option<&str>,
original_body: &Value,
) -> Value {
let model = model.map(str::trim).filter(|value| !value.is_empty());
let client_api_format = client_api_format.trim().to_ascii_lowercase();
if client_api_format.eq_ignore_ascii_case("openai:image")
&& grok_is_image_edit_request(original_body)
{
return build_grok_image_edit_body(original_body);
}
let is_image_generation =
grok_is_image_generation_request(client_api_format.as_str(), model, original_body);
let message = if is_image_generation {
extract_grok_image_prompt(original_body)
.map(|prompt| format!("Drawing: {prompt}"))
.unwrap_or_else(|| "Drawing:".to_string())
} else {
match client_api_format.as_str() {
"openai:responses" | "openai:responses:compact" => {
extract_grok_message_from_openai_responses_body(original_body)
}
"claude:messages" => extract_grok_message_from_claude_messages_body(original_body),
_ => extract_grok_message_from_openai_chat_body(original_body),
}
};
let mut payload = grok_base_app_chat_payload(message, grok_mode_id_for_model(model));
if is_image_generation {
payload.insert(
"imageGenerationCount".to_string(),
json!(extract_grok_image_count(original_body)
.unwrap_or(1)
.clamp(1, GROK_IMAGE_GENERATION_MAX_COUNT)),
);
if let Some(size) = extract_grok_image_option(
original_body,
&["aspect_ratio", "aspectRatio", "ratio", "size"],
) {
payload.insert("size".to_string(), Value::String(size));
}
}
Value::Object(payload)
}
fn grok_is_image_generation_request(
client_api_format: &str,
model: Option<&str>,
body: &Value,
) -> bool {
if client_api_format.eq_ignore_ascii_case("openai:image") {
return !grok_is_image_edit_request(body);
}
if !matches!(
client_api_format,
"openai:chat" | "openai:responses" | "openai:responses:compact"
) {
return false;
}
grok_is_image_generation_model(model) || grok_has_image_generation_tool(body)
}
fn grok_is_image_generation_model(model: Option<&str>) -> bool {
let model = model.unwrap_or_default().trim().to_ascii_lowercase();
model.contains("grok-imagine-image") && !model.contains("edit")
}
fn grok_has_image_generation_tool(body: &Value) -> bool {
body.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.any(|tool| {
tool.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("image_generation"))
&& !tool
.get("action")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("edit"))
})
}
pub fn grok_is_image_edit_request(body: &Value) -> bool {
body.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.any(|tool| {
tool.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("image_generation"))
&& tool
.get("action")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("edit"))
})
}
pub fn build_grok_image_edit_body(body: &Value) -> Value {
json!({
"temporary": true,
"modelName": GROK_IMAGE_EDIT_MODEL_NAME,
"message": extract_grok_image_prompt(body).unwrap_or_default(),
"enableImageGeneration": true,
"returnImageBytes": false,
"returnRawGrokInXaiRequest": false,
"enableImageStreaming": true,
"imageGenerationCount": extract_grok_image_count(body).unwrap_or(2).clamp(1, 2),
"forceConcise": false,
"enableSideBySide": true,
"sendFinalMetadata": true,
"isReasoning": false,
"disableTextFollowUps": true,
"responseMetadata": {
"modelConfigOverride": {
"modelMap": {
"imageEditModel": GROK_IMAGE_EDIT_MODEL_KIND,
"imageEditModelConfig": {
"imageReferences": [],
"parentPostId": ""
}
}
}
},
"disableMemory": true,
"forceSideBySide": false,
})
}
fn grok_base_app_chat_payload(message: String, mode_id: &'static str) -> Map<String, Value> {
let mut payload = Map::new();
payload.insert("collectionIds".to_string(), json!([]));
payload.insert("connectors".to_string(), json!([]));
payload.insert(
"deviceEnvInfo".to_string(),
json!({
"darkModeEnabled": false,
"devicePixelRatio": 2,
"screenHeight": 1329,
"screenWidth": 2056,
"viewportHeight": 1083,
"viewportWidth": 2056,
}),
);
payload.insert("disableMemory".to_string(), json!(true));
payload.insert("disableSearch".to_string(), json!(false));
payload.insert("disableSelfHarmShortCircuit".to_string(), json!(false));
payload.insert("disableTextFollowUps".to_string(), json!(false));
payload.insert("enableImageGeneration".to_string(), json!(true));
payload.insert("enableImageStreaming".to_string(), json!(true));
payload.insert("enableSideBySide".to_string(), json!(true));
payload.insert("fileAttachments".to_string(), json!([]));
payload.insert("forceConcise".to_string(), json!(false));
payload.insert("forceSideBySide".to_string(), json!(false));
payload.insert("imageAttachments".to_string(), json!([]));
payload.insert("imageGenerationCount".to_string(), json!(2));
payload.insert("isAsyncChat".to_string(), json!(false));
payload.insert("message".to_string(), Value::String(message));
payload.insert("modeId".to_string(), Value::String(mode_id.to_string()));
payload.insert("responseMetadata".to_string(), json!({}));
payload.insert("returnImageBytes".to_string(), json!(false));
payload.insert("returnRawGrokInXaiRequest".to_string(), json!(false));
payload.insert("searchAllConnectors".to_string(), json!(false));
payload.insert("sendFinalMetadata".to_string(), json!(true));
payload.insert("temporary".to_string(), json!(true));
payload.insert(
"toolOverrides".to_string(),
json!({
"gmailSearch": false,
"googleCalendarSearch": false,
"outlookSearch": false,
"outlookCalendarSearch": false,
"googleDriveSearch": false,
}),
);
payload
}
fn grok_mode_id_for_model(model: Option<&str>) -> &'static str {
let model = model.unwrap_or_default().to_ascii_lowercase();
if model.contains("4.3") || model.contains("computer") {
"grok-420-computer-use-sa"
} else if model.contains("multi-agent") {
"heavy"
} else if model.contains("non-reasoning") || model.contains("fast") || model.contains("lite") {
"fast"
} else if model.contains("expert") || model.contains("reasoning") {
"expert"
} else if model.contains("0309-heavy") {
"auto"
} else if model.contains("heavy") {
"heavy"
} else {
"auto"
}
}
fn extract_grok_message_from_openai_chat_body(body: &Value) -> String {
let parts = body
.get("messages")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(|message| {
let role = message
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("user");
grok_chat_message_text(message).map(|text| format!("[{role}]: {text}"))
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
(!parts.is_empty())
.then(|| parts.join("\n\n"))
.or_else(|| value_text(body.get("prompt")?))
.unwrap_or_default()
}
fn extract_grok_message_from_openai_responses_body(body: &Value) -> String {
let mut parts = Vec::new();
if let Some(instructions) = body
.get("instructions")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
parts.push(format!("system: {instructions}"));
}
if let Some(input) = body.get("input").and_then(value_text) {
if !input.trim().is_empty() {
parts.push(input);
}
}
parts.join("\n\n")
}
fn extract_grok_message_from_claude_messages_body(body: &Value) -> String {
let mut parts = Vec::new();
if let Some(system) = body.get("system").and_then(value_text) {
if !system.trim().is_empty() {
parts.push(format!("system: {system}"));
}
}
if let Some(messages) = body.get("messages").and_then(Value::as_array) {
for message in messages {
let role = message
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("user");
if let Some(text) = grok_chat_message_text(message) {
parts.push(format!("[{role}]: {text}"));
}
}
}
parts.join("\n\n")
}
fn grok_chat_message_text(message: &Value) -> Option<String> {
let role = message
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if role.eq_ignore_ascii_case("tool") {
let text = value_text(message.get("content")?)?;
let text = text.trim();
if text.is_empty() {
return None;
}
let label = message
.get("tool_call_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|id| format!("[tool result for {id}]"))
.unwrap_or_else(|| "[tool result]".to_string());
return Some(format!("{label}:\n{text}"));
}
if role.eq_ignore_ascii_case("assistant") {
if let Some(tool_calls) = message.get("tool_calls").and_then(Value::as_array) {
if !tool_calls.is_empty() {
let mut parts = Vec::new();
if let Some(text) = message.get("content").and_then(value_text) {
let text = text.trim();
if !text.is_empty() {
parts.push(text.to_string());
}
}
parts.push(grok_tool_calls_text(tool_calls));
return Some(parts.join("\n"));
}
}
}
let text = value_text(message.get("content")?)?;
let text = strip_grok_generated_artifacts(text.trim())
.trim()
.to_string();
(!text.is_empty()).then_some(text)
}
fn grok_tool_calls_text(tool_calls: &[Value]) -> String {
let mut parts = Vec::new();
for tool_call in tool_calls {
let function = tool_call.get("function").unwrap_or(tool_call);
let name = function
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("tool");
let arguments = function
.get("arguments")
.and_then(value_text)
.unwrap_or_else(|| "{}".to_string());
parts.push(format!(
"<tool_call name=\"{name}\">{arguments}</tool_call>"
));
}
parts.join("\n")
}
fn strip_grok_generated_artifacts(text: &str) -> String {
let marker = "[grok2api-sources]: #";
let Some(marker_index) = text.find(marker) else {
return text.to_string();
};
let prefix = &text[..marker_index];
if let Some(sources_index) = prefix.rfind("## Sources") {
prefix[..sources_index].trim_end().to_string()
} else {
text.to_string()
}
}
fn extract_grok_image_prompt(body: &Value) -> Option<String> {
body.get("prompt")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| body.get("input").and_then(value_text))
.or_else(|| extract_grok_last_user_text_from_openai_chat_body(body))
}
fn extract_grok_image_count(body: &Value) -> Option<u64> {
body.get("n")
.and_then(Value::as_u64)
.or_else(|| body.get("imageGenerationCount").and_then(Value::as_u64))
.or_else(|| {
body.get("image_config").and_then(|config| {
config
.get("n")
.or_else(|| config.get("imageGenerationCount"))
.and_then(Value::as_u64)
})
})
.or_else(|| {
body.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.find_map(|tool| {
tool.get("n")
.or_else(|| tool.get("imageGenerationCount"))
.and_then(Value::as_u64)
})
})
}
fn extract_grok_image_option(body: &Value, keys: &[&str]) -> Option<String> {
for key in keys {
if let Some(value) = body
.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(value.to_string());
}
}
if let Some(config) = body.get("image_config").and_then(Value::as_object) {
for key in keys {
if let Some(value) = config
.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(value.to_string());
}
}
}
body.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.find_map(|tool| {
keys.iter().find_map(|key| {
tool.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
})
}
fn extract_grok_last_user_text_from_openai_chat_body(body: &Value) -> Option<String> {
body.get("messages")
.and_then(Value::as_array)?
.iter()
.rev()
.find_map(|message| {
let role = message.get("role").and_then(Value::as_str)?;
if !role.eq_ignore_ascii_case("user") {
return None;
}
grok_chat_message_text(message)
})
}
fn value_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Array(items) => {
let mut parts = Vec::new();
for item in items {
if let Some(text) = value_text(item) {
if !text.trim().is_empty() {
parts.push(text);
}
}
}
Some(parts.join("\n"))
}
Value::Object(object) => {
if is_attachment_content_block(object) {
return None;
}
if let Some(text) = object
.get("text")
.or_else(|| object.get("input_text"))
.or_else(|| object.get("content"))
.and_then(value_text)
{
return Some(text);
}
None
}
_ => None,
}
}
fn is_attachment_content_block(object: &Map<String, Value>) -> bool {
let Some(block_type) = object.get("type").and_then(Value::as_str) else {
return object.contains_key("image_url")
|| object.contains_key("file")
|| object.contains_key("file_data")
|| object.contains_key("file_url");
};
matches!(
block_type,
"image_url" | "input_image" | "input_file" | "file" | "image" | "document"
)
}
fn parse_grok_auth_config(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 grok_cookie_from_transport(transport: &GatewayProviderTransportSnapshot) -> Option<String> {
let auth_config = parse_grok_auth_config(transport);
let token = auth_config
.as_ref()
.and_then(|config| {
grok_config_string(
Some(config),
&[
"sso_token",
"ssoToken",
"access_token",
"token",
"session_token",
"sessionToken",
],
)
})
.or_else(|| {
let raw = transport.key.decrypted_api_key.trim();
(!raw.is_empty()).then(|| raw.to_string())
})?;
let token = strip_cookie_prefix(token.trim(), "sso=");
if token.is_empty() {
return None;
}
let sso_rw = auth_config
.as_ref()
.and_then(|config| grok_config_string(Some(config), &["sso_rw_token", "ssoRwToken"]))
.map(|value| strip_cookie_prefix(value.trim(), "sso-rw="))
.filter(|value| !value.is_empty())
.unwrap_or_else(|| token.clone());
let mut parts = vec![format!("sso={token}"), format!("sso-rw={sso_rw}")];
if let Some(extra_cookies) = auth_config
.as_ref()
.and_then(|config| grok_config_string(Some(config), &["cf_cookies", "cfCookies", "cookie"]))
.and_then(|value| normalize_grok_extra_cookies(value.as_str()))
{
parts.push(extra_cookies);
}
if let Some(cf_clearance) = auth_config
.as_ref()
.and_then(|config| grok_config_string(Some(config), &["cf_clearance", "cfClearance"]))
.map(|value| strip_cookie_prefix(value.trim(), "cf_clearance="))
.filter(|value| !value.is_empty())
{
if !parts.iter().any(|part| part.contains("cf_clearance=")) {
parts.push(format!("cf_clearance={cf_clearance}"));
}
}
Some(parts.join("; "))
}
fn grok_config_string(config: Option<&Value>, fields: &[&str]) -> Option<String> {
grok_config_string_from_object(config.and_then(Value::as_object), fields)
}
fn grok_config_string_from_object(
object: Option<&Map<String, Value>>,
fields: &[&str],
) -> Option<String> {
let object = object?;
fields.iter().find_map(|field| {
object
.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn strip_cookie_prefix(value: &str, prefix: &str) -> String {
value
.strip_prefix(prefix)
.unwrap_or(value)
.trim()
.to_string()
}
fn normalize_grok_extra_cookies(value: &str) -> Option<String> {
let parts = value
.trim()
.trim_matches(';')
.split(';')
.filter_map(|segment| {
let (name, value) = segment.trim().split_once('=')?;
let name = name.trim();
let value = value.trim();
if name.is_empty()
|| value.is_empty()
|| name.eq_ignore_ascii_case("sso")
|| name.eq_ignore_ascii_case("sso-rw")
{
return None;
}
Some(format!("{name}={value}"))
})
.collect::<Vec<_>>();
(!parts.is_empty()).then(|| parts.join("; "))
}
#[cfg(test)]
mod tests {
use aether_contracts::{
ResolvedTransportProfile, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_HTTP_MODE_AUTO,
TRANSPORT_POOL_SCOPE_KEY,
};
use serde_json::json;
use super::{
build_grok_app_chat_body, build_grok_browser_headers, build_grok_upstream_url,
grok_browser_resolved_transport_profile,
grok_browser_resolved_transport_profile_from_auth_config, resolve_grok_session_auth,
GrokHeaderInput, GROK_CHAT_PATH, GROK_INTERNAL_HEADER,
};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
fn sample_transport(auth_config: serde_json::Value) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Grok".to_string(),
provider_type: "grok".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
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:chat".to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://grok.com/".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: "key".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,
decrypted_api_key: String::new(),
decrypted_auth_config: Some(auth_config.to_string()),
},
}
}
#[test]
fn resolves_sso_cookie_from_auth_config() {
let transport = sample_transport(json!({
"sso_token": "sso=abc",
"sso_rw_token": "rw",
"cf_clearance": "cf"
}));
let (_, value) = resolve_grok_session_auth(&transport).expect("cookie should resolve");
assert_eq!(value, "sso=abc; sso-rw=rw; cf_clearance=cf");
}
#[test]
fn resolves_sso_cookie_without_duplicate_session_cookies() {
let transport = sample_transport(json!({
"sso_token": "abc",
"sso_rw_token": "rw",
"cf_cookies": "i18nextLng=zh; sso=ignored; sso-rw=ignored-rw; cf_clearance=cf"
}));
let (_, value) = resolve_grok_session_auth(&transport).expect("cookie should resolve");
assert_eq!(value, "sso=abc; sso-rw=rw; i18nextLng=zh; cf_clearance=cf");
}
#[test]
fn builds_grok_browser_headers_with_internal_marker() {
let transport = sample_transport(json!({"sso_token": "abc"}));
let transport_profile =
grok_browser_resolved_transport_profile(None, "test").expect("profile should resolve");
let headers = build_grok_browser_headers(GrokHeaderInput {
transport: &transport,
transport_profile: Some(&transport_profile),
request_headers: None,
content_type: "application/json",
accept: "*/*",
header_rules: None,
provider_request_body: &json!({"message":"hi"}),
original_request_body: &json!({"messages":[]}),
})
.expect("headers should build");
assert_eq!(headers.get(GROK_INTERNAL_HEADER), Some(&"1".to_string()));
assert_eq!(
headers.get("cookie"),
Some(&"sso=abc; sso-rw=abc".to_string())
);
assert_eq!(headers.get("origin"), Some(&"https://grok.com".to_string()));
assert!(headers
.get("user-agent")
.is_some_and(|value| value.contains("Chrome/136.0.0.0")));
assert_eq!(
headers.get("sec-ch-ua"),
Some(
&r#""Google Chrome";v="136", "Chromium";v="136", "Not(A:Brand";v="24""#.to_string()
)
);
assert_eq!(
headers.get("sec-ch-ua-platform"),
Some(&r#""macOS""#.to_string())
);
}
#[test]
fn builds_grok_browser_headers_from_transport_profile() {
let mut transport = sample_transport(json!({
"sso_token": "abc",
"user_agent": "Mozilla/5.0 custom",
"browser_profile": "chrome136"
}));
transport.key.fingerprint = Some(json!({
"transport_profile": {
"profile_id": "chrome137",
"backend": "browser_wreq",
"extra": {"browser_profile": "chrome137"}
}
}));
let transport_profile = grok_browser_resolved_transport_profile(Some("chrome137"), "test")
.expect("profile should resolve");
let headers = build_grok_browser_headers(GrokHeaderInput {
transport: &transport,
transport_profile: Some(&transport_profile),
request_headers: None,
content_type: "application/json",
accept: "*/*",
header_rules: None,
provider_request_body: &json!({"message":"hi"}),
original_request_body: &json!({"messages":[]}),
})
.expect("headers should build");
assert!(headers
.get("user-agent")
.is_some_and(|value| value.contains("Chrome/137.0.0.0")));
assert_eq!(
headers.get("sec-ch-ua"),
Some(
&r#""Google Chrome";v="137", "Chromium";v="137", "Not(A:Brand";v="24""#.to_string()
)
);
}
#[test]
fn infers_grok_browser_profile_from_legacy_user_agent() {
let auth_config = json!({
"sso_token": "abc",
"user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36"
});
let auth_config = auth_config.as_object().expect("object");
let profile = grok_browser_resolved_transport_profile_from_auth_config(auth_config, "test")
.expect("profile should resolve from user agent");
assert_eq!(profile.profile_id, "chrome137");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(serde_json::Value::as_str),
Some("chrome137")
);
}
#[test]
fn rejects_unsupported_legacy_grok_user_agent_profile() {
let auth_config = json!({
"sso_token": "abc",
"user_agent": "Mozilla/5.0 Version/18.0 Safari/605.1.15"
});
let auth_config = auth_config.as_object().expect("object");
assert!(
grok_browser_resolved_transport_profile_from_auth_config(auth_config, "test").is_none()
);
}
#[test]
fn rejects_unsupported_grok_browser_header_profile() {
let mut transport = sample_transport(json!({"sso_token": "abc"}));
transport.key.fingerprint = Some(json!({
"transport_profile": {
"profile_id": "safari999",
"backend": "browser_wreq",
"extra": {"browser_profile": "safari999"}
}
}));
let transport_profile = ResolvedTransportProfile {
profile_id: "safari999".to_string(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: Some(json!({"browser_profile": "safari999"})),
};
let headers = build_grok_browser_headers(GrokHeaderInput {
transport: &transport,
transport_profile: Some(&transport_profile),
request_headers: None,
content_type: "application/json",
accept: "*/*",
header_rules: None,
provider_request_body: &json!({"message":"hi"}),
original_request_body: &json!({"messages":[]}),
});
assert!(headers.is_none());
}
#[test]
fn builds_rest_chat_url_from_endpoint_base() {
let transport = sample_transport(json!({"sso_token": "abc"}));
assert_eq!(
build_grok_upstream_url(&transport, GROK_CHAT_PATH),
"https://grok.com/rest/app-chat/conversations/new"
);
}
#[test]
fn builds_app_chat_body_from_openai_chat_messages() {
let body = build_grok_app_chat_body(
"openai:chat",
Some("grok-4.20-fast"),
&json!({
"messages": [
{"role": "system", "content": "be short"},
{"role": "user", "content": [{"type":"text", "text":"hello"}]}
]
}),
);
assert_eq!(body["modeId"], "fast");
assert_eq!(body["message"], "[system]: be short\n\n[user]: hello");
assert_eq!(body["enableImageGeneration"], true);
}
#[test]
fn builds_app_chat_body_omits_uploaded_attachment_blocks_from_prompt() {
let body = build_grok_app_chat_body(
"openai:chat",
Some("grok-4.20-fast"),
&json!({
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "describe this"},
{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
{"type": "file", "file": {"filename": "notes.txt", "file_data": "data:text/plain;base64,bm90ZXM="}}
]
}]
}),
);
assert_eq!(body["message"], "[user]: describe this");
}
#[test]
fn builds_app_chat_body_keeps_openai_chat_history() {
let body = build_grok_app_chat_body(
"openai:chat",
Some("grok-4.20-fast"),
&json!({
"messages": [
{"role": "system", "content": "remember the code word"},
{"role": "user", "content": "the code word is lantern"},
{"role": "assistant", "content": "Got it.\n\n## Sources\n[grok2api-sources]: #\n- [old](https://example.com)"},
{"role": "user", "content": "what is the code word?"}
]
}),
);
assert_eq!(
body["message"],
"[system]: remember the code word\n\n[user]: the code word is lantern\n\n[assistant]: Got it.\n\n[user]: what is the code word?"
);
}
#[test]
fn builds_app_chat_body_keeps_tool_context() {
let body = build_grok_app_chat_body(
"openai:chat",
Some("grok-4.20-fast"),
&json!({
"messages": [
{
"role": "assistant",
"content": null,
"tool_calls": [{
"id": "call_1",
"type": "function",
"function": {"name": "lookup", "arguments": "{\"q\":\"weather\"}"}
}]
},
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"}
]
}),
);
assert!(body["message"]
.as_str()
.expect("message should be string")
.contains("<tool_call name=\"lookup\">{\"q\":\"weather\"}</tool_call>"));
assert!(body["message"]
.as_str()
.expect("message should be string")
.contains("[tool]: [tool result for call_1]:\nsunny"));
}
#[test]
fn maps_grok_0309_model_variants_to_web_modes() {
let request = json!({"messages": [{"role": "user", "content": "hello"}]});
let fast = build_grok_app_chat_body(
"openai:chat",
Some("grok-4.20-0309-non-reasoning-heavy"),
&request,
);
let auto = build_grok_app_chat_body("openai:chat", Some("grok-4.20-0309-heavy"), &request);
let expert = build_grok_app_chat_body(
"openai:chat",
Some("grok-4.20-0309-reasoning-super"),
&request,
);
let heavy =
build_grok_app_chat_body("openai:chat", Some("grok-4.20-multi-agent-0309"), &request);
assert_eq!(fast["modeId"], "fast");
assert_eq!(auto["modeId"], "auto");
assert_eq!(expert["modeId"], "expert");
assert_eq!(heavy["modeId"], "heavy");
}
#[test]
fn builds_app_chat_body_from_openai_image_prompt() {
let body = build_grok_app_chat_body(
"openai:image",
Some("grok-imagine-image"),
&json!({"prompt":"a red chair", "n": 3}),
);
assert_eq!(body["modeId"], "auto");
assert_eq!(body["message"], "Drawing: a red chair");
assert_eq!(body["imageGenerationCount"], 3);
}
#[test]
fn builds_app_chat_body_from_openai_responses_image_generation_tool() {
let body = build_grok_app_chat_body(
"openai:responses",
Some("grok-imagine-image-lite"),
&json!({
"input": "a red chair",
"tools": [{
"type": "image_generation",
"n": 4,
"size": "1280x720"
}]
}),
);
assert_eq!(body["modeId"], "fast");
assert_eq!(body["message"], "Drawing: a red chair");
assert_eq!(body["imageGenerationCount"], 4);
assert_eq!(body["size"], "1280x720");
}
#[test]
fn app_chat_image_generation_count_uses_gateway_ceiling() {
let body = build_grok_app_chat_body(
"openai:responses",
Some("grok-imagine-image-lite"),
&json!({
"input": "a red chair",
"tools": [{
"type": "image_generation",
"n": 10
}]
}),
);
assert_eq!(body["imageGenerationCount"], 4);
}
#[test]
fn builds_app_chat_body_from_openai_chat_image_model_request() {
let body = build_grok_app_chat_body(
"openai:chat",
Some("grok-imagine-image-pro"),
&json!({
"messages": [
{"role": "system", "content": "ignore prior style"},
{"role": "user", "content": [{"type":"text", "text":"a blue sofa"}]}
],
"image_config": {
"n": 2,
"size": "720x1280"
}
}),
);
assert_eq!(body["message"], "Drawing: a blue sofa");
assert_eq!(body["imageGenerationCount"], 2);
assert_eq!(body["size"], "720x1280");
}
#[test]
fn builds_app_chat_body_from_openai_image_edit_request() {
let body = build_grok_app_chat_body(
"openai:image",
Some("grok-imagine-image-edit"),
&json!({
"input": [
{"type": "input_text", "text": "make the chair blue"},
{"type": "input_image", "image_url": "data:image/png;base64,aW1hZ2U="}
],
"tools": [{
"type": "image_generation",
"action": "edit"
}]
}),
);
assert_eq!(body["modelName"], "imagine-image-edit");
assert_eq!(body["message"], "make the chair blue");
assert_eq!(body["enableImageGeneration"], true);
assert_eq!(body["enableImageStreaming"], true);
assert_eq!(body["disableTextFollowUps"], true);
assert_eq!(body["imageGenerationCount"], 2);
assert_eq!(
body["responseMetadata"]["modelConfigOverride"]["modelMap"]["imageEditModel"],
"imagine"
);
}
}
Reference in New Issue View Git Blame Copy Permalink