2f11d69
use models::ai::ai_job;
use models::ai::ai_job::{ActiveModel, AiJobStatus};
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, PaginatorTrait, QueryFilter, Set};
use time::OffsetDateTime;
use uuid::Uuid;
/// 消费者侧的作业状态流转。
impl super::AgentService {
/// 标记作业开始执行。
pub async fn mark_running(&self, job_id: Uuid, kind_str: &str) -> Result<(), sea_orm::DbErr> {
let db = self.db_write();
let now = OffsetDateTime::now_utc();
if let Some(row) = ai_job::Entity::find_by_id(job_id).one(db).await? {
let attempts = row.attempts.saturating_add(1);
let mut model: ActiveModel = row.into();
model.status = Set(AiJobStatus::Running);
model.started_at = Set(Some(now));
model.updated_at = Set(Some(now));
model.attempts = Set(attempts);
model.update(db).await?;
return Ok(());
}
// 队列里存在但库中缺失：补一条 running 记录，保证任务可审计。
ActiveModel {
id: Set(job_id),
repo_id: Set(Uuid::nil()),
kind: Set(kind_from_str(kind_str)),
target_id: Set(None),
status: Set(AiJobStatus::Running),
attempts: Set(1),
model: Set(self.config().ai.model.clone()),
target_ref: Set(None),
language: Set(self.config().ai.output_language.clone()),
request_json: Set(None),
response_json: Set(None),
prompt_tokens: Set(None),
completion_tokens: Set(None),
error: Set(None),
created_at: Set(now),
started_at: Set(Some(now)),
finished_at: Set(None),
updated_at: Set(Some(now)),
}
.insert(db)
.await?;
Ok(())
}
pub async fn mark_success(
&self,
job_id: Uuid,
request: Option<serde_json::Value>,
response: Option<serde_json::Value>,
prompt_tokens: Option<i32>,
completion_tokens: Option<i32>,
) {
let db = self.db_write();
let now = OffsetDateTime::now_utc();
if let Ok(Some(row)) = ai_job::Entity::find_by_id(job_id).one(db).await {
let mut model: ActiveModel = row.into();
model.status = Set(AiJobStatus::Success);
model.request_json = Set(request);
model.response_json = Set(response);
model.prompt_tokens = Set(prompt_tokens);
model.completion_tokens = Set(completion_tokens);
model.error = Set(None);
model.finished_at = Set(Some(now));
model.updated_at = Set(Some(now));
if let Err(error) = model.update(db).await {
tracing::warn!(%error, %job_id, "failed to mark ai job success");
}
}
}
pub async fn mark_failed(&self, job_id: Uuid, error: &str) {
let db = self.db_write();
let now = OffsetDateTime::now_utc();
if let Ok(Some(row)) = ai_job::Entity::find_by_id(job_id).one(db).await {
let mut model: ActiveModel = row.into();
model.status = Set(AiJobStatus::Failed);
model.error = Set(Some(error.to_string()));
model.finished_at = Set(Some(now));
model.updated_at = Set(Some(now));
if let Err(update_error) = model.update(db).await {
tracing::warn!(%update_error, %job_id, "failed to mark ai job failure");
}
}
}
/// 统计某仓库仍未结束的作业数，用于限流与积压观测。
pub async fn pending_jobs(&self, repo_id: Uuid) -> Result<u64, sea_orm::DbErr> {
ai_job::Entity::find()
.filter(ai_job::Column::RepoId.eq(repo_id))
.filter(ai_job::Column::Status.is_in([AiJobStatus::Pending, AiJobStatus::Running]))
.count(self.db_read())
.await
}
}
fn kind_from_str(value: &str) -> ai_job::AiJobKind {
use ai_job::AiJobKind;
match value {
"commit_review" => AiJobKind::CommitReview,
"repo_summary" => AiJobKind::RepoSummary,
"release_notes" => AiJobKind::ReleaseNotes,
_ => AiJobKind::PullReview,
}
}