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 mut map = self.0.map.lock().await;
288        if map
289            .get(id)
290            .is_some_and(|job| created_at(job) + JOB_TTL <= SystemTime::now())
291        {
292            map.remove(id);
293            if let Some(store) = &self.0.store {
294                let _ = store.tasks().kv_delete(key(id));
295            }
296        }
297        let job = map
298            .get(id)
299            .filter(|j| j.owner == owner && created_at(j) + JOB_TTL > SystemTime::now())
300            .ok_or_else(|| not_found(id))?;
301        let mut body = match &job.outcome {
302            Some(o) if evidence => agent_body_with_evidence(o),
303            Some(o) => agent_body(o),
304            None => json!({
305                "success": true,
306                "status": job.status,
307                "creditsUsed": job.credits,
308            }),
309        };
310        body["expiresAt"] = json!(rfc3339(created_at(job) + JOB_TTL));
311        if let Some(e) = &job.effort {
312            body["effort"] = json!(e);
313        }
314        if !job.warnings.is_empty() {
315            let mut all = job.warnings.clone();
316            if let Some(w) = body["warnings"].as_array() {
317                all.extend(w.iter().filter_map(|v| v.as_str().map(str::to_string)));
318            }
319            body["warnings"] = json!(all);
320        }
321        Ok(body)
322    }
323
324    /// 取消：在两步之间生效，已经发出的模型调用或抓取会先跑完。
325    pub async fn cancel(&self, id: &str, owner: &str) -> Result<(), ApiError> {
326        let map = self.0.map.lock().await;
327        let job = map
328            .get(id)
329            .filter(|j| j.owner == owner && created_at(j) + JOB_TTL > SystemTime::now())
330            .ok_or_else(|| not_found(id))?;
331        job.cancel.store(true, Ordering::SeqCst);
332        Ok(())
333    }
334}
335
336fn not_found(id: &str) -> ApiError {
337    ApiError::new(ErrorCode::NotFound, format!("agent job {id} not found"))
338}
339
340/// 对外的任务状态：出错、取消、预算用尽而结果不完整都是 `failed`（和 Firecrawl 一样，`failed` 时没有 `data`）。
341fn status_of(o: &AgentOutcome) -> &'static str {
342    if o.error.is_some() || o.budget.exhausted.is_some() || o.finish_reason == "cancelled" {
343        "failed"
344    } else {
345        "completed"
346    }
347}
348
349/// 结束后的响应体：`GET /v2/agent/{id}` 和 CLI 输出的共同部分。
350pub fn agent_body(o: &AgentOutcome) -> Value {
351    let status = status_of(o);
352    let mut b = json!({
353        "success": status == "completed",
354        "status": status,
355        "creditsUsed": o.credits(),
356        "finishReason": o.finish_reason,
357        "complete": o.complete,
358        "missing": o.missing,
359        "sources": o.sources,
360        "fieldSources": o.field_sources,
361        "usage": o.usage,
362        "budget": o.budget,
363        "steps": o.steps,
364    });
365    if let Some(m) = &o.message {
366        b["message"] = json!(m);
367    }
368    if !o.warnings.is_empty() {
369        b["warnings"] = json!(o.warnings);
370    }
371    if status == "completed" {
372        b["data"] = o.data.clone().unwrap_or(Value::Null);
373        return b;
374    }
375    if let Some(partial) = o.data.as_ref().filter(|d| !d.is_null()) {
376        b["partial"] = partial.clone();
377        if let Some(valid) = o.schema_valid {
378            b["partialSchemaValid"] = json!(valid);
379        }
380    }
381    if let Some(kind) = &o.budget.exhausted {
382        b["stopReason"] = json!("credit_limit_reached");
383        b["error"] = json!(format!("the run reached its budget ({kind})"));
384    } else if let Some(e) = &o.error {
385        b["error"] = json!(format!("{}: {}", e.code, e.message));
386    } else {
387        b["error"] = json!("the agent run was cancelled");
388    }
389    b
390}
391
392/// [`agent_body`] 加上存档的证据全文。
393pub fn agent_body_with_evidence(o: &AgentOutcome) -> Value {
394    let mut b = agent_body(o);
395    b["evidence"] = json!(o.evidence);
396    b
397}
398
399struct JobSink {
400    agents: Agents,
401    id: String,
402    cancel: Arc<AtomicBool>,
403    caller: Option<Caller>,
404    webhook: Option<WebhookConfig>,
405    /// 发 webhook 用的任务快照。
406    metadata: AgentJob,
407}
408
409impl AgentSink for JobSink {
410    fn step<'a>(&'a self, step: &'a AgentStep, credits: u64) -> BoxFuture<'a, ()> {
411        Box::pin(async move {
412            if let Some(j) = self.agents.0.map.lock().await.get_mut(&self.id) {
413                j.credits = credits;
414            }
415            if self.webhook.is_some() {
416                let (agents, job, data) = (
417                    self.agents.clone(),
418                    self.metadata.clone(),
419                    json!([{ "creditsUsed": credits, "action": step }]),
420                );
421                // 收件方慢或者不通时不能拖住这次运行
422                tokio::spawn(async move {
423                    agents.notify(&job, "action", true, data, None).await;
424                });
425            }
426        })
427    }
428
429    fn cancelled(&self) -> bool {
430        self.cancel.load(Ordering::SeqCst)
431    }
432
433    fn charge_page(&self) -> BoxFuture<'_, ()> {
434        Box::pin(async move {
435            if let (Some(c), Some(l)) = (&self.caller, &self.agents.0.limiter) {
436                l.consume_pages(c, 1).await;
437            }
438        })
439    }
440
441    fn pages_left(&self) -> Option<u64> {
442        let (c, l) = (self.caller.as_ref()?, self.agents.0.limiter.as_ref()?);
443        l.remaining_pages(c)
444    }
445}
pagesift_api/agent.rs

pagesift_api/
agent.rs
