Skip to main content

rap_protocol/
lib.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::future::Future;
4use std::panic::AssertUnwindSafe;
5
6// ── Identifier kinds ──
7
8strkind::strkind! {
9    /// Identifies a conversation thread.
10    ///
11    /// This is the identifier RAP calls `group_id` on the wire: every
12    /// invocation and callback carries the thread it belongs to. In the
13    /// embedded daemon runtime thread IDs are UUIDs; in the Lambda runtime a
14    /// root thread ID is the caller-chosen conversation key (which doubles as
15    /// the SQS FIFO `MessageGroupId`). Treat the contents as opaque.
16    pub ThreadId;
17
18    /// Identifies one tool call within a thread.
19    ///
20    /// This is the `id` field on RAP invocations and callbacks. It
21    /// originates from the model provider's tool-call stream (e.g. Bedrock
22    /// `tooluse_...` IDs) or, for runtime-synthesized calls (thread spawns,
23    /// subscription events), from a minted UUID. Treat the contents as
24    /// opaque; formats are provider-controlled.
25    pub ToolCallId;
26
27    /// The provider-scoped companion ID of a tool call (`call_id` on RAP
28    /// invocations and callbacks). Some providers issue a second identifier
29    /// alongside the tool-call ID; it must be echoed back verbatim in tool
30    /// results. Absent for providers that do not use one.
31    pub ProviderCallId;
32}
33
34/// Identifies a pending user choice.
35///
36/// Currently an alias of [`ToolCallId`]: a choice is keyed by the tool call
37/// that requested it (`RapUserChoice.id` is copied from the invocation).
38/// Tracked as a candidate for becoming a distinct kind in
39/// <https://github.com/hydro-project/infinity/issues/109>.
40pub type ChoiceId<T = String> = ToolCallId<T>;
41
42// ── RAP protocol types ──
43
44#[derive(Debug, Serialize, Deserialize)]
45pub struct RapInvocation {
46    pub operation: String,
47    #[serde(default)]
48    pub arguments: serde_json::Value,
49    pub id: ToolCallId,
50    pub call_id: Option<ProviderCallId>,
51    pub callback_url: String,
52    pub group_id: ThreadId,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub user_id: Option<String>,
55    /// Ordered ancestor thread group IDs, from root to the parent of the current thread.
56    #[serde(default)]
57    pub thread_ancestors: Option<Vec<ThreadId>>,
58}
59
60/// A segment of display content for human-facing UIs.
61#[derive(Debug, Clone, Serialize, Deserialize)]
62#[serde(tag = "type", content = "content")]
63pub enum DisplaySegment {
64    /// Plain text display.
65    #[serde(rename = "text")]
66    Text(String),
67    /// A unified diff.
68    #[serde(rename = "diff")]
69    Diff(DiffContent),
70    /// An inline image (e.g. an image file read by a tool). Clients that
71    /// cannot render images should skip this segment and fall back to the
72    /// next supported one.
73    #[serde(rename = "image")]
74    Image(ImageContent),
75}
76
77/// Content for a diff display segment.
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct DiffContent {
80    /// File path the diff applies to.
81    pub path: String,
82    /// Unified diff string.
83    pub patch: String,
84}
85
86/// Content for an image display segment.
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct ImageContent {
89    /// Base64-encoded image bytes.
90    pub data: String,
91    /// MIME type of the image, e.g. `image/png`.
92    #[serde(rename = "mediaType")]
93    pub media_type: String,
94}
95
96/// Build the prioritized display segments list: tool-provided `display_as`
97/// segments first, then the raw `text` as a trailing `Text` fallback.
98pub fn build_display_segments(
99    display_as: Option<&[DisplaySegment]>,
100    text: &str,
101) -> Vec<DisplaySegment> {
102    let mut segments: Vec<DisplaySegment> = display_as.map(|s| s.to_vec()).unwrap_or_default();
103    segments.push(DisplaySegment::Text(text.to_owned()));
104    segments
105}
106
107/// A single item of structured tool-result content.
108///
109/// A [`RapToolResult`] carries either a plain `text` result or a structured
110/// `content` list (which may include images) for runtimes and models that
111/// support multimodal inputs.
112#[derive(Debug, Clone, Serialize, Deserialize)]
113#[serde(tag = "type")]
114pub enum RapToolResultContent {
115    /// Plain text content.
116    #[serde(rename = "text")]
117    Text { text: String },
118    /// An image, base64-encoded.
119    #[serde(rename = "image")]
120    Image {
121        /// Base64-encoded image bytes.
122        data: String,
123        /// MIME type of the image, e.g. `image/png`.
124        #[serde(rename = "mediaType")]
125        media_type: String,
126    },
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct RapToolResult {
131    pub group_id: ThreadId,
132    pub id: ToolCallId,
133    #[serde(default, skip_serializing_if = "Option::is_none")]
134    pub call_id: Option<ProviderCallId>,
135    /// Plain-text result — the shorthand for a text-only result. A tool MUST
136    /// provide either `text` or `content`. When `content` is also present it
137    /// supersedes `text` for the model-facing result.
138    #[serde(default, skip_serializing_if = "Option::is_none")]
139    pub text: Option<String>,
140    /// Structured content (text and images) for multimodal results. A tool
141    /// MUST provide either `text` or `content`. When present, the runtime
142    /// builds the model-facing tool result from these items (replacing each
143    /// `image` item with a text placeholder for models without image support).
144    #[serde(default, skip_serializing_if = "Option::is_none")]
145    pub content: Option<Vec<RapToolResultContent>>,
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub display_as: Option<Vec<DisplaySegment>>,
148    #[serde(default, skip_serializing_if = "Option::is_none")]
149    pub subscription: Option<bool>,
150}
151
152#[derive(Debug, Clone, Serialize, Deserialize)]
153pub struct RapUserChoice {
154    pub group_id: ThreadId,
155    pub id: ChoiceId,
156    #[serde(default, skip_serializing_if = "Option::is_none")]
157    pub call_id: Option<ProviderCallId>,
158    pub prompt: String,
159    pub choices: Vec<String>,
160    #[serde(default)]
161    pub default: usize,
162    pub response_url: String,
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize)]
166pub struct RapSubscriptionEvent {
167    pub group_id: ThreadId,
168    pub tool_call_id: ToolCallId,
169    pub text: String,
170    #[serde(default)]
171    pub associative: bool,
172    /// When `true`, this is the final event for this subscription. The runtime
173    /// SHOULD remove the subscription from its active tracking.
174    #[serde(default, skip_serializing_if = "Option::is_none")]
175    pub r#final: Option<bool>,
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct RapViewUpdate {
180    pub group_id: ThreadId,
181    /// The type of view being updated (e.g. "diff").
182    pub view_type: String,
183    /// View-specific payload. The runtime passes this through to clients without interpretation.
184    pub content: serde_json::Value,
185}
186
187#[derive(Debug, Clone, Serialize, Deserialize)]
188pub struct RapOAuth {
189    pub group_id: ThreadId,
190    pub id: ToolCallId,
191    #[serde(default, skip_serializing_if = "Option::is_none")]
192    pub call_id: Option<ProviderCallId>,
193    pub auth_url: String,
194}
195
196/// Tagged enum for all RAP callback message types.
197#[derive(Debug, Clone, Serialize, Deserialize)]
198#[serde(tag = "type")]
199pub enum RapCallback {
200    #[serde(rename = "tool_result")]
201    ToolResult(RapToolResult),
202    #[serde(rename = "subscription_event")]
203    SubscriptionEvent(RapSubscriptionEvent),
204    #[serde(rename = "oauth")]
205    OAuth(RapOAuth),
206    #[serde(rename = "user_choice")]
207    UserChoice(RapUserChoice),
208    #[serde(rename = "view_update")]
209    ViewUpdate(RapViewUpdate),
210}
211
212#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct ToolsetManifest {
214    pub name: String,
215    #[serde(default, skip_serializing_if = "Option::is_none")]
216    pub description: Option<String>,
217    pub endpoint: String,
218    pub tools: Vec<ToolDef>,
219    #[serde(default, rename = "needsMigration")]
220    pub needs_migration: bool,
221}
222
223#[derive(Debug, Clone, Serialize, Deserialize)]
224pub struct ToolDef {
225    pub name: String,
226    pub description: String,
227    #[serde(rename = "inputSchema")]
228    pub input_schema: serde_json::Value,
229    #[serde(default, skip_serializing_if = "Option::is_none")]
230    pub annotations: Option<serde_json::Value>,
231    #[serde(
232        default,
233        rename = "displayScript",
234        skip_serializing_if = "Option::is_none"
235    )]
236    pub display_script: Option<String>,
237}
238
239// ── Callback client ──
240
241/// Trait for sending tool results back to the RAP callback URL.
242#[async_trait]
243pub trait CallbackClient: Send + Sync + 'static {
244    async fn post_json(
245        &self,
246        url: &str,
247        body: &str,
248    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
249}
250
251/// Plain reqwest client for local mode (no auth needed).
252pub struct PlainCallbackClient {
253    client: reqwest::Client,
254}
255
256impl Default for PlainCallbackClient {
257    fn default() -> Self {
258        Self::new()
259    }
260}
261
262impl PlainCallbackClient {
263    pub fn new() -> Self {
264        Self {
265            client: reqwest::Client::new(),
266        }
267    }
268}
269
270#[async_trait]
271impl CallbackClient for PlainCallbackClient {
272    async fn post_json(
273        &self,
274        url: &str,
275        body: &str,
276    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
277        self.client
278            .post(url)
279            .header("content-type", "application/json")
280            .body(body.to_owned())
281            .send()
282            .await?;
283        Ok(())
284    }
285}
286
287// ── Helper functions for sending RAP callbacks ──
288
289/// Send a `tool_result` callback.
290pub async fn send_tool_result<C: CallbackClient>(
291    client: &C,
292    invocation: &RapInvocation,
293    text: &str,
294    display_as: Option<Vec<DisplaySegment>>,
295    subscription: bool,
296) {
297    let result = RapCallback::ToolResult(RapToolResult {
298        group_id: invocation.group_id.clone(),
299        id: invocation.id.clone(),
300        call_id: invocation.call_id.clone(),
301        text: Some(text.to_owned()),
302        content: None,
303        display_as,
304        subscription: if subscription { Some(true) } else { None },
305    });
306    let body = serde_json::to_string(&result).expect("bug: failed to serialize tool result");
307    if let Err(e) = client.post_json(&invocation.callback_url, &body).await {
308        tracing::error!("failed to send tool result: {e}");
309    }
310}
311
312/// Send a `user_choice` callback requesting user confirmation.
313pub async fn send_user_choice<C: CallbackClient>(
314    client: &C,
315    invocation: &RapInvocation,
316    prompt: &str,
317    choices: Vec<String>,
318    default: usize,
319    response_url: &str,
320) {
321    let msg = RapCallback::UserChoice(RapUserChoice {
322        group_id: invocation.group_id.clone(),
323        id: invocation.id.clone(),
324        call_id: invocation.call_id.clone(),
325        prompt: prompt.to_owned(),
326        choices,
327        default,
328        response_url: response_url.to_owned(),
329    });
330    let body = serde_json::to_string(&msg).expect("bug: failed to serialize user_choice");
331    if let Err(e) = client.post_json(&invocation.callback_url, &body).await {
332        tracing::error!("failed to send user_choice: {e}");
333    }
334}
335
336/// Send a `subscription_event` callback.
337pub async fn send_subscription_event<C: CallbackClient>(
338    client: &C,
339    callback_url: &str,
340    group_id: ThreadId,
341    tool_call_id: ToolCallId,
342    text: &str,
343    associative: bool,
344    r#final: bool,
345) {
346    let event = RapCallback::SubscriptionEvent(RapSubscriptionEvent {
347        group_id,
348        tool_call_id,
349        text: text.to_owned(),
350        associative,
351        r#final: if r#final { Some(true) } else { None },
352    });
353    let body = serde_json::to_string(&event).expect("bug: failed to serialize subscription event");
354    if let Err(e) = client.post_json(callback_url, &body).await {
355        tracing::error!("failed to send subscription event: {e}");
356    }
357}
358
359/// Send a `view_update` callback.
360pub async fn send_view_update<C: CallbackClient>(
361    client: &C,
362    callback_url: &str,
363    group_id: ThreadId,
364    view_type: &str,
365    content: serde_json::Value,
366) {
367    let msg = RapCallback::ViewUpdate(RapViewUpdate {
368        group_id,
369        view_type: view_type.to_owned(),
370        content,
371    });
372    let body = serde_json::to_string(&msg).expect("bug: failed to serialize view_update");
373    if let Err(e) = client.post_json(callback_url, &body).await {
374        tracing::error!("failed to send view_update: {e}");
375    }
376}
377
378/// Run a future and log if it panics, instead of silently swallowing the panic.
379/// Use this to wrap fire-and-forget spawned tasks.
380pub async fn log_panic<F: Future<Output = ()>>(name: &'static str, f: F) {
381    use futures_util::FutureExt;
382    if let Err(e) = AssertUnwindSafe(f).catch_unwind().await {
383        let msg = if let Some(s) = e.downcast_ref::<&str>() {
384            *s
385        } else if let Some(s) = e.downcast_ref::<String>() {
386            &**s
387        } else {
388            "unknown panic payload"
389        };
390        tracing::error!("task '{name}' panicked: {msg}");
391    }
392}