diff --git a/Cargo.lock b/Cargo.lock index 35d6cd46..2ea404c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -127,6 +127,7 @@ dependencies = [ name = "agentflare-backend" version = "0.1.0" dependencies = [ + "agentflare-db-kit", "hex", "hmac", "rusqlite", diff --git a/crates/agentflare-backend/Cargo.toml b/crates/agentflare-backend/Cargo.toml index 84b0a508..b1d488ea 100644 --- a/crates/agentflare-backend/Cargo.toml +++ b/crates/agentflare-backend/Cargo.toml @@ -8,6 +8,7 @@ description = "Universal agent backend: workspace/project/item schema + CRUD, po publish = false [dependencies] +db_kit = { package = "agentflare-db-kit", path = "../agentflare-db-kit" } rusqlite = { version = "0.40", features = ["bundled"] } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/crates/agentflare-backend/src/claim.rs b/crates/agentflare-backend/src/claim.rs new file mode 100644 index 00000000..99242ca2 --- /dev/null +++ b/crates/agentflare-backend/src/claim.rs @@ -0,0 +1,37 @@ +//! Item claim lease — a thin wrapper over agentflare-db-kit's generic +//! `ClaimLedger`, keyed by `item_id`. Pure lease primitive, no item-state +//! knowledge; `item::claim`/`item::claim_done` compose this with +//! `item::update_state` to make claiming actually mean something. +use db_kit::claim::ClaimLedger; +use rusqlite::Connection; + +pub use db_kit::claim::Acquire; + +const LEDGER: ClaimLedger = ClaimLedger::new("item_claims", &["item_id"]); + +pub fn acquire( + conn: &Connection, + item_id: &str, + owner: &str, + now: i64, + ttl_secs: i64, +) -> rusqlite::Result { + LEDGER.acquire(conn, &[item_id], owner, now, ttl_secs) +} + +pub fn heartbeat( + conn: &Connection, + item_id: &str, + owner: &str, + now: i64, +) -> rusqlite::Result { + LEDGER.heartbeat(conn, &[item_id], owner, now) +} + +pub fn release(conn: &Connection, item_id: &str, owner: &str) -> rusqlite::Result { + LEDGER.release(conn, &[item_id], owner) +} + +pub fn done(conn: &Connection, item_id: &str, owner: &str, now: i64) -> rusqlite::Result { + LEDGER.done(conn, &[item_id], owner, now) +} diff --git a/crates/agentflare-backend/src/item.rs b/crates/agentflare-backend/src/item.rs index 00fddcb3..ab619949 100644 --- a/crates/agentflare-backend/src/item.rs +++ b/crates/agentflare-backend/src/item.rs @@ -401,6 +401,56 @@ pub fn list_dependencies(conn: &Connection, item_id: &str) -> Result Ok(rows.collect::>()?) } +/// Claims an item so other agents don't duplicate the work: on a fresh +/// acquire, sets the assignee and moves state into the project's "started" +/// group (which sets `started_at`, via `update_state`). A live claim held by +/// someone else returns `Held` and leaves the item untouched. Acquisition, +/// the state transition, and the assignee update are one transaction — a +/// mid-sequence failure can't leave `item_claims` saying "claimed" while the +/// item itself never reflects it. +pub fn claim( + conn: &Connection, + item_id: &str, + owner: &str, + now: i64, + ttl_secs: i64, +) -> Result { + let tx = conn.unchecked_transaction()?; + let outcome = crate::claim::acquire(&tx, item_id, owner, now, ttl_secs)?; + if outcome == crate::claim::Acquire::Acquired { + let item = get(&tx, item_id)?; + let started_state = crate::state::first_in_group(&tx, &item.project_id, "started")?; + update_state(&tx, item_id, &started_state.id)?; + update( + &tx, + item_id, + UpdateItem { + assignee_agent: Some(owner.to_string()), + ..Default::default() + }, + )?; + } + tx.commit()?; + Ok(outcome) +} + +/// Marks a claimed item done: releases the claim (re-claimable by anyone +/// afterward) and moves state into the project's "completed" group. Always +/// "completed", never "cancelled" — those are distinct outcomes. The done +/// transition and the state update are one transaction, same reasoning as +/// `claim()`. +pub fn claim_done(conn: &Connection, item_id: &str, owner: &str, now: i64) -> Result { + let tx = conn.unchecked_transaction()?; + let done = crate::claim::done(&tx, item_id, owner, now)?; + if done { + let item = get(&tx, item_id)?; + let completed_state = crate::state::first_in_group(&tx, &item.project_id, "completed")?; + update_state(&tx, item_id, &completed_state.id)?; + } + tx.commit()?; + Ok(done) +} + #[cfg(test)] mod tests { use super::*; @@ -917,4 +967,99 @@ mod tests { Err(crate::error::Error::InvalidTransition(_)) )); } + + const TTL: i64 = 14400; + + fn make_item(conn: &Connection, pid: &str, sid: &str) -> Item { + create( + conn, + CreateItem { + project_id: pid.to_string(), + state_id: sid.to_string(), + name: "Test".into(), + description: None, + priority: None, + parent_id: None, + assignee_agent: None, + sort_order: None, + external_source: None, + external_id: None, + metadata: None, + label_ids: vec![], + assignee_ids: vec![], + dependency_ids: vec![], + }, + ) + .unwrap() + } + + #[test] + fn claim_acquires_sets_assignee_and_moves_to_started_state() { + let conn = db::open_in_memory().unwrap(); + let (pid, sid) = seed_project(&conn, ""); + let item = make_item(&conn, &pid, &sid); + let outcome = claim(&conn, &item.id, "agent:1", 1000, TTL).unwrap(); + assert_eq!(outcome, crate::claim::Acquire::Acquired); + let updated = get(&conn, &item.id).unwrap(); + assert_eq!(updated.assignee_agent.as_deref(), Some("agent:1")); + assert_eq!(updated.state_id, state_in_group(&conn, &pid, "started")); + assert!(updated.started_at.is_some()); + } + + #[test] + fn claim_on_already_held_item_returns_held_and_leaves_item_unchanged() { + let conn = db::open_in_memory().unwrap(); + let (pid, sid) = seed_project(&conn, ""); + let item = make_item(&conn, &pid, &sid); + claim(&conn, &item.id, "agent:1", 1000, TTL).unwrap(); + let outcome = claim(&conn, &item.id, "agent:2", 1001, TTL).unwrap(); + assert!(matches!( + outcome, + crate::claim::Acquire::Held { ref owner, .. } if owner == "agent:1" + )); + let unchanged = get(&conn, &item.id).unwrap(); + assert_eq!(unchanged.assignee_agent.as_deref(), Some("agent:1")); + } + + #[test] + fn stale_claim_is_stealable_by_a_different_owner() { + let conn = db::open_in_memory().unwrap(); + let (pid, sid) = seed_project(&conn, ""); + let item = make_item(&conn, &pid, &sid); + claim(&conn, &item.id, "agent:1", 1000, TTL).unwrap(); + let outcome = claim(&conn, &item.id, "agent:2", 1000 + TTL + 1, TTL).unwrap(); + assert_eq!(outcome, crate::claim::Acquire::Acquired); + let updated = get(&conn, &item.id).unwrap(); + assert_eq!(updated.assignee_agent.as_deref(), Some("agent:2")); + } + + #[test] + fn claim_done_moves_to_completed_state_and_is_reclaimable() { + let conn = db::open_in_memory().unwrap(); + let (pid, sid) = seed_project(&conn, ""); + let item = make_item(&conn, &pid, &sid); + claim(&conn, &item.id, "agent:1", 1000, TTL).unwrap(); + assert!(claim_done(&conn, &item.id, "agent:1", 1100).unwrap()); + let done_item = get(&conn, &item.id).unwrap(); + assert_eq!(done_item.state_id, state_in_group(&conn, &pid, "completed")); + assert!(done_item.completed_at.is_some()); + + let outcome = claim(&conn, &item.id, "agent:2", 1200, TTL).unwrap(); + assert_eq!(outcome, crate::claim::Acquire::Acquired); + } + + #[test] + fn heartbeat_release_done_are_owner_scoped() { + let conn = db::open_in_memory().unwrap(); + let (pid, sid) = seed_project(&conn, ""); + let item = make_item(&conn, &pid, &sid); + claim(&conn, &item.id, "agent:1", 1000, TTL).unwrap(); + + assert!(!crate::claim::heartbeat(&conn, &item.id, "agent:2", 1100).unwrap()); + assert!(!crate::claim::release(&conn, &item.id, "agent:2").unwrap()); + assert!(!claim_done(&conn, &item.id, "agent:2", 1100).unwrap()); + + assert!(crate::claim::heartbeat(&conn, &item.id, "agent:1", 1100).unwrap()); + assert!(claim_done(&conn, &item.id, "agent:1", 1200).unwrap()); + } } diff --git a/crates/agentflare-backend/src/lib.rs b/crates/agentflare-backend/src/lib.rs index 90149761..db57ebac 100644 --- a/crates/agentflare-backend/src/lib.rs +++ b/crates/agentflare-backend/src/lib.rs @@ -1,4 +1,5 @@ pub mod asset; +pub mod claim; pub mod db; pub mod error; pub mod events; diff --git a/crates/agentflare-backend/src/schema.sql b/crates/agentflare-backend/src/schema.sql index 873f7b69..f05df65c 100644 --- a/crates/agentflare-backend/src/schema.sql +++ b/crates/agentflare-backend/src/schema.sql @@ -153,3 +153,11 @@ CREATE TABLE IF NOT EXISTS webhook_logs ( ); CREATE INDEX IF NOT EXISTS idx_webhook_logs_webhook ON webhook_logs(webhook_id); CREATE INDEX IF NOT EXISTS idx_webhook_logs_workspace ON webhook_logs(workspace_id); + +CREATE TABLE IF NOT EXISTS item_claims ( + item_id TEXT PRIMARY KEY REFERENCES items(id), + owner TEXT NOT NULL, + status TEXT NOT NULL, + created_at INTEGER NOT NULL, + heartbeat_at INTEGER NOT NULL +); diff --git a/crates/agentflare-backend/src/state.rs b/crates/agentflare-backend/src/state.rs index 4df4a0ee..379504bb 100644 --- a/crates/agentflare-backend/src/state.rs +++ b/crates/agentflare-backend/src/state.rs @@ -134,6 +134,23 @@ pub fn get(conn: &Connection, id: &str) -> Result { }) } +/// First state (by sequence) in `group` for the project — used to resolve +/// the "Started"/"Completed" target when claiming or completing an item. +pub fn first_in_group(conn: &Connection, project_id: &str, group: &str) -> Result { + conn.query_row( + "SELECT id, project_id, name, group_name, sequence, is_default, color, created_at, updated_at, deleted_at + FROM states WHERE project_id = ?1 AND group_name = ?2 AND deleted_at IS NULL ORDER BY sequence LIMIT 1", + rusqlite::params![project_id, group], + row_to_state, + ) + .map_err(|e| match e { + rusqlite::Error::QueryReturnedNoRows => { + crate::error::Error::NotFound(format!("no '{group}' state for project {project_id}")) + } + other => other.into(), + }) +} + pub fn list_by_project(conn: &Connection, project_id: &str) -> Result> { let mut stmt = conn.prepare( "SELECT id, project_id, name, group_name, sequence, is_default, color, created_at, updated_at, deleted_at @@ -256,6 +273,25 @@ mod tests { ); } + #[test] + fn first_in_group_returns_lowest_sequence_match() { + let conn = db::open_in_memory().unwrap(); + let pid = seed_project(&conn); + let started = first_in_group(&conn, &pid, "started").unwrap(); + assert_eq!(started.name, "In Progress"); + assert_eq!(started.group_name, "started"); + } + + #[test] + fn first_in_group_errors_when_no_state_matches() { + let conn = db::open_in_memory().unwrap(); + let pid = seed_project(&conn); + assert!(matches!( + first_in_group(&conn, &pid, "no-such-group"), + Err(crate::error::Error::NotFound(_)) + )); + } + #[test] fn create_custom_state() { let conn = db::open_in_memory().unwrap(); diff --git a/src/mcp_server.rs b/src/mcp_server.rs index ee537765..e651724b 100644 --- a/src/mcp_server.rs +++ b/src/mcp_server.rs @@ -550,6 +550,16 @@ struct ProjectLink { identifier: String, } +/// Default 4h — item claims are plausibly longer-running than +/// `src/claims.rs`'s 30-min GitHub-issue-claim default, hence a separate env +/// var rather than sharing `AGENTFLARE_CLAIM_TTL_SECS`. +fn backend_claim_ttl_secs() -> i64 { + std::env::var("AGENTFLARE_BACKEND_CLAIM_TTL_SECS") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(14400) as i64 +} + /// NotFound/Duplicate/InvalidTransition are caller-fixable → invalid_params; /// a raw database error is ours to fix → internal_error. Same split as /// `skill_load`'s NotFound/Ambiguous handling above. @@ -647,6 +657,12 @@ struct BackendItemDeleteRequest { id: String, } +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct BackendItemClaimRequest { + #[schemars(description = "Item ID")] + item_id: String, +} + #[derive(Debug, Deserialize, schemars::JsonSchema)] struct BackendLabelCreateRequest { #[schemars(description = "Label name")] @@ -2286,6 +2302,90 @@ impl AgentflareMcp { })? } + #[tool( + description = "Claim a work item so other agents don't duplicate the work. Sets assignee + moves state to the project's Started state. Returns acquired, or held with the current owner if a live claim exists." + )] + fn backend_item_claim( + &self, + Parameters(BackendItemClaimRequest { item_id }): Parameters, + ) -> Result { + if item_id.trim().is_empty() { + return Err(ErrorData::invalid_params("item_id is required", None)); + } + let owner = crate::claims::owner_id(); + let now = crate::claims::now(); + let ttl = backend_claim_ttl_secs(); + self.with_backend_db(|conn| { + let outcome = agentflare_backend::item::claim(conn, &item_id, &owner, now, ttl) + .map_err(map_backend_err)?; + Ok(match outcome { + agentflare_backend::claim::Acquire::Acquired => { + serde_json::json!({"status": "acquired", "item_id": item_id, "owner": owner}) + } + agentflare_backend::claim::Acquire::Held { owner: holder, age_secs } => { + serde_json::json!({"status": "held", "item_id": item_id, "owner": holder, "age_secs": age_secs}) + } + } + .to_string()) + })? + } + + #[tool( + description = "Refresh the lease on a work item claim you own. Returns heartbeat=false if you don't hold a live claim." + )] + fn backend_item_heartbeat( + &self, + Parameters(BackendItemClaimRequest { item_id }): Parameters, + ) -> Result { + if item_id.trim().is_empty() { + return Err(ErrorData::invalid_params("item_id is required", None)); + } + let owner = crate::claims::owner_id(); + let now = crate::claims::now(); + self.with_backend_db(|conn| { + let ok = agentflare_backend::claim::heartbeat(conn, &item_id, &owner, now) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + Ok(serde_json::json!({"heartbeat": ok, "item_id": item_id}).to_string()) + })? + } + + #[tool( + description = "Release a work item claim you own, without changing its state. Returns released=false if you don't hold a live claim." + )] + fn backend_item_release( + &self, + Parameters(BackendItemClaimRequest { item_id }): Parameters, + ) -> Result { + if item_id.trim().is_empty() { + return Err(ErrorData::invalid_params("item_id is required", None)); + } + let owner = crate::claims::owner_id(); + self.with_backend_db(|conn| { + let ok = agentflare_backend::claim::release(conn, &item_id, &owner) + .map_err(|e| ErrorData::internal_error(e.to_string(), None))?; + Ok(serde_json::json!({"released": ok, "item_id": item_id}).to_string()) + })? + } + + #[tool( + description = "Mark a claimed work item done: releases the claim and moves state to the project's Completed state. Returns done=false if you don't hold a live claim." + )] + fn backend_item_done( + &self, + Parameters(BackendItemClaimRequest { item_id }): Parameters, + ) -> Result { + if item_id.trim().is_empty() { + return Err(ErrorData::invalid_params("item_id is required", None)); + } + let owner = crate::claims::owner_id(); + let now = crate::claims::now(); + self.with_backend_db(|conn| { + let done = agentflare_backend::item::claim_done(conn, &item_id, &owner, now) + .map_err(map_backend_err)?; + Ok(serde_json::json!({"done": done, "item_id": item_id}).to_string()) + })? + } + #[tool(description = "Create a label in the repo's linked project.")] fn backend_label_create( &self,