use std::collections::{BTreeMap, HashMap};
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Instant, SystemTime};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use crate::module::{Manifest, Module};
use crate::ops::{
APP_MIME, App, Ask, Body, Completing, Content, Ctx, JsonObject, Looks, Message, Ops, Progress,
Prompt, Resource, Tool, with_input,
};
use crate::seams::{
Change, ChangeFeed, Entitlements, Feature, Identity, InProcessChanges, Local, Metering, Open,
Space, Unmetered, Usage,
};
use crate::tasks::{InProcessTasks, StoredTask, TaskState, TaskStore};
pub struct Kit<M, I = Local, E = Open, T = Unmetered, S = InProcessTasks> {
core: Core<M>,
identity: I,
entitlements: E,
metering: T,
tasks: S,
}
struct Core<M> {
module: Arc<M>,
manifest: Manifest,
ops: Ops<M>,
routes: matchit::Router<usize>,
changes: Arc<dyn ChangeFeed>,
running: Mutex<HashMap<String, CancellationToken>>,
}
#[derive(Default)]
pub(crate) struct Caller<'a> {
pub bearer: Option<&'a str>,
pub cancel: Option<CancellationToken>,
pub progress: Option<Arc<dyn Progress>>,
pub answers: Option<BTreeMap<String, Value>>,
pub space: Option<Arc<Space>>,
}
impl<'a> Caller<'a> {
pub(crate) fn holding(bearer: Option<&'a str>) -> Self {
Self {
bearer,
..Self::default()
}
}
}
impl<M: Module, I: Identity> Kit<M, I> {
pub fn new(module: M, identity: I) -> Self {
let mut ops = Ops::new();
module.register(&mut ops);
let mut routes = matchit::Router::new();
for (index, readable) in ops.resources.iter().enumerate() {
let uri = readable.resource.uri;
if let Err(error) = routes.insert(uri, index) {
panic!("resource {uri} cannot be routed: {error}");
}
}
Self {
core: Core {
manifest: module.manifest(),
module: Arc::new(module),
ops,
routes,
changes: Arc::new(InProcessChanges::default()),
running: Mutex::default(),
},
identity,
entitlements: Open,
metering: Unmetered,
tasks: InProcessTasks::default(),
}
}
}
impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Kit<M, I, E, T, S> {
pub fn with_entitlements<E2: Entitlements>(self, entitlements: E2) -> Kit<M, I, E2, T, S> {
Kit {
core: self.core,
identity: self.identity,
entitlements,
metering: self.metering,
tasks: self.tasks,
}
}
pub fn with_metering<T2: Metering>(self, metering: T2) -> Kit<M, I, E, T2, S> {
Kit {
core: self.core,
identity: self.identity,
entitlements: self.entitlements,
metering,
tasks: self.tasks,
}
}
pub fn with_changes(mut self, changes: impl ChangeFeed) -> Self {
self.core.changes = Arc::new(changes);
self
}
pub fn with_tasks<S2: TaskStore>(self, tasks: S2) -> Kit<M, I, E, T, S2> {
Kit {
core: self.core,
identity: self.identity,
entitlements: self.entitlements,
metering: self.metering,
tasks,
}
}
pub fn manifest(&self) -> &Manifest {
&self.core.manifest
}
pub fn tools(&self) -> impl Iterator<Item = &Tool> {
self.core.ops.list.iter().map(|op| &op.tool)
}
pub fn instructions(&self) -> Option<&'static str> {
self.core.ops.instructions
}
pub fn looks(&self) -> &Looks {
&self.core.ops.server
}
pub fn apps(&self) -> &[App] {
&self.core.ops.apps
}
pub fn resources(&self) -> impl Iterator<Item = &Resource> {
self.core.ops.resources.iter().map(|r| &r.resource)
}
pub async fn read(&self, uri: &str, bearer: Option<&str>) -> Result<Body, CallError> {
let caller = Caller::holding(bearer);
self.read_as(uri, caller).await
}
pub(crate) async fn read_as(&self, uri: &str, caller: Caller<'_>) -> Result<Body, CallError> {
if let Some(app) = self.apps().iter().find(|app| app.uri == uri) {
return Ok(Body::Text {
mime: APP_MIME.into(),
text: app.html.to_owned(),
});
}
let space = self.identified(&caller).await?;
let found = self
.core
.routes
.at(uri)
.map_err(|_| CallError::Unknown(uri.to_owned()))?;
let readable = &self.core.ops.resources[*found.value];
let Resource {
uri: template,
feature,
..
} = readable.resource;
self.entitled(&space, feature)?;
let variables = if found.params.is_empty() {
Value::Null
} else {
found
.params
.iter()
.map(|(name, value)| {
let value = percent_encoding::percent_decode_str(value).decode_utf8_lossy();
(name.to_owned(), Value::String(value.into_owned()))
})
.collect::<JsonObject>()
.into()
};
let ctx = Ctx::new(space, caller.cancel, caller.progress);
self.metered(template, feature, ctx, caller.answers, |ctx| {
(readable.read)(self.core.module.clone(), ctx, variables)
})
.await
.map(|(value, _)| value)
.map_err(CallError::from)
}
pub(crate) fn has_completions(&self) -> bool {
!self.core.ops.completers.is_empty()
}
pub fn prompts(&self) -> impl Iterator<Item = &Prompt> {
self.core.ops.prompts.iter().map(|p| &p.prompt)
}
pub async fn prompt(
&self,
name: &str,
bearer: Option<&str>,
arguments: JsonObject,
) -> Result<Vec<Message>, CallError> {
let caller = Caller::holding(bearer);
self.prompt_as(name, caller, arguments).await
}
pub(crate) async fn prompt_as(
&self,
name: &str,
caller: Caller<'_>,
arguments: JsonObject,
) -> Result<Vec<Message>, CallError> {
let space = self.identified(&caller).await?;
let promptable = self
.core
.ops
.prompts
.iter()
.find(|p| p.prompt.name == name)
.ok_or_else(|| CallError::Unknown(name.to_owned()))?;
let Prompt { name, feature, .. } = promptable.prompt;
self.entitled(&space, feature)?;
let ctx = Ctx::new(space, caller.cancel, caller.progress);
self.metered(name, feature, ctx, caller.answers, |ctx| {
(promptable.get)(self.core.module.clone(), ctx, Value::Object(arguments))
})
.await
.map(|(value, _)| value)
.map_err(CallError::from)
}
pub(crate) async fn complete_as(
&self,
target: &str,
argument: &str,
completing: Completing,
caller: Caller<'_>,
) -> Result<Vec<String>, CallError> {
let space = self.identified(&caller).await?;
let Some(completer) = self
.core
.ops
.completers
.iter()
.find(|c| c.target == target && c.argument == argument)
else {
return Ok(Vec::new());
};
let feature = self
.prompts()
.find(|p| p.name == target)
.map(|p| p.feature)
.or_else(|| {
self.resources()
.find(|r| r.uri == target)
.map(|r| r.feature)
})
.expect("Ops::complete checked the target");
self.entitled(&space, feature)?;
let ctx = Ctx::new(space, caller.cancel, caller.progress);
self.metered(completer.target, feature, ctx, caller.answers, |ctx| {
let completing = serde_json::to_value(completing).unwrap_or_default();
(completer.complete)(self.core.module.clone(), ctx, completing)
})
.await
.map(|(value, _)| value)
.map_err(CallError::from)
}
pub fn tool(&self, name: &str) -> Option<&Tool> {
self.tools().find(|tool| tool.name == name)
}
pub async fn call(
&self,
name: &str,
bearer: Option<&str>,
arguments: JsonObject,
) -> Result<Reply, CallError> {
let caller = Caller::holding(bearer);
self.call_as(name, caller, arguments).await
}
pub(crate) async fn call_as(
&self,
name: &str,
caller: Caller<'_>,
arguments: JsonObject,
) -> Result<Reply, CallError> {
let (index, space) = self.admit(name, &caller).await?;
let ctx = Ctx::new(space, caller.cancel, caller.progress);
self.run(index, ctx, Value::Object(arguments), caller.answers)
.await
.map_err(CallError::from)
}
pub async fn start(
self: &Arc<Self>,
name: &str,
bearer: Option<&str>,
arguments: JsonObject,
) -> Result<StoredTask, CallError> {
self.start_as(name, &Caller::holding(bearer), arguments)
.await
}
pub(crate) async fn start_as(
self: &Arc<Self>,
name: &str,
caller: &Caller<'_>,
arguments: JsonObject,
) -> Result<StoredTask, CallError> {
let (index, space) = self.admit(name, caller).await?;
let now = SystemTime::now();
let arguments = Value::Object(arguments);
let task = StoredTask {
id: uuid::Uuid::new_v4().to_string(),
space: Space::clone(&space),
operation: self.core.ops.list[index].tool.name.to_owned(),
arguments: arguments.clone(),
state: TaskState::Working,
status: None,
created_at: now,
updated_at: now,
version: 0,
};
self.tasks
.put(task.clone(), None)
.await
.map_err(unavailable)?;
self.launch(index, space, task.id.clone(), arguments, None);
Ok(task)
}
pub async fn update(
self: &Arc<Self>,
bearer: Option<&str>,
id: &str,
answers: BTreeMap<String, Value>,
) -> Result<StoredTask, CallError> {
self.update_as(&Caller::holding(bearer), id, answers).await
}
pub(crate) async fn update_as(
self: &Arc<Self>,
caller: &Caller<'_>,
id: &str,
answers: BTreeMap<String, Value>,
) -> Result<StoredTask, CallError> {
self.task_as(caller, id).await?;
let (task, resumed) = self
.transition(id, |task| {
matches!(task.state, TaskState::InputRequired(_)).then(|| StoredTask {
state: TaskState::Working,
status: None,
..task.clone()
})
})
.await
.map_err(unavailable)?
.ok_or_else(|| CallError::NoTask(id.to_owned()))?;
if resumed {
let index = self
.core
.ops
.list
.iter()
.position(|op| op.tool.name == task.operation)
.ok_or_else(|| CallError::Unknown(task.operation.clone()))?;
let space = Arc::new(task.space.clone());
self.launch(
index,
space,
task.id.clone(),
task.arguments.clone(),
Some(answers),
);
}
Ok(task)
}
fn launch(
self: &Arc<Self>,
index: usize,
space: Arc<Space>,
id: String,
arguments: Value,
answers: Option<BTreeMap<String, Value>>,
) {
let cancel = CancellationToken::new();
self.running().insert(id.clone(), cancel.clone());
let progress = Arc::new(TaskProgress {
kit: self.clone(),
id: id.clone(),
});
let ctx = Ctx::new(space, Some(cancel), Some(progress));
let kit = self.clone();
tokio::spawn(tracing::Instrument::in_current_span(async move {
let state = match kit.run(index, ctx, arguments, answers).await {
Ok(result) => TaskState::Succeeded(result),
Err(Halt::Failed(message)) => TaskState::Errored(message),
Err(Halt::Input(asked)) => TaskState::InputRequired(asked),
};
kit.running().remove(&id);
if let Err(error) = kit.finish(&id, state).await {
tracing::error!(%error, "a finished task could not be stored");
}
}));
}
pub async fn task(&self, bearer: Option<&str>, id: &str) -> Result<StoredTask, CallError> {
self.task_as(&Caller::holding(bearer), id).await
}
pub(crate) async fn task_as(
&self,
caller: &Caller<'_>,
id: &str,
) -> Result<StoredTask, CallError> {
let space = self.identified(caller).await?;
match self.tasks.get(id).await.map_err(unavailable)? {
Some(task) if task.space == *space => Ok(task),
_ => Err(CallError::NoTask(id.to_owned())),
}
}
pub async fn cancel(&self, bearer: Option<&str>, id: &str) -> Result<StoredTask, CallError> {
self.cancel_as(&Caller::holding(bearer), id).await
}
pub(crate) async fn cancel_as(
&self,
caller: &Caller<'_>,
id: &str,
) -> Result<StoredTask, CallError> {
self.task_as(caller, id).await?;
if let Some(cancel) = self.running().get(id) {
cancel.cancel();
}
self.finish(id, TaskState::Cancelled)
.await
.map_err(unavailable)?
.ok_or_else(|| CallError::NoTask(id.to_owned()))
}
async fn status(&self, id: &str, status: String) -> std::io::Result<()> {
self.transition(id, |task| {
(!task.state.is_terminal()).then(|| StoredTask {
status: Some(status.clone()),
..task.clone()
})
})
.await
.map(drop)
}
fn running(&self) -> std::sync::MutexGuard<'_, HashMap<String, CancellationToken>> {
self.core
.running
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
fn entitled(&self, space: &Space, feature: Feature) -> Result<(), CallError> {
if self.entitlements.allows(space, feature) {
Ok(())
} else {
Err(CallError::NotEntitled(feature))
}
}
pub(crate) fn changes(&self) -> tokio::sync::broadcast::Receiver<Change> {
self.core.changes.subscribe()
}
pub(crate) async fn identified(&self, caller: &Caller<'_>) -> Result<Arc<Space>, CallError> {
match &caller.space {
Some(space) => Ok(space.clone()),
None => self.identify(caller.bearer).await,
}
}
pub(crate) async fn identify(&self, bearer: Option<&str>) -> Result<Arc<Space>, CallError> {
self.identity
.identify(bearer)
.await
.ok_or(CallError::Unidentified)
}
async fn admit(
&self,
name: &str,
caller: &Caller<'_>,
) -> Result<(usize, Arc<Space>), CallError> {
let space = self.identified(caller).await?;
let index = self
.core
.ops
.list
.iter()
.position(|op| op.tool.name == name)
.ok_or_else(|| CallError::Unknown(name.to_owned()))?;
self.entitled(&space, self.core.ops.list[index].tool.feature)?;
Ok((index, space))
}
async fn run(
&self,
index: usize,
ctx: Ctx,
arguments: Value,
answers: Option<BTreeMap<String, Value>>,
) -> Result<Reply, Halt> {
let op = &self.core.ops.list[index];
let (value, content) = self
.metered(op.tool.name, op.tool.feature, ctx, answers, |ctx| {
(op.call)(self.core.module.clone(), ctx, arguments)
})
.await?;
Ok(Reply { value, content })
}
async fn metered<V, Fut>(
&self,
operation: &'static str,
feature: Feature,
ctx: Ctx,
answers: Option<BTreeMap<String, Value>>,
work: impl FnOnce(Ctx) -> Fut,
) -> Result<(V, Vec<Content>), Halt>
where
Fut: Future<Output = Result<V, String>>,
{
let space = ctx.shared_space();
let started = Instant::now();
let (result, effects) = with_input(answers, work(ctx)).await;
for uri in effects.updated {
self.core.changes.publish(Change {
space: space.clone(),
uri,
});
}
let result = match result {
Ok(value) => Ok((value, effects.attached)),
Err(_) if !effects.asked.is_empty() => Err(Halt::Input(effects.asked)),
Err(message) => Err(Halt::Failed(message)),
};
self.metering.record(
&space,
Usage {
operation,
feature,
elapsed: started.elapsed(),
ok: result.is_ok(),
},
);
result
}
async fn finish(&self, id: &str, state: TaskState) -> std::io::Result<Option<StoredTask>> {
let changed = self
.transition(id, |task| {
(!task.state.is_terminal()).then(|| StoredTask {
state: state.clone(),
..task.clone()
})
})
.await?;
Ok(changed.map(|(task, _)| task))
}
async fn transition(
&self,
id: &str,
change: impl Fn(&StoredTask) -> Option<StoredTask>,
) -> std::io::Result<Option<(StoredTask, bool)>> {
loop {
let Some(current) = self.tasks.get(id).await? else {
return Ok(None);
};
let Some(next) = change(¤t) else {
return Ok(Some((current, false)));
};
let next = StoredTask {
version: current.version + 1,
updated_at: SystemTime::now(),
..next
};
if self.tasks.put(next.clone(), Some(current.version)).await? {
return Ok(Some((next, true)));
}
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Reply {
pub value: Value,
pub content: Vec<Content>,
}
pub(crate) enum Halt {
Failed(String),
Input(Vec<Ask>),
}
impl From<Halt> for CallError {
fn from(halt: Halt) -> Self {
match halt {
Halt::Failed(message) => Self::Failed(message),
Halt::Input(asked) => Self::InputRequired(asked),
}
}
}
struct TaskProgress<K> {
kit: Arc<K>,
id: String,
}
impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Progress
for TaskProgress<Kit<M, I, E, T, S>>
{
fn report(
&self,
done: f64,
total: Option<f64>,
message: Option<String>,
) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
let status = match (total, message) {
(Some(total), Some(message)) => format!("{done}/{total}: {message}"),
(Some(total), None) => format!("{done}/{total}"),
(None, Some(message)) => format!("{done}: {message}"),
(None, None) => format!("{done}"),
};
Box::pin(async move {
if let Err(error) = self.kit.status(&self.id, status).await {
tracing::warn!(%error, "a task's progress could not be stored");
}
})
}
}
fn unavailable(error: std::io::Error) -> CallError {
CallError::Unavailable(error.to_string())
}
#[derive(Debug, PartialEq)]
pub enum CallError {
Unidentified,
Unknown(String),
NotEntitled(Feature),
Failed(String),
InputRequired(Vec<Ask>),
NoTask(String),
Unavailable(String),
}
impl std::fmt::Display for CallError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Unidentified => f.write_str("the caller is not identified"),
Self::Unknown(name) => write!(f, "no operation is named {name}"),
Self::NotEntitled(Feature::Core) => f.write_str("this operation is not enabled"),
Self::NotEntitled(Feature::Pro(feature)) => {
write!(
f,
"this operation needs the {feature} feature, which is not enabled"
)
}
Self::Failed(message) => f.write_str(message),
Self::InputRequired(asked) => {
write!(f, "the operation asks the user {} question(s)", asked.len())
}
Self::NoTask(id) => write!(f, "no task {id} is kept for this caller"),
Self::Unavailable(message) => write!(f, "task store unavailable: {message}"),
}
}
}
impl std::error::Error for CallError {}