MCPcopy Create free account
hub / github.com/NVIDIA/OpenShell / handle_relay_stream

Function handle_relay_stream

crates/openshell-server/src/supervisor_session.rs:467–567  ·  view source on GitHub ↗

Handle a `RelayStream` RPC from a supervisor. The first inbound `RelayFrame` must carry a `RelayInit` identifying the pending relay; subsequent frames carry raw bytes forward to the gateway-side waiter. Bytes flowing the other way are chunked and sent as `RelayFrame::data` messages back over the response stream.

(
    registry: &SupervisorSessionRegistry,
    request: Request<tonic::Streaming<RelayFrame>>,
)

Source from the content-addressed store, hash-verified

465/// gateway-side waiter. Bytes flowing the other way are chunked and sent as
466/// `RelayFrame::data` messages back over the response stream.
467pub async fn handle_relay_stream(
468 registry: &SupervisorSessionRegistry,
469 request: Request<tonic::Streaming<RelayFrame>>,
470) -> Result<
471 Response<
472 Pin<Box<dyn tokio_stream::Stream<Item = Result<RelayFrame, Status>> + Send + 'static>>,
473 >,
474 Status,
475> {
476 let principal = request.extensions().get::<Principal>().cloned();
477 let mut inbound = request.into_inner();
478
479 // First frame must identify the channel.
480 let first = inbound
481 .message()
482 .await?
483 .ok_or_else(|| Status::invalid_argument("empty RelayStream"))?;
484 let channel_id = match first.payload {
485 Some(openshell_core::proto::relay_frame::Payload::Init(RelayInit { channel_id }))
486 if !channel_id.is_empty() =>
487 {
488 channel_id
489 }
490 _ => {
491 return Err(Status::invalid_argument(
492 "first RelayFrame must be init with non-empty channel_id",
493 ));
494 }
495 };
496
497 // Claim the pending relay. Consumes the entry — it cannot be reused.
498 let supervisor_side = registry.claim_relay(&channel_id, principal.as_ref())?;
499 info!(channel_id = %channel_id, "relay stream: claimed pending relay, bridging");
500
501 let (mut read_half, mut write_half) = tokio::io::split(supervisor_side);
502
503 // Supervisor → gateway: drain `inbound` and write to the DuplexStream.
504 let channel_id_in = channel_id.clone();
505 tokio::spawn(async move {
506 loop {
507 match inbound.message().await {
508 Ok(Some(frame)) => {
509 let Some(openshell_core::proto::relay_frame::Payload::Data(data)) =
510 frame.payload
511 else {
512 warn!(channel_id = %channel_id_in, "relay stream: received non-data frame after init");
513 break;
514 };
515 if data.is_empty() {
516 continue;
517 }
518 if let Err(e) =
519 tokio::io::AsyncWriteExt::write_all(&mut write_half, &data).await
520 {
521 warn!(channel_id = %channel_id_in, error = %e, "relay stream: write to duplex failed");
522 break;
523 }
524 }

Callers 2

relay_streamMethod · 0.85
relay_streamMethod · 0.85

Calls 5

write_allFunction · 0.85
messageMethod · 0.80
claim_relayMethod · 0.80
spawnFunction · 0.50
is_emptyMethod · 0.45

Tested by 1

relay_streamMethod · 0.68