pagesift_api/
agent.rs

1//! agent 任务：提交后异步跑，结果留 [`JOB_TTL`]。配了存储时写进 `agent/<id>`，重启后结束的任务照样能查；
2//! 重启时还没跑完的任务不续跑（模型的对话没有存），如实报 `failed`。
3
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::time::{Duration, SystemTime, UNIX_EPOCH};
8
9use pagesift_store::Store;
10use serde::{Deserialize, Serialize};
11use serde_json::{Value, json};
12use tokio::sync::Mutex;
13
14use crate::engine::{AgentSink, BoxFuture, Engine};
15use crate::error::{ApiError, ErrorCode};
16use crate::jobs::{JOB_TTL, rfc3339};
17use crate::limits::{Caller, Limiter};
18use crate::types::{AgentFailure, AgentOutcome, AgentRequest, AgentStep};
19use crate::validate::{check_url, check_url_count};
20use crate::webhook::{WebhookConfig, Webhooks};
21
22/// [协议] Firecrawl 文档：`prompt` 最长 10,000 个字符。
23const MAX_PROMPT_CHARS: usize = 10_000;
24
25#[derive(Clone, Serialize, Deserialize)]
26#[serde(rename_all = "camelCase")]
27struct AgentJob {
28    id: String,
29    owner: String,
30    /// `processing`、`completed`、`failed`。
31    status: String,
32    created_ms: u64,
33    effort: Option<String>,
34    /// 提交时的警告（`ignored_fields`）。
35    warnings: Vec<String>,
36    credits: u64,
37    webhook: Option<WebhookConfig>,
38    outcome: Option<AgentOutcome>,
39    #[serde(skip)]
40    cancel: Arc<AtomicBool>,
41}
42
43#[derive(Default)]
44struct Inner {
45    map: Mutex<HashMap<String, AgentJob>>,
46    store: Option<Store>,
47    webhooks: Option<Webhooks>,
48    limiter: Option<Arc<Limiter>>,
49}
50
51/// agent 任务表。HTTP 服务和 MCP 各有一份（MCP 的没有存储、webhook 和配额）。
52#[derive(Clone, Default)]
53pub struct Agents(Arc<Inner>);
54
55fn key(id: &str) -> String {
56    format!("agent/{id}")
57}
58
59fn millis(t: SystemTime) -> u64 {
60    t.duration_since(UNIX_EPOCH).unwrap_or_default().as_millis() as u64
61}
62
63fn created_at(job: &AgentJob) -> SystemTime {
64    UNIX_EPOCH + Duration::from_millis(job.created_ms)
65}
66
67/// 校验请求体（和具体入口无关）。
68pub fn check_agent_request(req: &AgentRequest) -> Result<(), ApiError> {
69    if req.prompt.trim().is_empty() {
70        return Err(ApiError::invalid("prompt must not be empty"));
71    }
72    if req.prompt.chars().count() > MAX_PROMPT_CHARS {
73        return Err(ApiError::invalid(format!(
74            "prompt is longer than {MAX_PROMPT_CHARS} characters"
75        )));
76    }
77    let urls = req.urls.as_deref().unwrap_or_default();
78    check_url_count(urls.len())?;
79    for u in urls {
80        check_url(u)?;
81    }
82    if req.strict_constrain_to_urls && urls.is_empty() {
83        return Err(ApiError::invalid(
84            "strictConstrainToURLs needs a non-empty urls list",
85        ));
86    }
87    if req.schema.as_ref().is_some_and(|v| !v.is_object()) {
88        return Err(ApiError::invalid("schema must be a JSON object"));
89    }
90    if req.max_credits == Some(0) {
91        return Err(ApiError::invalid("maxCredits must be at least 1"));
92    }
93    Ok(())
94}
95
96impl Agents {
97    pub(crate) fn new(
98        store: Option<Store>,
99        webhooks: Option<Webhooks>,
100        limiter: Option<Arc<Limiter>>,
101    ) -> Self {
102        let mut map = HashMap::new();
103        if let Some(s) = &store
104            && let Ok(keys) = s.tasks().kv_keys("agent/")
105        {
106            for k in keys {
107                let Ok(Some(v)) = s.tasks().kv_get(&k) else {
108                    continue;
109                };
110                let Ok(mut job) = serde_json::from_slice::<AgentJob>(&v) else {
111                    continue;
112                };
113                if created_at(&job) + JOB_TTL <= SystemTime::now() {
114                    let _ = s.tasks().kv_delete(&k);
115                    continue;
116                }
117                if job.status == "processing" {
118                    job.status = "failed".into();
119                    job.outcome = Some(AgentOutcome {
120                        finish_reason: "error".into(),
121                        error: Some(AgentFailure {
122                            code: ErrorCode::Internal.as_str().into(),
123                            message: "the server restarted while this agent run was in progress; \
124                                      the run was not resumed"
125                                .into(),
126                        }),
127                        ..Default::default()
128                    });
129                    if let Ok(v) = serde_json::to_vec(&job) {
130                        let _ = s.tasks().kv_put(&k, &v);
131                    }
132                }
133                map.insert(job.id.clone(), job);
134            }
135        }
136        Self(Arc::new(Inner {
137            map: Mutex::new(map),
138            store,
139            webhooks,
140            limiter,
141        }))
142    }
143
144    fn persist(&self, job: &AgentJob) {
145        if let (Some(store), Ok(v)) = (&self.0.store, serde_json::to_vec(job))
146            && let Err(e) = store.tasks().kv_put(key(&job.id), &v)
147        {
148            tracing::warn!(job = %job.id, error = %e, "could not save the agent run");
149        }
150    }
151
152    /// 没有调用方身份、webhook 和配额的入口（MCP）。
153    pub async fn start_local(
154        &self,
155        engine: Arc<dyn Engine>,
156        req: AgentRequest,
157    ) -> Result<String, ApiError> {
158        self.start(engine, req, Vec::new(), None, None).await
159    }
160
161    /// 校验、登记、开跑，返回任务 id。没配 LLM 等问题在这里就报错。
162    pub(crate) async fn start(
163        &self,
164        engine: Arc<dyn Engine>,
165        req: AgentRequest,
166        warnings: Vec<String>,
167        caller: Option<Caller>,
168        webhook: Option<WebhookConfig>,
169    ) -> Result<String, ApiError> {
170        check_agent_request(&req)?;
171        engine.agent_check(&req)?;
172        let id = uuid::Uuid::new_v4().to_string();
173        let owner = caller
174            .as_ref()
175            .map_or_else(|| "anonymous".to_string(), |c| c.account.to_string());
176        let cancel = Arc::new(AtomicBool::new(false));
177        let job = AgentJob {
178            id: id.clone(),
179            owner: owner.clone(),
180            status: "processing".into(),
181            created_ms: millis(SystemTime::now()),
182            effort: req.effort.map(|e| e.as_str().to_string()),
183            warnings,
184            credits: 0,
185            webhook,
186            outcome: None,
187            cancel: cancel.clone(),
188        };
189        {
190            let mut map = self.0.map.lock().await;
191            let now = SystemTime::now();
192            map.retain(|k, j| {
193                let live = created_at(j) + JOB_TTL > now;
194                if !live && let Some(s) = &self.0.store {
195                    let _ = s.tasks().kv_delete(key(k));
196                }
197                live
198            });
199            if let (Some(c), Some(l)) = (&caller, &self.0.limiter) {
200                if c.limits.max_concurrent_jobs.is_some() {
201                    let active = map
202                        .values()
203                        .filter(|j| j.status == "processing" && j.owner == owner)
204                        .count();
205                    l.check_jobs(c, active)?;
206                }
207                l.check_quota(c)?;
208            }
209            map.insert(id.clone(), job.clone());
210        }
211        self.persist(&job);
212        self.notify(&job, "started", true, json!([]), None).await;
213
214        let sink: Arc<dyn AgentSink> = Arc::new(JobSink {
215            agents: self.clone(),
216            id: id.clone(),
217            cancel,
218            caller,
219            webhook: job.webhook.clone(),
220            metadata: job.clone(),
221        });
222        let agents = self.clone();
223        let job_id = id.clone();
224        tokio::spawn(async move {
225            let outcome = engine.agent(&req, sink).await;
226            agents.finish(&job_id, outcome).await;
227        });
228        Ok(id)
229    }
230
231    async fn finish(&self, id: &str, outcome: AgentOutcome) {
232        let status = status_of(&outcome);
233        let job = {
234            let mut map = self.0.map.lock().await;
235            let Some(job) = map.get_mut(id) else { return };
236            job.status = status.into();
237            job.credits = outcome.credits();
238            job.outcome = Some(outcome);
239            job.clone()
240        };
241        self.persist(&job);
242        let o = job.outcome.as_ref().expect("just set");
243        let event = match (status, o.finish_reason.as_str()) {
244            ("completed", _) => "completed",
245            (_, "cancelled") => "cancelled",
246            _ => "failed",
247        };
248        let body = agent_body(o);
249        let error = body
250            .get("error")
251            .and_then(Value::as_str)
252            .map(str::to_string);
253        self.notify(&job, event, status == "completed", json!([body]), error)
254            .await;
255    }
256
257    /// 发一条 webhook（`agent.<event>`），载荷照 Firecrawl。
258    async fn notify(
259        &self,
260        job: &AgentJob,
261        event: &str,
262        success: bool,
263        data: Value,
264        error: Option<String>,
265    ) {
266        let (Some(cfg), Some(hooks)) = (&job.webhook, &self.0.webhooks) else {
267            return;
268        };
269        if !cfg.events.iter().any(|e| e == event) {
270            return;
271        }
272        let mut body = json!({
273            "success": success,
274            "type": format!("agent.{event}"),
275            "id": job.id,
276            "data": data,
277            "metadata": cfg.metadata,
278        });
279        if let Some(e) = error {
280            body["error"] = json!(e);
281        }
282        let _ = hooks.deliver(cfg, &body).await;
283    }
284
285    /// 查询：只有提交者能看。`evidence` 为真时带上存档的证据全文。
286    pub async fn status(&self, id: &str, owner: &str, evidence: bool) -> Result<Value, ApiError> {
287        let map = self.0.map.lock().await;
288        let job = map
289            .get(id)
290            .filter(|j| j.owner == owner)
291            .ok_or_else(|| not_found(id))?;
292        let mut body = match &job.outcome {
293            Some(o) if evidence => agent_body_with_evidence(o),
294            Some(o) => agent_body(o),
295            None => json!({
296                "success": true,
297                "status": job.status,
298                "creditsUsed": job.credits,
299            }),
300        };
301        body["expiresAt"] = json!(rfc3339(created_at(job) + JOB_TTL));
302        if let Some(e) = &job.effort {
303            body["effort"] = json!(e);
304        }
305        if !job.warnings.is_empty() {
306            let mut all = job.warnings.clone();
307            if let Some(w) = body["warnings"].as_array() {
308                all.extend(w.iter().filter_map(|v| v.as_str().map(str::to_string)));
309            }
310            body["warnings"] = json!(all);
311        }
312        Ok(body)
313    }
314
315    /// 取消：在两步之间生效，已经发出的模型调用或抓取会先跑完。
316    pub async fn cancel(&self, id: &str, owner: &str) -> Result<(), ApiError> {
317        let map = self.0.map.lock().await;
318        let job = map
319            .get(id)
320            .filter(|j| j.owner == owner)
321            .ok_or_else(|| not_found(id))?;
322        job.cancel.store(true, Ordering::SeqCst);
323        Ok(())
324    }
325}
326
327fn not_found(id: &str) -> ApiError {
328    ApiError::new(ErrorCode::NotFound, format!("agent job {id} not found"))
329}
330
331/// 对外的任务状态：出错、取消、预算用尽而结果不完整都是 `failed`（和 Firecrawl 一样，`failed` 时没有 `data`）。
332fn status_of(o: &AgentOutcome) -> &'static str {
333    if o.error.is_some() || o.budget.exhausted.is_some() || o.finish_reason == "cancelled" {
334        "failed"
335    } else {
336        "completed"
337    }
338}
339
340/// 结束后的响应体：`GET /v2/agent/{id}` 和 CLI 输出的共同部分。
341pub fn agent_body(o: &AgentOutcome) -> Value {
342    let status = status_of(o);
343    let mut b = json!({
344        "success": status == "completed",
345        "status": status,
346        "creditsUsed": o.credits(),
347        "finishReason": o.finish_reason,
348        "complete": o.complete,
349        "missing": o.missing,
350        "sources": o.sources,
351        "fieldSources": o.field_sources,
352        "usage": o.usage,
353        "budget": o.budget,
354        "steps": o.steps,
355    });
356    if let Some(m) = &o.message {
357        b["message"] = json!(m);
358    }
359    if !o.warnings.is_empty() {
360        b["warnings"] = json!(o.warnings);
361    }
362    if status == "completed" {
363        b["data"] = o.data.clone().unwrap_or(Value::Null);
364        return b;
365    }
366    if let Some(partial) = o.data.as_ref().filter(|d| !d.is_null()) {
367        b["partial"] = partial.clone();
368        if let Some(valid) = o.schema_valid {
369            b["partialSchemaValid"] = json!(valid);
370        }
371    }
372    if let Some(kind) = &o.budget.exhausted {
373        b["stopReason"] = json!("credit_limit_reached");
374        b["error"] = json!(format!("the run reached its budget ({kind})"));
375    } else if let Some(e) = &o.error {
376        b["error"] = json!(format!("{}: {}", e.code, e.message));
377    } else {
378        b["error"] = json!("the agent run was cancelled");
379    }
380    b
381}
382
383/// [`agent_body`] 加上存档的证据全文。
384pub fn agent_body_with_evidence(o: &AgentOutcome) -> Value {
385    let mut b = agent_body(o);
386    b["evidence"] = json!(o.evidence);
387    b
388}
389
390struct JobSink {
391    agents: Agents,
392    id: String,
393    cancel: Arc<AtomicBool>,
394    caller: Option<Caller>,
395    webhook: Option<WebhookConfig>,
396    /// 发 webhook 用的任务快照。
397    metadata: AgentJob,
398}
399
400impl AgentSink for JobSink {
401    fn step<'a>(&'a self, step: &'a AgentStep, credits: u64) -> BoxFuture<'a, ()> {
402        Box::pin(async move {
403            if let Some(j) = self.agents.0.map.lock().await.get_mut(&self.id) {
404                j.credits = credits;
405            }
406            if self.webhook.is_some() {
407                let (agents, job, data) = (
408                    self.agents.clone(),
409                    self.metadata.clone(),
410                    json!([{ "creditsUsed": credits, "action": step }]),
411                );
412                // 收件方慢或者不通时不能拖住这次运行
413                tokio::spawn(async move {
414                    agents.notify(&job, "action", true, data, None).await;
415                });
416            }
417        })
418    }
419
420    fn cancelled(&self) -> bool {
421        self.cancel.load(Ordering::SeqCst)
422    }
423
424    fn charge_page(&self) -> BoxFuture<'_, ()> {
425        Box::pin(async move {
426            if let (Some(c), Some(l)) = (&self.caller, &self.agents.0.limiter) {
427                l.consume_pages(c, 1).await;
428            }
429        })
430    }
431
432    fn pages_left(&self) -> Option<u64> {
433        let (c, l) = (self.caller.as_ref()?, self.agents.0.limiter.as_ref()?);
434        l.remaining_pages(c)
435    }
436}
pagesift_api/agent.rs

pagesift_api/
agent.rs
