use std::collections::{HashMap, VecDeque};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;
use crate::config::{AgentConfig, LlmHttpAuthMode, SharedAgentConfig};
use crate::llm::ChatCompletionsBackend;
use crate::memory::long_term_memory::LongTermMemoryRuntime;
use crate::per_turn_flight::PerTurnFlight;
use crate::request_audit::WebRequestAudit;
use crate::sse::SseStreamHub;
use crate::types::{CommandApprovalDecision, LlmSeedOverride, Message, Tool};
use crate::web::WebChatJobAppFacet;
use log::debug;
use tokio::sync::{Semaphore, mpsc, oneshot};
mod stream_finish;
mod worker;
#[cfg(test)]
mod tests;
const RECENT_CAP: usize = 32;
#[derive(Clone)]
pub(crate) struct WebChatQueueDeps {
pub cfg: SharedAgentConfig,
pub api_key: String,
pub client: reqwest::Client,
pub tools: Vec<Tool>,
pub chat_queue: ChatJobQueue,
pub long_term_memory: Option<Arc<LongTermMemoryRuntime>>,
pub sse_stream_hub: Arc<SseStreamHub>,
pub llm_backend: Option<&'static (dyn ChatCompletionsBackend + 'static)>,
pub turn_runner: Arc<dyn crate::TurnRunner>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WebClientLlmThinkingMode {
On,
Off,
}
#[derive(Clone, Debug, Default)]
pub struct WebChatLlmOverride {
pub api_base: Option<String>,
pub model: Option<String>,
pub api_key: Option<String>,
pub llm_context_tokens: Option<u32>,
pub llm_thinking_mode: Option<WebClientLlmThinkingMode>,
}
pub(super) fn resolve_web_llm_for_job(
deps: &WebChatQueueDeps,
cfg_snap: Arc<AgentConfig>,
ov: Option<&WebChatLlmOverride>,
) -> (Arc<AgentConfig>, String) {
let (mut cfg, key) = match ov {
None => (cfg_snap, deps.api_key.clone()),
Some(o) => {
let mut c = (*cfg_snap).clone();
let mut key = deps.api_key.clone();
if let Some(ref x) = o.api_base {
c.llm.api_base.clone_from(x);
}
if let Some(ref x) = o.api_key {
key.clone_from(x);
c.llm.llm_http_auth_mode = LlmHttpAuthMode::Bearer;
}
if let Some(n) = o.llm_context_tokens {
c.llm_sampling.llm_context_tokens = n;
}
if let Some(mode) = o.llm_thinking_mode {
match mode {
WebClientLlmThinkingMode::On => {
c.llm_vendor_flags.llm_bigmodel_thinking = true;
c.llm_vendor_flags.llm_kimi_thinking_disabled = false;
}
WebClientLlmThinkingMode::Off => {
c.llm_vendor_flags.llm_bigmodel_thinking = false;
c.llm_vendor_flags.llm_kimi_thinking_disabled = true;
}
}
}
(Arc::new(c), key)
}
};
{
let mut c = (*cfg).clone();
c.per_plan_policy.planner_executor_mode = crate::config::PlannerExecutorMode::SingleAgent;
c.per_plan_policy.orchestration_profile = crate::config::OrchestrationProfile::ReAct;
cfg = Arc::new(c);
}
(cfg, key)
}
pub(super) fn resolve_executor_llm_for_job(
deps: &WebChatQueueDeps,
cfg_snap: Arc<AgentConfig>,
ov: Option<&WebChatLlmOverride>,
) -> Option<(Arc<AgentConfig>, String)> {
let o = ov?;
let mut c = (*cfg_snap).clone();
let mut key = deps.api_key.clone();
let mut has_override = false;
if let Some(ref x) = o.api_base {
c.llm.api_base.clone_from(x);
has_override = true;
}
if let Some(ref x) = o.model {
c.llm.model.clone_from(x);
has_override = true;
}
if let Some(ref x) = o.api_key {
key.clone_from(x);
c.llm.llm_http_auth_mode = LlmHttpAuthMode::Bearer;
has_override = true;
}
if has_override {
Some((Arc::new(c), key))
} else {
None
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct PerFlightStatusEntry {
pub job_id: u64,
pub awaiting_plan_rewrite_model: bool,
pub plan_rewrite_attempts: usize,
pub require_plan_in_final_content: bool,
}
struct PerFlightJobGuard {
queue: ChatJobQueue,
job_id: u64,
}
impl Drop for PerFlightJobGuard {
fn drop(&mut self) {
self.queue.unregister_per_job_per_flight(self.job_id);
}
}
#[derive(Debug, Clone, Copy)]
pub struct ChatQueueFull {
pub max_pending: usize,
}
#[derive(Debug)]
pub enum ChatJsonJobFailure {
ConversationConflict,
Agent(crate::agent::agent_turn::RunAgentTurnError),
}
pub struct WebChatJobEnvelope {
pub job_id: u64,
pub queue_deps: Arc<WebChatQueueDeps>,
pub app: WebChatJobAppFacet,
pub conversation_id: String,
pub messages: Vec<Message>,
pub expected_revision: Option<u64>,
pub request_agent_role: Option<String>,
pub persisted_active_agent_role: Option<String>,
pub request_session_mode: Option<String>,
pub persisted_active_session_mode: Option<String>,
pub work_dir: PathBuf,
pub workspace_is_set: bool,
pub temperature_override: Option<f32>,
pub seed_override: LlmSeedOverride,
pub llm_override: Option<WebChatLlmOverride>,
pub executor_llm_override: Option<WebChatLlmOverride>,
pub readonly_tool_ttl_cache_secs: Option<u64>,
pub request_audit: WebRequestAudit,
pub request_id: Option<String>,
pub github_token: Option<String>,
#[allow(dead_code)]
pub client_sse_protocol: Option<u8>,
}
pub struct JsonSubmitParams {
pub envelope: WebChatJobEnvelope,
pub reply_tx: oneshot::Sender<Result<Vec<Message>, ChatJsonJobFailure>>,
}
pub struct StreamSubmitParams {
pub envelope: WebChatJobEnvelope,
pub stream_event_tx: mpsc::Sender<(u64, String)>,
pub web_approval_session: Option<WebApprovalSession>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct ChatJobRecord {
pub job_id: u64,
pub kind: String,
pub ok: bool,
pub cancelled: bool,
pub duration_ms: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_preview: Option<String>,
}
struct QueueMetrics {
running: AtomicUsize,
completed_ok: AtomicU64,
completed_cancelled: AtomicU64,
completed_err: AtomicU64,
}
impl Default for QueueMetrics {
fn default() -> Self {
Self {
running: AtomicUsize::new(0),
completed_ok: AtomicU64::new(0),
completed_cancelled: AtomicU64::new(0),
completed_err: AtomicU64::new(0),
}
}
}
pub(super) enum QueuedChatJob {
Stream {
envelope: WebChatJobEnvelope,
stream_event_tx: mpsc::Sender<(u64, String)>,
web_approval_session: Option<WebApprovalSession>,
},
Json {
envelope: WebChatJobEnvelope,
reply_tx: oneshot::Sender<Result<Vec<Message>, ChatJsonJobFailure>>,
},
}
pub struct WebApprovalSession {
pub session_id: String,
pub approval_rx: mpsc::Receiver<CommandApprovalDecision>,
}
impl QueuedChatJob {
fn job_id(&self) -> u64 {
match self {
QueuedChatJob::Stream { envelope, .. } | QueuedChatJob::Json { envelope, .. } => {
envelope.job_id
}
}
}
}
struct Inner {
submit_tx: mpsc::Sender<QueuedChatJob>,
max_concurrent: usize,
max_pending: usize,
next_job_id: AtomicU64,
metrics: Arc<QueueMetrics>,
recent: Arc<Mutex<VecDeque<ChatJobRecord>>>,
active_per_flights: Arc<Mutex<HashMap<u64, Arc<PerTurnFlight>>>>,
shutdown_triggered: Arc<AtomicBool>,
}
#[derive(Clone)]
pub struct ChatJobQueue {
inner: Arc<Inner>,
}
impl ChatJobQueue {
pub fn new(max_concurrent: usize, max_pending: usize) -> Self {
let max_concurrent = max_concurrent.max(1);
let max_pending = max_pending.max(1);
let (submit_tx, rx) = mpsc::channel::<QueuedChatJob>(max_pending);
let sem = Arc::new(Semaphore::new(max_concurrent));
let metrics = Arc::new(QueueMetrics::default());
let recent = Arc::new(Mutex::new(VecDeque::with_capacity(RECENT_CAP)));
let metrics_loop = metrics.clone();
let recent_loop = recent.clone();
tokio::spawn(dispatcher_loop(rx, sem, metrics_loop, recent_loop));
Self {
inner: Arc::new(Inner {
submit_tx,
max_concurrent,
max_pending,
next_job_id: AtomicU64::new(1),
metrics,
recent,
active_per_flights: Arc::new(Mutex::new(HashMap::new())),
shutdown_triggered: Arc::new(AtomicBool::new(false)),
}),
}
}
pub fn max_concurrent(&self) -> usize {
self.inner.max_concurrent
}
pub fn max_pending(&self) -> usize {
self.inner.max_pending
}
pub fn next_job_id(&self) -> u64 {
self.inner.next_job_id.fetch_add(1, Ordering::SeqCst)
}
pub fn running_count(&self) -> usize {
self.inner.metrics.running.load(Ordering::SeqCst)
}
pub fn completed_ok(&self) -> u64 {
self.inner.metrics.completed_ok.load(Ordering::SeqCst)
}
pub fn completed_err(&self) -> u64 {
self.inner.metrics.completed_err.load(Ordering::SeqCst)
}
pub fn completed_cancelled(&self) -> u64 {
self.inner
.metrics
.completed_cancelled
.load(Ordering::SeqCst)
}
pub fn recent_jobs(&self) -> Vec<ChatJobRecord> {
self.inner
.recent
.lock()
.ok()
.map(|g| g.iter().rev().cloned().collect())
.unwrap_or_default()
}
fn begin_per_flight_job(&self, job_id: u64, flight: Arc<PerTurnFlight>) -> PerFlightJobGuard {
if let Ok(mut g) = self.inner.active_per_flights.lock() {
g.insert(job_id, flight);
}
PerFlightJobGuard {
queue: self.clone(),
job_id,
}
}
fn unregister_per_job_per_flight(&self, job_id: u64) {
if let Ok(mut g) = self.inner.active_per_flights.lock() {
g.remove(&job_id);
}
}
pub fn active_per_jobs(&self) -> Vec<PerFlightStatusEntry> {
let Ok(g) = self.inner.active_per_flights.lock() else {
return Vec::new();
};
let mut v: Vec<PerFlightStatusEntry> = g
.iter()
.map(|(&job_id, flight)| PerFlightStatusEntry {
job_id,
awaiting_plan_rewrite_model: flight
.awaiting_plan_rewrite_model
.load(Ordering::Relaxed),
plan_rewrite_attempts: flight.plan_rewrite_attempts.load(Ordering::Relaxed),
require_plan_in_final_content: flight
.require_plan_in_final_content
.load(Ordering::Relaxed),
})
.collect();
v.sort_by_key(|e| e.job_id);
v
}
pub fn shutdown(&self) {
self.inner.shutdown_triggered.store(true, Ordering::Release);
}
fn reject_if_shutdown(&self) -> Result<(), ChatQueueFull> {
if self.inner.shutdown_triggered.load(Ordering::Acquire) {
return Err(ChatQueueFull {
max_pending: self.inner.max_pending,
});
}
Ok(())
}
pub fn try_submit_stream(&self, p: StreamSubmitParams) -> Result<(), ChatQueueFull> {
self.reject_if_shutdown()?;
let StreamSubmitParams {
envelope,
stream_event_tx,
web_approval_session,
} = p;
let job = QueuedChatJob::Stream {
envelope,
stream_event_tx,
web_approval_session,
};
self.inner
.submit_tx
.try_send(job)
.map_err(|_| ChatQueueFull {
max_pending: self.inner.max_pending,
})
}
pub fn try_submit_json(&self, p: JsonSubmitParams) -> Result<(), ChatQueueFull> {
self.reject_if_shutdown()?;
let JsonSubmitParams { envelope, reply_tx } = p;
let job = QueuedChatJob::Json { envelope, reply_tx };
self.inner
.submit_tx
.try_send(job)
.map_err(|_| ChatQueueFull {
max_pending: self.inner.max_pending,
})
}
}
async fn dispatcher_loop(
mut rx: mpsc::Receiver<QueuedChatJob>,
sem: Arc<Semaphore>,
metrics: Arc<QueueMetrics>,
recent: Arc<Mutex<VecDeque<ChatJobRecord>>>,
) {
while let Some(job) = rx.recv().await {
let permit = match sem.clone().acquire_owned().await {
Ok(p) => p,
Err(_) => {
log::warn!(
target: "crabmate",
"队列调度器:信号量已关闭,仍有待处理任务被丢弃。请确保进程正常关闭。"
);
return;
}
};
let metrics = metrics.clone();
let recent = recent.clone();
metrics.running.fetch_add(1, Ordering::SeqCst);
tokio::spawn(async move {
finish_dispatched_job(job, permit, metrics, recent).await;
});
}
log::info!(target: "crabmate", "队列调度器:所有任务处理完毕,dispatcher 退出。");
}
fn bump_job_completion_metrics(metrics: &QueueMetrics, ok: bool, cancelled: bool) {
if cancelled {
metrics.completed_cancelled.fetch_add(1, Ordering::SeqCst);
} else if ok {
metrics.completed_ok.fetch_add(1, Ordering::SeqCst);
} else {
metrics.completed_err.fetch_add(1, Ordering::SeqCst);
}
}
async fn finish_dispatched_job(
job: QueuedChatJob,
permit: tokio::sync::OwnedSemaphorePermit,
metrics: Arc<QueueMetrics>,
recent: Arc<Mutex<VecDeque<ChatJobRecord>>>,
) {
let job_id = job.job_id();
let _permit = permit;
let start = Instant::now();
let outcome = worker::run_queued_job(job).await;
let ms = start.elapsed().as_millis() as u64;
metrics.running.fetch_sub(1, Ordering::SeqCst);
let kind = outcome.kind_label().to_string();
let (ok, cancelled, err) = outcome.fields();
bump_job_completion_metrics(&metrics, ok, cancelled);
let record = ChatJobRecord {
job_id,
kind,
ok,
cancelled,
duration_ms: ms,
error_preview: err,
};
debug!(
target: "crabmate",
"chat 队列任务结束 job_id={} kind={} ok={} duration_ms={}",
job_id,
record.kind,
record.ok,
record.duration_ms
);
if record.cancelled {
debug!(
target: "crabmate",
"chat 队列任务结束 job_id={} kind={} cancelled=true",
job_id,
record.kind
);
}
if let Ok(mut g) = recent.lock() {
g.push_back(record);
while g.len() > RECENT_CAP {
g.pop_front();
}
}
}