1use std::collections::{BTreeMap, HashMap};
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::{Arc, Mutex};
8use std::time::{Instant, SystemTime};
9
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use tokio_util::sync::CancellationToken;
13
14use crate::module::{Manifest, Module};
15use crate::ops::{
16 APP_MIME, App, Ask, Body, Completing, Content, Ctx, JsonObject, Looks, Message, Ops, Progress,
17 Prompt, Resource, Tool, with_input,
18};
19use crate::seams::{
20 Change, ChangeFeed, Entitlements, Feature, Identity, InProcessChanges, Local, Metering, Open,
21 Space, Unmetered, Usage,
22};
23use crate::tasks::{InProcessTasks, StoredTask, TaskState, TaskStore};
24
25pub struct Kit<M, I = Local, E = Open, T = Unmetered, S = InProcessTasks> {
27 core: Core<M>,
28 identity: I,
29 entitlements: E,
30 metering: T,
31 tasks: S,
32}
33
34struct Core<M> {
37 module: Arc<M>,
38 manifest: Manifest,
39 ops: Ops<M>,
40 routes: matchit::Router<usize>,
42 changes: Arc<dyn ChangeFeed>,
44 running: Mutex<HashMap<String, CancellationToken>>,
46}
47
48#[derive(Default)]
50pub(crate) struct Caller<'a> {
51 pub bearer: Option<&'a str>,
53 pub cancel: Option<CancellationToken>,
56 pub progress: Option<Arc<dyn Progress>>,
58 pub answers: Option<BTreeMap<String, Value>>,
60 pub space: Option<Arc<Space>>,
63}
64
65impl<'a> Caller<'a> {
66 pub(crate) fn holding(bearer: Option<&'a str>) -> Self {
68 Self {
69 bearer,
70 ..Self::default()
71 }
72 }
73}
74
75impl<M: Module, I: Identity> Kit<M, I> {
76 pub fn new(module: M, identity: I) -> Self {
78 let mut ops = Ops::new();
79 module.register(&mut ops);
80 let mut routes = matchit::Router::new();
81 for (index, readable) in ops.resources.iter().enumerate() {
82 let uri = readable.resource.uri;
83 if let Err(error) = routes.insert(uri, index) {
84 panic!("resource {uri} cannot be routed: {error}");
85 }
86 }
87 Self {
88 core: Core {
89 manifest: module.manifest(),
90 module: Arc::new(module),
91 ops,
92 routes,
93 changes: Arc::new(InProcessChanges::default()),
94 running: Mutex::default(),
95 },
96 identity,
97 entitlements: Open,
98 metering: Unmetered,
99 tasks: InProcessTasks::default(),
100 }
101 }
102}
103
104impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Kit<M, I, E, T, S> {
105 pub fn with_entitlements<E2: Entitlements>(self, entitlements: E2) -> Kit<M, I, E2, T, S> {
106 Kit {
107 core: self.core,
108 identity: self.identity,
109 entitlements,
110 metering: self.metering,
111 tasks: self.tasks,
112 }
113 }
114
115 pub fn with_metering<T2: Metering>(self, metering: T2) -> Kit<M, I, E, T2, S> {
116 Kit {
117 core: self.core,
118 identity: self.identity,
119 entitlements: self.entitlements,
120 metering,
121 tasks: self.tasks,
122 }
123 }
124
125 pub fn with_changes(mut self, changes: impl ChangeFeed) -> Self {
128 self.core.changes = Arc::new(changes);
129 self
130 }
131
132 pub fn with_tasks<S2: TaskStore>(self, tasks: S2) -> Kit<M, I, E, T, S2> {
133 Kit {
134 core: self.core,
135 identity: self.identity,
136 entitlements: self.entitlements,
137 metering: self.metering,
138 tasks,
139 }
140 }
141
142 pub fn manifest(&self) -> &Manifest {
143 &self.core.manifest
144 }
145
146 pub fn tools(&self) -> impl Iterator<Item = &Tool> {
148 self.core.ops.list.iter().map(|op| &op.tool)
149 }
150
151 pub fn instructions(&self) -> Option<&'static str> {
153 self.core.ops.instructions
154 }
155
156 pub fn looks(&self) -> &Looks {
158 &self.core.ops.server
159 }
160
161 pub fn apps(&self) -> &[App] {
163 &self.core.ops.apps
164 }
165
166 pub fn resources(&self) -> impl Iterator<Item = &Resource> {
168 self.core.ops.resources.iter().map(|r| &r.resource)
169 }
170
171 pub async fn read(&self, uri: &str, bearer: Option<&str>) -> Result<Body, CallError> {
173 let caller = Caller::holding(bearer);
174 self.read_as(uri, caller).await
175 }
176
177 pub(crate) async fn read_as(&self, uri: &str, caller: Caller<'_>) -> Result<Body, CallError> {
179 if let Some(app) = self.apps().iter().find(|app| app.uri == uri) {
182 return Ok(Body::Text {
183 mime: APP_MIME.into(),
184 text: app.html.to_owned(),
185 });
186 }
187 let space = self.identified(&caller).await?;
188 let found = self
189 .core
190 .routes
191 .at(uri)
192 .map_err(|_| CallError::Unknown(uri.to_owned()))?;
193 let readable = &self.core.ops.resources[*found.value];
194 let Resource {
195 uri: template,
196 feature,
197 ..
198 } = readable.resource;
199 self.entitled(&space, feature)?;
200 let variables = if found.params.is_empty() {
201 Value::Null
202 } else {
203 found
204 .params
205 .iter()
206 .map(|(name, value)| {
207 let value = percent_encoding::percent_decode_str(value).decode_utf8_lossy();
208 (name.to_owned(), Value::String(value.into_owned()))
209 })
210 .collect::<JsonObject>()
211 .into()
212 };
213 let ctx = Ctx::new(space, caller.cancel, caller.progress);
214 self.metered(template, feature, ctx, caller.answers, |ctx| {
215 (readable.read)(self.core.module.clone(), ctx, variables)
216 })
217 .await
218 .map(|(value, _)| value)
219 .map_err(CallError::from)
220 }
221
222 pub(crate) fn has_completions(&self) -> bool {
224 !self.core.ops.completers.is_empty()
225 }
226
227 pub fn prompts(&self) -> impl Iterator<Item = &Prompt> {
229 self.core.ops.prompts.iter().map(|p| &p.prompt)
230 }
231
232 pub async fn prompt(
234 &self,
235 name: &str,
236 bearer: Option<&str>,
237 arguments: JsonObject,
238 ) -> Result<Vec<Message>, CallError> {
239 let caller = Caller::holding(bearer);
240 self.prompt_as(name, caller, arguments).await
241 }
242
243 pub(crate) async fn prompt_as(
245 &self,
246 name: &str,
247 caller: Caller<'_>,
248 arguments: JsonObject,
249 ) -> Result<Vec<Message>, CallError> {
250 let space = self.identified(&caller).await?;
251 let promptable = self
252 .core
253 .ops
254 .prompts
255 .iter()
256 .find(|p| p.prompt.name == name)
257 .ok_or_else(|| CallError::Unknown(name.to_owned()))?;
258 let Prompt { name, feature, .. } = promptable.prompt;
259 self.entitled(&space, feature)?;
260 let ctx = Ctx::new(space, caller.cancel, caller.progress);
261 self.metered(name, feature, ctx, caller.answers, |ctx| {
262 (promptable.get)(self.core.module.clone(), ctx, Value::Object(arguments))
263 })
264 .await
265 .map(|(value, _)| value)
266 .map_err(CallError::from)
267 }
268
269 pub(crate) async fn complete_as(
273 &self,
274 target: &str,
275 argument: &str,
276 completing: Completing,
277 caller: Caller<'_>,
278 ) -> Result<Vec<String>, CallError> {
279 let space = self.identified(&caller).await?;
280 let Some(completer) = self
281 .core
282 .ops
283 .completers
284 .iter()
285 .find(|c| c.target == target && c.argument == argument)
286 else {
287 return Ok(Vec::new());
288 };
289 let feature = self
290 .prompts()
291 .find(|p| p.name == target)
292 .map(|p| p.feature)
293 .or_else(|| {
294 self.resources()
295 .find(|r| r.uri == target)
296 .map(|r| r.feature)
297 })
298 .expect("Ops::complete checked the target");
299 self.entitled(&space, feature)?;
300 let ctx = Ctx::new(space, caller.cancel, caller.progress);
301 self.metered(completer.target, feature, ctx, caller.answers, |ctx| {
302 let completing = serde_json::to_value(completing).unwrap_or_default();
303 (completer.complete)(self.core.module.clone(), ctx, completing)
304 })
305 .await
306 .map(|(value, _)| value)
307 .map_err(CallError::from)
308 }
309
310 pub fn tool(&self, name: &str) -> Option<&Tool> {
312 self.tools().find(|tool| tool.name == name)
313 }
314
315 pub async fn call(
318 &self,
319 name: &str,
320 bearer: Option<&str>,
321 arguments: JsonObject,
322 ) -> Result<Reply, CallError> {
323 let caller = Caller::holding(bearer);
324 self.call_as(name, caller, arguments).await
325 }
326
327 pub(crate) async fn call_as(
329 &self,
330 name: &str,
331 caller: Caller<'_>,
332 arguments: JsonObject,
333 ) -> Result<Reply, CallError> {
334 let (index, space) = self.admit(name, &caller).await?;
335 let ctx = Ctx::new(space, caller.cancel, caller.progress);
336 self.run(index, ctx, Value::Object(arguments), caller.answers)
337 .await
338 .map_err(CallError::from)
339 }
340
341 pub async fn start(
344 self: &Arc<Self>,
345 name: &str,
346 bearer: Option<&str>,
347 arguments: JsonObject,
348 ) -> Result<StoredTask, CallError> {
349 self.start_as(name, &Caller::holding(bearer), arguments)
350 .await
351 }
352
353 pub(crate) async fn start_as(
355 self: &Arc<Self>,
356 name: &str,
357 caller: &Caller<'_>,
358 arguments: JsonObject,
359 ) -> Result<StoredTask, CallError> {
360 let (index, space) = self.admit(name, caller).await?;
361 let now = SystemTime::now();
362 let arguments = Value::Object(arguments);
363 let task = StoredTask {
364 id: uuid::Uuid::new_v4().to_string(),
365 space: Space::clone(&space),
366 operation: self.core.ops.list[index].tool.name.to_owned(),
367 arguments: arguments.clone(),
368 state: TaskState::Working,
369 status: None,
370 created_at: now,
371 updated_at: now,
372 version: 0,
373 };
374 self.tasks
376 .put(task.clone(), None)
377 .await
378 .map_err(unavailable)?;
379 self.launch(index, space, task.id.clone(), arguments, None);
380 Ok(task)
381 }
382
383 pub async fn update(
387 self: &Arc<Self>,
388 bearer: Option<&str>,
389 id: &str,
390 answers: BTreeMap<String, Value>,
391 ) -> Result<StoredTask, CallError> {
392 self.update_as(&Caller::holding(bearer), id, answers).await
393 }
394
395 pub(crate) async fn update_as(
397 self: &Arc<Self>,
398 caller: &Caller<'_>,
399 id: &str,
400 answers: BTreeMap<String, Value>,
401 ) -> Result<StoredTask, CallError> {
402 self.task_as(caller, id).await?;
403 let (task, resumed) = self
404 .transition(id, |task| {
405 matches!(task.state, TaskState::InputRequired(_)).then(|| StoredTask {
406 state: TaskState::Working,
407 status: None,
408 ..task.clone()
409 })
410 })
411 .await
412 .map_err(unavailable)?
413 .ok_or_else(|| CallError::NoTask(id.to_owned()))?;
414 if resumed {
417 let index = self
418 .core
419 .ops
420 .list
421 .iter()
422 .position(|op| op.tool.name == task.operation)
423 .ok_or_else(|| CallError::Unknown(task.operation.clone()))?;
424 let space = Arc::new(task.space.clone());
425 self.launch(
426 index,
427 space,
428 task.id.clone(),
429 task.arguments.clone(),
430 Some(answers),
431 );
432 }
433 Ok(task)
434 }
435
436 fn launch(
439 self: &Arc<Self>,
440 index: usize,
441 space: Arc<Space>,
442 id: String,
443 arguments: Value,
444 answers: Option<BTreeMap<String, Value>>,
445 ) {
446 let cancel = CancellationToken::new();
449 self.running().insert(id.clone(), cancel.clone());
450 let progress = Arc::new(TaskProgress {
451 kit: self.clone(),
452 id: id.clone(),
453 });
454 let ctx = Ctx::new(space, Some(cancel), Some(progress));
455 let kit = self.clone();
456 tokio::spawn(tracing::Instrument::in_current_span(async move {
458 let state = match kit.run(index, ctx, arguments, answers).await {
459 Ok(result) => TaskState::Succeeded(result),
460 Err(Halt::Failed(message)) => TaskState::Errored(message),
461 Err(Halt::Input(asked)) => TaskState::InputRequired(asked),
462 };
463 kit.running().remove(&id);
464 if let Err(error) = kit.finish(&id, state).await {
465 tracing::error!(%error, "a finished task could not be stored");
466 }
467 }));
468 }
469
470 pub async fn task(&self, bearer: Option<&str>, id: &str) -> Result<StoredTask, CallError> {
472 self.task_as(&Caller::holding(bearer), id).await
473 }
474
475 pub(crate) async fn task_as(
477 &self,
478 caller: &Caller<'_>,
479 id: &str,
480 ) -> Result<StoredTask, CallError> {
481 let space = self.identified(caller).await?;
482 match self.tasks.get(id).await.map_err(unavailable)? {
483 Some(task) if task.space == *space => Ok(task),
484 _ => Err(CallError::NoTask(id.to_owned())),
487 }
488 }
489
490 pub async fn cancel(&self, bearer: Option<&str>, id: &str) -> Result<StoredTask, CallError> {
495 self.cancel_as(&Caller::holding(bearer), id).await
496 }
497
498 pub(crate) async fn cancel_as(
500 &self,
501 caller: &Caller<'_>,
502 id: &str,
503 ) -> Result<StoredTask, CallError> {
504 self.task_as(caller, id).await?;
505 if let Some(cancel) = self.running().get(id) {
506 cancel.cancel();
507 }
508 self.finish(id, TaskState::Cancelled)
509 .await
510 .map_err(unavailable)?
511 .ok_or_else(|| CallError::NoTask(id.to_owned()))
512 }
513
514 async fn status(&self, id: &str, status: String) -> std::io::Result<()> {
516 self.transition(id, |task| {
517 (!task.state.is_terminal()).then(|| StoredTask {
518 status: Some(status.clone()),
519 ..task.clone()
520 })
521 })
522 .await
523 .map(drop)
524 }
525
526 fn running(&self) -> std::sync::MutexGuard<'_, HashMap<String, CancellationToken>> {
527 self.core
528 .running
529 .lock()
530 .unwrap_or_else(|poison| poison.into_inner())
531 }
532
533 fn entitled(&self, space: &Space, feature: Feature) -> Result<(), CallError> {
536 if self.entitlements.allows(space, feature) {
537 Ok(())
538 } else {
539 Err(CallError::NotEntitled(feature))
540 }
541 }
542
543 pub(crate) fn changes(&self) -> tokio::sync::broadcast::Receiver<Change> {
546 self.core.changes.subscribe()
547 }
548
549 pub(crate) async fn identified(&self, caller: &Caller<'_>) -> Result<Arc<Space>, CallError> {
552 match &caller.space {
553 Some(space) => Ok(space.clone()),
554 None => self.identify(caller.bearer).await,
555 }
556 }
557
558 pub(crate) async fn identify(&self, bearer: Option<&str>) -> Result<Arc<Space>, CallError> {
559 self.identity
560 .identify(bearer)
561 .await
562 .ok_or(CallError::Unidentified)
563 }
564
565 async fn admit(
568 &self,
569 name: &str,
570 caller: &Caller<'_>,
571 ) -> Result<(usize, Arc<Space>), CallError> {
572 let space = self.identified(caller).await?;
575 let index = self
577 .core
578 .ops
579 .list
580 .iter()
581 .position(|op| op.tool.name == name)
582 .ok_or_else(|| CallError::Unknown(name.to_owned()))?;
583 self.entitled(&space, self.core.ops.list[index].tool.feature)?;
584 Ok((index, space))
585 }
586
587 async fn run(
589 &self,
590 index: usize,
591 ctx: Ctx,
592 arguments: Value,
593 answers: Option<BTreeMap<String, Value>>,
594 ) -> Result<Reply, Halt> {
595 let op = &self.core.ops.list[index];
596 let (value, content) = self
597 .metered(op.tool.name, op.tool.feature, ctx, answers, |ctx| {
598 (op.call)(self.core.module.clone(), ctx, arguments)
599 })
600 .await?;
601 Ok(Reply { value, content })
602 }
603
604 async fn metered<V, Fut>(
607 &self,
608 operation: &'static str,
609 feature: Feature,
610 ctx: Ctx,
611 answers: Option<BTreeMap<String, Value>>,
612 work: impl FnOnce(Ctx) -> Fut,
613 ) -> Result<(V, Vec<Content>), Halt>
614 where
615 Fut: Future<Output = Result<V, String>>,
616 {
617 let space = ctx.shared_space();
618 let started = Instant::now();
619 let (result, effects) = with_input(answers, work(ctx)).await;
620 for uri in effects.updated {
623 self.core.changes.publish(Change {
624 space: space.clone(),
625 uri,
626 });
627 }
628 let result = match result {
629 Ok(value) => Ok((value, effects.attached)),
630 Err(_) if !effects.asked.is_empty() => Err(Halt::Input(effects.asked)),
632 Err(message) => Err(Halt::Failed(message)),
633 };
634 self.metering.record(
635 &space,
636 Usage {
637 operation,
638 feature,
639 elapsed: started.elapsed(),
640 ok: result.is_ok(),
641 },
642 );
643 result
644 }
645
646 async fn finish(&self, id: &str, state: TaskState) -> std::io::Result<Option<StoredTask>> {
650 let changed = self
651 .transition(id, |task| {
652 (!task.state.is_terminal()).then(|| StoredTask {
653 state: state.clone(),
654 ..task.clone()
655 })
656 })
657 .await?;
658 Ok(changed.map(|(task, _)| task))
659 }
660
661 async fn transition(
667 &self,
668 id: &str,
669 change: impl Fn(&StoredTask) -> Option<StoredTask>,
670 ) -> std::io::Result<Option<(StoredTask, bool)>> {
671 loop {
672 let Some(current) = self.tasks.get(id).await? else {
673 return Ok(None);
674 };
675 let Some(next) = change(¤t) else {
676 return Ok(Some((current, false)));
677 };
678 let next = StoredTask {
679 version: current.version + 1,
680 updated_at: SystemTime::now(),
681 ..next
682 };
683 if self.tasks.put(next.clone(), Some(current.version)).await? {
684 return Ok(Some((next, true)));
685 }
686 }
687 }
688}
689
690#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
692pub struct Reply {
693 pub value: Value,
694 pub content: Vec<Content>,
695}
696
697pub(crate) enum Halt {
699 Failed(String),
700 Input(Vec<Ask>),
702}
703
704impl From<Halt> for CallError {
705 fn from(halt: Halt) -> Self {
706 match halt {
707 Halt::Failed(message) => Self::Failed(message),
708 Halt::Input(asked) => Self::InputRequired(asked),
709 }
710 }
711}
712
713struct TaskProgress<K> {
715 kit: Arc<K>,
716 id: String,
717}
718
719impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Progress
720 for TaskProgress<Kit<M, I, E, T, S>>
721{
722 fn report(
723 &self,
724 done: f64,
725 total: Option<f64>,
726 message: Option<String>,
727 ) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
728 let status = match (total, message) {
729 (Some(total), Some(message)) => format!("{done}/{total}: {message}"),
730 (Some(total), None) => format!("{done}/{total}"),
731 (None, Some(message)) => format!("{done}: {message}"),
732 (None, None) => format!("{done}"),
733 };
734 Box::pin(async move {
735 if let Err(error) = self.kit.status(&self.id, status).await {
736 tracing::warn!(%error, "a task's progress could not be stored");
737 }
738 })
739 }
740}
741
742fn unavailable(error: std::io::Error) -> CallError {
743 CallError::Unavailable(error.to_string())
744}
745
746#[derive(Debug, PartialEq)]
748pub enum CallError {
749 Unidentified,
751 Unknown(String),
752 NotEntitled(Feature),
753 Failed(String),
756 InputRequired(Vec<Ask>),
759 NoTask(String),
761 Unavailable(String),
763}
764
765impl std::fmt::Display for CallError {
766 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
767 match self {
768 Self::Unidentified => f.write_str("the caller is not identified"),
769 Self::Unknown(name) => write!(f, "no operation is named {name}"),
770 Self::NotEntitled(Feature::Core) => f.write_str("this operation is not enabled"),
771 Self::NotEntitled(Feature::Pro(feature)) => {
772 write!(
773 f,
774 "this operation needs the {feature} feature, which is not enabled"
775 )
776 }
777 Self::Failed(message) => f.write_str(message),
778 Self::InputRequired(asked) => {
779 write!(f, "the operation asks the user {} question(s)", asked.len())
780 }
781 Self::NoTask(id) => write!(f, "no task {id} is kept for this caller"),
782 Self::Unavailable(message) => write!(f, "task store unavailable: {message}"),
783 }
784 }
785}
786
787impl std::error::Error for CallError {}