klieo-a2a 0.4.0

//! `A2aTaskStore` — durable Task persistence over [`klieo_core::KvStore`].
//!
//! Layout in bucket `a2a.tasks`:
//! - `task.<task_id>` → JSON-encoded [`Task`].
//! - `index.context.<context_id>` → JSON array of task ids in that context.
//!
//! Index entries are maintained on `put` / `delete`. `list(Some(ctx))`
//! reads the index, then fetches each task. `list(None)` is not
//! supported in v0.0.1 (would require a global index entry which is a
//! hot key on writes); callers must filter by context.

use crate::error::A2aError;
use crate::types::Task;
use bytes::Bytes;
use klieo_core::KvStore;
use std::sync::Arc;

/// Default KV bucket name for [`A2aTaskStore`] entries.
pub const DEFAULT_BUCKET: &str = "a2a.tasks";

/// Task store over a CAS-style KV.
pub struct A2aTaskStore {
    kv: Arc<dyn KvStore>,
    bucket: String,
}

impl A2aTaskStore {
    /// Build a new store backed by `kv` writing under `bucket` (typical:
    /// [`DEFAULT_BUCKET`]).
    pub fn new(kv: Arc<dyn KvStore>, bucket: String) -> Self {
        Self { kv, bucket }
    }

    fn task_key(&self, id: &str) -> String {
        format!("task.{id}")
    }

    fn index_key(&self, context_id: &str) -> String {
        format!("index.context.{context_id}")
    }

    /// Persist a task and update its context index.
    pub async fn put(&self, task: &Task) -> Result<(), A2aError> {
        let bytes = Bytes::from(serde_json::to_vec(task)?);
        self.kv
            .put(&self.bucket, &self.task_key(&task.id), bytes)
            .await?;
        // Update index.
        let key = self.index_key(&task.contextId);
        let current = self.kv.get(&self.bucket, &key).await?;
        let mut ids: Vec<String> = match current {
            Some(entry) => serde_json::from_slice(&entry.value)?,
            None => vec![],
        };
        if !ids.contains(&task.id) {
            ids.push(task.id.clone());
        }
        let updated = Bytes::from(serde_json::to_vec(&ids)?);
        self.kv.put(&self.bucket, &key, updated).await?;
        Ok(())
    }

    /// Fetch a task by id.
    pub async fn get(&self, id: &str) -> Result<Option<Task>, A2aError> {
        match self.kv.get(&self.bucket, &self.task_key(id)).await? {
            Some(entry) => Ok(Some(serde_json::from_slice(&entry.value)?)),
            None => Ok(None),
        }
    }

    /// List tasks. `context_id = Some(ctx)` reads the index; `None` is
    /// not supported in v0.0.1 and returns an empty vec — callers must
    /// supply a context.
    pub async fn list(&self, context_id: Option<&str>) -> Result<Vec<Task>, A2aError> {
        let Some(ctx) = context_id else {
            return Ok(vec![]);
        };
        let entry = self.kv.get(&self.bucket, &self.index_key(ctx)).await?;
        let ids: Vec<String> = match entry {
            Some(e) => serde_json::from_slice(&e.value)?,
            None => return Ok(vec![]),
        };
        let mut out = Vec::with_capacity(ids.len());
        for id in ids {
            if let Some(t) = self.get(&id).await? {
                out.push(t);
            }
        }
        Ok(out)
    }

    /// Delete a task and remove its id from the context index.
    pub async fn delete(&self, id: &str) -> Result<(), A2aError> {
        // Read existing to find context id.
        let task = self.get(id).await?;
        self.kv.delete(&self.bucket, &self.task_key(id)).await?;
        if let Some(t) = task {
            let key = self.index_key(&t.contextId);
            if let Some(entry) = self.kv.get(&self.bucket, &key).await? {
                let mut ids: Vec<String> = serde_json::from_slice(&entry.value)?;
                ids.retain(|x| x != id);
                let updated = Bytes::from(serde_json::to_vec(&ids)?);
                self.kv.put(&self.bucket, &key, updated).await?;
            }
        }
        Ok(())
    }
}