mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 08:27:46 +08:00
feat(providers): add xAI provider with device code OAuth
Add a separate `xai` provider type for xAI Grok CLI subscription accounts. It is independent of the existing `grok` provider, which reverse-proxies grok.com with browser cookies; behavior of `grok` is unchanged. Account binding uses the xAI device code flow, so no local callback listener is needed and headless deployments can bind accounts. Refresh tokens can also be imported individually or in batches, and are rotated on refresh. OAuth requests default to the cli-chat-proxy Responses API; API keys and compact stay on api.x.ai. Explicit custom gateways are preserved. Only `openai:responses` and `openai:responses:compact` are exposed; Chat, Claude and Gemini clients reach the provider through Aether's existing cross-format conversion rather than new native endpoints. Upstream Responses payloads are sanitized for what xAI actually rejects: `previous_response_id` and `metadata.user_id` are dropped, hosted `tool_choice` is rewritten, `web_search` is restored for converted clients, `image_generation` is stripped on older Grok conversation models, unsupported reasoning effort is removed, and requested `reasoning.encrypted_content` is preserved with a replay policy keyed on the configured provider type rather than the model name. Quota refresh reads /user and /billing?format=credits and stores a structured usage snapshot; a prepaid balance keeps an account selectable after the weekly allowance is exhausted. API-key accounts skip the subscription billing surface. The admin UI shows remaining weekly quota as a labeled bar in the provider drawer and the pool list. Co-Authored-By: Claude Opus 5 <[email protected]>
This commit is contained in:
@@ -6649,7 +6649,7 @@ mod tests {
|
||||
}
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
|
||||
@@ -7375,9 +7375,27 @@ mod tests {
|
||||
queue: Arc<dyn RuntimeQueueStore>,
|
||||
policy_started: Arc<tokio::sync::Notify>,
|
||||
release_policy: Arc<tokio::sync::Notify>,
|
||||
policy_released: Arc<AtomicBool>,
|
||||
policy_reads: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl BlockingPolicyQueueConfiguredUsageStore {
|
||||
fn new(queue: Arc<dyn RuntimeQueueStore>) -> Self {
|
||||
Self {
|
||||
queue,
|
||||
policy_started: Arc::new(tokio::sync::Notify::new()),
|
||||
release_policy: Arc::new(tokio::sync::Notify::new()),
|
||||
policy_released: Arc::new(AtomicBool::new(false)),
|
||||
policy_reads: Arc::new(AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
fn release_blocked_policy(&self) {
|
||||
self.policy_released.store(true, Ordering::Release);
|
||||
self.release_policy.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FailingPolicyUsageStore {
|
||||
inner: NoRedisUsageStore,
|
||||
@@ -8409,7 +8427,18 @@ mod tests {
|
||||
async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
|
||||
self.policy_reads.fetch_add(1, Ordering::AcqRel);
|
||||
self.policy_started.notify_one();
|
||||
self.release_policy.notified().await;
|
||||
// Latch the gate: Notify is edge-triggered, and later policy reads
|
||||
// (or a waiter that subscribed after a single notify) must not hang.
|
||||
loop {
|
||||
if self.policy_released.load(Ordering::Acquire) {
|
||||
break;
|
||||
}
|
||||
let notified = self.release_policy.notified();
|
||||
if self.policy_released.load(Ordering::Acquire) {
|
||||
break;
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
Ok(UsageBodyCapturePolicy::default())
|
||||
}
|
||||
}
|
||||
@@ -9043,15 +9072,39 @@ mod tests {
|
||||
.await
|
||||
.expect("a duplicate first-byte marker must release the terminal barrier");
|
||||
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(
|
||||
records.len(),
|
||||
2,
|
||||
"the duplicate first byte must be coalesced"
|
||||
);
|
||||
assert_eq!(records[0].status, "streaming");
|
||||
assert_eq!(records[1].status, "completed");
|
||||
drop(records);
|
||||
{
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(
|
||||
records.len(),
|
||||
2,
|
||||
"the duplicate first byte must be coalesced"
|
||||
);
|
||||
assert_eq!(records[0].status, "streaming");
|
||||
assert_eq!(records[1].status, "completed");
|
||||
}
|
||||
|
||||
// The terminal persistence notification can arrive before the submission
|
||||
// dispatcher accounts for its completed task and releases admission.
|
||||
timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
let snapshot = runtime.metrics_snapshot();
|
||||
if snapshot.lifecycle_submission_pending == 0
|
||||
&& snapshot.first_byte_persistence_pending == 0
|
||||
&& snapshot.ordered_lifecycle_pending == 0
|
||||
&& runtime
|
||||
.lifecycle_submission
|
||||
.state
|
||||
.admission
|
||||
.available_permits()
|
||||
== CAPACITY
|
||||
{
|
||||
break;
|
||||
}
|
||||
sleep(Duration::from_millis(1)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("duplicate first-byte submission accounting should drain");
|
||||
|
||||
let snapshot = runtime.metrics_snapshot();
|
||||
assert_eq!(snapshot.lifecycle_submission_pending, 0);
|
||||
@@ -12325,12 +12378,9 @@ mod tests {
|
||||
async fn event_capture_budget_bounds_blocked_policy_waiters_and_releases_on_cancel_or_basic() {
|
||||
for limit in [0, 64 * 1024] {
|
||||
let runtime = UsageRuntime::new(UsageRuntimeConfig::default()).expect("runtime");
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore {
|
||||
queue: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
|
||||
policy_started: Arc::new(tokio::sync::Notify::new()),
|
||||
release_policy: Arc::new(tokio::sync::Notify::new()),
|
||||
policy_reads: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore::new(Arc::new(
|
||||
RuntimeState::memory(MemoryRuntimeStateConfig::default()),
|
||||
));
|
||||
let budget = Arc::new(crate::event_capture_budget::EventCaptureMemoryBudget::new(
|
||||
limit,
|
||||
));
|
||||
@@ -12391,7 +12441,7 @@ mod tests {
|
||||
.await
|
||||
.expect("replacement policy read starts");
|
||||
assert_eq!(budget.retained_bytes(), retained);
|
||||
store.release_policy.notify_one();
|
||||
store.release_blocked_policy();
|
||||
let event = timeout(Duration::from_secs(2), completing)
|
||||
.await
|
||||
.expect("Basic policy completes")
|
||||
@@ -13398,12 +13448,7 @@ mod tests {
|
||||
Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
|
||||
let tracked_queue = Arc::new(FlakyAppendQueueStore::new(inner_queue, 0));
|
||||
let queue: Arc<dyn RuntimeQueueStore> = tracked_queue.clone();
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore {
|
||||
queue,
|
||||
policy_started: Arc::new(tokio::sync::Notify::new()),
|
||||
release_policy: Arc::new(tokio::sync::Notify::new()),
|
||||
policy_reads: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore::new(queue);
|
||||
let runtime = UsageRuntime::new(config).expect("usage runtime should build");
|
||||
let request_id = "req-terminal-seed-waits-for-turn";
|
||||
let plan = terminal_test_plan(request_id);
|
||||
@@ -13425,9 +13470,10 @@ mod tests {
|
||||
assert_eq!(blocked_snapshot.terminal_submission_in_flight, 0);
|
||||
assert!(blocked_snapshot.lifecycle_submission_pending >= 2);
|
||||
|
||||
store.release_policy.notify_waiters();
|
||||
store.release_blocked_policy();
|
||||
timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
store.release_blocked_policy();
|
||||
let snapshot = runtime.metrics_snapshot();
|
||||
if tracked_queue.successful_appends.load(Ordering::Acquire) == 1
|
||||
&& snapshot.lifecycle_submission_pending == 0
|
||||
@@ -13435,7 +13481,7 @@ mod tests {
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
sleep(Duration::from_millis(1)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
@@ -13468,12 +13514,7 @@ mod tests {
|
||||
Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
|
||||
let tracked_queue = Arc::new(FlakyAppendQueueStore::new(inner_queue, 0));
|
||||
let queue: Arc<dyn RuntimeQueueStore> = tracked_queue.clone();
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore {
|
||||
queue,
|
||||
policy_started: Arc::new(tokio::sync::Notify::new()),
|
||||
release_policy: Arc::new(tokio::sync::Notify::new()),
|
||||
policy_reads: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore::new(queue);
|
||||
let runtime = UsageRuntime::new(config).expect("usage runtime should build");
|
||||
let policy_started = store.policy_started.notified();
|
||||
|
||||
@@ -13520,9 +13561,10 @@ mod tests {
|
||||
assert_eq!(blocked_snapshot.terminal_submission_in_flight, 1);
|
||||
assert!(blocked_snapshot.lifecycle_submission_pending <= BACKLOG + 1);
|
||||
|
||||
store.release_policy.notify_waiters();
|
||||
store.release_blocked_policy();
|
||||
timeout(Duration::from_secs(5), async {
|
||||
loop {
|
||||
store.release_blocked_policy();
|
||||
let snapshot = runtime.metrics_snapshot();
|
||||
if tracked_queue.successful_appends.load(Ordering::Acquire) == BACKLOG + 1
|
||||
&& snapshot.lifecycle_submission_pending == 0
|
||||
@@ -13530,7 +13572,7 @@ mod tests {
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
sleep(Duration::from_millis(1)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
@@ -13828,12 +13870,7 @@ mod tests {
|
||||
Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
|
||||
let tracked_queue = Arc::new(FlakyAppendQueueStore::new(inner_queue, 0));
|
||||
let queue: Arc<dyn RuntimeQueueStore> = tracked_queue.clone();
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore {
|
||||
queue,
|
||||
policy_started: Arc::new(tokio::sync::Notify::new()),
|
||||
release_policy: Arc::new(tokio::sync::Notify::new()),
|
||||
policy_reads: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let store = BlockingPolicyQueueConfiguredUsageStore::new(queue);
|
||||
let runtime = UsageRuntime::new(config).expect("usage runtime should build");
|
||||
let policy_started = store.policy_started.notified();
|
||||
runtime
|
||||
@@ -13892,10 +13929,10 @@ mod tests {
|
||||
.expect("terminal submissions should reach the execution backlog");
|
||||
let saturated_snapshot = runtime.metrics_snapshot();
|
||||
|
||||
store.release_policy.notify_waiters();
|
||||
store.release_blocked_policy();
|
||||
let all_completed = timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
store.release_policy.notify_waiters();
|
||||
store.release_blocked_policy();
|
||||
if tracked_queue.successful_appends.load(Ordering::Acquire)
|
||||
== EXCESS_SUBMISSIONS + 1
|
||||
&& runtime.metrics_snapshot().terminal_submission_in_flight == 0
|
||||
|
||||
Reference in New Issue
Block a user