Skip to main content

alien_commands/server/
mod.rs

1//! Command Server Implementation
2//!
3//! The Command server implements the command lifecycle:
4//! - CommandRegistry is the SOURCE OF TRUTH for all metadata (state, timestamps, etc.)
5//! - KV stores ONLY operational data (params/response blobs, indices, leases)
6
7use std::sync::Arc;
8use std::time::Duration;
9
10use alien_bindings::presigned::PresignedRequest;
11use alien_bindings::traits::{Kv, PutCondition, PutOptions, Storage};
12use alien_error::{AlienError, Context, ContextError, IntoAlienError};
13use chrono::{DateTime, Utc};
14use hex;
15use hmac::{Hmac, Mac};
16use serde::{Deserialize, Serialize};
17use sha2::Sha256;
18use tracing::{debug, info, warn};
19use uuid::Uuid;
20
21use base64::{engine::general_purpose, Engine as _};
22use bytes::Bytes;
23use object_store::path::Path as StoragePath;
24
25use crate::error::{ErrorData, Result};
26use crate::types::*;
27use crate::INLINE_MAX_BYTES;
28
29/// Max serialized KV value size. Conservative threshold below the hard 24KB
30/// boundary (Azure Table Storage) to account for JSON wrapping overhead.
31const KV_VALUE_THRESHOLD: usize = 20_000;
32
33fn is_definite_dispatch_rejection(error: &crate::error::Error) -> bool {
34    matches!(
35        error.error.as_ref(),
36        Some(ErrorData::TransportDispatchRejected { .. })
37    )
38}
39
40pub mod axum_handlers;
41pub mod command_registry;
42pub mod storage;
43
44pub use crate::dispatchers::{CommandDispatcher, NullCommandDispatcher};
45pub use axum_handlers::{
46    create_axum_router, CommandPayloadResponse, HasCommandServer, StorePayloadRequest,
47};
48pub use command_registry::{
49    delivery_mode_for, select_command_target, validate_command_name, validate_command_target_id,
50    CommandAccessContext, CommandEnvelopeData, CommandMetadata, CommandRegistry, CommandStatus,
51    InMemoryCommandRegistry, ResolvedCommandTarget,
52};
53
54// =============================================================================
55// KV Data Structures (Operational Data Only)
56// =============================================================================
57
58/// Params stored in KV (just the blob)
59#[derive(Debug, Clone, Serialize, Deserialize)]
60struct CommandParamsData {
61    pub params: BodySpec,
62}
63
64/// Response stored in KV (just the blob)
65#[derive(Debug, Clone, Serialize, Deserialize)]
66struct CommandResponseData {
67    pub response: CommandResponse,
68}
69
70/// Lease record with TTL
71#[derive(Debug, Clone, Serialize, Deserialize)]
72struct LeaseData {
73    pub lease_id: String,
74    pub acquired_at: DateTime<Utc>,
75    pub expires_at: DateTime<Utc>,
76    pub owner: String,
77}
78
79/// Deadline index data
80#[derive(Debug, Clone, Serialize, Deserialize)]
81struct DeadlineIndexData {
82    pub command_id: String,
83    pub deadline: DateTime<Utc>,
84}
85
86// =============================================================================
87// Command Server
88// =============================================================================
89
90/// Core command server implementation.
91///
92/// Uses CommandRegistry as source of truth for metadata.
93/// Uses KV for operational data (params, responses, indices, leases).
94pub struct CommandServer {
95    kv: Arc<dyn Kv>,
96    storage: Arc<dyn Storage>,
97    command_dispatcher: Arc<dyn CommandDispatcher>,
98    command_registry: Arc<dyn CommandRegistry>,
99    inline_max_bytes: usize,
100    base_url: String,
101    response_signing_key: Vec<u8>,
102}
103
104impl CommandServer {
105    /// Create a new command server instance
106    pub fn new(
107        kv: Arc<dyn Kv>,
108        storage: Arc<dyn Storage>,
109        command_dispatcher: Arc<dyn CommandDispatcher>,
110        command_registry: Arc<dyn CommandRegistry>,
111        base_url: String,
112        response_signing_key: Vec<u8>,
113    ) -> Self {
114        Self {
115            kv,
116            storage,
117            command_dispatcher,
118            command_registry,
119            inline_max_bytes: INLINE_MAX_BYTES,
120            base_url,
121            response_signing_key,
122        }
123    }
124
125    /// Create a new command server with custom inline size limit
126    pub fn with_inline_limit(
127        kv: Arc<dyn Kv>,
128        storage: Arc<dyn Storage>,
129        command_dispatcher: Arc<dyn CommandDispatcher>,
130        command_registry: Arc<dyn CommandRegistry>,
131        base_url: String,
132        inline_max_bytes: usize,
133        response_signing_key: Vec<u8>,
134    ) -> Self {
135        Self {
136            kv,
137            storage,
138            command_dispatcher,
139            command_registry,
140            inline_max_bytes,
141            base_url,
142            response_signing_key,
143        }
144    }
145
146    /// Maximum allowed response token lifetime (2 hours).
147    const MAX_RESPONSE_TOKEN_LIFETIME_SECS: i64 = 7200;
148    /// Worker execution can run for 1 hour and the operator lease adds 60
149    /// seconds of response-submission headroom. Response credentials therefore
150    /// need to outlive both; 2 hours stays within the verifier cap.
151    const RESPONSE_CREDENTIAL_LIFETIME_SECS: u64 = 7200;
152
153    /// Sign a response URL for a specific command.
154    ///
155    /// Returns `(hmac_hex, expires_epoch)`. The HMAC is computed over
156    /// `"arc.v1:{command_id}:{expires}"` using the server's signing key.
157    fn sign_response_url(&self, command_id: &str) -> (String, i64) {
158        let expires = Utc::now().timestamp() + Self::RESPONSE_CREDENTIAL_LIFETIME_SECS as i64;
159        let message = format!("commands.v1:{}:{}", command_id, expires);
160
161        type HmacSha256 = Hmac<Sha256>;
162        let mut mac =
163            HmacSha256::new_from_slice(&self.response_signing_key).expect("HMAC accepts any key");
164        mac.update(message.as_bytes());
165        let result = mac.finalize();
166        let token = hex::encode(result.into_bytes());
167
168        (token, expires)
169    }
170
171    /// Verify a response token for a specific command.
172    ///
173    /// Performs HMAC verification first (constant-time), then checks expiration,
174    /// to avoid leaking timing information about token validity windows.
175    pub fn verify_response_token(&self, command_id: &str, token: &str, expires: i64) -> bool {
176        // Validate token format: SHA-256 HMAC = 32 bytes = 64 hex chars.
177        if token.len() != 64 {
178            return false;
179        }
180
181        let message = format!("commands.v1:{}:{}", command_id, expires);
182
183        type HmacSha256 = Hmac<Sha256>;
184        let mut mac =
185            HmacSha256::new_from_slice(&self.response_signing_key).expect("HMAC accepts any key");
186        mac.update(message.as_bytes());
187
188        // Decode hex and verify HMAC (constant-time comparison).
189        let Ok(token_bytes) = hex::decode(token) else {
190            return false;
191        };
192        let hmac_valid = mac.verify_slice(&token_bytes).is_ok();
193
194        // Check expiration and max lifetime AFTER HMAC to avoid timing leaks.
195        let now = Utc::now().timestamp();
196        let not_expired = now <= expires;
197        let within_max_lifetime = expires <= now + Self::MAX_RESPONSE_TOKEN_LIFETIME_SECS;
198
199        hmac_valid && not_expired && within_max_lifetime
200    }
201
202    // =========================================================================
203    // Public API Methods
204    // =========================================================================
205
206    /// Create a new command.
207    ///
208    /// Flow:
209    /// 1. Validate request
210    /// 2. Registry creates command metadata (source of truth)
211    /// 3. KV stores params blob
212    /// 4. KV creates pending index (for Pull) or dispatch (for Push)
213    pub async fn create_command(
214        &self,
215        request: CreateCommandRequest,
216    ) -> Result<CreateCommandResponse> {
217        // Validate the request
218        self.validate_create_command(&request).await?;
219
220        // Resolve which command-capable resource this command targets
221        // (explicit targetResourceId, or single-target shorthand).
222        let resolved_target = self
223            .command_registry
224            .resolve_target(
225                &request.deployment_id,
226                request.target_resource_id.as_deref(),
227            )
228            .await?;
229
230        // Compose the target-scoped idempotency key once, from the resolved
231        // target, and reuse it for both the pre-create check and the
232        // post-create mapping. Both must derive from the same target, so a
233        // single composition is the source of truth (idempotency is scoped per
234        // target: the same key addressed to two different targets is two
235        // commands).
236        let composed_idempotency_key = request.idempotency_key.as_ref().map(|idem_key| {
237            Self::compose_idempotency_key(
238                &request.deployment_id,
239                &resolved_target.target.resource_id,
240                &request.command,
241                idem_key,
242            )
243        });
244
245        // Check idempotency if key provided.
246        if let Some(ref composed_key) = composed_idempotency_key {
247            if let Some(existing_id) = self.check_idempotency(composed_key).await? {
248                // Return existing command status
249                let status = self
250                    .command_registry
251                    .get_command_status(&existing_id)
252                    .await?;
253                if let Some(s) = status {
254                    return Ok(CreateCommandResponse {
255                        command_id: existing_id,
256                        state: s.state,
257                        storage_upload: None,
258                        inline_allowed_up_to: self.inline_max_bytes as u64,
259                        next: "poll".to_string(),
260                    });
261                }
262            }
263        }
264
265        // Determine initial state and request size
266        let (initial_state, request_size_bytes) = match &request.params {
267            BodySpec::Inline { inline_base64 } => {
268                let size = inline_base64.len() as u64;
269                (CommandState::Pending, Some(size))
270            }
271            BodySpec::Storage { size, .. } => {
272                if size.unwrap_or(0) > self.inline_max_bytes as u64 {
273                    (CommandState::PendingUpload, *size)
274                } else {
275                    (CommandState::Pending, *size)
276                }
277            }
278        };
279
280        // 1. Registry creates command metadata (SOURCE OF TRUTH)
281        let metadata = self
282            .command_registry
283            .create_command(
284                &request.deployment_id,
285                &request.command,
286                &resolved_target,
287                initial_state,
288                request.deadline,
289                request_size_bytes,
290            )
291            .await?;
292
293        let command_id = metadata.command_id;
294        let delivery_mode = metadata.delivery_mode;
295
296        // 2. Store idempotency mapping in KV, reusing the key composed above
297        // from the same resolved target. Losing the conditional put means a
298        // concurrent create with the same key raced past the pre-create check
299        // together with this one; dedupe by failing the command created above
300        // and answering with the winner's, so the caller never observes two
301        // live commands for one idempotency key.
302        if let Some(ref composed_key) = composed_idempotency_key {
303            if let Some(winner_id) = self.store_idempotency(composed_key, &command_id).await? {
304                let _ = self
305                    .command_registry
306                    .complete_command(
307                        &command_id,
308                        CommandState::Failed,
309                        Utc::now(),
310                        None,
311                        Some(serde_json::json!({
312                            "code": "IDEMPOTENT_DUPLICATE",
313                            "message": format!(
314                                "Superseded by concurrent create '{}' with the same idempotency key",
315                                winner_id
316                            ),
317                        })),
318                    )
319                    .await?;
320                let status = self.command_registry.get_command_status(&winner_id).await?;
321                let state = status.map(|s| s.state).unwrap_or(CommandState::Pending);
322                return Ok(CreateCommandResponse {
323                    command_id: winner_id,
324                    state,
325                    storage_upload: None,
326                    inline_allowed_up_to: self.inline_max_bytes as u64,
327                    next: "poll".to_string(),
328                });
329            }
330        }
331
332        // 3. Store params in KV
333        self.store_params(&command_id, &request.params).await?;
334
335        // 4. Generate storage upload URL if needed
336        let storage_upload = if initial_state == CommandState::PendingUpload {
337            Some(self.generate_params_upload(&command_id).await?)
338        } else {
339            None
340        };
341
342        // 5. Handle dispatch based on state and the target's delivery mode
343        let (final_state, next_action) = if initial_state == CommandState::Pending {
344            match delivery_mode {
345                CommandDeliveryMode::Push => {
346                    // Push delivery: dispatch immediately
347                    let state = self
348                        .dispatch_command_push(&command_id, &request.deployment_id)
349                        .await?;
350                    (state, "poll")
351                }
352                CommandDeliveryMode::Pull => {
353                    // Pull delivery: create a pending index for a receiver or
354                    // environment-local operator relay.
355                    self.create_pending_index(
356                        &request.deployment_id,
357                        &metadata.target.resource_id,
358                        &command_id,
359                    )
360                    .await?;
361                    debug!("Command {} ready for target-scoped lease", command_id);
362                    (CommandState::Pending, "poll")
363                }
364            }
365        } else {
366            // PendingUpload - need upload first
367            (initial_state, "upload")
368        };
369
370        // 6. Create deadline index if deadline provided
371        if let Some(deadline) = request.deadline {
372            self.create_deadline_index(&command_id, deadline).await?;
373        }
374
375        Ok(CreateCommandResponse {
376            command_id,
377            state: final_state,
378            storage_upload,
379            inline_allowed_up_to: self.inline_max_bytes as u64,
380            next: next_action.to_string(),
381        })
382    }
383
384    /// Mark upload as complete and dispatch command.
385    pub async fn upload_complete(
386        &self,
387        command_id: &str,
388        upload_request: UploadCompleteRequest,
389    ) -> Result<UploadCompleteResponse> {
390        // 1. Get current status from registry (source of truth)
391        let status = self
392            .command_registry
393            .get_command_status(command_id)
394            .await?
395            .ok_or_else(|| {
396                AlienError::new(ErrorData::CommandNotFound {
397                    command_id: command_id.to_string(),
398                })
399            })?;
400
401        // 2. Validate current state
402        if status.state != CommandState::PendingUpload {
403            return Err(AlienError::new(ErrorData::InvalidStateTransition {
404                from: status.state.as_ref().to_string(),
405                to: CommandState::Pending.as_ref().to_string(),
406            }));
407        }
408
409        // 3. Update params in KV with storage reference
410        let storage_get_request = self.generate_storage_get_request(command_id).await?;
411        let params = BodySpec::Storage {
412            size: Some(upload_request.size),
413            storage_get_request: Some(storage_get_request),
414            storage_put_used: None,
415        };
416        self.store_params(command_id, &params).await?;
417
418        // 4. Update registry to Pending state
419        let transitioned = self
420            .command_registry
421            .update_command_state(command_id, CommandState::Pending, None, None, None, None)
422            .await?;
423        if !transitioned {
424            let current_state = self
425                .command_registry
426                .get_command_status(command_id)
427                .await?
428                .map(|current| current.state)
429                .unwrap_or(status.state);
430            return Err(AlienError::new(ErrorData::InvalidStateTransition {
431                from: current_state.as_ref().to_string(),
432                to: CommandState::Pending.as_ref().to_string(),
433            }));
434        }
435
436        // 5. Get delivery mode from registry and handle dispatch
437        let metadata = self
438            .command_registry
439            .get_command_metadata(command_id)
440            .await?
441            .ok_or_else(|| {
442                AlienError::new(ErrorData::CommandNotFound {
443                    command_id: command_id.to_string(),
444                })
445            })?;
446
447        let final_state = match metadata.delivery_mode {
448            CommandDeliveryMode::Push => {
449                self.dispatch_command_push(command_id, &status.deployment_id)
450                    .await?
451            }
452            CommandDeliveryMode::Pull => {
453                self.create_pending_index(
454                    &status.deployment_id,
455                    &metadata.target.resource_id,
456                    command_id,
457                )
458                .await?;
459                debug!(
460                    "Command {} ready for pull after upload (target will poll)",
461                    command_id
462                );
463                CommandState::Pending
464            }
465        };
466
467        Ok(UploadCompleteResponse {
468            command_id: command_id.to_string(),
469            state: final_state,
470        })
471    }
472
473    /// Get command status.
474    ///
475    /// Queries registry for metadata (source of truth), KV for response blob.
476    pub async fn get_command_status(&self, command_id: &str) -> Result<CommandStatusResponse> {
477        // 1. Get status from registry (SOURCE OF TRUTH)
478        let status = self
479            .command_registry
480            .get_command_status(command_id)
481            .await?
482            .ok_or_else(|| {
483                AlienError::new(ErrorData::CommandNotFound {
484                    command_id: command_id.to_string(),
485                })
486            })?;
487
488        // 2. Check deadline expiry inline
489        if let Some(deadline) = status.deadline {
490            if Utc::now() > deadline && !status.state.is_terminal() {
491                // Expire the command — CONDITIONALLY: a submit can land
492                // between the status read above and this write, and an
493                // unconditional write would stomp its terminal state (the
494                // torn-record class the conditional transition exists for).
495                let won = self
496                    .command_registry
497                    .complete_command(command_id, CommandState::Expired, Utc::now(), None, None)
498                    .await?;
499                if won {
500                    // Clean up pending index
501                    self.delete_pending_index(
502                        &status.deployment_id,
503                        &status.target.resource_id,
504                        command_id,
505                    )
506                    .await?;
507
508                    // Return expired status directly (avoid recursion)
509                    return Ok(CommandStatusResponse {
510                        command_id: command_id.to_string(),
511                        state: CommandState::Expired,
512                        attempt: status.attempt,
513                        target: status.target,
514                        response: None,
515                    });
516                }
517                // Lost to a concurrent submit: fall through and serve the
518                // freshly-terminal state below.
519                let status = self
520                    .command_registry
521                    .get_command_status(command_id)
522                    .await?
523                    .ok_or_else(|| {
524                        AlienError::new(ErrorData::CommandNotFound {
525                            command_id: command_id.to_string(),
526                        })
527                    })?;
528                let response = if status.state.is_terminal() {
529                    self.get_response(command_id).await?
530                } else {
531                    None
532                };
533                return Ok(CommandStatusResponse {
534                    command_id: command_id.to_string(),
535                    state: status.state,
536                    attempt: status.attempt,
537                    target: status.target,
538                    response,
539                });
540            }
541        }
542
543        // 3. Get response blob from KV if terminal state
544        let response = if status.state.is_terminal() {
545            self.get_response(command_id).await?
546        } else {
547            None
548        };
549
550        Ok(CommandStatusResponse {
551            command_id: command_id.to_string(),
552            state: status.state,
553            attempt: status.attempt,
554            target: status.target,
555            response,
556        })
557    }
558
559    /// Submit response from deployment.
560    ///
561    /// Stores response blob in KV, updates state in registry.
562    pub async fn submit_command_response(
563        &self,
564        command_id: &str,
565        mut response: CommandResponse,
566    ) -> Result<()> {
567        // 1. Get current status from registry
568        let status = self
569            .command_registry
570            .get_command_status(command_id)
571            .await?
572            .ok_or_else(|| {
573                AlienError::new(ErrorData::CommandNotFound {
574                    command_id: command_id.to_string(),
575                })
576            })?;
577
578        // 2. Handle duplicate responses gracefully
579        if status.state.is_terminal() {
580            debug!(
581                "Ignoring duplicate response for terminal command {}",
582                command_id
583            );
584            return Ok(());
585        }
586
587        // 3. Validate state transition
588        if status.state != CommandState::Dispatched {
589            return Err(AlienError::new(ErrorData::InvalidStateTransition {
590                from: status.state.as_ref().to_string(),
591                to: CommandState::Succeeded.as_ref().to_string(),
592            }));
593        }
594
595        // 4. If response was uploaded to storage, generate download URL
596        if let CommandResponse::Success {
597            response: ref mut body,
598        } = response
599        {
600            if let BodySpec::Storage {
601                size,
602                storage_get_request,
603                storage_put_used,
604            } = body
605            {
606                if storage_get_request.is_none() && storage_put_used.unwrap_or(false) {
607                    let get_request = self
608                        .generate_response_storage_get_request(command_id)
609                        .await?;
610                    *body = BodySpec::Storage {
611                        size: *size,
612                        storage_get_request: Some(get_request),
613                        storage_put_used: *storage_put_used,
614                    };
615                }
616            }
617        }
618
619        // 5. Store response blob in KV
620        self.store_response(command_id, &response).await?;
621
622        // 6. Update registry state (SOURCE OF TRUTH) BEFORE cleaning up the lease
623        // and pending index. Committing the terminal state first guarantees the
624        // command can never be stranded as Dispatched-with-no-lease-no-index: if
625        // the process dies (or a cleanup step errors) after this point, get/poll
626        // sees the terminal state and returns the already-stored response, and any
627        // orphaned lease/pending-index entry is reaped by `acquire_lease`'s
628        // terminal-state check. The old order (cleanup first, state last) left a
629        // crash window in which the response was stored but permanently invisible.
630        let (new_state, error) = if response.is_success() {
631            (CommandState::Succeeded, None)
632        } else if let CommandResponse::Error { code, message, .. } = &response {
633            (
634                CommandState::Failed,
635                Some(serde_json::json!({ "code": code, "message": message })),
636            )
637        } else {
638            (CommandState::Failed, None)
639        };
640
641        let response_size = match &response {
642            CommandResponse::Success {
643                response: BodySpec::Inline { inline_base64 },
644            } => Some(inline_base64.len() as u64),
645            CommandResponse::Success {
646                response: BodySpec::Storage { size, .. },
647            } => *size,
648            _ => None,
649        };
650
651        // Conditional terminal transition: exactly one of two racing
652        // submitters (a redelivered execution racing the original whose lease
653        // expired) wins; a terminal record is never overwritten. The blob was
654        // pre-stored above for crash-safety (state never flips terminal with
655        // no blob on disk); the winner re-stores its own response below so
656        // the recorded state and the served blob agree. Residual: a loser
657        // whose pre-store lands after the winner's re-store can still leave
658        // its blob — that requires the loser's KV write to outlast the
659        // winner's entire transition+re-store, a pathological schedule.
660        let won = self
661            .command_registry
662            .complete_command(command_id, new_state, Utc::now(), response_size, error)
663            .await?;
664        if !won {
665            debug!(
666                "Ignoring duplicate response for command {} (lost the terminal transition race)",
667                command_id
668            );
669            return Ok(());
670        }
671        self.store_response(command_id, &response).await?;
672
673        // 7. Clean up lease from KV (best-effort; the terminal state is already
674        // committed above, so a failure here cannot strand the command). Log a
675        // warning and continue instead of failing the call — a leftover entry
676        // is reaped by `acquire_lease`'s terminal-state check on its next scan.
677        if let Err(e) = self.delete_lease(command_id).await {
678            warn!(
679                command_id,
680                error = %e,
681                "Failed to clean up lease after terminal response; will be reaped on next lease scan"
682            );
683        }
684
685        // 8. Clean up pending index from KV (best-effort; terminal state).
686        // Same reasoning as the lease cleanup above.
687        if let Err(e) = self
688            .delete_pending_index(
689                &status.deployment_id,
690                &status.target.resource_id,
691                command_id,
692            )
693            .await
694        {
695            warn!(
696                command_id,
697                error = %e,
698                "Failed to clean up pending index after terminal response; will be reaped on next lease scan"
699            );
700        }
701
702        info!(
703            "Command {} completed with state {:?}",
704            command_id, new_state
705        );
706        Ok(())
707    }
708
709    /// Acquire leases for polling deployments.
710    ///
711    /// Scans KV pending index, queries registry for metadata, creates leases in KV.
712    pub async fn acquire_lease(
713        &self,
714        deployment_id: &str,
715        lease_request: &LeaseRequest,
716    ) -> Result<LeaseResponse> {
717        let mut leases = Vec::new();
718
719        // 1. Scan KV pending index — ONLY the requesting target's prefix.
720        // Commands for other targets in the same deployment are invisible here.
721        let target_prefix = format!(
722            "target:{}:{}:pending:",
723            deployment_id, lease_request.target.resource_id
724        );
725        let scan_result = self
726            .kv
727            .scan_prefix(&target_prefix, Some(lease_request.max_leases * 2), None)
728            .await
729            .into_alien_error()
730            .context(ErrorData::KvOperationFailed {
731                operation: "scan_prefix".to_string(),
732                key: target_prefix.clone(),
733                message: "Failed to scan for pending commands".to_string(),
734            })?;
735
736        for entry in scan_result.items {
737            let index_key = entry.key;
738            if leases.len() >= lease_request.max_leases {
739                break;
740            }
741
742            let command_id = self.extract_command_id_from_index_key(&index_key)?;
743
744            // 2. Try to acquire lease atomically in KV
745            let lease_id = format!("lease_{}", Uuid::new_v4());
746            let lease_duration = Duration::from_secs(lease_request.lease_seconds);
747            let expires_at =
748                Utc::now() + chrono::Duration::seconds(lease_request.lease_seconds as i64);
749
750            let lease_data = LeaseData {
751                lease_id: lease_id.clone(),
752                acquired_at: Utc::now(),
753                expires_at,
754                owner: deployment_id.to_string(),
755            };
756
757            let lease_key = format!("cmd:{}:lease", command_id);
758            let lease_value = serde_json::to_vec(&lease_data).into_alien_error().context(
759                ErrorData::SerializationFailed {
760                    message: "Failed to serialize lease data".to_string(),
761                    data_type: Some("LeaseData".to_string()),
762                },
763            )?;
764
765            let options = Some(PutOptions {
766                ttl: Some(lease_duration),
767                condition: PutCondition::Absent,
768            });
769
770            let success = self
771                .kv
772                .put(&lease_key, lease_value, options)
773                .await
774                .context(ErrorData::KvOperationFailed {
775                    operation: "put".to_string(),
776                    key: lease_key.clone(),
777                    message: "Failed to create lease".to_string(),
778                })?;
779
780            if !success {
781                // Lease already exists, skip
782                continue;
783            }
784
785            // Reverse index for O(1) release-by-lease-id lookups. Shares the
786            // lease's TTL so it self-cleans on expiry; a stale entry is
787            // harmless because every reader re-verifies the lease_id against
788            // the live `cmd:{id}:lease` record.
789            let reverse_key = format!("lease:{}", lease_id);
790            self.kv
791                .put(
792                    &reverse_key,
793                    command_id.as_bytes().to_vec(),
794                    Some(PutOptions {
795                        ttl: Some(lease_duration),
796                        condition: PutCondition::None,
797                    }),
798                )
799                .await
800                .context(ErrorData::KvOperationFailed {
801                    operation: "put".to_string(),
802                    key: reverse_key,
803                    message: "Failed to create lease reverse index".to_string(),
804                })?;
805
806            // 3. Get metadata from registry
807            let mut metadata = match self
808                .command_registry
809                .get_command_metadata(&command_id)
810                .await?
811            {
812                Some(m) => m,
813                None => {
814                    // Command doesn't exist in registry, clean up
815                    self.delete_lease(&command_id).await?;
816                    let _ = self.kv.delete(&index_key, None).await;
817                    continue;
818                }
819            };
820
821            // 3.1 Expiry-driven redelivery: this command was already
822            // dispatched, yet the lease slot was free (the conditional put
823            // above succeeded) — the previous lease TTL-expired with no
824            // response and no explicit release. Increment the attempt so the
825            // redelivered envelope carries `attempt > 1`, the at-least-once
826            // redelivery signal both receiver twins document. The explicit
827            // release path (`release_lease`) increments the same counter.
828            if metadata.state == CommandState::Dispatched {
829                self.command_registry.increment_attempt(&command_id).await?;
830                metadata.attempt += 1;
831            }
832
833            // 3.5 Defense-in-depth: the pending index key said this command
834            // belongs to the requesting target — verify the registry agrees.
835            // A mismatch means the index is corrupt; fail loudly rather than
836            // delivering a command to the wrong resource. The corrupt index
837            // key is deliberately retained (only the lease is cleaned up):
838            // genuine corruption should never occur, and the key is the
839            // evidence an operator needs — do not "fix" this by deleting it.
840            if metadata.target != lease_request.target {
841                self.delete_lease(&command_id).await?;
842                return Err(AlienError::new(ErrorData::Other {
843                    message: format!(
844                        "Pending index corruption: command '{}' is indexed under target '{}' \
845                         but the registry says it belongs to target '{}' — refusing to deliver",
846                        command_id, lease_request.target.resource_id, metadata.target.resource_id,
847                    ),
848                }));
849            }
850
851            // 4. Check if command is in terminal state (stale index)
852            if metadata.state.is_terminal() {
853                // Clean up stale data
854                self.delete_lease(&command_id).await?;
855                let _ = self.kv.delete(&index_key, None).await;
856                continue;
857            }
858
859            // 5. Check deadline expiry — conditionally: a racing submit's
860            // terminal state must never be stomped back to Expired.
861            if let Some(deadline) = metadata.deadline {
862                if Utc::now() > deadline {
863                    let _ = self
864                        .command_registry
865                        .complete_command(
866                            &command_id,
867                            CommandState::Expired,
868                            Utc::now(),
869                            None,
870                            None,
871                        )
872                        .await?;
873                    self.delete_lease(&command_id).await?;
874                    let _ = self.kv.delete(&index_key, None).await;
875                    continue;
876                }
877            }
878
879            // 6. Get params from KV
880            let params = match self.get_params(&command_id).await? {
881                Some(p) => p,
882                None => {
883                    // No params, something went wrong
884                    self.delete_lease(&command_id).await?;
885                    continue;
886                }
887            };
888
889            // 7. Mark Dispatched — conditionally. Between the terminal
890            // check in step 3.1 and here, the ORIGINAL holder of a
891            // TTL-expired lease can still submit; handing out this lease
892            // would re-execute a completed command and stomp its state.
893            let dispatched = self
894                .command_registry
895                .mark_dispatched_if_not_terminal(&command_id, Utc::now())
896                .await?;
897            if !dispatched {
898                debug!(
899                    command_id = %command_id,
900                    "Command turned terminal while leasing; releasing the lease"
901                );
902                self.delete_lease(&command_id).await?;
903                let _ = self.kv.delete(&index_key, None).await;
904                continue;
905            }
906
907            // 8. Build envelope. Lease-served envelopes carry manager URLs
908            // as path-relative references from the lease endpoint: a pull
909            // consumer resolves them against the exact endpoint it reached,
910            // preserving both a network-corrected origin and any reverse-
911            // proxy path prefix the manager cannot know. Push envelopes keep
912            // absolute URLs: push transports have no configured base, and
913            // reaching the manager's public address is inherent to push
914            // delivery.
915            let mut envelope = self.build_envelope(&command_id, &metadata, params).await?;
916            Self::relativize_manager_urls(&mut envelope, &self.base_url);
917
918            leases.push(LeaseInfo {
919                lease_id,
920                lease_expires_at: expires_at,
921                command_id: command_id.clone(),
922                attempt: metadata.attempt,
923                envelope,
924            });
925        }
926
927        Ok(LeaseResponse { leases })
928    }
929
930    /// Rewrite manager-origin URLs in a lease-served envelope to references
931    /// relative to that lease endpoint (see [`Self::acquire_lease`]). Only
932    /// same-origin URLs can be made relative, so cloud-presigned storage URLs
933    /// pass through byte-for-byte.
934    fn relativize_manager_urls(envelope: &mut Envelope, base_url: &str) {
935        let Ok(mut lease_endpoint) = reqwest::Url::parse(base_url) else {
936            // An unparseable base cannot produce relative references; leave
937            // the envelope absolute (the pre-relative behavior).
938            return;
939        };
940        let Ok(mut segments) = lease_endpoint.path_segments_mut() else {
941            return;
942        };
943        segments.pop_if_empty().push("commands").push("leases");
944        drop(segments);
945        lease_endpoint.set_query(None);
946        lease_endpoint.set_fragment(None);
947
948        let relativize = |target: &mut String| {
949            let Ok(mut target_url) = reqwest::Url::parse(target.as_str()) else {
950                return;
951            };
952            let suffix = target
953                .find(['?', '#'])
954                .map(|index| target[index..].to_string())
955                .unwrap_or_default();
956            target_url.set_query(None);
957            target_url.set_fragment(None);
958            if let Some(relative) = lease_endpoint.make_relative(&target_url) {
959                *target = format!("{relative}{suffix}");
960            }
961        };
962
963        relativize(&mut envelope.response_handling.submit_response_url);
964        if let alien_core::presigned::PresignedRequestBackend::Http { url, .. } =
965            &mut envelope.response_handling.storage_upload_request.backend
966        {
967            relativize(url);
968        }
969        if let BodySpec::Storage {
970            storage_get_request: Some(request),
971            ..
972        } = &mut envelope.params
973        {
974            if let alien_core::presigned::PresignedRequestBackend::Http { url, .. } =
975                &mut request.backend
976            {
977                relativize(url);
978            }
979        }
980    }
981
982    /// Release a lease manually.
983    ///
984    /// Increments attempt count in registry, returns command to Pending state.
985    pub async fn release_lease(&self, command_id: &str, lease_id: &str) -> Result<()> {
986        let lease_key = format!("cmd:{}:lease", command_id);
987
988        // 1. Verify lease ownership
989        if let Ok(Some(lease_data)) = self.kv.get(&lease_key).await {
990            let lease: LeaseData = serde_json::from_slice(&lease_data.value)
991                .into_alien_error()
992                .context(ErrorData::SerializationFailed {
993                    message: "Failed to deserialize lease data".to_string(),
994                    data_type: Some("LeaseData".to_string()),
995                })?;
996
997            if lease.lease_id != lease_id {
998                return Err(AlienError::new(ErrorData::LeaseNotFound {
999                    lease_id: lease_id.to_string(),
1000                }));
1001            }
1002
1003            // 2. Delete lease from KV (and its reverse index)
1004            self.delete_lease(command_id).await?;
1005            let _ = self.kv.delete(&format!("lease:{}", lease_id), None).await;
1006
1007            // 3. Return the command to Pending only while it is still
1008            // non-terminal. A racing response/deadline completion wins.
1009            let released = self
1010                .command_registry
1011                .update_command_state(command_id, CommandState::Pending, None, None, None, None)
1012                .await?;
1013            if released {
1014                // 4. Increment the attempt only for a command that was
1015                // actually released for redelivery.
1016                self.command_registry.increment_attempt(command_id).await?;
1017            }
1018
1019            // Note: Pending index is NOT removed on lease, so command is still there
1020            debug!(
1021                released,
1022                "Lease {} released for command {}", lease_id, command_id
1023            );
1024        }
1025
1026        Ok(())
1027    }
1028
1029    /// Get the deployment_id that owns a command.
1030    ///
1031    /// Used by the manager's auth layer to check whether the caller has access
1032    /// to a specific command without fetching the full status.
1033    pub async fn get_command_deployment_id(&self, command_id: &str) -> Result<Option<String>> {
1034        let status = self.command_registry.get_command_status(command_id).await?;
1035        Ok(status.map(|s| s.deployment_id))
1036    }
1037
1038    /// Get the command's canonical ownership fields for route authorization.
1039    pub async fn get_command_access_context(
1040        &self,
1041        command_id: &str,
1042    ) -> Result<Option<CommandAccessContext>> {
1043        self.command_registry
1044            .get_command_access_context(command_id)
1045            .await
1046    }
1047
1048    /// Resolve a lease_id to `(command_id, owner_deployment_id)` via the
1049    /// `lease:{lease_id}` reverse index, re-verifying against the live lease
1050    /// record (the reverse entry can outlive a released lease briefly, and a
1051    /// command can have been re-leased under a new lease_id since).
1052    ///
1053    /// Used by the manager's auth layer to check that the caller may act on
1054    /// the deployment that holds the lease.
1055    pub async fn get_lease_owner(&self, lease_id: &str) -> Result<Option<(String, String)>> {
1056        let reverse_key = format!("lease:{}", lease_id);
1057        let Some(command_id_bytes) =
1058            self.kv
1059                .get(&reverse_key)
1060                .await
1061                .context(ErrorData::KvOperationFailed {
1062                    operation: "get".to_string(),
1063                    key: reverse_key.clone(),
1064                    message: "Failed to look up lease reverse index".to_string(),
1065                })?
1066        else {
1067            return Ok(None);
1068        };
1069        let command_id = String::from_utf8(command_id_bytes.value).map_err(|_| {
1070            AlienError::new(ErrorData::Other {
1071                message: format!("Lease reverse index '{}' is not valid UTF-8", reverse_key),
1072            })
1073        })?;
1074
1075        let lease_key = format!("cmd:{}:lease", command_id);
1076        let Some(lease_data) =
1077            self.kv
1078                .get(&lease_key)
1079                .await
1080                .context(ErrorData::KvOperationFailed {
1081                    operation: "get".to_string(),
1082                    key: lease_key,
1083                    message: "Failed to look up lease".to_string(),
1084                })?
1085        else {
1086            return Ok(None);
1087        };
1088        let lease: LeaseData = serde_json::from_slice(&lease_data.value)
1089            .into_alien_error()
1090            .context(ErrorData::SerializationFailed {
1091                message: "Failed to deserialize lease data".to_string(),
1092                data_type: Some("LeaseData".to_string()),
1093            })?;
1094        if lease.lease_id != lease_id {
1095            return Ok(None);
1096        }
1097
1098        Ok(Some((command_id, lease.owner)))
1099    }
1100
1101    /// Release a lease by lease_id only (for the API).
1102    pub async fn release_lease_by_id(&self, lease_id: &str) -> Result<()> {
1103        match self.get_lease_owner(lease_id).await? {
1104            Some((command_id, _owner)) => self.release_lease(&command_id, lease_id).await,
1105            None => Err(AlienError::new(ErrorData::LeaseNotFound {
1106                lease_id: lease_id.to_string(),
1107            })),
1108        }
1109    }
1110
1111    // =========================================================================
1112    // Internal Helper Methods
1113    // =========================================================================
1114
1115    async fn validate_create_command(&self, request: &CreateCommandRequest) -> Result<()> {
1116        if request.command.is_empty() {
1117            return Err(AlienError::new(ErrorData::InvalidCommand {
1118                message: "Command name cannot be empty".to_string(),
1119            }));
1120        }
1121
1122        // The command name occupies one segment of the `:`-delimited
1123        // idempotency key (`{dep}:{rid}:{command}:{key}`). A ':' in the name
1124        // would let it bleed into the client-key segment — which routinely
1125        // contains ':' — so (command="a:b", key="c") and (command="a",
1126        // key="b:c") would forge the same key and be treated as the same
1127        // command. Reject ':' here, mirroring the target-id colon guard, so the
1128        // command segment is always unambiguous.
1129        validate_command_name(&request.command)?;
1130
1131        if request.deployment_id.is_empty() {
1132            return Err(AlienError::new(ErrorData::InvalidCommand {
1133                message: "Deployment ID cannot be empty".to_string(),
1134            }));
1135        }
1136
1137        if let Some(deadline) = request.deadline {
1138            if deadline <= Utc::now() {
1139                return Err(AlienError::new(ErrorData::InvalidCommand {
1140                    message: "Deadline must be in the future".to_string(),
1141                }));
1142            }
1143        }
1144
1145        Ok(())
1146    }
1147
1148    // --- Idempotency ---
1149
1150    /// Compose the target-scoped idempotency key:
1151    /// `{deploymentId}:{targetResourceId}:{commandName}:{key}`.
1152    ///
1153    /// Scoping by target means the same client key addressed to two different
1154    /// targets creates two distinct commands.
1155    fn compose_idempotency_key(
1156        deployment_id: &str,
1157        target_resource_id: &str,
1158        command_name: &str,
1159        idem_key: &str,
1160    ) -> String {
1161        // Invariant: at every real call site the target id and command name are
1162        // both `:`-free (the former enforced at resolution via
1163        // `validate_command_target_id`, the latter via `validate_command_name`
1164        // in `validate_create_command`), so each occupies exactly one segment of
1165        // `{dep}:{rid}:{command}:{key}` and only the trailing client key may
1166        // carry ':'. Those guards live upstream; this is a pure formatter, so it
1167        // is not asserted on `command_name` here (the collision tests below
1168        // deliberately format a ':'-bearing command to demonstrate what the
1169        // upstream guard prevents).
1170        debug_assert!(
1171            !target_resource_id.contains(':'),
1172            "target_resource_id must be ':'-free before key composition: {target_resource_id}"
1173        );
1174        format!(
1175            "{}:{}:{}:{}",
1176            deployment_id, target_resource_id, command_name, idem_key
1177        )
1178    }
1179
1180    async fn check_idempotency(&self, idem_key: &str) -> Result<Option<String>> {
1181        let key = format!("idem:{}", idem_key);
1182        if let Some(data) = self
1183            .kv
1184            .get(&key)
1185            .await
1186            .context(ErrorData::KvOperationFailed {
1187                operation: "get".to_string(),
1188                key: key.clone(),
1189                message: "Failed to check idempotency".to_string(),
1190            })?
1191        {
1192            let command_id = String::from_utf8(data.value).into_alien_error().context(
1193                ErrorData::SerializationFailed {
1194                    message: "Invalid idempotency data".to_string(),
1195                    data_type: Some("String".to_string()),
1196                },
1197            )?;
1198            return Ok(Some(command_id));
1199        }
1200        Ok(None)
1201    }
1202
1203    /// Claim the idempotency key for `command_id`.
1204    ///
1205    /// Returns `None` when this command won the key. Returns
1206    /// `Some(winner_id)` when a concurrent create with the same key won the
1207    /// conditional put first — both requests passed the pre-create
1208    /// `check_idempotency` before either stored, so the loser must be
1209    /// detected here, after its command was already created.
1210    async fn store_idempotency(&self, idem_key: &str, command_id: &str) -> Result<Option<String>> {
1211        let key = format!("idem:{}", idem_key);
1212        let ttl = Duration::from_secs(24 * 60 * 60); // 24 hours
1213        let won = self
1214            .kv
1215            .put(
1216                &key,
1217                command_id.as_bytes().to_vec(),
1218                Some(PutOptions {
1219                    ttl: Some(ttl),
1220                    condition: PutCondition::Absent,
1221                }),
1222            )
1223            .await
1224            .context(ErrorData::KvOperationFailed {
1225                operation: "put".to_string(),
1226                key: key.clone(),
1227                message: "Failed to store idempotency".to_string(),
1228            })?;
1229        if won {
1230            return Ok(None);
1231        }
1232        // Lost the conditional put: read back who won. The winner entry can
1233        // only be absent if it TTL-expired in the microseconds since — treat
1234        // that as an inconsistency rather than silently duplicating.
1235        match self.check_idempotency(idem_key).await? {
1236            Some(winner_id) => Ok(Some(winner_id)),
1237            None => Err(AlienError::new(ErrorData::Other {
1238                message: format!(
1239                    "Idempotency key '{}' was concurrently claimed but has no winner entry",
1240                    key
1241                ),
1242            })),
1243        }
1244    }
1245
1246    // --- Params ---
1247
1248    pub async fn store_params(&self, command_id: &str, params: &BodySpec) -> Result<()> {
1249        let key = format!("cmd:{}:params", command_id);
1250
1251        // Try serializing as-is first
1252        let data = CommandParamsData {
1253            params: params.clone(),
1254        };
1255        let value = serde_json::to_vec(&data).into_alien_error().context(
1256            ErrorData::SerializationFailed {
1257                message: "Failed to serialize params".to_string(),
1258                data_type: Some("CommandParamsData".to_string()),
1259            },
1260        )?;
1261
1262        // If it fits in KV, store directly (fast path)
1263        if value.len() <= KV_VALUE_THRESHOLD {
1264            self.kv
1265                .put(&key, value, None)
1266                .await
1267                .context(ErrorData::KvOperationFailed {
1268                    operation: "put".to_string(),
1269                    key: key.clone(),
1270                    message: "Failed to store params".to_string(),
1271                })?;
1272            return Ok(());
1273        }
1274
1275        // Auto-promote: inline data exceeds KV limit, store raw bytes in blob
1276        if let BodySpec::Inline { inline_base64 } = params {
1277            let raw_bytes = general_purpose::STANDARD
1278                .decode(inline_base64)
1279                .into_alien_error()
1280                .context(ErrorData::SerializationFailed {
1281                    message: "Failed to decode inline base64 params for auto-promotion".to_string(),
1282                    data_type: Some("base64".to_string()),
1283                })?;
1284
1285            let raw_len = raw_bytes.len() as u64;
1286            let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1287
1288            self.storage
1289                .put(&blob_path, Bytes::from(raw_bytes).into())
1290                .await
1291                .into_alien_error()
1292                .context(ErrorData::StorageOperationFailed {
1293                    message: "Failed to auto-promote params to blob storage".to_string(),
1294                    operation: Some("put".to_string()),
1295                    path: Some(blob_path.to_string()),
1296                })?;
1297
1298            debug!(
1299                "Auto-promoted params for command {} to blob ({} bytes raw)",
1300                command_id, raw_len
1301            );
1302
1303            // Store tiny reference in KV instead
1304            let promoted = CommandParamsData {
1305                params: BodySpec::Storage {
1306                    size: Some(raw_len),
1307                    storage_get_request: None,
1308                    storage_put_used: Some(true),
1309                },
1310            };
1311            let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1312                ErrorData::SerializationFailed {
1313                    message: "Failed to serialize promoted params reference".to_string(),
1314                    data_type: Some("CommandParamsData".to_string()),
1315                },
1316            )?;
1317            self.kv.put(&key, promoted_value, None).await.context(
1318                ErrorData::KvOperationFailed {
1319                    operation: "put".to_string(),
1320                    key: key.clone(),
1321                    message: "Failed to store promoted params reference".to_string(),
1322                },
1323            )?;
1324            return Ok(());
1325        }
1326
1327        // Storage references are always tiny, store as-is
1328        self.kv
1329            .put(&key, value, None)
1330            .await
1331            .context(ErrorData::KvOperationFailed {
1332                operation: "put".to_string(),
1333                key: key.clone(),
1334                message: "Failed to store params".to_string(),
1335            })?;
1336        Ok(())
1337    }
1338
1339    pub async fn get_params(&self, command_id: &str) -> Result<Option<BodySpec>> {
1340        let key = format!("cmd:{}:params", command_id);
1341        if let Some(value) = self
1342            .kv
1343            .get(&key)
1344            .await
1345            .context(ErrorData::KvOperationFailed {
1346                operation: "get".to_string(),
1347                key: key.clone(),
1348                message: "Failed to get params".to_string(),
1349            })?
1350        {
1351            let data: CommandParamsData = serde_json::from_slice(&value.value)
1352                .into_alien_error()
1353                .context(ErrorData::SerializationFailed {
1354                message: "Failed to deserialize params".to_string(),
1355                data_type: Some("CommandParamsData".to_string()),
1356            })?;
1357            return Ok(Some(data.params));
1358        }
1359        Ok(None)
1360    }
1361
1362    // --- Response ---
1363
1364    pub async fn store_response(&self, command_id: &str, response: &CommandResponse) -> Result<()> {
1365        let key = format!("cmd:{}:response", command_id);
1366        let data = CommandResponseData {
1367            response: response.clone(),
1368        };
1369        let value = serde_json::to_vec(&data).into_alien_error().context(
1370            ErrorData::SerializationFailed {
1371                message: "Failed to serialize response".to_string(),
1372                data_type: Some("CommandResponseData".to_string()),
1373            },
1374        )?;
1375
1376        // If it fits in KV, store directly (fast path)
1377        if value.len() <= KV_VALUE_THRESHOLD {
1378            self.kv
1379                .put(&key, value, None)
1380                .await
1381                .context(ErrorData::KvOperationFailed {
1382                    operation: "put".to_string(),
1383                    key: key.clone(),
1384                    message: "Failed to store response".to_string(),
1385                })?;
1386            return Ok(());
1387        }
1388
1389        // Auto-promote: inline response exceeds KV limit
1390        if let CommandResponse::Success {
1391            response: BodySpec::Inline { inline_base64 },
1392        } = response
1393        {
1394            let raw_bytes = general_purpose::STANDARD
1395                .decode(inline_base64)
1396                .into_alien_error()
1397                .context(ErrorData::SerializationFailed {
1398                    message: "Failed to decode inline base64 response for auto-promotion"
1399                        .to_string(),
1400                    data_type: Some("base64".to_string()),
1401                })?;
1402
1403            let raw_len = raw_bytes.len() as u64;
1404            let blob_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1405
1406            self.storage
1407                .put(&blob_path, Bytes::from(raw_bytes).into())
1408                .await
1409                .into_alien_error()
1410                .context(ErrorData::StorageOperationFailed {
1411                    message: "Failed to auto-promote response to blob storage".to_string(),
1412                    operation: Some("put".to_string()),
1413                    path: Some(blob_path.to_string()),
1414                })?;
1415
1416            // Generate presigned GET URL for the caller
1417            let get_request = self
1418                .generate_response_storage_get_request(command_id)
1419                .await?;
1420
1421            debug!(
1422                "Auto-promoted response for command {} to blob ({} bytes raw)",
1423                command_id, raw_len
1424            );
1425
1426            // Store tiny reference in KV
1427            let promoted = CommandResponseData {
1428                response: CommandResponse::Success {
1429                    response: BodySpec::Storage {
1430                        size: Some(raw_len),
1431                        storage_get_request: Some(get_request),
1432                        storage_put_used: Some(true),
1433                    },
1434                },
1435            };
1436            let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1437                ErrorData::SerializationFailed {
1438                    message: "Failed to serialize promoted response reference".to_string(),
1439                    data_type: Some("CommandResponseData".to_string()),
1440                },
1441            )?;
1442            self.kv.put(&key, promoted_value, None).await.context(
1443                ErrorData::KvOperationFailed {
1444                    operation: "put".to_string(),
1445                    key: key.clone(),
1446                    message: "Failed to store promoted response reference".to_string(),
1447                },
1448            )?;
1449            return Ok(());
1450        }
1451
1452        // Error responses or storage references are always small, store as-is
1453        self.kv
1454            .put(&key, value, None)
1455            .await
1456            .context(ErrorData::KvOperationFailed {
1457                operation: "put".to_string(),
1458                key: key.clone(),
1459                message: "Failed to store response".to_string(),
1460            })?;
1461        Ok(())
1462    }
1463
1464    pub async fn get_response(&self, command_id: &str) -> Result<Option<CommandResponse>> {
1465        let key = format!("cmd:{}:response", command_id);
1466        if let Some(value) = self
1467            .kv
1468            .get(&key)
1469            .await
1470            .context(ErrorData::KvOperationFailed {
1471                operation: "get".to_string(),
1472                key: key.clone(),
1473                message: "Failed to get response".to_string(),
1474            })?
1475        {
1476            let data: CommandResponseData = serde_json::from_slice(&value.value)
1477                .into_alien_error()
1478                .context(ErrorData::SerializationFailed {
1479                    message: "Failed to deserialize response".to_string(),
1480                    data_type: Some("CommandResponseData".to_string()),
1481                })?;
1482            return Ok(Some(data.response));
1483        }
1484        Ok(None)
1485    }
1486
1487    // --- Pending Index ---
1488
1489    async fn create_pending_index(
1490        &self,
1491        deployment_id: &str,
1492        target_resource_id: &str,
1493        command_id: &str,
1494    ) -> Result<()> {
1495        // Invariant: the target id is `:`-free (enforced at resolution via
1496        // `validate_command_target_id`), so its prefix `target:{dep}:{rid}:`
1497        // cannot overlap another target's pending keys.
1498        debug_assert!(
1499            !target_resource_id.contains(':'),
1500            "target_resource_id must be ':'-free in the pending index: {target_resource_id}"
1501        );
1502        let timestamp = Utc::now().timestamp_nanos_opt().unwrap_or(0);
1503        let key = format!(
1504            "target:{}:{}:pending:{}:{}",
1505            deployment_id, target_resource_id, timestamp, command_id
1506        );
1507
1508        // Store empty value - just for ordering
1509        self.kv
1510            .put(&key, vec![], None)
1511            .await
1512            .context(ErrorData::KvOperationFailed {
1513                operation: "put".to_string(),
1514                key: key.clone(),
1515                message: "Failed to create pending index".to_string(),
1516            })?;
1517        Ok(())
1518    }
1519
1520    async fn delete_pending_index(
1521        &self,
1522        deployment_id: &str,
1523        target_resource_id: &str,
1524        command_id: &str,
1525    ) -> Result<()> {
1526        // We need to scan to find the exact key since we don't know the timestamp
1527        let prefix = format!("target:{}:{}:pending:", deployment_id, target_resource_id);
1528        let scan_result = self
1529            .kv
1530            .scan_prefix(&prefix, Some(100), None)
1531            .await
1532            .into_alien_error()
1533            .context(ErrorData::KvOperationFailed {
1534                operation: "scan_prefix".to_string(),
1535                key: prefix.clone(),
1536                message: "Failed to scan pending index".to_string(),
1537            })?;
1538
1539        for entry in scan_result.items {
1540            if entry.key.ends_with(&format!(":{}", command_id)) {
1541                let _ = self.kv.delete(&entry.key, None).await;
1542                break;
1543            }
1544        }
1545        Ok(())
1546    }
1547
1548    // --- Lease ---
1549
1550    async fn delete_lease(&self, command_id: &str) -> Result<()> {
1551        let key = format!("cmd:{}:lease", command_id);
1552        let _ = self.kv.delete(&key, None).await;
1553        Ok(())
1554    }
1555
1556    /// Expire every overdue non-terminal command recorded in the deadline
1557    /// index. Intended to run periodically from the hosting process.
1558    ///
1559    /// Deadlines are otherwise only enforced lazily (status polls and lease
1560    /// scans), which never reaches a command nobody polls — most notably a
1561    /// `PendingUpload` whose params upload never completed, which has no
1562    /// pending-index entry and would otherwise live forever.
1563    ///
1564    /// Uses the conditional terminal transition, so racing a concurrent
1565    /// submit is safe: whoever wins, the record stays consistent. Returns
1566    /// the number of commands expired.
1567    pub async fn reap_expired_commands(&self) -> Result<u32> {
1568        let now = Utc::now();
1569        let mut expired = 0u32;
1570        let mut cursor: Option<String> = None;
1571        // Paginate the whole index (bounded page count as a runaway guard):
1572        // keys are not numerically ordered — the timestamp segment isn't
1573        // zero-padded — so due entries can sit behind any number of
1574        // future-dated ones and a single capped scan would starve them.
1575        for _ in 0..64 {
1576            let scan = self
1577                .kv
1578                .scan_prefix("deadline:", Some(256), cursor.clone())
1579                .await
1580                .into_alien_error()
1581                .context(ErrorData::KvOperationFailed {
1582                    operation: "scan_prefix".to_string(),
1583                    key: "deadline:".to_string(),
1584                    message: "Failed to scan the deadline index".to_string(),
1585                })?;
1586            let next_cursor = scan.next_cursor.clone();
1587            for entry in scan.items {
1588                let Ok(data) = serde_json::from_slice::<DeadlineIndexData>(&entry.value) else {
1589                    warn!(key = %entry.key, "Unparseable deadline index entry; deleting");
1590                    let _ = self.kv.delete(&entry.key, None).await;
1591                    continue;
1592                };
1593                if data.deadline > now {
1594                    // Not due yet. (Keys are not sortable numerically — the
1595                    // timestamp segment isn't zero-padded — so keep scanning.)
1596                    continue;
1597                }
1598
1599                let status = self
1600                    .command_registry
1601                    .get_command_status(&data.command_id)
1602                    .await?;
1603                match status {
1604                    None => {
1605                        let _ = self.kv.delete(&entry.key, None).await;
1606                    }
1607                    Some(status) if status.state.is_terminal() => {
1608                        let _ = self.kv.delete(&entry.key, None).await;
1609                    }
1610                    Some(status) => {
1611                        let won = self
1612                        .command_registry
1613                        .complete_command(
1614                            &data.command_id,
1615                            CommandState::Expired,
1616                            now,
1617                            None,
1618                            Some(serde_json::json!({
1619                                "code": "COMMAND_EXPIRED",
1620                                "message": format!("Deadline {} elapsed", data.deadline.to_rfc3339()),
1621                            })),
1622                        )
1623                        .await?;
1624                        if won {
1625                            expired += 1;
1626                            info!(command_id = %data.command_id, "Expired overdue command");
1627                            let _ = self.delete_lease(&data.command_id).await;
1628                            let _ = self
1629                                .delete_pending_index(
1630                                    &status.deployment_id,
1631                                    &status.target.resource_id,
1632                                    &data.command_id,
1633                                )
1634                                .await;
1635                        }
1636                        let _ = self.kv.delete(&entry.key, None).await;
1637                    }
1638                }
1639            }
1640            match next_cursor {
1641                Some(next) => cursor = Some(next),
1642                None => break,
1643            }
1644        }
1645        Ok(expired)
1646    }
1647
1648    // --- Deadline Index ---
1649
1650    async fn create_deadline_index(&self, command_id: &str, deadline: DateTime<Utc>) -> Result<()> {
1651        let key = format!(
1652            "deadline:{}:{}",
1653            deadline.timestamp_nanos_opt().unwrap_or(0),
1654            command_id
1655        );
1656
1657        let data = DeadlineIndexData {
1658            command_id: command_id.to_string(),
1659            deadline,
1660        };
1661        let value = serde_json::to_vec(&data).into_alien_error().context(
1662            ErrorData::SerializationFailed {
1663                message: "Failed to serialize deadline index".to_string(),
1664                data_type: Some("DeadlineIndexData".to_string()),
1665            },
1666        )?;
1667
1668        // The entry must remain VISIBLE well past the deadline: scan_prefix
1669        // treats logically-expired keys as absent on every provider, so a
1670        // TTL equal to the deadline would hide the entry at exactly the
1671        // moment the reaper needs it (the bug that made the reaper inert).
1672        // The reaper deletes entries as it processes them; the 7-day grace
1673        // is only self-cleaning for processes that never run a reaper.
1674        const DEADLINE_INDEX_GRACE: chrono::Duration = chrono::Duration::days(7);
1675        let ttl = deadline
1676            .signed_duration_since(Utc::now())
1677            .checked_add(&DEADLINE_INDEX_GRACE)
1678            .unwrap_or(DEADLINE_INDEX_GRACE);
1679        let options = (ttl.num_seconds() > 0).then(|| PutOptions {
1680            ttl: Some(Duration::from_secs(ttl.num_seconds() as u64)),
1681            condition: PutCondition::None,
1682        });
1683
1684        self.kv
1685            .put(&key, value, options)
1686            .await
1687            .context(ErrorData::KvOperationFailed {
1688                operation: "put".to_string(),
1689                key: key.clone(),
1690                message: "Failed to create deadline index".to_string(),
1691            })?;
1692        Ok(())
1693    }
1694
1695    // --- Dispatch ---
1696
1697    async fn dispatch_command_push(
1698        &self,
1699        command_id: &str,
1700        deployment_id: &str,
1701    ) -> Result<CommandState> {
1702        // Get metadata from registry
1703        let metadata = self
1704            .command_registry
1705            .get_command_metadata(command_id)
1706            .await?
1707            .ok_or_else(|| {
1708                AlienError::new(ErrorData::CommandNotFound {
1709                    command_id: command_id.to_string(),
1710                })
1711            })?;
1712
1713        // Get params from KV
1714        let params = self.get_params(command_id).await?.ok_or_else(|| {
1715            AlienError::new(ErrorData::CommandNotFound {
1716                command_id: command_id.to_string(),
1717            })
1718        })?;
1719
1720        // Build envelope
1721        let envelope = self.build_envelope(command_id, &metadata, params).await?;
1722
1723        // Mark Dispatched BEFORE invoking the transport: a fast worker can
1724        // execute and submit before this function resumes, and submit
1725        // validates state == Dispatched — the old dispatch-then-mark order
1726        // could reject that submit (losing the response) or stomp its
1727        // terminal state back to Dispatched.
1728        if !self
1729            .command_registry
1730            .mark_dispatched_if_not_terminal(command_id, Utc::now())
1731            .await?
1732        {
1733            return Ok(self
1734                .command_registry
1735                .get_command_status(command_id)
1736                .await?
1737                .map(|status| status.state)
1738                .unwrap_or(CommandState::Dispatched));
1739        }
1740
1741        // A definite pre-delivery rejection (connection refusal, request
1742        // builder failure, or any HTTP status other than the runtime's exact
1743        // 202 acceptance) is safe to record as terminal DELIVERY_FAILED.
1744        // Other transport errors are ambiguous: the target may have accepted
1745        // the envelope before the response was lost. Keep those Dispatched so
1746        // a late response remains valid and return the durable ID for polling;
1747        // reverting/retrying could execute the command twice.
1748        if let Err(error) = self.command_dispatcher.dispatch(&envelope).await {
1749            if is_definite_dispatch_rejection(&error) {
1750                let delivery_failure = CommandResponse::error(
1751                    "DELIVERY_FAILED",
1752                    "Worker runtime did not accept command delivery",
1753                );
1754                self.submit_command_response(command_id, delivery_failure)
1755                    .await?;
1756                warn!(
1757                    command_id,
1758                    deployment_id,
1759                    error = %error,
1760                    "Push dispatch was definitely rejected; command marked Failed"
1761                );
1762                return Ok(CommandState::Failed);
1763            }
1764
1765            let error = error.context(ErrorData::TransportDispatchFailed {
1766                message: "Failed to dispatch command".to_string(),
1767                transport_type: None,
1768                target: Some(deployment_id.to_string()),
1769            });
1770            warn!(
1771                command_id,
1772                deployment_id,
1773                error = %error,
1774                "Push dispatch acknowledgement failed; command remains Dispatched for a possible late response"
1775            );
1776            return Ok(CommandState::Dispatched);
1777        }
1778
1779        info!("Command {} dispatched via push", envelope.command_id);
1780        Ok(CommandState::Dispatched)
1781    }
1782
1783    async fn build_envelope(
1784        &self,
1785        command_id: &str,
1786        metadata: &CommandEnvelopeData,
1787        mut params: BodySpec,
1788    ) -> Result<Envelope> {
1789        let response_handling = self.create_response_handling(command_id).await?;
1790
1791        // Re-inline: if params are in blob but fit in transport limit, read and embed inline.
1792        // This avoids unnecessary storage downloads for medium-sized params (18KB–150KB).
1793        if let BodySpec::Storage { size, .. } = &params {
1794            let raw_size = size.unwrap_or(0) as usize;
1795            if raw_size > 0 && raw_size <= self.inline_max_bytes {
1796                let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1797                match self.storage.get(&blob_path).await {
1798                    Ok(get_result) => match get_result.bytes().await {
1799                        Ok(raw_bytes) => {
1800                            params = BodySpec::inline(&raw_bytes);
1801                            debug!(
1802                                "Re-inlined params for command {} ({} bytes) into envelope",
1803                                command_id, raw_size
1804                            );
1805                        }
1806                        Err(e) => {
1807                            debug!(
1808                                    "Failed to read blob bytes for re-inline (command {}), falling back to presigned URL: {}",
1809                                    command_id, e
1810                                );
1811                        }
1812                    },
1813                    Err(e) => {
1814                        debug!(
1815                            "Failed to read blob for re-inline (command {}), falling back to presigned URL: {}",
1816                            command_id, e
1817                        );
1818                    }
1819                }
1820            }
1821        }
1822
1823        // If params are still Storage (either too large or re-inline failed),
1824        // always mint a fresh GET request. A command may remain Pending longer
1825        // than the URL minted at upload completion, and leasing must never hand
1826        // a Worker an already-expired storage credential.
1827        if let BodySpec::Storage {
1828            size,
1829            storage_get_request: _,
1830            storage_put_used,
1831        } = &params
1832        {
1833            let get_request = self.generate_storage_get_request(command_id).await?;
1834            params = BodySpec::Storage {
1835                size: *size,
1836                storage_get_request: Some(get_request),
1837                storage_put_used: *storage_put_used,
1838            };
1839        }
1840
1841        Ok(Envelope::new(
1842            metadata.deployment_id.clone(),
1843            metadata.target.clone(),
1844            command_id.to_string(),
1845            metadata.attempt,
1846            metadata.deadline,
1847            metadata.command.clone(),
1848            params,
1849            response_handling,
1850        ))
1851    }
1852
1853    async fn create_response_handling(&self, command_id: &str) -> Result<ResponseHandling> {
1854        let upload_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1855        let expires_in = Duration::from_secs(Self::RESPONSE_CREDENTIAL_LIFETIME_SECS);
1856        let presigned = self
1857            .storage
1858            .presigned_put(&upload_path, expires_in)
1859            .await
1860            .context(ErrorData::StorageOperationFailed {
1861                message: "Failed to create response upload URL".to_string(),
1862                operation: Some("presigned_put".to_string()),
1863                path: Some(upload_path.to_string()),
1864            })?;
1865
1866        let (response_token, expires) = self.sign_response_url(command_id);
1867
1868        Ok(ResponseHandling {
1869            max_inline_bytes: self.inline_max_bytes as u64,
1870            submit_response_url: format!(
1871                "{}/commands/{}/response?response_token={}&expires={}",
1872                self.base_url.trim_end_matches('/'),
1873                command_id,
1874                response_token,
1875                expires,
1876            ),
1877            storage_upload_request: presigned,
1878        })
1879    }
1880
1881    async fn generate_params_upload(&self, command_id: &str) -> Result<StorageUpload> {
1882        let upload_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1883        let expires_in = Duration::from_secs(3600);
1884        let presigned = self
1885            .storage
1886            .presigned_put(&upload_path, expires_in)
1887            .await
1888            .into_alien_error()
1889            .context(ErrorData::StorageOperationFailed {
1890                message: "Failed to create presigned URL".to_string(),
1891                operation: Some("presigned_put".to_string()),
1892                path: Some(upload_path.to_string()),
1893            })?;
1894
1895        Ok(StorageUpload {
1896            put_request: presigned.clone(),
1897            expires_at: presigned.expiration,
1898        })
1899    }
1900
1901    async fn generate_storage_get_request(&self, command_id: &str) -> Result<PresignedRequest> {
1902        let path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1903        let expires_in = Duration::from_secs(3600);
1904        self.storage.presigned_get(&path, expires_in).await.context(
1905            ErrorData::StorageOperationFailed {
1906                message: "Failed to create storage get request".to_string(),
1907                operation: Some("presigned_get".to_string()),
1908                path: Some(path.to_string()),
1909            },
1910        )
1911    }
1912
1913    async fn generate_response_storage_get_request(
1914        &self,
1915        command_id: &str,
1916    ) -> Result<PresignedRequest> {
1917        let path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1918        let expires_in = Duration::from_secs(3600);
1919        self.storage.presigned_get(&path, expires_in).await.context(
1920            ErrorData::StorageOperationFailed {
1921                message: "Failed to create response storage get request".to_string(),
1922                operation: Some("presigned_get".to_string()),
1923                path: Some(path.to_string()),
1924            },
1925        )
1926    }
1927
1928    fn extract_command_id_from_index_key(&self, index_key: &str) -> Result<String> {
1929        index_key
1930            .split(':')
1931            .last()
1932            .ok_or_else(|| {
1933                AlienError::new(ErrorData::Other {
1934                    message: format!("Invalid index key format: {}", index_key),
1935                })
1936            })
1937            .map(|s| s.to_string())
1938    }
1939}
1940
1941#[cfg(test)]
1942mod relative_url_tests {
1943    use super::*;
1944    use alien_core::presigned::PresignedOperation;
1945    use std::collections::HashMap;
1946
1947    fn http_request(url: &str, operation: PresignedOperation) -> PresignedRequest {
1948        PresignedRequest::new_http(
1949            url.to_string(),
1950            match operation {
1951                PresignedOperation::Get => "GET",
1952                PresignedOperation::Put => "PUT",
1953                PresignedOperation::Delete => "DELETE",
1954            }
1955            .to_string(),
1956            HashMap::new(),
1957            operation,
1958            "commands/test".to_string(),
1959            Utc::now() + chrono::Duration::minutes(5),
1960        )
1961    }
1962
1963    #[test]
1964    fn leased_manager_urls_are_relative_to_lease_endpoint() {
1965        let mut envelope = Envelope::new(
1966            "deployment",
1967            CommandTarget::new("daemon", CommandTargetType::Daemon),
1968            "command",
1969            1,
1970            None,
1971            "run",
1972            BodySpec::Storage {
1973                size: Some(2048),
1974                storage_get_request: Some(http_request(
1975                    "http://manager.internal/storage/params?signature=params",
1976                    PresignedOperation::Get,
1977                )),
1978                storage_put_used: Some(true),
1979            },
1980            ResponseHandling {
1981                max_inline_bytes: 1024,
1982                submit_response_url:
1983                    "http://manager.internal/v1/commands/command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1984                        .to_string(),
1985                storage_upload_request: http_request(
1986                    "http://manager.internal/v1/storage/response?signature=DoNotCanonicalize%2FUpload",
1987                    PresignedOperation::Put,
1988                ),
1989            },
1990        );
1991
1992        CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
1993
1994        assert_eq!(
1995            envelope.response_handling.submit_response_url,
1996            "command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1997        );
1998        assert_eq!(
1999            envelope.response_handling.storage_upload_request.url(),
2000            "../storage/response?signature=DoNotCanonicalize%2FUpload"
2001        );
2002        let BodySpec::Storage {
2003            storage_get_request: Some(params),
2004            ..
2005        } = &envelope.params
2006        else {
2007            panic!("storage params request");
2008        };
2009        assert_eq!(params.url(), "../../storage/params?signature=params");
2010
2011        envelope.response_handling.storage_upload_request = http_request(
2012            "https://storage.example.com/result?signature=cloud",
2013            PresignedOperation::Put,
2014        );
2015        CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
2016        assert_eq!(
2017            envelope.response_handling.storage_upload_request.url(),
2018            "https://storage.example.com/result?signature=cloud",
2019            "cloud-presigned URLs must remain byte-for-byte absolute"
2020        );
2021    }
2022}
2023
2024#[cfg(test)]
2025mod idempotency_key_tests {
2026    use super::*;
2027    use crate::server::{validate_command_name, validate_command_target_id};
2028
2029    #[test]
2030    fn definite_dispatch_rejection_is_classified_by_typed_error_data() {
2031        let rejected = AlienError::new(ErrorData::TransportDispatchRejected {
2032            message: "not accepted".to_string(),
2033            transport_type: Some("http".to_string()),
2034            target: Some("command-id".to_string()),
2035        });
2036        assert!(is_definite_dispatch_rejection(&rejected));
2037
2038        let ambiguous = AlienError::new(ErrorData::TransportDispatchFailed {
2039            message: "acknowledgement lost".to_string(),
2040            transport_type: Some("http".to_string()),
2041            target: Some("command-id".to_string()),
2042        });
2043        assert!(!is_definite_dispatch_rejection(&ambiguous));
2044    }
2045
2046    /// Idempotency keys are `{dep}:{rid}:{command}:{key}`. If a target id could
2047    /// contain ':', the `rid` segment would be ambiguous: the two triples
2048    ///   (rid="svc",   command="a:b", key="k")
2049    ///   (rid="svc:a", command="b",   key="k")
2050    /// both compose to `dep:svc:a:b:k`. The shared guard forbids ':' in a target
2051    /// id, so the second can never be a resolved target — closing the collision
2052    /// at the rid boundary. Only the colon-free composition is exercised here;
2053    /// the colliding one is proven unreachable via the guard.
2054    #[test]
2055    fn target_id_colon_guard_prevents_idempotency_key_collision() {
2056        let colliding = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "k");
2057        assert_eq!(colliding, "dep:svc:a:b:k");
2058
2059        // The alternate triple that would collide needs target id "svc:a".
2060        assert!(
2061            validate_command_target_id("svc:a").is_err(),
2062            "a ':'-bearing target id must be rejected so it cannot forge the rid segment"
2063        );
2064        assert!(validate_command_target_id("svc").is_ok());
2065    }
2066
2067    /// The command name is the other forgeable segment. Client keys routinely
2068    /// contain ':', so without a command-name guard these two distinct inputs
2069    ///   (command="a:b", key="c")
2070    ///   (command="a",   key="b:c")
2071    /// both compose to `dep:svc:a:b:c` — a cross-command idempotency collision.
2072    /// The guard rejects the ':'-bearing command name, so only the second input
2073    /// is ever composed; the first can never reach key composition.
2074    #[test]
2075    fn command_name_colon_guard_prevents_idempotency_key_collision() {
2076        // Without the guard, both inputs compose to the identical string.
2077        let forged = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "c");
2078        let legitimate = CommandServer::compose_idempotency_key("dep", "svc", "a", "b:c");
2079        assert_eq!(
2080            forged, legitimate,
2081            "these inputs are exactly the colliding pair the guard must separate"
2082        );
2083        assert_eq!(legitimate, "dep:svc:a:b:c");
2084
2085        // The guard makes the colliding input unreachable: a ':'-bearing command
2086        // name is rejected, while the legitimate command name is accepted. With
2087        // the forger blocked, `a:b`+`c` can never compose the shared key — only
2088        // the distinct `a`+`b:c` command can.
2089        let err = validate_command_name("a:b").expect_err("':'-bearing command must be rejected");
2090        assert_eq!(err.code, "INVALID_COMMAND");
2091        assert!(validate_command_name("a").is_ok());
2092    }
2093}