atman_runtime/
rendezvous.rs1use std::sync::Arc;
2
3use tokio::sync::oneshot;
4use uuid::Uuid;
5
6use crate::error::RuntimeError;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub struct PromptId(pub Uuid);
12
13impl PromptId {
14 pub fn now() -> Self {
15 Self(Uuid::now_v7())
16 }
17}
18
19impl std::fmt::Display for PromptId {
20 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21 self.0.fmt(f)
22 }
23}
24
25pub trait PromptResolver: Send + Sync {
26 fn register(&self, id: PromptId) -> oneshot::Receiver<serde_json::Value>;
27 fn drop_pending(&self, id: &PromptId);
28
29 fn expire_pending(&self, id: &PromptId) -> bool {
30 self.drop_pending(id);
31 true
32 }
33
34 fn register_with_payload(
35 &self,
36 id: PromptId,
37 _kind: &str,
38 _payload: serde_json::Value,
39 ) -> oneshot::Receiver<serde_json::Value> {
40 self.register(id)
41 }
42}
43
44pub struct AutoResolveResolver {
47 pub default: serde_json::Value,
48}
49
50impl PromptResolver for AutoResolveResolver {
51 fn register(&self, _id: PromptId) -> oneshot::Receiver<serde_json::Value> {
52 let (tx, rx) = oneshot::channel();
53 let _ = tx.send(self.default.clone());
54 rx
55 }
56 fn drop_pending(&self, _id: &PromptId) {}
57}
58
59pub async fn await_prompt(
60 resolver: &Arc<dyn PromptResolver>,
61 id: PromptId,
62 timeout: std::time::Duration,
63) -> Result<serde_json::Value, RuntimeError> {
64 await_prompt_inner(resolver.register(id), resolver, id, timeout).await
65}
66
67pub async fn await_prompt_with_payload(
68 resolver: &Arc<dyn PromptResolver>,
69 id: PromptId,
70 kind: &str,
71 payload: serde_json::Value,
72 timeout: std::time::Duration,
73) -> Result<serde_json::Value, RuntimeError> {
74 await_prompt_inner(
75 resolver.register_with_payload(id, kind, payload),
76 resolver,
77 id,
78 timeout,
79 )
80 .await
81}
82
83pub async fn await_expirable_prompt_with_payload(
84 resolver: &Arc<dyn PromptResolver>,
85 id: PromptId,
86 kind: &str,
87 payload: serde_json::Value,
88 timeout: std::time::Duration,
89) -> Result<serde_json::Value, RuntimeError> {
90 let mut rx = resolver.register_with_payload(id, kind, payload);
91 match tokio::time::timeout(timeout, &mut rx).await {
92 Ok(Ok(value)) => Ok(value),
93 Ok(Err(_)) => {
94 resolver.drop_pending(&id);
95 Err(RuntimeError::ToolFailed(format!(
96 "prompt {id} channel closed before answer"
97 )))
98 }
99 Err(_) => {
100 if resolver.expire_pending(&id) {
101 Err(RuntimeError::ToolFailed(format!(
102 "prompt {id} timed out after {}s",
103 timeout.as_secs()
104 )))
105 } else {
106 rx.await.map_err(|_| {
107 RuntimeError::ToolFailed(format!("prompt {id} channel closed before answer"))
108 })
109 }
110 }
111 }
112}
113
114async fn await_prompt_inner(
115 rx: oneshot::Receiver<serde_json::Value>,
116 resolver: &Arc<dyn PromptResolver>,
117 id: PromptId,
118 timeout: std::time::Duration,
119) -> Result<serde_json::Value, RuntimeError> {
120 match tokio::time::timeout(timeout, rx).await {
121 Ok(Ok(v)) => Ok(v),
122 Ok(Err(_)) => {
123 resolver.drop_pending(&id);
124 Err(RuntimeError::ToolFailed(format!(
125 "prompt {id} channel closed before answer"
126 )))
127 }
128 Err(_) => {
129 resolver.drop_pending(&id);
130 Err(RuntimeError::ToolFailed(format!(
131 "prompt {id} timed out after {}s",
132 timeout.as_secs()
133 )))
134 }
135 }
136}