194c44c
use std::net::SocketAddr;
use std::time::Duration;
use agent::{AgentService, AgentTask};
use anyhow::Context;
use config::scoped::{AiHostConfig, ScopedConfig};
use git::sync::queue::SyncQueue;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
use tracing::Instrument as _;
use tracing::{info, warn};
#[tokio::main]
async fn main() -> anyhow::Result<()> {
export::telemetry::init_tracing();
let cfg = AiHostConfig::load().context("loading configuration")?;
if !cfg.aihost.enabled() {
info!("aihost is disabled by config; exiting");
return Ok(());
}
if !cfg.ai.enabled {
info!("ai is disabled by config; aihost exits without a model endpoint");
return Ok(());
}
info!("connecting to database / redis / storage");
let state = export::AppState::with_config(cfg.clone().into_app_config())
.await
.context("building AppState")?;
if let Err(err) = export::telemetry::init_meter_provider() {
warn!(error = %err, "failed to init meter provider; metrics disabled");
}
export::metrics::spawn_queue_depth_reporter(
state.redis.clone(),
&[
("ai.jobs", cfg.ai.queue_topic.clone()),
("ai.jobs.dlq", cfg.ai.dlq_topic.clone()),
],
);
let service = AgentService::new(state).context("building ai service")?;
info!(model = service.client().model(), "dependencies ready");
let cancel = CancellationToken::new();
requeue_stale_jobs(&service).await;
let mut workers = Vec::with_capacity(cfg.aihost.concurrency);
for slot in 0..cfg.aihost.concurrency.max(1) {
workers.push(tokio::spawn(consume(
service.clone(),
cfg.ai.queue_topic.clone(),
cfg.ai.dlq_topic.clone(),
cfg.aihost.retry_backoff_base_secs,
slot,
cancel.clone(),
)));
}
let health = tokio::spawn(run_health(
cfg.health.bind_addr().to_string(),
cfg.health.aihost_port(),
service.state.database.clone(),
service.state.redis.clone(),
cancel.clone(),
));
export::telemetry::shutdown_signal().await;
info!("aihost shutting down");
cancel.cancel();
for worker in workers {
let _ = worker.await;
}
let _ = health.await;
info!("aihost shut down");
Ok(())
}
/// 消费 AI 队列；失败按退避重投，超过 `max_attempts` 进 DLQ。
async fn consume(
service: AgentService,
topic: String,
dlq: String,
backoff_base_secs: u64,
slot: usize,
cancel: CancellationToken,
) {
let queue = match SyncQueue::with_consumer_conns(service.state.redis.clone(), slot + 1).await {
Ok(queue) => queue,
Err(error) => {
warn!(%error, slot, "failed to open ai queue consumer connection");
return;
}
};
let mut attempts: std::collections::HashMap<uuid::Uuid, i32> = std::collections::HashMap::new();
let max_attempts = service.config().ai.max_attempts.max(1);
loop {
let payload = tokio::select! {
_ = cancel.cancelled() => break,
res = queue.consume(&topic, 5, slot) => match res {
Ok(Some(payload)) => payload,
Ok(None) => continue,
Err(error) => {
warn!(%error, slot, "ai queue consume failed");
tokio::time::sleep(Duration::from_secs(2)).await;
continue;
}
},
};
let task = match AgentTask::decode(&payload) {
Ok(task) => task,
Err(error) => {
warn!(%error, "dropping malformed ai task");
let _ = queue.requeue(&dlq, &payload).await;
continue;
}
};
let job_id = task.job_id();
let span = tracing::info_span!(
"ai.task",
job_id = %job_id,
kind = task.kind_str(),
request_id = tracing::field::Empty
);
if let Some(request_id) = task.request_id() {
span.record("request_id", request_id);
}
export::telemetry::set_parent_from_traceparent(&span, task.trace());
let started = std::time::Instant::now();
match export::middleware::with_request_id(
task.request_id(),
service.run_task(&task).instrument(span),
)
.await
{
Ok(()) => {
attempts.remove(&job_id);
export::metrics::ai_job(task.kind_str(), "ok");
export::metrics::ai_job_duration(task.kind_str(), started.elapsed());
info!(%job_id, kind = task.kind_str(), "ai task completed");
}
Err(error) => {
if !error.is_retryable() {
attempts.remove(&job_id);
export::metrics::ai_job(task.kind_str(), "failed");
warn!(%error, %job_id, "ai task is not retryable; sending to DLQ");
service.mark_failed(job_id, &error.to_string()).await;
let _ = queue.requeue(&dlq, &payload).await;
continue;
}
let count = attempts.entry(job_id).or_insert(0);
*count += 1;
if *count >= max_attempts {
attempts.remove(&job_id);
export::metrics::ai_job(task.kind_str(), "failed");
warn!(%error, %job_id, "ai task failed permanently; sending to DLQ");
service.mark_failed(job_id, &error.to_string()).await;
let _ = queue.requeue(&dlq, &payload).await;
} else {
warn!(%error, %job_id, attempt = *count, "ai task failed; retrying");
let backoff = backoff_base_secs.max(1) * u64::from(*count as u32);
tokio::time::sleep(Duration::from_secs(backoff)).await;
let _ = queue.requeue(&topic, &payload).await;
}
}
}
}
info!(slot, "ai consumer stopped");
}
/// 启动时收敛上次进程留下的 running 作业。
///
/// 队列是 Redis 列表且没有投递确认：进程在 `BRPOP` 之后崩溃，载荷已经离开
/// 队列，仅把数据库状态改回 pending 不会有任何消费者再取到它。这里用提交时
/// 持久化的载荷重新投递，投不出去的保持 pending，等待管理端手动重试。
async fn requeue_stale_jobs(service: &AgentService) {
use models::ai::ai_job::{self, AiJobStatus};
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
let stale = ai_job::Entity::find()
.filter(ai_job::Column::Status.eq(AiJobStatus::Running))
.all(service.db_write())
.await;
let rows = match stale {
Ok(rows) => rows,
Err(error) => {
warn!(%error, "stale ai job scan failed");
return;
}
};
for row in rows {
match git::sync::retry_ai_job(
service.db_write(),
&service.state.redis,
&service.config().ai,
&row,
)
.await
{
Ok(()) => info!(job_id = %row.id, "stale ai job requeued"),
Err(error) => warn!(%error, job_id = %row.id, "stale ai job requeue failed"),
}
}
}
async fn run_health(
bind: String,
port: u16,
database: std::sync::Arc<export::AppDatabase>,
redis: export::AppRedis,
cancel: CancellationToken,
) {
let addr: SocketAddr = match format!("{bind}:{port}").parse() {
Ok(addr) => addr,
Err(error) => {
warn!(%error, "invalid aihost health address");
return;
}
};
let listener = match TcpListener::bind(addr).await {
Ok(listener) => listener,
Err(error) => {
warn!(%error, %addr, "aihost health bind failed");
return;
}
};
info!(%addr, "aihost health listening");
let router = export::health::router_with(Some(database), redis, None);
let _ = axum::serve(listener, router)
.with_graceful_shutdown(async move { cancel.cancelled().await })
.await;
}