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, PageHold};
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            holds: Mutex::new(HashMap::new()),
222        });
223        let agents = self.clone();
224        let job_id = id.clone();
225        tokio::spawn(async move {
226            let outcome = engine.agent(&req, sink).await;
227            agents.finish(&job_id, outcome).await;
228        });
229        Ok(id)
230    }
231
232    async fn finish(&self, id: &str, outcome: AgentOutcome) {
233        let status = status_of(&outcome);
234        let job = {
235            let mut map = self.0.map.lock().await;
236            let Some(job) = map.get_mut(id) else { return };
237            job.status = status.into();
238            job.credits = outcome.credits();
239            job.outcome = Some(outcome);
240            job.clone()
241        };
242        self.persist(&job);
243        let o = job.outcome.as_ref().expect("just set");
244        let event = match (status, o.finish_reason.as_str()) {
245            ("completed", _) => "completed",
246            (_, "cancelled") => "cancelled",
247            _ => "failed",
248        };
249        let body = agent_body(o);
250        let error = body
251            .get("error")
252            .and_then(Value::as_str)
253            .map(str::to_string);
254        self.notify(&job, event, status == "completed", json!([body]), error)
255            .await;
256    }
257
258    /// 发一条 webhook（`agent.<event>`），载荷照 Firecrawl。
259    async fn notify(
260        &self,
261        job: &AgentJob,
262        event: &str,
263        success: bool,
264        data: Value,
265        error: Option<String>,
266    ) {
267        let (Some(cfg), Some(hooks)) = (&job.webhook, &self.0.webhooks) else {
268            return;
269        };
270        if !cfg.events.iter().any(|e| e == event) {
271            return;
272        }
273        let mut body = json!({
274            "success": success,
275            "type": format!("agent.{event}"),
276            "id": job.id,
277            "data": data,
278            "metadata": cfg.metadata,
279        });
280        if let Some(e) = error {
281            body["error"] = json!(e);
282        }
283        let _ = hooks.deliver(cfg, &body).await;
284    }
285
286    /// 查询：只有提交者能看。`evidence` 为真时带上存档的证据全文。
287    pub async fn status(&self, id: &str, owner: &str, evidence: bool) -> Result<Value, ApiError> {
288        let mut map = self.0.map.lock().await;
289        if map
290            .get(id)
291            .is_some_and(|job| created_at(job) + JOB_TTL <= SystemTime::now())
292        {
293            map.remove(id);
294            if let Some(store) = &self.0.store {
295                let _ = store.tasks().kv_delete(key(id));
296            }
297        }
298        let job = map
299            .get(id)
300            .filter(|j| j.owner == owner && created_at(j) + JOB_TTL > SystemTime::now())
301            .ok_or_else(|| not_found(id))?;
302        let mut body = match &job.outcome {
303            Some(o) if evidence => agent_body_with_evidence(o),
304            Some(o) => agent_body(o),
305            None => json!({
306                "success": true,
307                "status": job.status,
308                "creditsUsed": job.credits,
309            }),
310        };
311        body["expiresAt"] = json!(rfc3339(created_at(job) + JOB_TTL));
312        if let Some(e) = &job.effort {
313            body["effort"] = json!(e);
314        }
315        if !job.warnings.is_empty() {
316            let mut all = job.warnings.clone();
317            if let Some(w) = body["warnings"].as_array() {
318                all.extend(w.iter().filter_map(|v| v.as_str().map(str::to_string)));
319            }
320            body["warnings"] = json!(all);
321        }
322        Ok(body)
323    }
324
325    /// 取消：在两步之间生效，已经发出的模型调用或抓取会先跑完。
326    pub async fn cancel(&self, id: &str, owner: &str) -> Result<(), ApiError> {
327        let map = self.0.map.lock().await;
328        let job = map
329            .get(id)
330            .filter(|j| j.owner == owner && created_at(j) + JOB_TTL > SystemTime::now())
331            .ok_or_else(|| not_found(id))?;
332        job.cancel.store(true, Ordering::SeqCst);
333        Ok(())
334    }
335}
336
337fn not_found(id: &str) -> ApiError {
338    ApiError::new(ErrorCode::NotFound, format!("agent job {id} not found"))
339}
340
341/// 对外的任务状态：出错、取消、预算用尽而结果不完整都是 `failed`（和 Firecrawl 一样，`failed` 时没有 `data`）。
342fn status_of(o: &AgentOutcome) -> &'static str {
343    if o.error.is_some() || o.budget.exhausted.is_some() || o.finish_reason == "cancelled" {
344        "failed"
345    } else {
346        "completed"
347    }
348}
349
350/// 结束后的响应体：`GET /v2/agent/{id}` 和 CLI 输出的共同部分。
351pub fn agent_body(o: &AgentOutcome) -> Value {
352    let status = status_of(o);
353    let mut b = json!({
354        "success": status == "completed",
355        "status": status,
356        "creditsUsed": o.credits(),
357        "finishReason": o.finish_reason,
358        "complete": o.complete,
359        "missing": o.missing,
360        "sources": o.sources,
361        "fieldSources": o.field_sources,
362        "usage": o.usage,
363        "budget": o.budget,
364        "steps": o.steps,
365    });
366    if let Some(m) = &o.message {
367        b["message"] = json!(m);
368    }
369    if !o.warnings.is_empty() {
370        b["warnings"] = json!(o.warnings);
371    }
372    if status == "completed" {
373        b["data"] = o.data.clone().unwrap_or(Value::Null);
374        return b;
375    }
376    if let Some(partial) = o.data.as_ref().filter(|d| !d.is_null()) {
377        b["partial"] = partial.clone();
378        if let Some(valid) = o.schema_valid {
379            b["partialSchemaValid"] = json!(valid);
380        }
381    }
382    if let Some(kind) = &o.budget.exhausted {
383        b["stopReason"] = json!("credit_limit_reached");
384        b["error"] = json!(format!("the run reached its budget ({kind})"));
385    } else if let Some(e) = &o.error {
386        b["error"] = json!(format!("{}: {}", e.code, e.message));
387    } else {
388        b["error"] = json!("the agent run was cancelled");
389    }
390    b
391}
392
393/// [`agent_body`] 加上存档的证据全文。
394pub fn agent_body_with_evidence(o: &AgentOutcome) -> Value {
395    let mut b = agent_body(o);
396    b["evidence"] = json!(o.evidence);
397    b
398}
399
400struct JobSink {
401    agents: Agents,
402    id: String,
403    cancel: Arc<AtomicBool>,
404    caller: Option<Caller>,
405    webhook: Option<WebhookConfig>,
406    /// 发 webhook 用的任务快照。
407    metadata: AgentJob,
408    /// 已预留的页面配额，按 URL 索引。
409    holds: Mutex<HashMap<String, PageHold>>,
410}
411
412impl AgentSink for JobSink {
413    fn step<'a>(&'a self, step: &'a AgentStep, credits: u64) -> BoxFuture<'a, ()> {
414        Box::pin(async move {
415            if let Some(j) = self.agents.0.map.lock().await.get_mut(&self.id) {
416                j.credits = credits;
417            }
418            if self.webhook.is_some() {
419                let (agents, job, data) = (
420                    self.agents.clone(),
421                    self.metadata.clone(),
422                    json!([{ "creditsUsed": credits, "action": step }]),
423                );
424                // 收件方慢或者不通时不能拖住这次运行
425                tokio::spawn(async move {
426                    agents.notify(&job, "action", true, data, None).await;
427                });
428            }
429        })
430    }
431
432    fn cancelled(&self) -> bool {
433        self.cancel.load(Ordering::SeqCst)
434    }
435
436    fn reserve_page<'a>(&'a self, url: &'a str) -> BoxFuture<'a, Result<(), ApiError>> {
437        Box::pin(async move {
438            let mut holds = self.holds.lock().await;
439            if holds.contains_key(url) {
440                return Ok(());
441            }
442            let (Some(c), Some(l)) = (&self.caller, &self.agents.0.limiter) else {
443                return Ok(());
444            };
445            let hold = l.reserve_page(c).await?;
446            holds.insert(url.to_string(), hold);
447            Ok(())
448        })
449    }
450
451    fn charge_page<'a>(&'a self, url: &'a str) -> BoxFuture<'a, ()> {
452        Box::pin(async move {
453            let hold = { self.holds.lock().await.remove(url) };
454            if let Some(hold) = hold {
455                hold.commit().await;
456            }
457        })
458    }
459
460    fn release_page<'a>(&'a self, url: &'a str) -> BoxFuture<'a, ()> {
461        Box::pin(async move {
462            drop(self.holds.lock().await.remove(url));
463        })
464    }
465
466    fn pages_left(&self) -> Option<u64> {
467        let (c, l) = (self.caller.as_ref()?, self.agents.0.limiter.as_ref()?);
468        l.remaining_pages(c)
469    }
470}
pagesift_api/agent.rs

pagesift_api/
agent.rs
