1use std::marker::PhantomData;
4use std::sync::Arc;
5
6use dashmap::DashMap;
7use tokio::sync::{mpsc, oneshot, Mutex};
8
9use arcp_core::envelope::Envelope;
10use arcp_core::error::{ARCPError, ErrorCode};
11use arcp_core::ids::{ArtifactId, JobId, MessageId, SessionId, SubscriptionId};
12use arcp_core::messages::{
13 ArtifactFetchPayload, ArtifactPutPayload, ArtifactRef, ArtifactReleasePayload, CancelPayload,
14 CancelTargetKind, Capabilities, ClientIdentity, Credentials, JobAcceptedPayload,
15 JobCompletedPayload, JobFailedPayload, MessageType, NackPayload, SessionAcceptedPayload,
16 SessionOpenPayload, SubscribePayload, SubscriptionFilter, SubscriptionSince, ToolInvokePayload,
17 UnsubscribePayload,
18};
19use arcp_core::transport::Transport;
20
21mod sealed {
24 pub trait State {}
25 impl State for super::Unauthenticated {}
26 impl State for super::Authenticated {}
27}
28
29#[derive(Debug)]
31pub struct Unauthenticated;
32
33#[derive(Debug)]
35pub struct Authenticated;
36
37pub struct Session<S: sealed::State, T: Transport + 'static> {
43 inner: Arc<SessionInner<T>>,
44 _state: PhantomData<S>,
45}
46
47struct SessionInner<T: Transport + 'static> {
48 transport: Arc<dyn Transport>,
49 session_id: Mutex<Option<SessionId>>,
50 capabilities: Mutex<Capabilities>,
51 pending_jobs: DashMap<MessageId, JobNotifier>,
53 pending_accepted: DashMap<MessageId, oneshot::Sender<Result<JobId, ARCPError>>>,
57 pending_artifact: DashMap<MessageId, oneshot::Sender<ArtifactReply>>,
59 active_subscriptions: Arc<DashMap<SubscriptionId, mpsc::UnboundedSender<Envelope>>>,
63 pending_subscribe: DashMap<MessageId, oneshot::Sender<SubscriptionId>>,
65 reader: Mutex<Option<tokio::task::JoinHandle<()>>>,
66 _transport_kind: PhantomData<T>,
67}
68
69#[derive(Debug)]
70enum JobNotifier {
71 Pending(oneshot::Sender<Result<serde_json::Value, ARCPError>>),
72 Taken,
74}
75
76impl JobNotifier {
77 fn take(&mut self) -> Option<oneshot::Sender<Result<serde_json::Value, ARCPError>>> {
78 match std::mem::replace(self, Self::Taken) {
79 Self::Pending(tx) => Some(tx),
80 Self::Taken => None,
81 }
82 }
83}
84
85#[derive(Debug)]
87enum ArtifactReply {
88 Ref(ArtifactRef),
90 Inline { data: String, media_type: String },
93 Nack(NackPayload),
95}
96
97impl<S: sealed::State, T: Transport + 'static> std::fmt::Debug for Session<S, T> {
98 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 f.debug_struct("Session")
100 .field("state", &std::any::type_name::<S>())
101 .finish_non_exhaustive()
102 }
103}
104
105impl<T: Transport + 'static> Session<Unauthenticated, T> {
106 pub async fn authenticate(
120 self,
121 creds: Credentials,
122 client: ClientIdentity,
123 caps: Capabilities,
124 ) -> Result<Session<Authenticated, T>, ARCPError> {
125 let open = Envelope::new(MessageType::SessionOpen(SessionOpenPayload {
126 auth: creds.clone(),
127 client,
128 capabilities: caps,
129 }));
130 let open_id = open.id.clone();
131 self.inner.transport.send(open).await?;
132
133 let env = self
134 .inner
135 .transport
136 .recv()
137 .await?
138 .ok_or_else(|| ARCPError::Unavailable {
139 detail: "transport closed during handshake".into(),
140 })?;
141 match env.payload {
142 MessageType::SessionAccepted(SessionAcceptedPayload {
143 session_id,
144 capabilities,
145 ..
146 }) => {
147 *self.inner.session_id.lock().await = Some(session_id);
148 *self.inner.capabilities.lock().await = capabilities;
149
150 let inner_for_reader = Arc::clone(&self.inner);
152 let reader = tokio::spawn(async move {
153 Self::reader_loop(inner_for_reader).await;
154 });
155 *self.inner.reader.lock().await = Some(reader);
156
157 Ok(Session {
158 inner: self.inner.clone(),
159 _state: PhantomData,
160 })
161 }
162 MessageType::SessionRejected(p) => Err(ARCPError::Unauthenticated {
163 detail: format!("session.rejected ({}): {}", p.code, p.message),
164 }),
165 MessageType::SessionUnauthenticated(p) => Err(ARCPError::Unauthenticated {
166 detail: format!("session.unauthenticated ({}): {}", p.code, p.message),
167 }),
168 MessageType::SessionChallenge(p) => Err(ARCPError::Unauthenticated {
169 detail: format!(
170 "runtime issued a challenge (\"{}\") but Phase 2 client cannot respond; \
171 correlation_id={}",
172 p.challenge, open_id
173 ),
174 }),
175 other => Err(ARCPError::Internal {
176 detail: format!("unexpected handshake response: type={}", other.type_name()),
177 }),
178 }
179 }
180
181 fn dispatch_envelope(inner: &SessionInner<T>, env: Envelope) {
182 if let MessageType::SubscribeEvent(p) = &env.payload {
185 if let Some(sub_id) = env.subscription_id.as_ref() {
186 if let Some(forwarder) = inner.active_subscriptions.get(sub_id) {
187 if let Ok(inner_env) = serde_json::from_value::<Envelope>(p.event.clone()) {
190 let _ = forwarder.send(inner_env);
191 }
192 }
193 }
194 return;
195 }
196
197 let Some(corr) = env.correlation_id.clone() else {
198 return;
199 };
200 match env.payload {
201 MessageType::JobAccepted(JobAcceptedPayload { job_id, .. }) => {
202 if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
203 let _ = tx.send(Ok(job_id));
204 }
205 }
206 MessageType::JobCompleted(JobCompletedPayload { value, .. }) => {
207 if let Some(mut entry) = inner.pending_jobs.get_mut(&corr) {
208 if let Some(tx) = entry.take() {
209 let _ = tx.send(Ok(value.unwrap_or(serde_json::Value::Null)));
210 }
211 }
212 }
213 MessageType::JobFailed(JobFailedPayload { code, message, .. }) => {
214 let err_for_accepted = ARCPError::Unknown {
220 detail: format!("job failed before accept ({code}): {message}"),
221 };
222 if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
223 let _ = tx.send(Err(err_for_accepted));
224 }
225 if let Some(mut entry) = inner.pending_jobs.get_mut(&corr) {
226 if let Some(tx) = entry.take() {
227 let _ = tx.send(Err(ARCPError::Unknown {
228 detail: format!("job failed ({code}): {message}"),
229 }));
230 }
231 }
232 }
233 MessageType::JobCancelled(p) => {
234 let reason = p.reason.unwrap_or_default();
235 if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
238 let _ = tx.send(Err(ARCPError::Cancelled {
239 reason: reason.clone(),
240 }));
241 }
242 if let Some(mut entry) = inner.pending_jobs.get_mut(&corr) {
243 if let Some(tx) = entry.take() {
244 let _ = tx.send(Err(ARCPError::Cancelled { reason }));
245 }
246 }
247 }
248 MessageType::ArtifactRef(arcp_core::messages::ArtifactRefPayload { artifact }) => {
249 if let Some((_, tx)) = inner.pending_artifact.remove(&corr) {
250 let _ = tx.send(ArtifactReply::Ref(artifact));
251 }
252 }
253 MessageType::ArtifactPut(ArtifactPutPayload {
254 media_type, data, ..
255 }) => {
256 if let Some((_, tx)) = inner.pending_artifact.remove(&corr) {
257 let _ = tx.send(ArtifactReply::Inline { data, media_type });
258 }
259 }
260 MessageType::Nack(payload) => {
261 if let Some((_, tx)) = inner.pending_accepted.remove(&corr) {
265 let _ = tx.send(Err(ARCPError::Unknown {
266 detail: format!("nack ({}): {}", payload.code, payload.message),
267 }));
268 }
269 if let Some((_, tx)) = inner.pending_artifact.remove(&corr) {
270 let _ = tx.send(ArtifactReply::Nack(payload));
271 }
272 }
273 MessageType::SubscribeAccepted(p) => {
274 if let Some((_, tx)) = inner.pending_subscribe.remove(&corr) {
275 let _ = tx.send(p.subscription_id);
276 }
277 }
278 _ => { }
279 }
280 }
281
282 async fn reader_loop(inner: Arc<SessionInner<T>>) {
283 while let Ok(Some(env)) = inner.transport.recv().await {
284 Self::dispatch_envelope(&inner, env);
285 }
286 let accepted_keys: Vec<MessageId> = inner
291 .pending_accepted
292 .iter()
293 .map(|r| r.key().clone())
294 .collect();
295 for k in accepted_keys {
296 if let Some((_, tx)) = inner.pending_accepted.remove(&k) {
297 let _ = tx.send(Err(ARCPError::Unavailable {
298 detail: "transport closed before job.accepted".into(),
299 }));
300 }
301 }
302 let artifact_keys: Vec<MessageId> = inner
303 .pending_artifact
304 .iter()
305 .map(|r| r.key().clone())
306 .collect();
307 for k in artifact_keys {
308 if let Some((_, tx)) = inner.pending_artifact.remove(&k) {
309 let _ = tx.send(ArtifactReply::Nack(NackPayload {
310 code: ErrorCode::Unavailable,
311 message: "transport closed before artifact response".into(),
312 details: None,
313 }));
314 }
315 }
316 let subscribe_keys: Vec<MessageId> = inner
317 .pending_subscribe
318 .iter()
319 .map(|r| r.key().clone())
320 .collect();
321 for k in subscribe_keys {
322 if let Some((_, tx)) = inner.pending_subscribe.remove(&k) {
323 drop(tx); }
325 }
326 let job_keys: Vec<MessageId> = inner.pending_jobs.iter().map(|r| r.key().clone()).collect();
327 for k in job_keys {
328 if let Some(mut entry) = inner.pending_jobs.get_mut(&k) {
329 if let Some(tx) = entry.take() {
330 let _ = tx.send(Err(ARCPError::Unavailable {
331 detail: "transport closed".into(),
332 }));
333 }
334 }
335 }
336 }
337}
338
339impl<T: Transport + 'static> Session<Authenticated, T> {
340 pub async fn id(&self) -> Result<SessionId, ARCPError> {
348 self.inner
349 .session_id
350 .lock()
351 .await
352 .clone()
353 .ok_or_else(|| ARCPError::Internal {
354 detail: "authenticated session missing id".into(),
355 })
356 }
357
358 pub async fn capabilities(&self) -> Capabilities {
360 self.inner.capabilities.lock().await.clone()
361 }
362
363 pub async fn invoke(
372 &self,
373 tool: impl Into<String>,
374 arguments: serde_json::Value,
375 ) -> Result<JobHandle, ARCPError> {
376 let session_id = self.id().await?;
377 let mut env = Envelope::new(MessageType::ToolInvoke(ToolInvokePayload::new(
378 tool, arguments,
379 )));
380 env.session_id = Some(session_id);
381 let correlation_id = env.id.clone();
382
383 let (acc_tx, acc_rx) = oneshot::channel::<Result<JobId, ARCPError>>();
384 let (term_tx, term_rx) = oneshot::channel::<Result<serde_json::Value, ARCPError>>();
385 self.inner
386 .pending_accepted
387 .insert(correlation_id.clone(), acc_tx);
388 self.inner
389 .pending_jobs
390 .insert(correlation_id.clone(), JobNotifier::Pending(term_tx));
391
392 if let Err(e) = self.inner.transport.send(env).await {
393 self.inner.pending_accepted.remove(&correlation_id);
396 self.inner.pending_jobs.remove(&correlation_id);
397 return Err(e);
398 }
399
400 let job_id = if let Ok(result) = acc_rx.await {
404 result?
405 } else {
406 self.inner.pending_jobs.remove(&correlation_id);
410 return Err(ARCPError::Unavailable {
411 detail: "runtime closed before job.accepted".into(),
412 });
413 };
414
415 Ok(JobHandle {
416 job_id,
417 correlation_id,
418 terminal: Mutex::new(Some(term_rx)),
419 transport: Arc::clone(&self.inner.transport),
420 session_id: self.id().await?,
421 })
422 }
423
424 pub async fn put_artifact(
436 &self,
437 media_type: impl Into<String>,
438 data: impl Into<String>,
439 retain_seconds: Option<u64>,
440 ) -> Result<ArtifactRef, ARCPError> {
441 let session_id = self.id().await?;
442 let mut env = Envelope::new(MessageType::ArtifactPut(ArtifactPutPayload {
443 media_type: media_type.into(),
444 data: data.into(),
445 sha256: None,
446 retain_seconds,
447 }));
448 env.session_id = Some(session_id);
449 let correlation_id = env.id.clone();
450
451 let (tx, rx) = oneshot::channel::<ArtifactReply>();
452 self.inner
453 .pending_artifact
454 .insert(correlation_id.clone(), tx);
455 if let Err(e) = self.inner.transport.send(env).await {
456 self.inner.pending_artifact.remove(&correlation_id);
457 return Err(e);
458 }
459
460 match rx.await {
461 Ok(ArtifactReply::Ref(reference)) => Ok(reference),
462 Ok(ArtifactReply::Inline { .. }) => Err(ARCPError::Internal {
463 detail: "expected artifact.ref, got inline body".into(),
464 }),
465 Ok(ArtifactReply::Nack(p)) => Err(map_nack(p)),
466 Err(_) => Err(ARCPError::Unavailable {
467 detail: "artifact.put response channel dropped".into(),
468 }),
469 }
470 }
471
472 pub async fn fetch_artifact(
479 &self,
480 artifact_id: ArtifactId,
481 ) -> Result<(String, String), ARCPError> {
482 let session_id = self.id().await?;
483 let mut env = Envelope::new(MessageType::ArtifactFetch(ArtifactFetchPayload {
484 artifact_id,
485 }));
486 env.session_id = Some(session_id);
487 let correlation_id = env.id.clone();
488
489 let (tx, rx) = oneshot::channel::<ArtifactReply>();
490 self.inner
491 .pending_artifact
492 .insert(correlation_id.clone(), tx);
493 if let Err(e) = self.inner.transport.send(env).await {
494 self.inner.pending_artifact.remove(&correlation_id);
495 return Err(e);
496 }
497
498 match rx.await {
499 Ok(ArtifactReply::Inline { data, media_type }) => Ok((data, media_type)),
500 Ok(ArtifactReply::Ref(_)) => Err(ARCPError::Internal {
501 detail: "expected inline body, got artifact.ref".into(),
502 }),
503 Ok(ArtifactReply::Nack(p)) => Err(map_nack(p)),
504 Err(_) => Err(ARCPError::Unavailable {
505 detail: "artifact.fetch response channel dropped".into(),
506 }),
507 }
508 }
509
510 pub async fn release_artifact(&self, artifact_id: ArtifactId) -> Result<(), ARCPError> {
517 let session_id = self.id().await?;
518 let mut env = Envelope::new(MessageType::ArtifactRelease(ArtifactReleasePayload {
519 artifact_id,
520 }));
521 env.session_id = Some(session_id);
522 self.inner.transport.send(env).await
523 }
524
525 pub async fn subscribe(
533 &self,
534 filter: SubscriptionFilter,
535 ) -> Result<SubscriptionHandle, ARCPError> {
536 let session_id = self.id().await?;
537 let mut env = Envelope::new(MessageType::Subscribe(SubscribePayload {
538 filter,
539 since: None,
540 }));
541 env.session_id = Some(session_id.clone());
542 let correlation_id = env.id.clone();
543
544 let (acc_tx, acc_rx) = oneshot::channel::<SubscriptionId>();
545 self.inner
546 .pending_subscribe
547 .insert(correlation_id.clone(), acc_tx);
548 if let Err(e) = self.inner.transport.send(env).await {
549 self.inner.pending_subscribe.remove(&correlation_id);
550 return Err(e);
551 }
552
553 let subscription_id = acc_rx.await.map_err(|_| ARCPError::Unavailable {
554 detail: "runtime closed before subscribe.accepted".into(),
555 })?;
556
557 let (fwd_tx, fwd_rx) = mpsc::unbounded_channel::<Envelope>();
558 self.inner
559 .active_subscriptions
560 .insert(subscription_id.clone(), fwd_tx);
561 Ok(SubscriptionHandle {
562 subscription_id,
563 session_id,
564 transport: Arc::clone(&self.inner.transport),
565 inbox: Mutex::new(fwd_rx),
566 forwarders: Arc::clone(&self.inner.active_subscriptions),
567 })
568 }
569}
570
571pub struct SubscriptionHandle {
578 pub subscription_id: SubscriptionId,
580 session_id: SessionId,
581 transport: Arc<dyn Transport>,
582 inbox: Mutex<mpsc::UnboundedReceiver<Envelope>>,
583 forwarders: Arc<DashMap<SubscriptionId, mpsc::UnboundedSender<Envelope>>>,
584}
585
586impl std::fmt::Debug for SubscriptionHandle {
587 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
588 f.debug_struct("SubscriptionHandle")
589 .field("subscription_id", &self.subscription_id)
590 .finish_non_exhaustive()
591 }
592}
593
594impl SubscriptionHandle {
595 pub async fn next(&self) -> Option<Envelope> {
598 self.inbox.lock().await.recv().await
599 }
600
601 pub async fn unsubscribe(self) -> Result<(), ARCPError> {
610 let mut env = Envelope::new(MessageType::Unsubscribe(UnsubscribePayload {
611 subscription_id: self.subscription_id.clone(),
612 }));
613 env.session_id = Some(self.session_id.clone());
614 let result = self.transport.send(env).await;
615 result
617 }
618}
619
620impl Drop for SubscriptionHandle {
621 fn drop(&mut self) {
622 self.forwarders.remove(&self.subscription_id);
623 }
624}
625
626#[allow(dead_code)] fn _since_marker(_x: SubscriptionSince) {}
628
629fn map_nack(p: NackPayload) -> ARCPError {
630 match p.code {
631 ErrorCode::NotFound => ARCPError::NotFound {
632 kind: "artifact",
633 id: p.message,
634 },
635 ErrorCode::InvalidArgument => ARCPError::InvalidArgument { detail: p.message },
636 other => ARCPError::Unknown {
637 detail: format!("nack ({other}): {}", p.message),
638 },
639 }
640}
641
642pub struct JobHandle {
644 pub job_id: JobId,
646 pub correlation_id: MessageId,
648 terminal: Mutex<Option<oneshot::Receiver<Result<serde_json::Value, ARCPError>>>>,
649 transport: Arc<dyn Transport>,
650 session_id: SessionId,
651}
652
653impl std::fmt::Debug for JobHandle {
654 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
655 f.debug_struct("JobHandle")
656 .field("job_id", &self.job_id)
657 .field("correlation_id", &self.correlation_id)
658 .finish_non_exhaustive()
659 }
660}
661
662impl JobHandle {
663 pub async fn join(&self) -> Result<serde_json::Value, ARCPError> {
672 let rx =
673 self.terminal
674 .lock()
675 .await
676 .take()
677 .ok_or_else(|| ARCPError::FailedPrecondition {
678 detail: "JobHandle::join called twice".into(),
679 })?;
680 rx.await.unwrap_or_else(|_| {
681 Err(ARCPError::Unavailable {
682 detail: "runtime channel closed before terminal event".into(),
683 })
684 })
685 }
686
687 pub async fn cancel(&self, reason: impl Into<String>) -> Result<(), ARCPError> {
696 let mut env = Envelope::new(MessageType::Cancel(CancelPayload {
697 target: CancelTargetKind::Job,
698 target_id: self.job_id.to_string(),
699 reason: Some(reason.into()),
700 deadline_ms: Some(5000),
701 }));
702 env.session_id = Some(self.session_id.clone());
703 self.transport.send(env).await
704 }
705}
706
707pub struct ARCPClient<T: Transport + 'static> {
709 transport: Option<T>,
710}
711
712impl<T: Transport + 'static> std::fmt::Debug for ARCPClient<T> {
713 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
714 f.debug_struct("ARCPClient")
715 .field("attached", &self.transport.is_some())
716 .finish_non_exhaustive()
717 }
718}
719
720impl<T: Transport + 'static> ARCPClient<T> {
721 #[must_use]
723 pub const fn new(transport: T) -> Self {
724 Self {
725 transport: Some(transport),
726 }
727 }
728
729 pub fn open(mut self) -> Result<Session<Unauthenticated, T>, ARCPError> {
737 let transport = self
738 .transport
739 .take()
740 .ok_or_else(|| ARCPError::FailedPrecondition {
741 detail: "client transport has already been consumed".into(),
742 })?;
743 let _ = ErrorCode::FailedPrecondition;
744 Ok(Session {
745 inner: Arc::new(SessionInner {
746 transport: Arc::new(transport),
747 session_id: Mutex::new(None),
748 capabilities: Mutex::new(Capabilities::default()),
749 pending_jobs: DashMap::new(),
750 pending_accepted: DashMap::new(),
751 pending_artifact: DashMap::new(),
752 active_subscriptions: Arc::new(DashMap::new()),
753 pending_subscribe: DashMap::new(),
754 reader: Mutex::new(None),
755 _transport_kind: PhantomData,
756 }),
757 _state: PhantomData,
758 })
759 }
760}