1use std::collections::BTreeSet;
8use std::fmt;
9use std::sync::Mutex;
10
11use serde::{Deserialize, Serialize};
12use serde_json::{json, Value};
13
14use crate::protocol::{
15 decode_protocol_frame, encode_protocol_frame, ErrorShape, EventFrame, EventFrameType,
16 ProtocolFrame, RequestFrame, ResponseFrame, ResponseFrameType,
17};
18use crate::{
19 InvocationEvent, InvocationEventKind, InvocationHandle, InvocationInput, InvocationOutput,
20 InvocationStatus, RuntimeError, VifuRuntime,
21};
22
23pub const VIFU_RUNTIME_BRIDGE_PROTOCOL_VERSION: &str = "vifu.runtime-bridge/1";
24
25pub const RUNTIME_BRIDGE_HELLO_METHOD: &str = "runtime.hello";
26pub const RUNTIME_BRIDGE_INVOKE_METHOD: &str = "runtime.invoke";
27pub const RUNTIME_BRIDGE_CANCEL_METHOD: &str = "runtime.cancel";
28
29pub const RUNTIME_BRIDGE_STARTED_EVENT: &str = "runtime.invocation.started";
30pub const RUNTIME_BRIDGE_OUTPUT_DELTA_EVENT: &str = "runtime.invocation.outputDelta";
31pub const RUNTIME_BRIDGE_COMPLETED_EVENT: &str = "runtime.invocation.completed";
32pub const RUNTIME_BRIDGE_FAILED_EVENT: &str = "runtime.invocation.failed";
33pub const RUNTIME_BRIDGE_CANCELLED_EVENT: &str = "runtime.invocation.cancelled";
34
35#[derive(Debug)]
36pub enum RuntimeBridgeError {
37 Protocol(String),
38 Runtime(RuntimeError),
39 StateUnavailable,
40}
41
42impl fmt::Display for RuntimeBridgeError {
43 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
44 match self {
45 Self::Protocol(message) => formatter.write_str(message),
46 Self::Runtime(error) => error.fmt(formatter),
47 Self::StateUnavailable => formatter.write_str("runtime bridge state is unavailable"),
48 }
49 }
50}
51
52impl std::error::Error for RuntimeBridgeError {}
53
54impl From<RuntimeError> for RuntimeBridgeError {
55 fn from(error: RuntimeError) -> Self {
56 Self::Runtime(error)
57 }
58}
59
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61#[serde(rename_all = "camelCase", deny_unknown_fields)]
62pub struct RuntimeBridgeHelloParams {
63 pub protocol: String,
64}
65
66#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
67#[serde(rename_all = "camelCase", deny_unknown_fields)]
68pub struct RuntimeBridgeHelloPayload {
69 pub protocol: String,
70 pub project_id: String,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
74#[serde(rename_all = "camelCase", deny_unknown_fields)]
75pub struct RuntimeBridgeInvokePayload {
76 pub handle: String,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
80#[serde(rename_all = "camelCase", deny_unknown_fields)]
81pub struct RuntimeBridgeCancelParams {
82 pub handle: String,
83}
84
85#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
86#[serde(rename_all = "camelCase", deny_unknown_fields)]
87pub struct RuntimeBridgeInvocationEvent {
88 pub handle: String,
89 pub event: InvocationEvent,
90 #[serde(default, skip_serializing_if = "Option::is_none")]
91 pub output: Option<InvocationOutput>,
92}
93
94pub struct RuntimeBridge {
96 runtime: VifuRuntime,
97 active_invocations: Mutex<BTreeSet<String>>,
98}
99
100impl RuntimeBridge {
101 pub fn new(runtime: VifuRuntime) -> Self {
102 Self {
103 runtime,
104 active_invocations: Mutex::new(BTreeSet::new()),
105 }
106 }
107
108 pub fn runtime(&self) -> &VifuRuntime {
109 &self.runtime
110 }
111
112 pub fn handle_encoded(&self, source: &str) -> Result<Vec<String>, RuntimeBridgeError> {
113 let frame = decode_protocol_frame(source).map_err(RuntimeBridgeError::Protocol)?;
114 self.handle_frame(frame)?
115 .iter()
116 .map(|frame| encode_protocol_frame(frame).map_err(RuntimeBridgeError::Protocol))
117 .collect()
118 }
119
120 pub fn handle_frame(
121 &self,
122 frame: ProtocolFrame,
123 ) -> Result<Vec<ProtocolFrame>, RuntimeBridgeError> {
124 let ProtocolFrame::Request(request) = frame else {
125 return Err(RuntimeBridgeError::Protocol(
126 "runtime bridge accepts request frames from the host".to_string(),
127 ));
128 };
129 Ok(vec![self.handle_request(request)])
130 }
131
132 pub fn drain_encoded(&self) -> Result<Vec<String>, RuntimeBridgeError> {
133 self.drain_events()?
134 .iter()
135 .map(|frame| encode_protocol_frame(frame).map_err(RuntimeBridgeError::Protocol))
136 .collect()
137 }
138
139 pub fn drain_events(&self) -> Result<Vec<ProtocolFrame>, RuntimeBridgeError> {
140 let handles = self
141 .active_invocations
142 .lock()
143 .map_err(|_| RuntimeBridgeError::StateUnavailable)?
144 .iter()
145 .cloned()
146 .collect::<Vec<_>>();
147 let mut frames = Vec::new();
148 let mut completed = Vec::new();
149
150 for handle in handles {
151 let invocation_handle = InvocationHandle(handle.clone());
152 let events = match self.runtime.drain_invocation_events(&invocation_handle) {
153 Ok(events) => events,
154 Err(RuntimeError::InvocationNotFound(_)) => {
155 completed.push(handle);
156 continue;
157 }
158 Err(error) => return Err(error.into()),
159 };
160 let poll = self.runtime.poll_invocation(&invocation_handle)?;
161 let output = poll.output.clone();
162 for event in events {
163 let event_name = invocation_event_name(event.kind);
164 let payload = RuntimeBridgeInvocationEvent {
165 handle: handle.clone(),
166 output: if event.kind == InvocationEventKind::Completed {
167 output.clone()
168 } else {
169 None
170 },
171 event,
172 };
173 frames.push(ProtocolFrame::Event(EventFrame {
174 frame_type: EventFrameType::Event,
175 event: event_name.to_string(),
176 payload: Some(
177 serde_json::to_value(payload)
178 .map_err(|error| RuntimeBridgeError::Protocol(error.to_string()))?,
179 ),
180 seq: None,
181 state_version: None,
182 }));
183 }
184 if is_terminal(poll.status) {
185 let _ = self.runtime.take_invocation(&invocation_handle)?;
186 completed.push(handle);
187 }
188 }
189
190 if !completed.is_empty() {
191 let mut active = self
192 .active_invocations
193 .lock()
194 .map_err(|_| RuntimeBridgeError::StateUnavailable)?;
195 for handle in completed {
196 active.remove(&handle);
197 }
198 }
199 Ok(frames)
200 }
201
202 fn handle_request(&self, request: RequestFrame) -> ProtocolFrame {
203 match request.method.as_str() {
204 RUNTIME_BRIDGE_HELLO_METHOD => self.handle_hello(request),
205 RUNTIME_BRIDGE_INVOKE_METHOD => self.handle_invoke(request),
206 RUNTIME_BRIDGE_CANCEL_METHOD => self.handle_cancel(request),
207 _ => error_response(
208 request.id,
209 "method_not_found",
210 "runtime bridge method is not supported",
211 None,
212 ),
213 }
214 }
215
216 fn handle_hello(&self, request: RequestFrame) -> ProtocolFrame {
217 let params = match decode_params::<RuntimeBridgeHelloParams>(&request) {
218 Ok(params) => params,
219 Err(error) => return invalid_params_response(request.id, error),
220 };
221 if params.protocol != VIFU_RUNTIME_BRIDGE_PROTOCOL_VERSION {
222 return error_response(
223 request.id,
224 "protocol_mismatch",
225 "runtime bridge protocol is not supported",
226 Some(json!({
227 "supported": VIFU_RUNTIME_BRIDGE_PROTOCOL_VERSION,
228 })),
229 );
230 }
231 success_response(
232 request.id,
233 RuntimeBridgeHelloPayload {
234 protocol: VIFU_RUNTIME_BRIDGE_PROTOCOL_VERSION.to_string(),
235 project_id: self.runtime.project_id().to_string(),
236 },
237 )
238 }
239
240 fn handle_invoke(&self, request: RequestFrame) -> ProtocolFrame {
241 let input = match decode_params::<InvocationInput>(&request) {
242 Ok(input) => input,
243 Err(error) => return invalid_params_response(request.id, error),
244 };
245 match self.runtime.start_invoke(input) {
246 Ok(handle) => {
247 let mut active = match self.active_invocations.lock() {
248 Ok(active) => active,
249 Err(_) => {
250 let _ = self.runtime.cancel_invocation(&handle);
251 return internal_error_response(request.id);
252 }
253 };
254 active.insert(handle.0.clone());
255 success_response(request.id, RuntimeBridgeInvokePayload { handle: handle.0 })
256 }
257 Err(error) => runtime_error_response(request.id, error),
258 }
259 }
260
261 fn handle_cancel(&self, request: RequestFrame) -> ProtocolFrame {
262 let params = match decode_params::<RuntimeBridgeCancelParams>(&request) {
263 Ok(params) => params,
264 Err(error) => return invalid_params_response(request.id, error),
265 };
266 match self
267 .runtime
268 .cancel_invocation(&InvocationHandle(params.handle.clone()))
269 {
270 Ok(()) => success_response(request.id, json!({"handle": params.handle})),
271 Err(error) => runtime_error_response(request.id, error),
272 }
273 }
274}
275
276fn decode_params<T>(request: &RequestFrame) -> Result<T, String>
277where
278 T: for<'de> Deserialize<'de>,
279{
280 let params = request
281 .params
282 .clone()
283 .ok_or_else(|| "request params are required".to_string())?;
284 serde_json::from_value(params).map_err(|error| error.to_string())
285}
286
287fn success_response(payload_id: String, payload: impl Serialize) -> ProtocolFrame {
288 match serde_json::to_value(payload) {
289 Ok(payload) => ProtocolFrame::Response(ResponseFrame {
290 frame_type: ResponseFrameType::Res,
291 id: payload_id,
292 ok: true,
293 payload: Some(payload),
294 error: None,
295 }),
296 Err(_) => internal_error_response(payload_id),
297 }
298}
299
300fn invalid_params_response(id: String, error: String) -> ProtocolFrame {
301 error_response(
302 id,
303 "invalid_params",
304 "runtime bridge request parameters are invalid",
305 Some(json!({"reason": error})),
306 )
307}
308
309fn runtime_error_response(id: String, error: RuntimeError) -> ProtocolFrame {
310 error_response(id, "runtime_error", &error.public_message(), None)
311}
312
313fn internal_error_response(id: String) -> ProtocolFrame {
314 error_response(
315 id,
316 "internal_error",
317 "runtime bridge could not complete the request",
318 None,
319 )
320}
321
322fn error_response(id: String, code: &str, message: &str, details: Option<Value>) -> ProtocolFrame {
323 ProtocolFrame::Response(ResponseFrame {
324 frame_type: ResponseFrameType::Res,
325 id,
326 ok: false,
327 payload: None,
328 error: Some(ErrorShape {
329 code: code.to_string(),
330 message: message.to_string(),
331 details,
332 retryable: None,
333 retry_after_ms: None,
334 }),
335 })
336}
337
338fn invocation_event_name(kind: InvocationEventKind) -> &'static str {
339 match kind {
340 InvocationEventKind::Started => RUNTIME_BRIDGE_STARTED_EVENT,
341 InvocationEventKind::OutputDelta => RUNTIME_BRIDGE_OUTPUT_DELTA_EVENT,
342 InvocationEventKind::Completed => RUNTIME_BRIDGE_COMPLETED_EVENT,
343 InvocationEventKind::Failed => RUNTIME_BRIDGE_FAILED_EVENT,
344 InvocationEventKind::Cancelled => RUNTIME_BRIDGE_CANCELLED_EVENT,
345 }
346}
347
348fn is_terminal(status: InvocationStatus) -> bool {
349 matches!(
350 status,
351 InvocationStatus::Completed | InvocationStatus::Failed | InvocationStatus::Cancelled
352 )
353}
354
355#[cfg(test)]
356mod tests {
357 use std::sync::Arc;
358 use std::time::{Duration, Instant};
359
360 use crate::{
361 AgentDefinition, AgentProvider, CancellationToken, EndpointDefinition, InvocationData,
362 ProviderFuture, ProviderRequest, ProviderResponse,
363 };
364
365 use super::*;
366
367 struct EchoProvider;
368
369 impl AgentProvider for EchoProvider {
370 fn supports(&self, capability: &str) -> bool {
371 capability == "chat"
372 }
373
374 fn invoke<'a>(
375 &'a self,
376 request: ProviderRequest,
377 _cancellation: CancellationToken,
378 ) -> ProviderFuture<'a> {
379 Box::pin(async move {
380 Ok(ProviderResponse {
381 data: request.data,
382 metadata: json!({"contentType": "application/json"}),
383 state: None,
384 })
385 })
386 }
387 }
388
389 fn configured_bridge() -> RuntimeBridge {
390 let runtime = VifuRuntime::new("bridge-project").unwrap();
391 runtime
392 .register_provider("echo", Arc::new(EchoProvider))
393 .unwrap();
394 runtime
395 .register_agent(AgentDefinition {
396 id: "guide".to_string(),
397 name: "Guide".to_string(),
398 provider: "echo".to_string(),
399 capabilities: vec!["chat".to_string()],
400 metadata: json!({}),
401 })
402 .unwrap();
403 runtime
404 .register_endpoint(EndpointDefinition {
405 name: "guide".to_string(),
406 agent: "guide".to_string(),
407 capability: "chat".to_string(),
408 timeout_ms: 1_000,
409 })
410 .unwrap();
411 RuntimeBridge::new(runtime)
412 }
413
414 #[test]
415 fn bridge_invokes_runtime_and_streams_terminal_event() {
416 let bridge = configured_bridge();
417 let response = bridge
418 .handle_frame(ProtocolFrame::Request(RequestFrame {
419 frame_type: crate::protocol::RequestFrameType::Req,
420 id: "invoke-1".to_string(),
421 method: RUNTIME_BRIDGE_INVOKE_METHOD.to_string(),
422 params: Some(json!({
423 "endpoint": "guide",
424 "sessionId": "player-one",
425 "data": {
426 "format": "json",
427 "value": {"message": "hello"}
428 },
429 "metadata": {}
430 })),
431 }))
432 .unwrap();
433 assert!(matches!(
434 response.as_slice(),
435 [ProtocolFrame::Response(ResponseFrame { ok: true, .. })]
436 ));
437
438 let deadline = Instant::now() + Duration::from_secs(1);
439 loop {
440 let events = bridge.drain_events().unwrap();
441 if events.iter().any(|frame| {
442 matches!(
443 frame,
444 ProtocolFrame::Event(EventFrame { event, .. })
445 if event == RUNTIME_BRIDGE_COMPLETED_EVENT
446 )
447 }) {
448 let completed = events
449 .into_iter()
450 .find(|frame| {
451 matches!(
452 frame,
453 ProtocolFrame::Event(EventFrame { event, .. })
454 if event == RUNTIME_BRIDGE_COMPLETED_EVENT
455 )
456 })
457 .unwrap();
458 let ProtocolFrame::Event(event) = completed else {
459 unreachable!()
460 };
461 let payload: RuntimeBridgeInvocationEvent =
462 serde_json::from_value(event.payload.unwrap()).unwrap();
463 assert_eq!(
464 payload.output.unwrap().data,
465 InvocationData::Json(json!({"message": "hello"}))
466 );
467 break;
468 }
469 assert!(Instant::now() < deadline);
470 std::thread::sleep(Duration::from_millis(5));
471 }
472 }
473
474 #[test]
475 fn bridge_rejects_unknown_methods_with_protocol_error() {
476 let bridge = configured_bridge();
477 let frames = bridge
478 .handle_frame(ProtocolFrame::Request(RequestFrame {
479 frame_type: crate::protocol::RequestFrameType::Req,
480 id: "unknown-1".to_string(),
481 method: "runtime.unknown".to_string(),
482 params: None,
483 }))
484 .unwrap();
485 assert!(matches!(
486 frames.as_slice(),
487 [ProtocolFrame::Response(ResponseFrame {
488 ok: false,
489 error: Some(ErrorShape { code, .. }),
490 ..
491 })] if code == "method_not_found"
492 ));
493 }
494}