1use crate::error::{ErrorData, Result};
13use alien_core::{CommandDeliveryMode, CommandState, CommandTarget, CommandTargetType};
14use alien_error::AlienError;
15use async_trait::async_trait;
16use chrono::{DateTime, Utc};
17use serde::{Deserialize, Serialize};
18use std::collections::HashMap;
19use std::sync::Arc;
20use tokio::sync::RwLock;
21use uuid::Uuid;
22
23#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct ResolvedCommandTarget {
26 pub target: CommandTarget,
28 pub delivery_mode: CommandDeliveryMode,
30}
31
32#[derive(Debug, Clone)]
34pub struct CommandMetadata {
35 pub command_id: String,
37 pub target: CommandTarget,
39 pub delivery_mode: CommandDeliveryMode,
41 pub project_id: String,
43}
44
45#[derive(Debug, Clone)]
47pub struct CommandEnvelopeData {
48 pub command_id: String,
49 pub deployment_id: String,
50 pub command: String, pub attempt: u32,
52 pub deadline: Option<DateTime<Utc>>,
53 pub state: CommandState,
54 pub target: CommandTarget,
55 pub delivery_mode: CommandDeliveryMode,
56}
57
58#[derive(Debug, Clone)]
60pub struct CommandStatus {
61 pub command_id: String,
62 pub workspace_id: String,
63 pub project_id: String,
64 pub deployment_id: String,
65 pub command: String, pub state: CommandState,
67 pub attempt: u32,
68 pub deadline: Option<DateTime<Utc>>,
69 pub created_at: DateTime<Utc>,
70 pub dispatched_at: Option<DateTime<Utc>>,
71 pub completed_at: Option<DateTime<Utc>>,
72 pub error: Option<serde_json::Value>,
73 pub request_size_bytes: Option<u64>,
74 pub response_size_bytes: Option<u64>,
75 pub target: CommandTarget,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
80pub struct CommandAccessContext {
81 pub workspace_id: String,
82 pub project_id: String,
83 pub deployment_id: String,
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
88#[serde(rename_all = "camelCase")]
89struct CommandRecord {
90 id: String,
91 deployment_id: String,
92 command: String,
93 state: CommandState,
94 attempt: u32,
95 deadline: Option<DateTime<Utc>>,
96 created_at: DateTime<Utc>,
97 dispatched_at: Option<DateTime<Utc>>,
98 completed_at: Option<DateTime<Utc>>,
99 request_size_bytes: Option<u64>,
100 response_size_bytes: Option<u64>,
101 error: Option<serde_json::Value>,
102 target: CommandTarget,
103 delivery_mode: CommandDeliveryMode,
104 project_id: String,
105}
106
107pub fn validate_command_target_id(resource_id: &str) -> Result<()> {
115 if resource_id.contains(':') {
116 return Err(AlienError::new(ErrorData::CommandTargetIdInvalid {
117 resource_id: resource_id.to_string(),
118 }));
119 }
120 Ok(())
121}
122
123pub fn validate_command_name(command: &str) -> Result<()> {
132 if command.contains(':') {
133 return Err(AlienError::new(ErrorData::InvalidCommand {
134 message: format!("Command name '{command}' must not contain ':'"),
135 }));
136 }
137 Ok(())
138}
139
140pub fn select_command_target(
153 deployment_id: &str,
154 targets: &[CommandTarget],
155 requested: Option<&str>,
156) -> Result<CommandTarget> {
157 let target = match requested {
158 Some(resource_id) => {
159 validate_command_target_id(resource_id)?;
161 let found = if resource_id.is_empty() {
164 None
165 } else {
166 targets.iter().find(|t| t.resource_id == resource_id)
167 };
168 found
169 .ok_or_else(|| {
170 AlienError::new(ErrorData::CommandTargetNotFound {
171 resource_id: resource_id.to_string(),
172 deployment_id: deployment_id.to_string(),
173 })
174 })?
175 .clone()
176 }
177 None => match targets {
178 [] => {
179 return Err(AlienError::new(ErrorData::NoCommandTargets {
180 deployment_id: deployment_id.to_string(),
181 }))
182 }
183 [single] => single.clone(),
184 _ => {
185 return Err(AlienError::new(ErrorData::CommandTargetAmbiguous {
186 deployment_id: deployment_id.to_string(),
187 }))
188 }
189 },
190 };
191
192 validate_command_target_id(&target.resource_id)?;
194 Ok(target)
195}
196
197pub fn delivery_mode_for(
203 resource_type: CommandTargetType,
204 worker_mode: CommandDeliveryMode,
205) -> CommandDeliveryMode {
206 match resource_type {
207 CommandTargetType::Container | CommandTargetType::Daemon => CommandDeliveryMode::Pull,
208 CommandTargetType::Worker => worker_mode,
209 }
210}
211
212#[async_trait]
217pub trait CommandRegistry: Send + Sync {
218 async fn resolve_target(
232 &self,
233 deployment_id: &str,
234 requested: Option<&str>,
235 ) -> Result<ResolvedCommandTarget>;
236
237 async fn create_command(
243 &self,
244 deployment_id: &str,
245 command_name: &str,
246 target: &ResolvedCommandTarget,
247 initial_state: CommandState,
248 deadline: Option<DateTime<Utc>>,
249 request_size_bytes: Option<u64>,
250 ) -> Result<CommandMetadata>;
251
252 async fn get_command_metadata(&self, command_id: &str) -> Result<Option<CommandEnvelopeData>>;
256
257 async fn get_command_status(&self, command_id: &str) -> Result<Option<CommandStatus>>;
261
262 async fn get_command_access_context(
267 &self,
268 command_id: &str,
269 ) -> Result<Option<CommandAccessContext>> {
270 Ok(self
271 .get_command_status(command_id)
272 .await?
273 .map(|status| CommandAccessContext {
274 workspace_id: status.workspace_id,
275 project_id: status.project_id,
276 deployment_id: status.deployment_id,
277 }))
278 }
279
280 async fn update_command_state(
285 &self,
286 command_id: &str,
287 state: CommandState,
288 dispatched_at: Option<DateTime<Utc>>,
289 completed_at: Option<DateTime<Utc>>,
290 response_size_bytes: Option<u64>,
291 error: Option<serde_json::Value>,
292 ) -> Result<bool>;
293
294 async fn complete_command(
302 &self,
303 command_id: &str,
304 state: CommandState,
305 completed_at: DateTime<Utc>,
306 response_size_bytes: Option<u64>,
307 error: Option<serde_json::Value>,
308 ) -> Result<bool>;
309
310 async fn mark_dispatched_if_not_terminal(
318 &self,
319 command_id: &str,
320 dispatched_at: DateTime<Utc>,
321 ) -> Result<bool>;
322
323 async fn increment_attempt(&self, command_id: &str) -> Result<u32>;
327}
328
329pub struct InMemoryCommandRegistry {
335 commands: Arc<RwLock<HashMap<String, CommandRecord>>>,
336 targets: Arc<RwLock<Vec<CommandTarget>>>,
342 worker_delivery_mode: CommandDeliveryMode,
351}
352
353impl InMemoryCommandRegistry {
354 pub fn new() -> Self {
357 Self::with_worker_delivery_mode(CommandDeliveryMode::Pull)
358 }
359
360 pub fn with_worker_delivery_mode(worker_delivery_mode: CommandDeliveryMode) -> Self {
364 Self {
365 commands: Arc::new(RwLock::new(HashMap::new())),
366 targets: Arc::new(RwLock::new(Vec::new())),
367 worker_delivery_mode,
368 }
369 }
370
371 pub async fn register_target(
379 &self,
380 resource_id: impl Into<String>,
381 resource_type: CommandTargetType,
382 ) -> Result<()> {
383 let resource_id = resource_id.into();
384 validate_command_target_id(&resource_id)?;
385 self.targets
386 .write()
387 .await
388 .push(CommandTarget::new(resource_id, resource_type));
389 Ok(())
390 }
391
392 #[allow(dead_code)]
394 pub async fn list_command_ids(&self) -> Vec<String> {
395 let commands = self.commands.read().await;
396 commands.keys().cloned().collect()
397 }
398}
399
400impl Default for InMemoryCommandRegistry {
401 fn default() -> Self {
402 Self::new()
403 }
404}
405
406#[async_trait]
407impl CommandRegistry for InMemoryCommandRegistry {
408 async fn resolve_target(
409 &self,
410 deployment_id: &str,
411 requested: Option<&str>,
412 ) -> Result<ResolvedCommandTarget> {
413 let targets = self.targets.read().await;
414 let target = select_command_target(deployment_id, &targets, requested)?;
415 let delivery_mode = delivery_mode_for(target.resource_type, self.worker_delivery_mode);
416 Ok(ResolvedCommandTarget {
417 target,
418 delivery_mode,
419 })
420 }
421
422 async fn create_command(
423 &self,
424 deployment_id: &str,
425 command_name: &str,
426 target: &ResolvedCommandTarget,
427 initial_state: CommandState,
428 deadline: Option<DateTime<Utc>>,
429 request_size_bytes: Option<u64>,
430 ) -> Result<CommandMetadata> {
431 let command_id = format!("cmd_{}", Uuid::new_v4());
432
433 let record = CommandRecord {
434 id: command_id.clone(),
435 deployment_id: deployment_id.to_string(),
436 command: command_name.to_string(),
437 state: initial_state,
438 attempt: 1,
439 deadline,
440 created_at: Utc::now(),
441 dispatched_at: None,
442 completed_at: None,
443 request_size_bytes,
444 response_size_bytes: None,
445 error: None,
446 target: target.target.clone(),
447 delivery_mode: target.delivery_mode,
448 project_id: "default".to_string(),
449 };
450
451 self.commands
452 .write()
453 .await
454 .insert(command_id.clone(), record);
455
456 Ok(CommandMetadata {
457 command_id,
458 target: target.target.clone(),
459 delivery_mode: target.delivery_mode,
460 project_id: "default".to_string(),
461 })
462 }
463
464 async fn get_command_metadata(&self, command_id: &str) -> Result<Option<CommandEnvelopeData>> {
465 let commands = self.commands.read().await;
466
467 Ok(commands.get(command_id).map(|r| CommandEnvelopeData {
468 command_id: r.id.clone(),
469 deployment_id: r.deployment_id.clone(),
470 command: r.command.clone(),
471 attempt: r.attempt,
472 deadline: r.deadline,
473 state: r.state,
474 target: r.target.clone(),
475 delivery_mode: r.delivery_mode,
476 }))
477 }
478
479 async fn get_command_status(&self, command_id: &str) -> Result<Option<CommandStatus>> {
480 let commands = self.commands.read().await;
481
482 Ok(commands.get(command_id).map(|r| CommandStatus {
483 command_id: r.id.clone(),
484 workspace_id: "default".to_string(),
485 project_id: r.project_id.clone(),
486 deployment_id: r.deployment_id.clone(),
487 command: r.command.clone(),
488 state: r.state,
489 attempt: r.attempt,
490 deadline: r.deadline,
491 created_at: r.created_at,
492 dispatched_at: r.dispatched_at,
493 completed_at: r.completed_at,
494 error: r.error.clone(),
495 request_size_bytes: r.request_size_bytes,
496 response_size_bytes: r.response_size_bytes,
497 target: r.target.clone(),
498 }))
499 }
500
501 async fn update_command_state(
502 &self,
503 command_id: &str,
504 state: CommandState,
505 dispatched_at: Option<DateTime<Utc>>,
506 completed_at: Option<DateTime<Utc>>,
507 response_size_bytes: Option<u64>,
508 error: Option<serde_json::Value>,
509 ) -> Result<bool> {
510 let mut commands = self.commands.write().await;
511
512 if let Some(record) = commands.get_mut(command_id) {
513 if record.state.is_terminal() {
514 return Ok(false);
515 }
516 record.state = state;
517
518 if let Some(ts) = dispatched_at {
519 record.dispatched_at = Some(ts);
520 }
521
522 if let Some(ts) = completed_at {
523 record.completed_at = Some(ts);
524 }
525
526 if let Some(size) = response_size_bytes {
527 record.response_size_bytes = Some(size);
528 }
529
530 if let Some(err) = error {
531 record.error = Some(err);
532 }
533 return Ok(true);
534 }
535
536 Ok(false)
537 }
538
539 async fn complete_command(
540 &self,
541 command_id: &str,
542 state: CommandState,
543 completed_at: DateTime<Utc>,
544 response_size_bytes: Option<u64>,
545 error: Option<serde_json::Value>,
546 ) -> Result<bool> {
547 let mut commands = self.commands.write().await;
548 let Some(record) = commands.get_mut(command_id) else {
549 return Ok(false);
550 };
551 if record.state.is_terminal() {
552 return Ok(false);
553 }
554 record.state = state;
555 record.completed_at = Some(completed_at);
556 if let Some(size) = response_size_bytes {
557 record.response_size_bytes = Some(size);
558 }
559 if let Some(err) = error {
560 record.error = Some(err);
561 }
562 Ok(true)
563 }
564
565 async fn mark_dispatched_if_not_terminal(
566 &self,
567 command_id: &str,
568 dispatched_at: DateTime<Utc>,
569 ) -> Result<bool> {
570 let mut commands = self.commands.write().await;
571 let Some(record) = commands.get_mut(command_id) else {
572 return Ok(false);
573 };
574 if record.state.is_terminal() {
575 return Ok(false);
576 }
577 record.state = CommandState::Dispatched;
578 record.dispatched_at = Some(dispatched_at);
579 Ok(true)
580 }
581
582 async fn increment_attempt(&self, command_id: &str) -> Result<u32> {
583 let mut commands = self.commands.write().await;
584
585 if let Some(record) = commands.get_mut(command_id) {
586 record.attempt += 1;
587 Ok(record.attempt)
588 } else {
589 Ok(1) }
591 }
592}
593
594#[cfg(test)]
595mod tests {
596 use super::*;
597 use alien_core::{CommandDeliveryMode, CommandTarget, CommandTargetType};
598
599 async fn resolved(
600 registry: &InMemoryCommandRegistry,
601 requested: Option<&str>,
602 ) -> Result<ResolvedCommandTarget> {
603 registry.resolve_target("dep-1", requested).await
604 }
605
606 #[tokio::test]
607 async fn test_resolve_explicit_target_found() {
608 let registry = InMemoryCommandRegistry::new();
609 registry
610 .register_target("worker-a", CommandTargetType::Worker)
611 .await
612 .unwrap();
613 registry
614 .register_target("daemon-b", CommandTargetType::Daemon)
615 .await
616 .unwrap();
617
618 let result = resolved(®istry, Some("daemon-b")).await.unwrap();
619 assert_eq!(
620 result.target,
621 CommandTarget::new("daemon-b", CommandTargetType::Daemon)
622 );
623 }
624
625 #[tokio::test]
626 async fn test_resolve_explicit_unknown_target_is_not_found() {
627 let registry = InMemoryCommandRegistry::new();
628 registry
629 .register_target("worker-a", CommandTargetType::Worker)
630 .await
631 .unwrap();
632
633 let err = resolved(®istry, Some("no-such-resource"))
634 .await
635 .unwrap_err();
636 assert_eq!(err.code, "COMMAND_TARGET_NOT_FOUND");
637 assert_eq!(err.http_status_code, Some(404));
638 }
639
640 #[tokio::test]
641 async fn test_resolve_explicit_empty_string_is_not_found() {
642 let registry = InMemoryCommandRegistry::new();
643 registry
644 .register_target("worker-a", CommandTargetType::Worker)
645 .await
646 .unwrap();
647
648 let err = resolved(®istry, Some("")).await.unwrap_err();
651 assert_eq!(err.code, "COMMAND_TARGET_NOT_FOUND");
652 }
653
654 #[tokio::test]
655 async fn test_resolve_shorthand_single_target_resolves() {
656 let registry = InMemoryCommandRegistry::new();
657 registry
658 .register_target("container-1", CommandTargetType::Container)
659 .await
660 .unwrap();
661
662 let result = resolved(®istry, None).await.unwrap();
663 assert_eq!(
664 result.target,
665 CommandTarget::new("container-1", CommandTargetType::Container)
666 );
667 }
668
669 #[tokio::test]
670 async fn test_resolve_shorthand_two_targets_is_ambiguous() {
671 let registry = InMemoryCommandRegistry::new();
672 registry
673 .register_target("worker-a", CommandTargetType::Worker)
674 .await
675 .unwrap();
676 registry
677 .register_target("worker-b", CommandTargetType::Worker)
678 .await
679 .unwrap();
680
681 let err = resolved(®istry, None).await.unwrap_err();
682 assert_eq!(err.code, "COMMAND_TARGET_AMBIGUOUS");
683 assert_eq!(err.http_status_code, Some(409));
684 }
685
686 #[tokio::test]
687 async fn test_resolve_shorthand_zero_targets_is_no_targets() {
688 let registry = InMemoryCommandRegistry::new();
689
690 let err = resolved(®istry, None).await.unwrap_err();
691 assert_eq!(err.code, "NO_COMMAND_TARGETS");
692 assert_eq!(err.http_status_code, Some(422));
693 }
694
695 #[tokio::test]
696 async fn test_delivery_mode_container_and_daemon_always_pull() {
697 let registry =
699 InMemoryCommandRegistry::with_worker_delivery_mode(CommandDeliveryMode::Push);
700 registry
701 .register_target("container-1", CommandTargetType::Container)
702 .await
703 .unwrap();
704 registry
705 .register_target("daemon-1", CommandTargetType::Daemon)
706 .await
707 .unwrap();
708
709 let container = resolved(®istry, Some("container-1")).await.unwrap();
710 assert_eq!(container.delivery_mode, CommandDeliveryMode::Pull);
711
712 let daemon = resolved(®istry, Some("daemon-1")).await.unwrap();
713 assert_eq!(daemon.delivery_mode, CommandDeliveryMode::Pull);
714 }
715
716 #[tokio::test]
717 async fn test_delivery_mode_worker_follows_registered_context() {
718 let push_registry =
719 InMemoryCommandRegistry::with_worker_delivery_mode(CommandDeliveryMode::Push);
720 push_registry
721 .register_target("worker-1", CommandTargetType::Worker)
722 .await
723 .unwrap();
724 let push_worker = push_registry
725 .resolve_target("dep-1", Some("worker-1"))
726 .await
727 .unwrap();
728 assert_eq!(push_worker.delivery_mode, CommandDeliveryMode::Push);
729
730 let pull_registry =
733 InMemoryCommandRegistry::with_worker_delivery_mode(CommandDeliveryMode::Pull);
734 pull_registry
735 .register_target("worker-1", CommandTargetType::Worker)
736 .await
737 .unwrap();
738 let pull_worker = pull_registry
739 .resolve_target("dep-1", Some("worker-1"))
740 .await
741 .unwrap();
742 assert_eq!(pull_worker.delivery_mode, CommandDeliveryMode::Pull);
743 }
744
745 #[tokio::test]
746 async fn test_create_command_stores_target_in_status_and_envelope_data() {
747 let registry = InMemoryCommandRegistry::new();
748 registry
749 .register_target("daemon-1", CommandTargetType::Daemon)
750 .await
751 .unwrap();
752
753 let resolved_target = registry.resolve_target("dep-1", None).await.unwrap();
754 let metadata = registry
755 .create_command(
756 "dep-1",
757 "sync-data",
758 &resolved_target,
759 CommandState::Pending,
760 None,
761 None,
762 )
763 .await
764 .unwrap();
765
766 let expected = CommandTarget::new("daemon-1", CommandTargetType::Daemon);
767 assert_eq!(metadata.target, expected);
768 assert_eq!(metadata.delivery_mode, CommandDeliveryMode::Pull);
769
770 let status = registry
771 .get_command_status(&metadata.command_id)
772 .await
773 .unwrap()
774 .unwrap();
775 assert_eq!(status.target, expected);
776
777 let envelope_data = registry
778 .get_command_metadata(&metadata.command_id)
779 .await
780 .unwrap()
781 .unwrap();
782 assert_eq!(envelope_data.target, expected);
783 assert_eq!(envelope_data.delivery_mode, CommandDeliveryMode::Pull);
784 }
785
786 #[tokio::test]
787 async fn non_terminal_update_cannot_resurrect_completed_command() {
788 let registry = InMemoryCommandRegistry::new();
789 registry
790 .register_target("daemon-1", CommandTargetType::Daemon)
791 .await
792 .unwrap();
793 let target = registry.resolve_target("dep-1", None).await.unwrap();
794 let command = registry
795 .create_command(
796 "dep-1",
797 "run",
798 &target,
799 CommandState::Dispatched,
800 None,
801 None,
802 )
803 .await
804 .unwrap();
805
806 assert!(registry
807 .complete_command(
808 &command.command_id,
809 CommandState::Succeeded,
810 Utc::now(),
811 None,
812 None,
813 )
814 .await
815 .unwrap());
816 assert!(!registry
817 .update_command_state(
818 &command.command_id,
819 CommandState::Pending,
820 None,
821 None,
822 None,
823 None,
824 )
825 .await
826 .unwrap());
827 assert_eq!(
828 registry
829 .get_command_status(&command.command_id)
830 .await
831 .unwrap()
832 .unwrap()
833 .state,
834 CommandState::Succeeded
835 );
836 }
837
838 #[tokio::test]
839 async fn in_memory_command_access_uses_oss_tenant_context() {
840 let registry = InMemoryCommandRegistry::new();
841 registry
842 .register_target("worker-1", CommandTargetType::Worker)
843 .await
844 .unwrap();
845 let target = registry.resolve_target("dep-1", None).await.unwrap();
846 let command = registry
847 .create_command("dep-1", "run", &target, CommandState::Pending, None, None)
848 .await
849 .unwrap();
850
851 assert_eq!(command.project_id, "default");
852 assert_eq!(
853 registry
854 .get_command_access_context(&command.command_id)
855 .await
856 .unwrap()
857 .unwrap(),
858 CommandAccessContext {
859 workspace_id: "default".to_string(),
860 project_id: "default".to_string(),
861 deployment_id: "dep-1".to_string(),
862 }
863 );
864 }
865
866 #[tokio::test]
867 async fn test_register_target_rejects_colon_in_id() {
868 let registry = InMemoryCommandRegistry::new();
869 let err = registry
870 .register_target("evil:pending:x", CommandTargetType::Worker)
871 .await
872 .unwrap_err();
873 assert_eq!(err.code, "COMMAND_TARGET_ID_INVALID");
874 assert_eq!(err.http_status_code, Some(400));
875 }
876
877 #[tokio::test]
878 async fn test_resolve_explicit_colon_id_is_invalid() {
879 let registry = InMemoryCommandRegistry::new();
880 registry
881 .register_target("worker-a", CommandTargetType::Worker)
882 .await
883 .unwrap();
884
885 let err = resolved(®istry, Some("worker-a:pending:1"))
888 .await
889 .unwrap_err();
890 assert_eq!(err.code, "COMMAND_TARGET_ID_INVALID");
891 }
892
893 #[test]
894 fn test_select_command_target_rejects_registered_colon_id() {
895 let targets = vec![CommandTarget::new("a:pending:x", CommandTargetType::Worker)];
898 let err = select_command_target("dep-1", &targets, None).unwrap_err();
899 assert_eq!(err.code, "COMMAND_TARGET_ID_INVALID");
900 }
901}