Skip to main content

atman_runtime/
rendezvous.rs

1use std::sync::Arc;
2
3use tokio::sync::oneshot;
4use uuid::Uuid;
5
6use crate::error::RuntimeError;
7
8// PromptId lives in atman-proto but runtime can't depend on proto (proto has no runtime deps).
9// Represent as Uuid at this trait boundary; adapter in daemon maps to proto::PromptId.
10#[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
44// In-proc fallback: prompts are auto-answered by a caller-provided default. Used when
45// no daemon is present; the flow author owns the auto-answer contract via tool args.
46pub 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}