context_harness/
agent_task_store.rs

1//! Durable accepted agent work. Tasks own scheduling facts; linked runs own all
2//! execution lifecycle, outcome, checkpoint, usage and effect history.
3
4use crate::agent_store::{AgentRun, AgentRunStore, FailureReason, RunBudgets, RunOutcome};
5use anyhow::{ensure, Context, Result};
6use chrono::Utc;
7use serde::{Deserialize, Serialize};
8use serde_json::{json, Value};
9use sha2::{Digest, Sha256};
10use sqlx::{FromRow, Sqlite, SqlitePool, Transaction};
11use uuid::Uuid;
12
13const REQUEST_DOMAIN: &[u8] = b"context-harness.agent-task.request.v1\0";
14pub const DEFAULT_QUEUE_LIMIT: u32 = 1_000;
15pub const MAX_QUEUE_LIMIT: u32 = 10_000;
16pub const MAX_TASK_INPUT_BYTES: usize = 64 * 1024;
17pub const MAX_PAYLOAD_IDENTITY_BYTES: usize = 8 * 1024;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, sqlx::Type)]
20#[serde(rename_all = "snake_case")]
21#[sqlx(type_name = "TEXT", rename_all = "snake_case")]
22pub enum AgentTaskStatus {
23    Queued,
24    Claimed,
25    CancelRequested,
26    Terminal,
27}
28
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub enum TaskSchedulingReason {
31    CancelledBeforeRun,
32    RunStopped,
33    AcceptedIdentityChanged,
34    ClaimAttemptsExhausted,
35    RestartRequired,
36    ReconciliationRequired,
37    Unknown(String),
38}
39
40impl TaskSchedulingReason {
41    pub fn from_stored(value: &str) -> Self {
42        match value {
43            "cancelled_before_run" => Self::CancelledBeforeRun,
44            "run_stopped" => Self::RunStopped,
45            "accepted_identity_changed" => Self::AcceptedIdentityChanged,
46            "claim_attempts_exhausted" => Self::ClaimAttemptsExhausted,
47            "restart_required" => Self::RestartRequired,
48            "reconciliation_required" => Self::ReconciliationRequired,
49            other => Self::Unknown(other.to_owned()),
50        }
51    }
52}
53
54/// Non-secret, statically resolved execution identity accepted at submission.
55#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
56pub struct AcceptedTaskIdentity {
57    value: Value,
58    digest: String,
59}
60
61impl AcceptedTaskIdentity {
62    pub fn new(value: Value) -> Result<Self> {
63        ensure!(
64            value.is_object(),
65            "accepted task identity must be an object"
66        );
67        ensure!(
68            value.get("schema_version").and_then(Value::as_u64) == Some(1),
69            "unsupported accepted task identity schema"
70        );
71        let value = canonical_json(value);
72        let encoded = serde_json::to_vec(&value)?;
73        ensure!(
74            encoded.len() <= 64 * 1024,
75            "accepted task identity is too large"
76        );
77        Ok(Self {
78            digest: format!("{:x}", Sha256::digest(&encoded)),
79            value,
80        })
81    }
82
83    pub fn value(&self) -> &Value {
84        &self.value
85    }
86
87    pub fn digest(&self) -> &str {
88        &self.digest
89    }
90}
91
92fn canonical_json(value: Value) -> Value {
93    match value {
94        Value::Object(object) => {
95            let mut entries = object.into_iter().collect::<Vec<_>>();
96            entries.sort_by(|left, right| left.0.cmp(&right.0));
97            Value::Object(
98                entries
99                    .into_iter()
100                    .map(|(key, value)| (key, canonical_json(value)))
101                    .collect(),
102            )
103        }
104        Value::Array(values) => Value::Array(values.into_iter().map(canonical_json).collect()),
105        other => other,
106    }
107}
108
109#[derive(Debug, Clone)]
110pub struct AgentTaskSubmission {
111    pub request_key: String,
112    pub agent_name: String,
113    pub agent_version: String,
114    pub input: String,
115    pub payload_identity: Vec<u8>,
116    pub accepted_identity: AcceptedTaskIdentity,
117    pub queue_limit: u32,
118}
119
120#[derive(Debug, Serialize, FromRow)]
121pub struct AgentTask {
122    pub id: String,
123    pub workspace_id: String,
124    pub request_key: String,
125    pub request_digest: String,
126    pub agent_name: String,
127    pub agent_version: String,
128    pub accepted_identity: sqlx::types::Json<Value>,
129    pub accepted_identity_digest: String,
130    pub input: String,
131    pub payload_identity: Vec<u8>,
132    pub status: AgentTaskStatus,
133    /// Raw stored value is retained so older readers preserve future reasons.
134    pub scheduling_reason: Option<String>,
135    pub scheduling_detail: Option<sqlx::types::Json<Value>>,
136    pub claim_owner: Option<String>,
137    pub claim_token: Option<String>,
138    pub lease_expires_at: Option<i64>,
139    pub claim_attempts: i64,
140    pub run_id: Option<String>,
141    pub created_at: i64,
142    pub updated_at: i64,
143    pub claimed_at: Option<i64>,
144    pub completed_at: Option<i64>,
145    pub last_sequence: i64,
146}
147
148impl AgentTask {
149    pub fn scheduling_reason_kind(&self) -> Option<TaskSchedulingReason> {
150        self.scheduling_reason
151            .as_deref()
152            .map(TaskSchedulingReason::from_stored)
153    }
154}
155
156#[derive(Debug, Serialize, FromRow)]
157pub struct AgentTaskEvent {
158    pub task_id: String,
159    pub sequence: i64,
160    pub timestamp: i64,
161    pub event_type: String,
162    pub payload: sqlx::types::Json<Value>,
163}
164
165#[derive(Debug, Serialize)]
166pub struct AgentTaskClaim {
167    pub task: AgentTask,
168    pub worker_id: String,
169    pub claim_token: String,
170}
171
172#[derive(Debug, Serialize)]
173pub struct ExpiredAgentTaskClaim {
174    pub claim: AgentTaskClaim,
175}
176
177/// Task scheduling state together with the authoritative linked run, when any.
178#[derive(Debug, Serialize)]
179pub struct AgentTaskInspection {
180    pub task: AgentTask,
181    pub run: Option<AgentRun>,
182}
183
184#[derive(Debug, Serialize)]
185#[serde(tag = "disposition", content = "task", rename_all = "snake_case")]
186pub enum AgentTaskSubmissionResult {
187    Created(AgentTask),
188    ExistingIdentical(AgentTask),
189    RequestConflict,
190    QueueFull,
191}
192
193#[derive(Clone)]
194pub struct AgentTaskStore {
195    pool: SqlitePool,
196    workspace_id: String,
197}
198
199impl AgentTaskStore {
200    pub fn new(pool: SqlitePool, workspace_id: &str) -> Result<Self> {
201        ensure!(!workspace_id.trim().is_empty(), "workspace ID is required");
202        Ok(Self {
203            pool,
204            workspace_id: workspace_id.to_owned(),
205        })
206    }
207
208    pub(crate) fn run_store(&self) -> Result<AgentRunStore> {
209        AgentRunStore::new(self.pool.clone(), &self.workspace_id)
210    }
211
212    pub async fn submit(
213        &self,
214        submission: AgentTaskSubmission,
215    ) -> Result<AgentTaskSubmissionResult> {
216        validate_submission(&submission)?;
217        let request_digest = request_digest(
218            &submission.agent_name,
219            &submission.input,
220            &submission.payload_identity,
221        )?;
222        let id = Uuid::new_v4().to_string();
223        let now = Utc::now().timestamp_millis();
224        let mut tx = self.pool.begin().await?;
225        let inserted: Option<String> = sqlx::query_scalar(
226            "INSERT INTO agent_tasks (id, workspace_id, request_key, request_digest, agent_name, agent_version, accepted_identity, accepted_identity_digest, input, payload_identity, status, created_at, updated_at) SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'queued', ?, ? WHERE (SELECT COUNT(*) FROM agent_tasks WHERE workspace_id = ? AND status != 'terminal') < ? ON CONFLICT(workspace_id, request_key) DO NOTHING RETURNING id",
227        )
228        .bind(&id)
229        .bind(&self.workspace_id)
230        .bind(&submission.request_key)
231        .bind(&request_digest)
232        .bind(&submission.agent_name)
233        .bind(&submission.agent_version)
234        .bind(sqlx::types::Json(submission.accepted_identity.value()))
235        .bind(submission.accepted_identity.digest())
236        .bind(&submission.input)
237        .bind(&submission.payload_identity)
238        .bind(now)
239        .bind(now)
240        .bind(&self.workspace_id)
241        .bind(i64::from(submission.queue_limit))
242        .fetch_optional(&mut *tx)
243        .await?;
244
245        if inserted.is_some() {
246            append_event(
247                &mut tx,
248                &id,
249                now,
250                "task.submitted",
251                &json!({
252                    "request_digest": request_digest,
253                    "accepted_identity_digest": submission.accepted_identity.digest(),
254                    "agent": submission.agent_name,
255                }),
256            )
257            .await?;
258            tx.commit().await?;
259            return Ok(AgentTaskSubmissionResult::Created(
260                self.get(&id).await?.context("created task missing")?,
261            ));
262        }
263
264        let existing: Option<(String, String, String)> = sqlx::query_as(
265            "SELECT id, request_digest, accepted_identity_digest FROM agent_tasks WHERE workspace_id = ? AND request_key = ?",
266        )
267        .bind(&self.workspace_id)
268        .bind(&submission.request_key)
269        .fetch_optional(&mut *tx)
270        .await?;
271        tx.commit().await?;
272        match existing {
273            Some((id, stored_request, stored_identity))
274                if stored_request == request_digest
275                    && stored_identity == submission.accepted_identity.digest() =>
276            {
277                Ok(AgentTaskSubmissionResult::ExistingIdentical(
278                    self.get(&id).await?.context("existing task missing")?,
279                ))
280            }
281            Some(_) => Ok(AgentTaskSubmissionResult::RequestConflict),
282            None => Ok(AgentTaskSubmissionResult::QueueFull),
283        }
284    }
285
286    pub async fn get(&self, id: &str) -> Result<Option<AgentTask>> {
287        Ok(
288            sqlx::query_as("SELECT * FROM agent_tasks WHERE id = ? AND workspace_id = ?")
289                .bind(id)
290                .bind(&self.workspace_id)
291                .fetch_optional(&self.pool)
292                .await?,
293        )
294    }
295
296    /// Inspect scheduling state without copying execution state into the task.
297    pub async fn inspect(&self, id: &str) -> Result<Option<AgentTaskInspection>> {
298        let Some(task) = self.get(id).await? else {
299            return Ok(None);
300        };
301        let run = match &task.run_id {
302            Some(run_id) => {
303                AgentRunStore::new(self.pool.clone(), &self.workspace_id)?
304                    .get_run(run_id)
305                    .await?
306            }
307            None => None,
308        };
309        Ok(Some(AgentTaskInspection { task, run }))
310    }
311
312    /// List bounded scheduling records with authoritative linked run projections.
313    pub async fn list_inspections(&self, limit: u32) -> Result<Vec<AgentTaskInspection>> {
314        let tasks = self.list(limit).await?;
315        let runs = self.run_store()?;
316        let mut inspections = Vec::with_capacity(tasks.len());
317        for task in tasks {
318            let run = match &task.run_id {
319                Some(run_id) => runs.get_run(run_id).await?,
320                None => None,
321            };
322            inspections.push(AgentTaskInspection { task, run });
323        }
324        Ok(inspections)
325    }
326
327    pub async fn list(&self, limit: u32) -> Result<Vec<AgentTask>> {
328        ensure!(
329            (1..=1_000).contains(&limit),
330            "task list limit must be 1-1000"
331        );
332        Ok(sqlx::query_as(
333            "SELECT * FROM agent_tasks WHERE workspace_id = ? ORDER BY created_at DESC, id LIMIT ?",
334        )
335        .bind(&self.workspace_id)
336        .bind(limit)
337        .fetch_all(&self.pool)
338        .await?)
339    }
340
341    pub async fn events(
342        &self,
343        id: &str,
344        after_sequence: i64,
345        limit: u32,
346    ) -> Result<Vec<AgentTaskEvent>> {
347        ensure!(after_sequence >= 0, "invalid task event cursor");
348        ensure!(
349            (1..=1_000).contains(&limit),
350            "task event limit must be 1-1000"
351        );
352        self.get(id).await?.context("task not found in workspace")?;
353        Ok(sqlx::query_as("SELECT e.* FROM agent_task_events e JOIN agent_tasks t ON t.id = e.task_id WHERE e.task_id = ? AND t.workspace_id = ? AND e.sequence > ? ORDER BY e.sequence LIMIT ?")
354            .bind(id)
355            .bind(&self.workspace_id)
356            .bind(after_sequence)
357            .bind(limit)
358            .fetch_all(&self.pool)
359            .await?)
360    }
361
362    /// Claim the oldest queued task using the database clock for its lease.
363    pub async fn claim_next(
364        &self,
365        worker_id: &str,
366        lease_seconds: u64,
367        max_attempts: u32,
368    ) -> Result<Option<AgentTaskClaim>> {
369        validate_worker_id(worker_id)?;
370        ensure!(
371            (5..=3_600).contains(&lease_seconds),
372            "lease duration must be 5-3600 seconds"
373        );
374        ensure!(
375            (1..=100).contains(&max_attempts),
376            "maximum claim attempts must be 1-100"
377        );
378        let claim_token = Uuid::new_v4().to_string();
379        let mut tx = self.pool.begin().await?;
380        let task: Option<AgentTask> = sqlx::query_as(
381            "WITH candidate AS (SELECT id FROM agent_tasks WHERE workspace_id = ? AND status = 'queued' AND claim_attempts < ? ORDER BY created_at, id LIMIT 1) UPDATE agent_tasks SET status = 'claimed', claim_owner = ?, claim_token = ?, lease_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + ?, claim_attempts = claim_attempts + 1, claimed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) WHERE id = (SELECT id FROM candidate) AND workspace_id = ? AND status = 'queued' RETURNING *",
382        )
383        .bind(&self.workspace_id)
384        .bind(i64::from(max_attempts))
385        .bind(worker_id)
386        .bind(&claim_token)
387        .bind(i64::try_from(lease_seconds)?.saturating_mul(1_000))
388        .bind(&self.workspace_id)
389        .fetch_optional(&mut *tx)
390        .await?;
391        let Some(task) = task else {
392            tx.commit().await?;
393            return Ok(None);
394        };
395        append_event(
396            &mut tx,
397            &task.id,
398            task.updated_at,
399            "task.claimed",
400            &json!({"worker_id": worker_id, "claim_attempt": task.claim_attempts}),
401        )
402        .await?;
403        tx.commit().await?;
404        let task = self.get(&task.id).await?.context("claimed task missing")?;
405        Ok(Some(AgentTaskClaim {
406            task,
407            worker_id: worker_id.to_owned(),
408            claim_token,
409        }))
410    }
411
412    /// Take scheduling ownership of the oldest expired claim for reconciliation.
413    pub async fn claim_next_expired(
414        &self,
415        worker_id: &str,
416        lease_seconds: u64,
417    ) -> Result<Option<ExpiredAgentTaskClaim>> {
418        validate_worker_id(worker_id)?;
419        ensure!(
420            (5..=3_600).contains(&lease_seconds),
421            "lease duration must be 5-3600 seconds"
422        );
423        let claim_token = Uuid::new_v4().to_string();
424        let mut tx = self.pool.begin().await?;
425        let task: Option<AgentTask> = sqlx::query_as(
426            "WITH candidate AS (SELECT id FROM agent_tasks WHERE workspace_id = ? AND status IN ('claimed', 'cancel_requested') AND lease_expires_at <= CAST(unixepoch('subsec') * 1000 AS INTEGER) ORDER BY lease_expires_at, created_at, id LIMIT 1) UPDATE agent_tasks SET claim_owner = ?, claim_token = ?, lease_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + ?, updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), last_sequence = last_sequence + 1 WHERE id = (SELECT id FROM candidate) AND workspace_id = ? AND status IN ('claimed', 'cancel_requested') AND lease_expires_at <= CAST(unixepoch('subsec') * 1000 AS INTEGER) RETURNING *",
427        )
428        .bind(&self.workspace_id)
429        .bind(worker_id)
430        .bind(&claim_token)
431        .bind(i64::try_from(lease_seconds)?.saturating_mul(1_000))
432        .bind(&self.workspace_id)
433        .fetch_optional(&mut *tx)
434        .await?;
435        let Some(task) = task else {
436            tx.commit().await?;
437            return Ok(None);
438        };
439        sqlx::query("INSERT INTO agent_task_events (task_id, sequence, timestamp, event_type, payload) VALUES (?, ?, ?, 'task.claimed', ?)")
440            .bind(&task.id)
441            .bind(task.last_sequence)
442            .bind(task.updated_at)
443            .bind(sqlx::types::Json(json!({"worker_id": worker_id, "reconciliation": true})))
444            .execute(&mut *tx)
445            .await?;
446        tx.commit().await?;
447        Ok(Some(ExpiredAgentTaskClaim {
448            claim: AgentTaskClaim {
449                task,
450                worker_id: worker_id.to_owned(),
451                claim_token,
452            },
453        }))
454    }
455
456    /// Extend an owned claim from the database's current time.
457    pub async fn heartbeat(&self, claim: &AgentTaskClaim, lease_seconds: u64) -> Result<AgentTask> {
458        ensure!(
459            (5..=3_600).contains(&lease_seconds),
460            "lease duration must be 5-3600 seconds"
461        );
462        let mut tx = self.pool.begin().await?;
463        let timestamp: Option<i64> = sqlx::query_scalar(
464            "UPDATE agent_tasks SET lease_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + ?, updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) WHERE id = ? AND workspace_id = ? AND claim_owner = ? AND claim_token = ? AND status IN ('claimed', 'cancel_requested') RETURNING updated_at",
465        )
466        .bind(i64::try_from(lease_seconds)?.saturating_mul(1_000))
467        .bind(&claim.task.id)
468        .bind(&self.workspace_id)
469        .bind(&claim.worker_id)
470        .bind(&claim.claim_token)
471        .fetch_optional(&mut *tx)
472        .await?;
473        let timestamp = timestamp.context("task claim is no longer owned by this worker")?;
474        append_event(
475            &mut tx,
476            &claim.task.id,
477            timestamp,
478            "task.heartbeat",
479            &json!({}),
480        )
481        .await?;
482        tx.commit().await?;
483        self.get(&claim.task.id)
484            .await?
485            .context("heartbeat task missing")
486    }
487
488    /// Atomically create the task's sole root run and persist its immutable link.
489    #[allow(clippy::too_many_arguments)]
490    pub async fn link_run(
491        &self,
492        claim: &AgentTaskClaim,
493        agent_version: &str,
494        model: &str,
495        budgets: RunBudgets,
496    ) -> Result<AgentRun> {
497        let mut tx = self.pool.begin().await?;
498        let task: AgentTask = sqlx::query_as(
499            "UPDATE agent_tasks SET updated_at = updated_at WHERE id = ? AND workspace_id = ? AND status = 'claimed' AND claim_owner = ? AND claim_token = ? AND run_id IS NULL RETURNING *",
500        )
501        .bind(&claim.task.id)
502        .bind(&self.workspace_id)
503        .bind(&claim.worker_id)
504        .bind(&claim.claim_token)
505        .fetch_optional(&mut *tx)
506        .await?
507        .context("task claim cannot create a run")?;
508        let runs = AgentRunStore::new(self.pool.clone(), &self.workspace_id)?;
509        let run_id = runs
510            .create_root_run_in(
511                &mut tx,
512                &task.agent_name,
513                agent_version,
514                model,
515                &task.input,
516                &budgets,
517            )
518            .await?;
519        let timestamp: i64 = sqlx::query_scalar(
520            "UPDATE agent_tasks SET run_id = ?, updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) WHERE id = ? AND workspace_id = ? AND status = 'claimed' AND claim_owner = ? AND claim_token = ? AND run_id IS NULL RETURNING updated_at",
521        )
522        .bind(&run_id)
523        .bind(&claim.task.id)
524        .bind(&self.workspace_id)
525        .bind(&claim.worker_id)
526        .bind(&claim.claim_token)
527        .fetch_one(&mut *tx)
528        .await?;
529        append_event(
530            &mut tx,
531            &claim.task.id,
532            timestamp,
533            "task.run_linked",
534            &json!({"run_id": run_id}),
535        )
536        .await?;
537        tx.commit().await?;
538        runs.get_run(&run_id).await?.context("linked run missing")
539    }
540
541    pub async fn finish_identity_changed(&self, claim: &AgentTaskClaim) -> Result<AgentTask> {
542        self.finish_claim(claim, "accepted_identity_changed", false)
543            .await
544    }
545
546    pub async fn finish_run_stopped(&self, claim: &AgentTaskClaim) -> Result<AgentTask> {
547        self.finish_claim(claim, "run_stopped", true).await
548    }
549
550    pub async fn finish_cancelled_before_run(&self, claim: &AgentTaskClaim) -> Result<AgentTask> {
551        self.finish_claim(claim, "cancelled_before_run", false)
552            .await
553    }
554
555    pub async fn finish_claim_attempts_exhausted(
556        &self,
557        claim: &AgentTaskClaim,
558    ) -> Result<AgentTask> {
559        self.finish_claim(claim, "claim_attempts_exhausted", false)
560            .await
561    }
562
563    pub async fn finish_restart_required(&self, claim: &AgentTaskClaim) -> Result<AgentTask> {
564        self.finish_claim(claim, "restart_required", true).await
565    }
566
567    pub async fn finish_reconciliation_required(
568        &self,
569        claim: &AgentTaskClaim,
570    ) -> Result<AgentTask> {
571        self.finish_claim(claim, "reconciliation_required", true)
572            .await
573    }
574
575    /// Return an expired pre-run claim to the queue without changing attempts.
576    pub async fn requeue(&self, claim: &AgentTaskClaim) -> Result<AgentTask> {
577        let mut tx = self.pool.begin().await?;
578        let changed: Option<(i64, i64)> = sqlx::query_as(
579            "UPDATE agent_tasks SET status = 'queued', claim_owner = NULL, claim_token = NULL, lease_expires_at = NULL, updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), last_sequence = last_sequence + 1 WHERE id = ? AND workspace_id = ? AND status = 'claimed' AND claim_owner = ? AND claim_token = ? AND run_id IS NULL RETURNING last_sequence, updated_at",
580        )
581        .bind(&claim.task.id)
582        .bind(&self.workspace_id)
583        .bind(&claim.worker_id)
584        .bind(&claim.claim_token)
585        .fetch_optional(&mut *tx)
586        .await?;
587        let (sequence, timestamp) = changed.context("task claim cannot be requeued")?;
588        sqlx::query("INSERT INTO agent_task_events (task_id, sequence, timestamp, event_type, payload) VALUES (?, ?, ?, 'task.requeued', '{}')")
589            .bind(&claim.task.id)
590            .bind(sequence)
591            .bind(timestamp)
592            .execute(&mut *tx)
593            .await?;
594        tx.commit().await?;
595        self.get(&claim.task.id)
596            .await?
597            .context("requeued task missing")
598    }
599
600    pub(crate) async fn fail_linked_run_setup(
601        &self,
602        claim: &AgentTaskClaim,
603        run_id: &str,
604    ) -> Result<AgentTask> {
605        let owned: bool = sqlx::query_scalar(
606            "SELECT EXISTS(SELECT 1 FROM agent_tasks WHERE id = ? AND workspace_id = ? AND status IN ('claimed', 'cancel_requested') AND claim_owner = ? AND claim_token = ? AND run_id = ?)",
607        )
608        .bind(&claim.task.id)
609        .bind(&self.workspace_id)
610        .bind(&claim.worker_id)
611        .bind(&claim.claim_token)
612        .bind(run_id)
613        .fetch_one(&self.pool)
614        .await?;
615        ensure!(owned, "task claim no longer owns the linked run");
616        AgentRunStore::new(self.pool.clone(), &self.workspace_id)?
617            .finish_run(
618                run_id,
619                RunOutcome::FailedWithReason {
620                    code: FailureReason::RecoveryRejected,
621                    error: "runtime assembly failed".into(),
622                },
623            )
624            .await?;
625        self.finish_run_stopped(claim).await
626    }
627
628    async fn finish_claim(
629        &self,
630        claim: &AgentTaskClaim,
631        reason: &str,
632        require_run: bool,
633    ) -> Result<AgentTask> {
634        let mut tx = self.pool.begin().await?;
635        let run_predicate = if require_run {
636            "run_id IS NOT NULL"
637        } else {
638            "run_id IS NULL"
639        };
640        let statement = format!(
641            "UPDATE agent_tasks SET status = 'terminal', scheduling_reason = ?, scheduling_detail = '{{}}', claim_owner = NULL, claim_token = NULL, lease_expires_at = NULL, updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), completed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), last_sequence = last_sequence + 1 WHERE id = ? AND workspace_id = ? AND status IN ('claimed', 'cancel_requested') AND claim_owner = ? AND claim_token = ? AND {run_predicate} RETURNING last_sequence, updated_at"
642        );
643        let changed: Option<(i64, i64)> = sqlx::query_as(&statement)
644            .bind(reason)
645            .bind(&claim.task.id)
646            .bind(&self.workspace_id)
647            .bind(&claim.worker_id)
648            .bind(&claim.claim_token)
649            .fetch_optional(&mut *tx)
650            .await?;
651        let (sequence, timestamp) = changed.context("task claim cannot be finalized")?;
652        sqlx::query("INSERT INTO agent_task_events (task_id, sequence, timestamp, event_type, payload) VALUES (?, ?, ?, 'task.terminal', ?)")
653            .bind(&claim.task.id)
654            .bind(sequence)
655            .bind(timestamp)
656            .bind(sqlx::types::Json(json!({"scheduling_reason": reason})))
657            .execute(&mut *tx)
658            .await?;
659        tx.commit().await?;
660        self.get(&claim.task.id)
661            .await?
662            .context("finalized task missing")
663    }
664
665    /// Request cancellation without rewriting linked run state.
666    pub async fn cancel(&self, id: &str) -> Result<AgentTask> {
667        let mut tx = self.pool.begin().await?;
668        let changed: Option<(String, i64, i64)> = sqlx::query_as("UPDATE agent_tasks SET status = CASE WHEN status = 'queued' THEN 'terminal' ELSE 'cancel_requested' END, scheduling_reason = CASE WHEN status = 'queued' THEN 'cancelled_before_run' ELSE scheduling_reason END, scheduling_detail = CASE WHEN status = 'queued' THEN '{}' ELSE scheduling_detail END, updated_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), completed_at = CASE WHEN status = 'queued' THEN CAST(unixepoch('subsec') * 1000 AS INTEGER) ELSE completed_at END, last_sequence = last_sequence + 1 WHERE id = ? AND workspace_id = ? AND status IN ('queued', 'claimed') AND (status != 'queued' OR run_id IS NULL) RETURNING status, last_sequence, updated_at")
669            .bind(id)
670            .bind(&self.workspace_id)
671            .fetch_optional(&mut *tx)
672            .await?;
673        if let Some((status, sequence, timestamp)) = changed {
674            let (event_type, payload) = if status == "terminal" {
675                (
676                    "task.terminal",
677                    json!({"scheduling_reason":"cancelled_before_run"}),
678                )
679            } else {
680                ("task.cancel_requested", json!({}))
681            };
682            sqlx::query("INSERT INTO agent_task_events (task_id, sequence, timestamp, event_type, payload) VALUES (?, ?, ?, ?, ?)")
683                .bind(id)
684                .bind(sequence)
685                .bind(timestamp)
686                .bind(event_type)
687                .bind(sqlx::types::Json(payload))
688                .execute(&mut *tx)
689                .await?;
690            tx.commit().await?;
691            return self.get(id).await?.context("cancelled task missing");
692        }
693        let task: AgentTask =
694            sqlx::query_as("SELECT * FROM agent_tasks WHERE id = ? AND workspace_id = ?")
695                .bind(id)
696                .bind(&self.workspace_id)
697                .fetch_optional(&mut *tx)
698                .await?
699                .context("task not found in workspace")?;
700        tx.commit().await?;
701        Ok(task)
702    }
703
704    /// Backward-compatible queued-only cancellation surface.
705    pub async fn cancel_queued(&self, id: &str) -> Result<AgentTask> {
706        let now = Utc::now().timestamp_millis();
707        let mut tx = self.pool.begin().await?;
708        let sequence: Option<i64> = sqlx::query_scalar("UPDATE agent_tasks SET status = 'terminal', scheduling_reason = 'cancelled_before_run', scheduling_detail = '{}', updated_at = ?, completed_at = ?, last_sequence = last_sequence + 1 WHERE id = ? AND workspace_id = ? AND status = 'queued' AND run_id IS NULL RETURNING last_sequence")
709            .bind(now)
710            .bind(now)
711            .bind(id)
712            .bind(&self.workspace_id)
713            .fetch_optional(&mut *tx)
714            .await?;
715        if let Some(sequence) = sequence {
716            sqlx::query("INSERT INTO agent_task_events (task_id, sequence, timestamp, event_type, payload) VALUES (?, ?, ?, 'task.terminal', ?)")
717                .bind(id)
718                .bind(sequence)
719                .bind(now)
720                .bind(sqlx::types::Json(json!({"scheduling_reason":"cancelled_before_run"})))
721                .execute(&mut *tx)
722                .await?;
723            tx.commit().await?;
724            return self.get(id).await?.context("cancelled task missing");
725        }
726        let task: AgentTask =
727            sqlx::query_as("SELECT * FROM agent_tasks WHERE id = ? AND workspace_id = ?")
728                .bind(id)
729                .bind(&self.workspace_id)
730                .fetch_optional(&mut *tx)
731                .await?
732                .context("task not found in workspace")?;
733        ensure!(
734            task.status == AgentTaskStatus::Terminal,
735            "task is no longer queued"
736        );
737        tx.commit().await?;
738        Ok(task)
739    }
740}
741
742fn validate_worker_id(worker_id: &str) -> Result<()> {
743    ensure!(
744        !worker_id.is_empty() && worker_id.len() <= 128 && !worker_id.chars().any(char::is_control),
745        "worker ID must contain 1-128 non-control UTF-8 bytes"
746    );
747    Ok(())
748}
749
750fn validate_submission(submission: &AgentTaskSubmission) -> Result<()> {
751    ensure!(
752        !submission.request_key.is_empty()
753            && submission.request_key.len() <= 256
754            && !submission.request_key.chars().any(char::is_control),
755        "request key must contain 1-256 non-control UTF-8 bytes"
756    );
757    ensure!(
758        !submission.agent_name.trim().is_empty(),
759        "agent name is required"
760    );
761    ensure!(
762        !submission.agent_version.trim().is_empty(),
763        "agent version is required"
764    );
765    ensure!(
766        !submission.input.trim().is_empty() && submission.input.len() <= MAX_TASK_INPUT_BYTES,
767        "task input must contain 1-65536 bytes"
768    );
769    ensure!(
770        submission.payload_identity.len() <= MAX_PAYLOAD_IDENTITY_BYTES,
771        "task payload identity exceeds 8 KiB"
772    );
773    ensure!(
774        (1..=MAX_QUEUE_LIMIT).contains(&submission.queue_limit),
775        "queue limit must be 1-10000"
776    );
777    Ok(())
778}
779
780fn request_digest(agent: &str, input: &str, payload: &[u8]) -> Result<String> {
781    let mut hash = Sha256::new();
782    hash.update(REQUEST_DOMAIN);
783    for bytes in [agent.as_bytes(), input.as_bytes(), payload] {
784        hash.update(u64::try_from(bytes.len())?.to_be_bytes());
785        hash.update(bytes);
786    }
787    Ok(format!("{:x}", hash.finalize()))
788}
789
790async fn append_event(
791    tx: &mut Transaction<'_, Sqlite>,
792    id: &str,
793    now: i64,
794    event_type: &str,
795    payload: &Value,
796) -> Result<i64> {
797    let sequence: i64 = sqlx::query_scalar(
798        "UPDATE agent_tasks SET last_sequence = last_sequence + 1, updated_at = ? WHERE id = ? RETURNING last_sequence",
799    )
800    .bind(now)
801    .bind(id)
802    .fetch_one(&mut **tx)
803    .await?;
804    sqlx::query("INSERT INTO agent_task_events (task_id, sequence, timestamp, event_type, payload) VALUES (?, ?, ?, ?, ?)")
805        .bind(id)
806        .bind(sequence)
807        .bind(now)
808        .bind(event_type)
809        .bind(sqlx::types::Json(payload))
810        .execute(&mut **tx)
811        .await?;
812    Ok(sequence)
813}
context_harness/agent_task_store.rs

context_harness/
agent_task_store.rs
