use std::marker::PhantomData;
use std::sync::Arc;
use dashmap::DashMap;
use tokio::sync::{mpsc, oneshot, Mutex};
use arcp_core::envelope::Envelope;
use arcp_core::error::{ARCPError, ErrorCode};
use arcp_core::ids::{ArtifactId, JobId, MessageId, SessionId, SubscriptionId};
use arcp_core::messages::{
ArtifactFetchPayload, ArtifactPutPayload, ArtifactRef, ArtifactReleasePayload, CancelPayload,
CancelTargetKind, Capabilities, ClientIdentity, Credentials, JobAcceptedPayload,
JobCompletedPayload, JobFailedPayload, MessageType, NackPayload, SessionAcceptedPayload,
SessionOpenPayload, SubscribePayload, SubscriptionFilter, SubscriptionSince, ToolInvokePayload,
UnsubscribePayload,
};
use arcp_core::transport::Transport;
mod sealed {
pub trait State {}
impl State for super::Unauthenticated {}
impl State for super::Authenticated {}
}
#[derive(Debug)]
pub struct Unauthenticated;
#[derive(Debug)]
pub struct Authenticated;
pub struct Session<S: sealed::State, T: Transport + 'static> {
inner: Arc<SessionInner<T>>,
_state: PhantomData<S>,
}
struct SessionInner<T: Transport + 'static> {
transport: Arc<dyn Transport>,
session_id: Mutex<Option<SessionId>>,
capabilities: Mutex<Capabilities>,
pending_jobs: DashMap<MessageId, JobNotifier>,
pending_accepted: DashMap<MessageId, oneshot::Sender<Result<JobId, ARCPError>>>,
pending_artifact: DashMap<MessageId, oneshot::Sender<ArtifactReply>>,
active_subscriptions: Arc<DashMap<SubscriptionId, mpsc::UnboundedSender<Envelope>>>,
pending_subscribe: DashMap<MessageId, oneshot::Sender<SubscriptionId>>,
reader: Mutex<Option<tokio::task::JoinHandle<()>>>,
_transport_kind: PhantomData<T>,
}
#[derive(Debug)]
enum JobNotifier {
Pending(oneshot::Sender<Result<serde_json::Value, ARCPError>>),
Taken,
}
impl JobNotifier {
fn take(&mut self) -> Option<oneshot::Sender<Result<serde_json::Value, ARCPError>>> {
match std::mem::replace(self, Self::Taken) {
Self::Pending(tx) => Some(tx),
Self::Taken => None,
}
}
}
#[derive(Debug)]
enum ArtifactReply {
Ref(ArtifactRef),
Inline { data: String, media_type: String },
Nack(NackPayload),
}
impl<S: sealed::State, T: Transport + 'static> std::fmt::Debug for Session<S, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Session")
.field("state", &std::any::type_name::<S>())
.finish_non_exhaustive()
}
}
impl<T: Transport + 'static> Session<Unauthenticated, T> {
pub async fn authenticate(
self,
creds: Credentials,
client: ClientIdentity,
caps: Capabilities,
) -> Result<Session<Authenticated, T>, ARCPError> {
let open = Envelope::new(MessageType::SessionOpen(SessionOpenPayload {
auth: creds.clone(),
client,
capabilities: caps,
}));
let open_id = open.id.clone();
self.inner.transport.send(open).await?;
let env = self
.inner
.transport
.recv()
.await?
.ok_or_else(|| ARCPError::Unavailable {
detail: "transport closed during handshake".into(),
})?;
match env.payload {
MessageType::SessionAccepted(SessionAcceptedPayload {
session_id,
capabilities,
..
}) => {
*self.inner.session_id.lock().await = Some(session_id);
*self.inner.capabilities.lock().await = capabilities;
let inner_for_reader = Arc::clone(&self.inner);
let reader = tokio::spawn(async move {
Self::reader_loop(inner_for_reader).await;
});
*self.inner.reader.lock().await = Some(reader);
Ok(Session {
inner: self.inner.clone(),
_state: PhantomData,
})
}
MessageType::SessionRejected(p) => Err(ARCPError::Unauthenticated {
detail: format!("session.rejected ({}): {}", p.code, p.message),
}),
MessageType::SessionUnauthenticated(p) => Err(ARCPError::Unauthenticated {
detail: format!("session.unauthenticated ({}): {}", p.code, p.message),
}),
MessageType::SessionChallenge(p) => Err(ARCPError::Unauthenticated {
detail: format!(
"runtime issued a challenge (\"{}\") but Phase 2 client cannot respond; \
correlation_id={}",
p.challenge, open_id
),
}),
other => Err(ARCPError::Internal {
detail: format!("unexpected handshake response: type={}", other.type_name()),
}),
}
}
fn dispatch_envelope(inner: &SessionInner<T>, env: Envelope) {
if let MessageType::SubscribeEvent(p) = &env.payload {
if let Some(sub_id) = env.subscription_id.as_ref() {
if let Some(forwarder) = inner.active_subscriptions.get(sub_id) {
if let Ok(inner_env) = serde_json::from_value::<Envelope>(p.event.clone()) {
let _ = forwarder.send(inner_env);
}
}
}
return;
}
let Some(corr) = env.correlation_id.clone() else {
return;
};
match env.payload {
MessageType::JobAccepted(JobAcceptedPayload { job_id, .. }) => {
if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
let _ = tx.send(Ok(job_id));
}
}
MessageType::JobCompleted(JobCompletedPayload { value, .. }) => {
if let Some(mut entry) = inner.pending_jobs.get_mut(&corr) {
if let Some(tx) = entry.take() {
let _ = tx.send(Ok(value.unwrap_or(serde_json::Value::Null)));
}
}
}
MessageType::JobFailed(JobFailedPayload { code, message, .. }) => {
let err_for_accepted = ARCPError::Unknown {
detail: format!("job failed before accept ({code}): {message}"),
};
if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
let _ = tx.send(Err(err_for_accepted));
}
if let Some(mut entry) = inner.pending_jobs.get_mut(&corr) {
if let Some(tx) = entry.take() {
let _ = tx.send(Err(ARCPError::Unknown {
detail: format!("job failed ({code}): {message}"),
}));
}
}
}
MessageType::JobCancelled(p) => {
let reason = p.reason.unwrap_or_default();
if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
let _ = tx.send(Err(ARCPError::Cancelled {
reason: reason.clone(),
}));
}
if let Some(mut entry) = inner.pending_jobs.get_mut(&corr) {
if let Some(tx) = entry.take() {
let _ = tx.send(Err(ARCPError::Cancelled { reason }));
}
}
}
MessageType::ArtifactRef(arcp_core::messages::ArtifactRefPayload { artifact }) => {
if let Some((_, tx)) = inner.pending_artifact.remove(&corr) {
let _ = tx.send(ArtifactReply::Ref(artifact));
}
}
MessageType::ArtifactPut(ArtifactPutPayload {
media_type, data, ..
}) => {
if let Some((_, tx)) = inner.pending_artifact.remove(&corr) {
let _ = tx.send(ArtifactReply::Inline { data, media_type });
}
}
MessageType::Nack(payload) => {
if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
let _ = tx.send(Err(ARCPError::Unknown {
detail: format!("nack ({}): {}", payload.code, payload.message),
}));
}
if let Some((_, tx)) = inner.pending_artifact.remove(&corr) {
let _ = tx.send(ArtifactReply::Nack(payload));
}
}
MessageType::SubscribeAccepted(p) => {
if let Some((_, tx)) = inner.pending_subscribe.remove(&corr) {
let _ = tx.send(p.subscription_id);
}
}
_ => { }
}
}
async fn reader_loop(inner: Arc<SessionInner<T>>) {
while let Ok(Some(env)) = inner.transport.recv().await {
Self::dispatch_envelope(&inner, env);
}
let accepted_keys: Vec<MessageId> = inner
.pending_accepted
.iter()
.map(|r| r.key().clone())
.collect();
for k in accepted_keys {
if let Some((_, tx)) = inner.pending_accepted.remove(&k) {
let _ = tx.send(Err(ARCPError::Unavailable {
detail: "transport closed before job.accepted".into(),
}));
}
}
let artifact_keys: Vec<MessageId> = inner
.pending_artifact
.iter()
.map(|r| r.key().clone())
.collect();
for k in artifact_keys {
if let Some((_, tx)) = inner.pending_artifact.remove(&k) {
let _ = tx.send(ArtifactReply::Nack(NackPayload {
code: ErrorCode::Unavailable,
message: "transport closed before artifact response".into(),
details: None,
}));
}
}
let subscribe_keys: Vec<MessageId> = inner
.pending_subscribe
.iter()
.map(|r| r.key().clone())
.collect();
for k in subscribe_keys {
if let Some((_, tx)) = inner.pending_subscribe.remove(&k) {
drop(tx); }
}
let job_keys: Vec<MessageId> = inner.pending_jobs.iter().map(|r| r.key().clone()).collect();
for k in job_keys {
if let Some(mut entry) = inner.pending_jobs.get_mut(&k) {
if let Some(tx) = entry.take() {
let _ = tx.send(Err(ARCPError::Unavailable {
detail: "transport closed".into(),
}));
}
}
}
}
}
impl<T: Transport + 'static> Session<Authenticated, T> {
pub async fn id(&self) -> Result<SessionId, ARCPError> {
self.inner
.session_id
.lock()
.await
.clone()
.ok_or_else(|| ARCPError::Internal {
detail: "authenticated session missing id".into(),
})
}
pub async fn capabilities(&self) -> Capabilities {
self.inner.capabilities.lock().await.clone()
}
pub async fn invoke(
&self,
tool: impl Into<String>,
arguments: serde_json::Value,
) -> Result<JobHandle, ARCPError> {
let session_id = self.id().await?;
let mut env = Envelope::new(MessageType::ToolInvoke(ToolInvokePayload::new(
tool, arguments,
)));
env.session_id = Some(session_id);
let correlation_id = env.id.clone();
let (acc_tx, acc_rx) = oneshot::channel::<Result<JobId, ARCPError>>();
let (term_tx, term_rx) = oneshot::channel::<Result<serde_json::Value, ARCPError>>();
self.inner
.pending_accepted
.insert(correlation_id.clone(), acc_tx);
self.inner
.pending_jobs
.insert(correlation_id.clone(), JobNotifier::Pending(term_tx));
if let Err(e) = self.inner.transport.send(env).await {
self.inner.pending_accepted.remove(&correlation_id);
self.inner.pending_jobs.remove(&correlation_id);
return Err(e);
}
let job_id = if let Ok(result) = acc_rx.await {
result?
} else {
self.inner.pending_jobs.remove(&correlation_id);
return Err(ARCPError::Unavailable {
detail: "runtime closed before job.accepted".into(),
});
};
Ok(JobHandle {
job_id,
correlation_id,
terminal: Mutex::new(Some(term_rx)),
transport: Arc::clone(&self.inner.transport),
session_id: self.id().await?,
})
}
pub async fn put_artifact(
&self,
media_type: impl Into<String>,
data: impl Into<String>,
retain_seconds: Option<u64>,
) -> Result<ArtifactRef, ARCPError> {
let session_id = self.id().await?;
let mut env = Envelope::new(MessageType::ArtifactPut(ArtifactPutPayload {
media_type: media_type.into(),
data: data.into(),
sha256: None,
retain_seconds,
}));
env.session_id = Some(session_id);
let correlation_id = env.id.clone();
let (tx, rx) = oneshot::channel::<ArtifactReply>();
self.inner
.pending_artifact
.insert(correlation_id.clone(), tx);
if let Err(e) = self.inner.transport.send(env).await {
self.inner.pending_artifact.remove(&correlation_id);
return Err(e);
}
match rx.await {
Ok(ArtifactReply::Ref(reference)) => Ok(reference),
Ok(ArtifactReply::Inline { .. }) => Err(ARCPError::Internal {
detail: "expected artifact.ref, got inline body".into(),
}),
Ok(ArtifactReply::Nack(p)) => Err(map_nack(p)),
Err(_) => Err(ARCPError::Unavailable {
detail: "artifact.put response channel dropped".into(),
}),
}
}
pub async fn fetch_artifact(
&self,
artifact_id: ArtifactId,
) -> Result<(String, String), ARCPError> {
let session_id = self.id().await?;
let mut env = Envelope::new(MessageType::ArtifactFetch(ArtifactFetchPayload {
artifact_id,
}));
env.session_id = Some(session_id);
let correlation_id = env.id.clone();
let (tx, rx) = oneshot::channel::<ArtifactReply>();
self.inner
.pending_artifact
.insert(correlation_id.clone(), tx);
if let Err(e) = self.inner.transport.send(env).await {
self.inner.pending_artifact.remove(&correlation_id);
return Err(e);
}
match rx.await {
Ok(ArtifactReply::Inline { data, media_type }) => Ok((data, media_type)),
Ok(ArtifactReply::Ref(_)) => Err(ARCPError::Internal {
detail: "expected inline body, got artifact.ref".into(),
}),
Ok(ArtifactReply::Nack(p)) => Err(map_nack(p)),
Err(_) => Err(ARCPError::Unavailable {
detail: "artifact.fetch response channel dropped".into(),
}),
}
}
pub async fn release_artifact(&self, artifact_id: ArtifactId) -> Result<(), ARCPError> {
let session_id = self.id().await?;
let mut env = Envelope::new(MessageType::ArtifactRelease(ArtifactReleasePayload {
artifact_id,
}));
env.session_id = Some(session_id);
self.inner.transport.send(env).await
}
pub async fn subscribe(
&self,
filter: SubscriptionFilter,
) -> Result<SubscriptionHandle, ARCPError> {
let session_id = self.id().await?;
let mut env = Envelope::new(MessageType::Subscribe(SubscribePayload {
filter,
since: None,
}));
env.session_id = Some(session_id.clone());
let correlation_id = env.id.clone();
let (acc_tx, acc_rx) = oneshot::channel::<SubscriptionId>();
self.inner
.pending_subscribe
.insert(correlation_id.clone(), acc_tx);
if let Err(e) = self.inner.transport.send(env).await {
self.inner.pending_subscribe.remove(&correlation_id);
return Err(e);
}
let subscription_id = acc_rx.await.map_err(|_| ARCPError::Unavailable {
detail: "runtime closed before subscribe.accepted".into(),
})?;
let (fwd_tx, fwd_rx) = mpsc::unbounded_channel::<Envelope>();
self.inner
.active_subscriptions
.insert(subscription_id.clone(), fwd_tx);
Ok(SubscriptionHandle {
subscription_id,
session_id,
transport: Arc::clone(&self.inner.transport),
inbox: Mutex::new(fwd_rx),
forwarders: Arc::clone(&self.inner.active_subscriptions),
})
}
}
pub struct SubscriptionHandle {
pub subscription_id: SubscriptionId,
session_id: SessionId,
transport: Arc<dyn Transport>,
inbox: Mutex<mpsc::UnboundedReceiver<Envelope>>,
forwarders: Arc<DashMap<SubscriptionId, mpsc::UnboundedSender<Envelope>>>,
}
impl std::fmt::Debug for SubscriptionHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SubscriptionHandle")
.field("subscription_id", &self.subscription_id)
.finish_non_exhaustive()
}
}
impl SubscriptionHandle {
pub async fn next(&self) -> Option<Envelope> {
self.inbox.lock().await.recv().await
}
pub async fn unsubscribe(self) -> Result<(), ARCPError> {
let mut env = Envelope::new(MessageType::Unsubscribe(UnsubscribePayload {
subscription_id: self.subscription_id.clone(),
}));
env.session_id = Some(self.session_id.clone());
let result = self.transport.send(env).await;
result
}
}
impl Drop for SubscriptionHandle {
fn drop(&mut self) {
self.forwarders.remove(&self.subscription_id);
}
}
#[allow(dead_code)] fn _since_marker(_x: SubscriptionSince) {}
fn map_nack(p: NackPayload) -> ARCPError {
match p.code {
ErrorCode::NotFound => ARCPError::NotFound {
kind: "artifact",
id: p.message,
},
ErrorCode::InvalidArgument => ARCPError::InvalidArgument { detail: p.message },
other => ARCPError::Unknown {
detail: format!("nack ({other}): {}", p.message),
},
}
}
pub struct JobHandle {
pub job_id: JobId,
pub correlation_id: MessageId,
terminal: Mutex<Option<oneshot::Receiver<Result<serde_json::Value, ARCPError>>>>,
transport: Arc<dyn Transport>,
session_id: SessionId,
}
impl std::fmt::Debug for JobHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("JobHandle")
.field("job_id", &self.job_id)
.field("correlation_id", &self.correlation_id)
.finish_non_exhaustive()
}
}
impl JobHandle {
pub async fn join(&self) -> Result<serde_json::Value, ARCPError> {
let rx =
self.terminal
.lock()
.await
.take()
.ok_or_else(|| ARCPError::FailedPrecondition {
detail: "JobHandle::join called twice".into(),
})?;
rx.await.unwrap_or_else(|_| {
Err(ARCPError::Unavailable {
detail: "runtime channel closed before terminal event".into(),
})
})
}
pub async fn cancel(&self, reason: impl Into<String>) -> Result<(), ARCPError> {
let mut env = Envelope::new(MessageType::Cancel(CancelPayload {
target: CancelTargetKind::Job,
target_id: self.job_id.to_string(),
reason: Some(reason.into()),
deadline_ms: Some(5000),
}));
env.session_id = Some(self.session_id.clone());
self.transport.send(env).await
}
}
pub struct ARCPClient<T: Transport + 'static> {
transport: Option<T>,
}
impl<T: Transport + 'static> std::fmt::Debug for ARCPClient<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ARCPClient")
.field("attached", &self.transport.is_some())
.finish_non_exhaustive()
}
}
impl<T: Transport + 'static> ARCPClient<T> {
#[must_use]
pub const fn new(transport: T) -> Self {
Self {
transport: Some(transport),
}
}
pub fn open(mut self) -> Result<Session<Unauthenticated, T>, ARCPError> {
let transport = self
.transport
.take()
.ok_or_else(|| ARCPError::FailedPrecondition {
detail: "client transport has already been consumed".into(),
})?;
let _ = ErrorCode::FailedPrecondition;
Ok(Session {
inner: Arc::new(SessionInner {
transport: Arc::new(transport),
session_id: Mutex::new(None),
capabilities: Mutex::new(Capabilities::default()),
pending_jobs: DashMap::new(),
pending_accepted: DashMap::new(),
pending_artifact: DashMap::new(),
active_subscriptions: Arc::new(DashMap::new()),
pending_subscribe: DashMap::new(),
reader: Mutex::new(None),
_transport_kind: PhantomData,
}),
_state: PhantomData,
})
}
}