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, 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 (index_key, _) in scan_result.items {
737            if leases.len() >= lease_request.max_leases {
738                break;
739            }
740
741            let command_id = self.extract_command_id_from_index_key(&index_key)?;
742
743            // 2. Try to acquire lease atomically in KV
744            let lease_id = format!("lease_{}", Uuid::new_v4());
745            let lease_duration = Duration::from_secs(lease_request.lease_seconds);
746            let expires_at =
747                Utc::now() + chrono::Duration::seconds(lease_request.lease_seconds as i64);
748
749            let lease_data = LeaseData {
750                lease_id: lease_id.clone(),
751                acquired_at: Utc::now(),
752                expires_at,
753                owner: deployment_id.to_string(),
754            };
755
756            let lease_key = format!("cmd:{}:lease", command_id);
757            let lease_value = serde_json::to_vec(&lease_data).into_alien_error().context(
758                ErrorData::SerializationFailed {
759                    message: "Failed to serialize lease data".to_string(),
760                    data_type: Some("LeaseData".to_string()),
761                },
762            )?;
763
764            let options = Some(PutOptions {
765                ttl: Some(lease_duration),
766                if_not_exists: true,
767            });
768
769            let success = self
770                .kv
771                .put(&lease_key, lease_value, options)
772                .await
773                .context(ErrorData::KvOperationFailed {
774                    operation: "put".to_string(),
775                    key: lease_key.clone(),
776                    message: "Failed to create lease".to_string(),
777                })?;
778
779            if !success {
780                // Lease already exists, skip
781                continue;
782            }
783
784            // Reverse index for O(1) release-by-lease-id lookups. Shares the
785            // lease's TTL so it self-cleans on expiry; a stale entry is
786            // harmless because every reader re-verifies the lease_id against
787            // the live `cmd:{id}:lease` record.
788            let reverse_key = format!("lease:{}", lease_id);
789            self.kv
790                .put(
791                    &reverse_key,
792                    command_id.as_bytes().to_vec(),
793                    Some(PutOptions {
794                        ttl: Some(lease_duration),
795                        if_not_exists: false,
796                    }),
797                )
798                .await
799                .context(ErrorData::KvOperationFailed {
800                    operation: "put".to_string(),
801                    key: reverse_key,
802                    message: "Failed to create lease reverse index".to_string(),
803                })?;
804
805            // 3. Get metadata from registry
806            let mut metadata = match self
807                .command_registry
808                .get_command_metadata(&command_id)
809                .await?
810            {
811                Some(m) => m,
812                None => {
813                    // Command doesn't exist in registry, clean up
814                    self.delete_lease(&command_id).await?;
815                    let _ = self.kv.delete(&index_key).await;
816                    continue;
817                }
818            };
819
820            // 3.1 Expiry-driven redelivery: this command was already
821            // dispatched, yet the lease slot was free (the conditional put
822            // above succeeded) — the previous lease TTL-expired with no
823            // response and no explicit release. Increment the attempt so the
824            // redelivered envelope carries `attempt > 1`, the at-least-once
825            // redelivery signal both receiver twins document. The explicit
826            // release path (`release_lease`) increments the same counter.
827            if metadata.state == CommandState::Dispatched {
828                self.command_registry.increment_attempt(&command_id).await?;
829                metadata.attempt += 1;
830            }
831
832            // 3.5 Defense-in-depth: the pending index key said this command
833            // belongs to the requesting target — verify the registry agrees.
834            // A mismatch means the index is corrupt; fail loudly rather than
835            // delivering a command to the wrong resource. The corrupt index
836            // key is deliberately retained (only the lease is cleaned up):
837            // genuine corruption should never occur, and the key is the
838            // evidence an operator needs — do not "fix" this by deleting it.
839            if metadata.target != lease_request.target {
840                self.delete_lease(&command_id).await?;
841                return Err(AlienError::new(ErrorData::Other {
842                    message: format!(
843                        "Pending index corruption: command '{}' is indexed under target '{}' \
844                         but the registry says it belongs to target '{}' — refusing to deliver",
845                        command_id, lease_request.target.resource_id, metadata.target.resource_id,
846                    ),
847                }));
848            }
849
850            // 4. Check if command is in terminal state (stale index)
851            if metadata.state.is_terminal() {
852                // Clean up stale data
853                self.delete_lease(&command_id).await?;
854                let _ = self.kv.delete(&index_key).await;
855                continue;
856            }
857
858            // 5. Check deadline expiry — conditionally: a racing submit's
859            // terminal state must never be stomped back to Expired.
860            if let Some(deadline) = metadata.deadline {
861                if Utc::now() > deadline {
862                    let _ = self
863                        .command_registry
864                        .complete_command(
865                            &command_id,
866                            CommandState::Expired,
867                            Utc::now(),
868                            None,
869                            None,
870                        )
871                        .await?;
872                    self.delete_lease(&command_id).await?;
873                    let _ = self.kv.delete(&index_key).await;
874                    continue;
875                }
876            }
877
878            // 6. Get params from KV
879            let params = match self.get_params(&command_id).await? {
880                Some(p) => p,
881                None => {
882                    // No params, something went wrong
883                    self.delete_lease(&command_id).await?;
884                    continue;
885                }
886            };
887
888            // 7. Mark Dispatched — conditionally. Between the terminal
889            // check in step 3.1 and here, the ORIGINAL holder of a
890            // TTL-expired lease can still submit; handing out this lease
891            // would re-execute a completed command and stomp its state.
892            let dispatched = self
893                .command_registry
894                .mark_dispatched_if_not_terminal(&command_id, Utc::now())
895                .await?;
896            if !dispatched {
897                debug!(
898                    command_id = %command_id,
899                    "Command turned terminal while leasing; releasing the lease"
900                );
901                self.delete_lease(&command_id).await?;
902                let _ = self.kv.delete(&index_key).await;
903                continue;
904            }
905
906            // 8. Build envelope. Lease-served envelopes carry manager URLs
907            // as path-relative references from the lease endpoint: a pull
908            // consumer resolves them against the exact endpoint it reached,
909            // preserving both a network-corrected origin and any reverse-
910            // proxy path prefix the manager cannot know. Push envelopes keep
911            // absolute URLs: push transports have no configured base, and
912            // reaching the manager's public address is inherent to push
913            // delivery.
914            let mut envelope = self.build_envelope(&command_id, &metadata, params).await?;
915            Self::relativize_manager_urls(&mut envelope, &self.base_url);
916
917            leases.push(LeaseInfo {
918                lease_id,
919                lease_expires_at: expires_at,
920                command_id: command_id.clone(),
921                attempt: metadata.attempt,
922                envelope,
923            });
924        }
925
926        Ok(LeaseResponse { leases })
927    }
928
929    /// Rewrite manager-origin URLs in a lease-served envelope to references
930    /// relative to that lease endpoint (see [`Self::acquire_lease`]). Only
931    /// same-origin URLs can be made relative, so cloud-presigned storage URLs
932    /// pass through byte-for-byte.
933    fn relativize_manager_urls(envelope: &mut Envelope, base_url: &str) {
934        let Ok(mut lease_endpoint) = reqwest::Url::parse(base_url) else {
935            // An unparseable base cannot produce relative references; leave
936            // the envelope absolute (the pre-relative behavior).
937            return;
938        };
939        let Ok(mut segments) = lease_endpoint.path_segments_mut() else {
940            return;
941        };
942        segments.pop_if_empty().push("commands").push("leases");
943        drop(segments);
944        lease_endpoint.set_query(None);
945        lease_endpoint.set_fragment(None);
946
947        let relativize = |target: &mut String| {
948            let Ok(mut target_url) = reqwest::Url::parse(target.as_str()) else {
949                return;
950            };
951            let suffix = target
952                .find(['?', '#'])
953                .map(|index| target[index..].to_string())
954                .unwrap_or_default();
955            target_url.set_query(None);
956            target_url.set_fragment(None);
957            if let Some(relative) = lease_endpoint.make_relative(&target_url) {
958                *target = format!("{relative}{suffix}");
959            }
960        };
961
962        relativize(&mut envelope.response_handling.submit_response_url);
963        if let alien_core::presigned::PresignedRequestBackend::Http { url, .. } =
964            &mut envelope.response_handling.storage_upload_request.backend
965        {
966            relativize(url);
967        }
968        if let BodySpec::Storage {
969            storage_get_request: Some(request),
970            ..
971        } = &mut envelope.params
972        {
973            if let alien_core::presigned::PresignedRequestBackend::Http { url, .. } =
974                &mut request.backend
975            {
976                relativize(url);
977            }
978        }
979    }
980
981    /// Release a lease manually.
982    ///
983    /// Increments attempt count in registry, returns command to Pending state.
984    pub async fn release_lease(&self, command_id: &str, lease_id: &str) -> Result<()> {
985        let lease_key = format!("cmd:{}:lease", command_id);
986
987        // 1. Verify lease ownership
988        if let Ok(Some(lease_data)) = self.kv.get(&lease_key).await {
989            let lease: LeaseData = serde_json::from_slice(&lease_data)
990                .into_alien_error()
991                .context(ErrorData::SerializationFailed {
992                    message: "Failed to deserialize lease data".to_string(),
993                    data_type: Some("LeaseData".to_string()),
994                })?;
995
996            if lease.lease_id != lease_id {
997                return Err(AlienError::new(ErrorData::LeaseNotFound {
998                    lease_id: lease_id.to_string(),
999                }));
1000            }
1001
1002            // 2. Delete lease from KV (and its reverse index)
1003            self.delete_lease(command_id).await?;
1004            let _ = self.kv.delete(&format!("lease:{}", lease_id)).await;
1005
1006            // 3. Return the command to Pending only while it is still
1007            // non-terminal. A racing response/deadline completion wins.
1008            let released = self
1009                .command_registry
1010                .update_command_state(command_id, CommandState::Pending, None, None, None, None)
1011                .await?;
1012            if released {
1013                // 4. Increment the attempt only for a command that was
1014                // actually released for redelivery.
1015                self.command_registry.increment_attempt(command_id).await?;
1016            }
1017
1018            // Note: Pending index is NOT removed on lease, so command is still there
1019            debug!(
1020                released,
1021                "Lease {} released for command {}", lease_id, command_id
1022            );
1023        }
1024
1025        Ok(())
1026    }
1027
1028    /// Get the deployment_id that owns a command.
1029    ///
1030    /// Used by the manager's auth layer to check whether the caller has access
1031    /// to a specific command without fetching the full status.
1032    pub async fn get_command_deployment_id(&self, command_id: &str) -> Result<Option<String>> {
1033        let status = self.command_registry.get_command_status(command_id).await?;
1034        Ok(status.map(|s| s.deployment_id))
1035    }
1036
1037    /// Get the command's canonical ownership fields for route authorization.
1038    pub async fn get_command_access_context(
1039        &self,
1040        command_id: &str,
1041    ) -> Result<Option<CommandAccessContext>> {
1042        self.command_registry
1043            .get_command_access_context(command_id)
1044            .await
1045    }
1046
1047    /// Resolve a lease_id to `(command_id, owner_deployment_id)` via the
1048    /// `lease:{lease_id}` reverse index, re-verifying against the live lease
1049    /// record (the reverse entry can outlive a released lease briefly, and a
1050    /// command can have been re-leased under a new lease_id since).
1051    ///
1052    /// Used by the manager's auth layer to check that the caller may act on
1053    /// the deployment that holds the lease.
1054    pub async fn get_lease_owner(&self, lease_id: &str) -> Result<Option<(String, String)>> {
1055        let reverse_key = format!("lease:{}", lease_id);
1056        let Some(command_id_bytes) =
1057            self.kv
1058                .get(&reverse_key)
1059                .await
1060                .context(ErrorData::KvOperationFailed {
1061                    operation: "get".to_string(),
1062                    key: reverse_key.clone(),
1063                    message: "Failed to look up lease reverse index".to_string(),
1064                })?
1065        else {
1066            return Ok(None);
1067        };
1068        let command_id = String::from_utf8(command_id_bytes).map_err(|_| {
1069            AlienError::new(ErrorData::Other {
1070                message: format!("Lease reverse index '{}' is not valid UTF-8", reverse_key),
1071            })
1072        })?;
1073
1074        let lease_key = format!("cmd:{}:lease", command_id);
1075        let Some(lease_data) =
1076            self.kv
1077                .get(&lease_key)
1078                .await
1079                .context(ErrorData::KvOperationFailed {
1080                    operation: "get".to_string(),
1081                    key: lease_key,
1082                    message: "Failed to look up lease".to_string(),
1083                })?
1084        else {
1085            return Ok(None);
1086        };
1087        let lease: LeaseData = serde_json::from_slice(&lease_data)
1088            .into_alien_error()
1089            .context(ErrorData::SerializationFailed {
1090                message: "Failed to deserialize lease data".to_string(),
1091                data_type: Some("LeaseData".to_string()),
1092            })?;
1093        if lease.lease_id != lease_id {
1094            return Ok(None);
1095        }
1096
1097        Ok(Some((command_id, lease.owner)))
1098    }
1099
1100    /// Release a lease by lease_id only (for the API).
1101    pub async fn release_lease_by_id(&self, lease_id: &str) -> Result<()> {
1102        match self.get_lease_owner(lease_id).await? {
1103            Some((command_id, _owner)) => self.release_lease(&command_id, lease_id).await,
1104            None => Err(AlienError::new(ErrorData::LeaseNotFound {
1105                lease_id: lease_id.to_string(),
1106            })),
1107        }
1108    }
1109
1110    // =========================================================================
1111    // Internal Helper Methods
1112    // =========================================================================
1113
1114    async fn validate_create_command(&self, request: &CreateCommandRequest) -> Result<()> {
1115        if request.command.is_empty() {
1116            return Err(AlienError::new(ErrorData::InvalidCommand {
1117                message: "Command name cannot be empty".to_string(),
1118            }));
1119        }
1120
1121        // The command name occupies one segment of the `:`-delimited
1122        // idempotency key (`{dep}:{rid}:{command}:{key}`). A ':' in the name
1123        // would let it bleed into the client-key segment — which routinely
1124        // contains ':' — so (command="a:b", key="c") and (command="a",
1125        // key="b:c") would forge the same key and be treated as the same
1126        // command. Reject ':' here, mirroring the target-id colon guard, so the
1127        // command segment is always unambiguous.
1128        validate_command_name(&request.command)?;
1129
1130        if request.deployment_id.is_empty() {
1131            return Err(AlienError::new(ErrorData::InvalidCommand {
1132                message: "Deployment ID cannot be empty".to_string(),
1133            }));
1134        }
1135
1136        if let Some(deadline) = request.deadline {
1137            if deadline <= Utc::now() {
1138                return Err(AlienError::new(ErrorData::InvalidCommand {
1139                    message: "Deadline must be in the future".to_string(),
1140                }));
1141            }
1142        }
1143
1144        Ok(())
1145    }
1146
1147    // --- Idempotency ---
1148
1149    /// Compose the target-scoped idempotency key:
1150    /// `{deploymentId}:{targetResourceId}:{commandName}:{key}`.
1151    ///
1152    /// Scoping by target means the same client key addressed to two different
1153    /// targets creates two distinct commands.
1154    fn compose_idempotency_key(
1155        deployment_id: &str,
1156        target_resource_id: &str,
1157        command_name: &str,
1158        idem_key: &str,
1159    ) -> String {
1160        // Invariant: at every real call site the target id and command name are
1161        // both `:`-free (the former enforced at resolution via
1162        // `validate_command_target_id`, the latter via `validate_command_name`
1163        // in `validate_create_command`), so each occupies exactly one segment of
1164        // `{dep}:{rid}:{command}:{key}` and only the trailing client key may
1165        // carry ':'. Those guards live upstream; this is a pure formatter, so it
1166        // is not asserted on `command_name` here (the collision tests below
1167        // deliberately format a ':'-bearing command to demonstrate what the
1168        // upstream guard prevents).
1169        debug_assert!(
1170            !target_resource_id.contains(':'),
1171            "target_resource_id must be ':'-free before key composition: {target_resource_id}"
1172        );
1173        format!(
1174            "{}:{}:{}:{}",
1175            deployment_id, target_resource_id, command_name, idem_key
1176        )
1177    }
1178
1179    async fn check_idempotency(&self, idem_key: &str) -> Result<Option<String>> {
1180        let key = format!("idem:{}", idem_key);
1181        if let Some(data) = self
1182            .kv
1183            .get(&key)
1184            .await
1185            .context(ErrorData::KvOperationFailed {
1186                operation: "get".to_string(),
1187                key: key.clone(),
1188                message: "Failed to check idempotency".to_string(),
1189            })?
1190        {
1191            let command_id = String::from_utf8(data).into_alien_error().context(
1192                ErrorData::SerializationFailed {
1193                    message: "Invalid idempotency data".to_string(),
1194                    data_type: Some("String".to_string()),
1195                },
1196            )?;
1197            return Ok(Some(command_id));
1198        }
1199        Ok(None)
1200    }
1201
1202    /// Claim the idempotency key for `command_id`.
1203    ///
1204    /// Returns `None` when this command won the key. Returns
1205    /// `Some(winner_id)` when a concurrent create with the same key won the
1206    /// conditional put first — both requests passed the pre-create
1207    /// `check_idempotency` before either stored, so the loser must be
1208    /// detected here, after its command was already created.
1209    async fn store_idempotency(&self, idem_key: &str, command_id: &str) -> Result<Option<String>> {
1210        let key = format!("idem:{}", idem_key);
1211        let ttl = Duration::from_secs(24 * 60 * 60); // 24 hours
1212        let won = self
1213            .kv
1214            .put(
1215                &key,
1216                command_id.as_bytes().to_vec(),
1217                Some(PutOptions {
1218                    ttl: Some(ttl),
1219                    if_not_exists: true,
1220                }),
1221            )
1222            .await
1223            .context(ErrorData::KvOperationFailed {
1224                operation: "put".to_string(),
1225                key: key.clone(),
1226                message: "Failed to store idempotency".to_string(),
1227            })?;
1228        if won {
1229            return Ok(None);
1230        }
1231        // Lost the conditional put: read back who won. The winner entry can
1232        // only be absent if it TTL-expired in the microseconds since — treat
1233        // that as an inconsistency rather than silently duplicating.
1234        match self.check_idempotency(idem_key).await? {
1235            Some(winner_id) => Ok(Some(winner_id)),
1236            None => Err(AlienError::new(ErrorData::Other {
1237                message: format!(
1238                    "Idempotency key '{}' was concurrently claimed but has no winner entry",
1239                    key
1240                ),
1241            })),
1242        }
1243    }
1244
1245    // --- Params ---
1246
1247    pub async fn store_params(&self, command_id: &str, params: &BodySpec) -> Result<()> {
1248        let key = format!("cmd:{}:params", command_id);
1249
1250        // Try serializing as-is first
1251        let data = CommandParamsData {
1252            params: params.clone(),
1253        };
1254        let value = serde_json::to_vec(&data).into_alien_error().context(
1255            ErrorData::SerializationFailed {
1256                message: "Failed to serialize params".to_string(),
1257                data_type: Some("CommandParamsData".to_string()),
1258            },
1259        )?;
1260
1261        // If it fits in KV, store directly (fast path)
1262        if value.len() <= KV_VALUE_THRESHOLD {
1263            self.kv
1264                .put(&key, value, None)
1265                .await
1266                .context(ErrorData::KvOperationFailed {
1267                    operation: "put".to_string(),
1268                    key: key.clone(),
1269                    message: "Failed to store params".to_string(),
1270                })?;
1271            return Ok(());
1272        }
1273
1274        // Auto-promote: inline data exceeds KV limit, store raw bytes in blob
1275        if let BodySpec::Inline { inline_base64 } = params {
1276            let raw_bytes = general_purpose::STANDARD
1277                .decode(inline_base64)
1278                .into_alien_error()
1279                .context(ErrorData::SerializationFailed {
1280                    message: "Failed to decode inline base64 params for auto-promotion".to_string(),
1281                    data_type: Some("base64".to_string()),
1282                })?;
1283
1284            let raw_len = raw_bytes.len() as u64;
1285            let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1286
1287            self.storage
1288                .put(&blob_path, Bytes::from(raw_bytes).into())
1289                .await
1290                .into_alien_error()
1291                .context(ErrorData::StorageOperationFailed {
1292                    message: "Failed to auto-promote params to blob storage".to_string(),
1293                    operation: Some("put".to_string()),
1294                    path: Some(blob_path.to_string()),
1295                })?;
1296
1297            debug!(
1298                "Auto-promoted params for command {} to blob ({} bytes raw)",
1299                command_id, raw_len
1300            );
1301
1302            // Store tiny reference in KV instead
1303            let promoted = CommandParamsData {
1304                params: BodySpec::Storage {
1305                    size: Some(raw_len),
1306                    storage_get_request: None,
1307                    storage_put_used: Some(true),
1308                },
1309            };
1310            let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1311                ErrorData::SerializationFailed {
1312                    message: "Failed to serialize promoted params reference".to_string(),
1313                    data_type: Some("CommandParamsData".to_string()),
1314                },
1315            )?;
1316            self.kv.put(&key, promoted_value, None).await.context(
1317                ErrorData::KvOperationFailed {
1318                    operation: "put".to_string(),
1319                    key: key.clone(),
1320                    message: "Failed to store promoted params reference".to_string(),
1321                },
1322            )?;
1323            return Ok(());
1324        }
1325
1326        // Storage references are always tiny, store as-is
1327        self.kv
1328            .put(&key, value, None)
1329            .await
1330            .context(ErrorData::KvOperationFailed {
1331                operation: "put".to_string(),
1332                key: key.clone(),
1333                message: "Failed to store params".to_string(),
1334            })?;
1335        Ok(())
1336    }
1337
1338    pub async fn get_params(&self, command_id: &str) -> Result<Option<BodySpec>> {
1339        let key = format!("cmd:{}:params", command_id);
1340        if let Some(value) = self
1341            .kv
1342            .get(&key)
1343            .await
1344            .context(ErrorData::KvOperationFailed {
1345                operation: "get".to_string(),
1346                key: key.clone(),
1347                message: "Failed to get params".to_string(),
1348            })?
1349        {
1350            let data: CommandParamsData = serde_json::from_slice(&value)
1351                .into_alien_error()
1352                .context(ErrorData::SerializationFailed {
1353                    message: "Failed to deserialize params".to_string(),
1354                    data_type: Some("CommandParamsData".to_string()),
1355                })?;
1356            return Ok(Some(data.params));
1357        }
1358        Ok(None)
1359    }
1360
1361    // --- Response ---
1362
1363    pub async fn store_response(&self, command_id: &str, response: &CommandResponse) -> Result<()> {
1364        let key = format!("cmd:{}:response", command_id);
1365        let data = CommandResponseData {
1366            response: response.clone(),
1367        };
1368        let value = serde_json::to_vec(&data).into_alien_error().context(
1369            ErrorData::SerializationFailed {
1370                message: "Failed to serialize response".to_string(),
1371                data_type: Some("CommandResponseData".to_string()),
1372            },
1373        )?;
1374
1375        // If it fits in KV, store directly (fast path)
1376        if value.len() <= KV_VALUE_THRESHOLD {
1377            self.kv
1378                .put(&key, value, None)
1379                .await
1380                .context(ErrorData::KvOperationFailed {
1381                    operation: "put".to_string(),
1382                    key: key.clone(),
1383                    message: "Failed to store response".to_string(),
1384                })?;
1385            return Ok(());
1386        }
1387
1388        // Auto-promote: inline response exceeds KV limit
1389        if let CommandResponse::Success {
1390            response: BodySpec::Inline { inline_base64 },
1391        } = response
1392        {
1393            let raw_bytes = general_purpose::STANDARD
1394                .decode(inline_base64)
1395                .into_alien_error()
1396                .context(ErrorData::SerializationFailed {
1397                    message: "Failed to decode inline base64 response for auto-promotion"
1398                        .to_string(),
1399                    data_type: Some("base64".to_string()),
1400                })?;
1401
1402            let raw_len = raw_bytes.len() as u64;
1403            let blob_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1404
1405            self.storage
1406                .put(&blob_path, Bytes::from(raw_bytes).into())
1407                .await
1408                .into_alien_error()
1409                .context(ErrorData::StorageOperationFailed {
1410                    message: "Failed to auto-promote response to blob storage".to_string(),
1411                    operation: Some("put".to_string()),
1412                    path: Some(blob_path.to_string()),
1413                })?;
1414
1415            // Generate presigned GET URL for the caller
1416            let get_request = self
1417                .generate_response_storage_get_request(command_id)
1418                .await?;
1419
1420            debug!(
1421                "Auto-promoted response for command {} to blob ({} bytes raw)",
1422                command_id, raw_len
1423            );
1424
1425            // Store tiny reference in KV
1426            let promoted = CommandResponseData {
1427                response: CommandResponse::Success {
1428                    response: BodySpec::Storage {
1429                        size: Some(raw_len),
1430                        storage_get_request: Some(get_request),
1431                        storage_put_used: Some(true),
1432                    },
1433                },
1434            };
1435            let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1436                ErrorData::SerializationFailed {
1437                    message: "Failed to serialize promoted response reference".to_string(),
1438                    data_type: Some("CommandResponseData".to_string()),
1439                },
1440            )?;
1441            self.kv.put(&key, promoted_value, None).await.context(
1442                ErrorData::KvOperationFailed {
1443                    operation: "put".to_string(),
1444                    key: key.clone(),
1445                    message: "Failed to store promoted response reference".to_string(),
1446                },
1447            )?;
1448            return Ok(());
1449        }
1450
1451        // Error responses or storage references are always small, store as-is
1452        self.kv
1453            .put(&key, value, None)
1454            .await
1455            .context(ErrorData::KvOperationFailed {
1456                operation: "put".to_string(),
1457                key: key.clone(),
1458                message: "Failed to store response".to_string(),
1459            })?;
1460        Ok(())
1461    }
1462
1463    pub async fn get_response(&self, command_id: &str) -> Result<Option<CommandResponse>> {
1464        let key = format!("cmd:{}:response", command_id);
1465        if let Some(value) = self
1466            .kv
1467            .get(&key)
1468            .await
1469            .context(ErrorData::KvOperationFailed {
1470                operation: "get".to_string(),
1471                key: key.clone(),
1472                message: "Failed to get response".to_string(),
1473            })?
1474        {
1475            let data: CommandResponseData = serde_json::from_slice(&value)
1476                .into_alien_error()
1477                .context(ErrorData::SerializationFailed {
1478                    message: "Failed to deserialize response".to_string(),
1479                    data_type: Some("CommandResponseData".to_string()),
1480                })?;
1481            return Ok(Some(data.response));
1482        }
1483        Ok(None)
1484    }
1485
1486    // --- Pending Index ---
1487
1488    async fn create_pending_index(
1489        &self,
1490        deployment_id: &str,
1491        target_resource_id: &str,
1492        command_id: &str,
1493    ) -> Result<()> {
1494        // Invariant: the target id is `:`-free (enforced at resolution via
1495        // `validate_command_target_id`), so its prefix `target:{dep}:{rid}:`
1496        // cannot overlap another target's pending keys.
1497        debug_assert!(
1498            !target_resource_id.contains(':'),
1499            "target_resource_id must be ':'-free in the pending index: {target_resource_id}"
1500        );
1501        let timestamp = Utc::now().timestamp_nanos_opt().unwrap_or(0);
1502        let key = format!(
1503            "target:{}:{}:pending:{}:{}",
1504            deployment_id, target_resource_id, timestamp, command_id
1505        );
1506
1507        // Store empty value - just for ordering
1508        self.kv
1509            .put(&key, vec![], None)
1510            .await
1511            .context(ErrorData::KvOperationFailed {
1512                operation: "put".to_string(),
1513                key: key.clone(),
1514                message: "Failed to create pending index".to_string(),
1515            })?;
1516        Ok(())
1517    }
1518
1519    async fn delete_pending_index(
1520        &self,
1521        deployment_id: &str,
1522        target_resource_id: &str,
1523        command_id: &str,
1524    ) -> Result<()> {
1525        // We need to scan to find the exact key since we don't know the timestamp
1526        let prefix = format!("target:{}:{}:pending:", deployment_id, target_resource_id);
1527        let scan_result = self
1528            .kv
1529            .scan_prefix(&prefix, Some(100), None)
1530            .await
1531            .into_alien_error()
1532            .context(ErrorData::KvOperationFailed {
1533                operation: "scan_prefix".to_string(),
1534                key: prefix.clone(),
1535                message: "Failed to scan pending index".to_string(),
1536            })?;
1537
1538        for (key, _) in scan_result.items {
1539            if key.ends_with(&format!(":{}", command_id)) {
1540                let _ = self.kv.delete(&key).await;
1541                break;
1542            }
1543        }
1544        Ok(())
1545    }
1546
1547    // --- Lease ---
1548
1549    async fn delete_lease(&self, command_id: &str) -> Result<()> {
1550        let key = format!("cmd:{}:lease", command_id);
1551        let _ = self.kv.delete(&key).await;
1552        Ok(())
1553    }
1554
1555    /// Expire every overdue non-terminal command recorded in the deadline
1556    /// index. Intended to run periodically from the hosting process.
1557    ///
1558    /// Deadlines are otherwise only enforced lazily (status polls and lease
1559    /// scans), which never reaches a command nobody polls — most notably a
1560    /// `PendingUpload` whose params upload never completed, which has no
1561    /// pending-index entry and would otherwise live forever.
1562    ///
1563    /// Uses the conditional terminal transition, so racing a concurrent
1564    /// submit is safe: whoever wins, the record stays consistent. Returns
1565    /// the number of commands expired.
1566    pub async fn reap_expired_commands(&self) -> Result<u32> {
1567        let now = Utc::now();
1568        let mut expired = 0u32;
1569        let mut cursor: Option<String> = None;
1570        // Paginate the whole index (bounded page count as a runaway guard):
1571        // keys are not numerically ordered — the timestamp segment isn't
1572        // zero-padded — so due entries can sit behind any number of
1573        // future-dated ones and a single capped scan would starve them.
1574        for _ in 0..64 {
1575            let scan = self
1576                .kv
1577                .scan_prefix("deadline:", Some(256), cursor.clone())
1578                .await
1579                .into_alien_error()
1580                .context(ErrorData::KvOperationFailed {
1581                    operation: "scan_prefix".to_string(),
1582                    key: "deadline:".to_string(),
1583                    message: "Failed to scan the deadline index".to_string(),
1584                })?;
1585            let next_cursor = scan.next_cursor.clone();
1586            for (key, value) in scan.items {
1587                let Ok(data) = serde_json::from_slice::<DeadlineIndexData>(&value) else {
1588                    warn!(key = %key, "Unparseable deadline index entry; deleting");
1589                    let _ = self.kv.delete(&key).await;
1590                    continue;
1591                };
1592                if data.deadline > now {
1593                    // Not due yet. (Keys are not sortable numerically — the
1594                    // timestamp segment isn't zero-padded — so keep scanning.)
1595                    continue;
1596                }
1597
1598                let status = self
1599                    .command_registry
1600                    .get_command_status(&data.command_id)
1601                    .await?;
1602                match status {
1603                    None => {
1604                        let _ = self.kv.delete(&key).await;
1605                    }
1606                    Some(status) if status.state.is_terminal() => {
1607                        let _ = self.kv.delete(&key).await;
1608                    }
1609                    Some(status) => {
1610                        let won = self
1611                        .command_registry
1612                        .complete_command(
1613                            &data.command_id,
1614                            CommandState::Expired,
1615                            now,
1616                            None,
1617                            Some(serde_json::json!({
1618                                "code": "COMMAND_EXPIRED",
1619                                "message": format!("Deadline {} elapsed", data.deadline.to_rfc3339()),
1620                            })),
1621                        )
1622                        .await?;
1623                        if won {
1624                            expired += 1;
1625                            info!(command_id = %data.command_id, "Expired overdue command");
1626                            let _ = self.delete_lease(&data.command_id).await;
1627                            let _ = self
1628                                .delete_pending_index(
1629                                    &status.deployment_id,
1630                                    &status.target.resource_id,
1631                                    &data.command_id,
1632                                )
1633                                .await;
1634                        }
1635                        let _ = self.kv.delete(&key).await;
1636                    }
1637                }
1638            }
1639            match next_cursor {
1640                Some(next) => cursor = Some(next),
1641                None => break,
1642            }
1643        }
1644        Ok(expired)
1645    }
1646
1647    // --- Deadline Index ---
1648
1649    async fn create_deadline_index(&self, command_id: &str, deadline: DateTime<Utc>) -> Result<()> {
1650        let key = format!(
1651            "deadline:{}:{}",
1652            deadline.timestamp_nanos_opt().unwrap_or(0),
1653            command_id
1654        );
1655
1656        let data = DeadlineIndexData {
1657            command_id: command_id.to_string(),
1658            deadline,
1659        };
1660        let value = serde_json::to_vec(&data).into_alien_error().context(
1661            ErrorData::SerializationFailed {
1662                message: "Failed to serialize deadline index".to_string(),
1663                data_type: Some("DeadlineIndexData".to_string()),
1664            },
1665        )?;
1666
1667        // The entry must remain VISIBLE well past the deadline: scan_prefix
1668        // treats logically-expired keys as absent on every provider, so a
1669        // TTL equal to the deadline would hide the entry at exactly the
1670        // moment the reaper needs it (the bug that made the reaper inert).
1671        // The reaper deletes entries as it processes them; the 7-day grace
1672        // is only self-cleaning for processes that never run a reaper.
1673        const DEADLINE_INDEX_GRACE: chrono::Duration = chrono::Duration::days(7);
1674        let ttl = deadline
1675            .signed_duration_since(Utc::now())
1676            .checked_add(&DEADLINE_INDEX_GRACE)
1677            .unwrap_or(DEADLINE_INDEX_GRACE);
1678        let options = (ttl.num_seconds() > 0).then(|| PutOptions {
1679            ttl: Some(Duration::from_secs(ttl.num_seconds() as u64)),
1680            if_not_exists: false,
1681        });
1682
1683        self.kv
1684            .put(&key, value, options)
1685            .await
1686            .context(ErrorData::KvOperationFailed {
1687                operation: "put".to_string(),
1688                key: key.clone(),
1689                message: "Failed to create deadline index".to_string(),
1690            })?;
1691        Ok(())
1692    }
1693
1694    // --- Dispatch ---
1695
1696    async fn dispatch_command_push(
1697        &self,
1698        command_id: &str,
1699        deployment_id: &str,
1700    ) -> Result<CommandState> {
1701        // Get metadata from registry
1702        let metadata = self
1703            .command_registry
1704            .get_command_metadata(command_id)
1705            .await?
1706            .ok_or_else(|| {
1707                AlienError::new(ErrorData::CommandNotFound {
1708                    command_id: command_id.to_string(),
1709                })
1710            })?;
1711
1712        // Get params from KV
1713        let params = self.get_params(command_id).await?.ok_or_else(|| {
1714            AlienError::new(ErrorData::CommandNotFound {
1715                command_id: command_id.to_string(),
1716            })
1717        })?;
1718
1719        // Build envelope
1720        let envelope = self.build_envelope(command_id, &metadata, params).await?;
1721
1722        // Mark Dispatched BEFORE invoking the transport: a fast worker can
1723        // execute and submit before this function resumes, and submit
1724        // validates state == Dispatched — the old dispatch-then-mark order
1725        // could reject that submit (losing the response) or stomp its
1726        // terminal state back to Dispatched.
1727        if !self
1728            .command_registry
1729            .mark_dispatched_if_not_terminal(command_id, Utc::now())
1730            .await?
1731        {
1732            return Ok(self
1733                .command_registry
1734                .get_command_status(command_id)
1735                .await?
1736                .map(|status| status.state)
1737                .unwrap_or(CommandState::Dispatched));
1738        }
1739
1740        // A definite pre-delivery rejection (connection refusal, request
1741        // builder failure, or any HTTP status other than the runtime's exact
1742        // 202 acceptance) is safe to record as terminal DELIVERY_FAILED.
1743        // Other transport errors are ambiguous: the target may have accepted
1744        // the envelope before the response was lost. Keep those Dispatched so
1745        // a late response remains valid and return the durable ID for polling;
1746        // reverting/retrying could execute the command twice.
1747        if let Err(error) = self.command_dispatcher.dispatch(&envelope).await {
1748            if is_definite_dispatch_rejection(&error) {
1749                let delivery_failure = CommandResponse::error(
1750                    "DELIVERY_FAILED",
1751                    "Worker runtime did not accept command delivery",
1752                );
1753                self.submit_command_response(command_id, delivery_failure)
1754                    .await?;
1755                warn!(
1756                    command_id,
1757                    deployment_id,
1758                    error = %error,
1759                    "Push dispatch was definitely rejected; command marked Failed"
1760                );
1761                return Ok(CommandState::Failed);
1762            }
1763
1764            let error = error.context(ErrorData::TransportDispatchFailed {
1765                message: "Failed to dispatch command".to_string(),
1766                transport_type: None,
1767                target: Some(deployment_id.to_string()),
1768            });
1769            warn!(
1770                command_id,
1771                deployment_id,
1772                error = %error,
1773                "Push dispatch acknowledgement failed; command remains Dispatched for a possible late response"
1774            );
1775            return Ok(CommandState::Dispatched);
1776        }
1777
1778        info!("Command {} dispatched via push", envelope.command_id);
1779        Ok(CommandState::Dispatched)
1780    }
1781
1782    async fn build_envelope(
1783        &self,
1784        command_id: &str,
1785        metadata: &CommandEnvelopeData,
1786        mut params: BodySpec,
1787    ) -> Result<Envelope> {
1788        let response_handling = self.create_response_handling(command_id).await?;
1789
1790        // Re-inline: if params are in blob but fit in transport limit, read and embed inline.
1791        // This avoids unnecessary storage downloads for medium-sized params (18KB–150KB).
1792        if let BodySpec::Storage { size, .. } = &params {
1793            let raw_size = size.unwrap_or(0) as usize;
1794            if raw_size > 0 && raw_size <= self.inline_max_bytes {
1795                let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1796                match self.storage.get(&blob_path).await {
1797                    Ok(get_result) => match get_result.bytes().await {
1798                        Ok(raw_bytes) => {
1799                            params = BodySpec::inline(&raw_bytes);
1800                            debug!(
1801                                "Re-inlined params for command {} ({} bytes) into envelope",
1802                                command_id, raw_size
1803                            );
1804                        }
1805                        Err(e) => {
1806                            debug!(
1807                                    "Failed to read blob bytes for re-inline (command {}), falling back to presigned URL: {}",
1808                                    command_id, e
1809                                );
1810                        }
1811                    },
1812                    Err(e) => {
1813                        debug!(
1814                            "Failed to read blob for re-inline (command {}), falling back to presigned URL: {}",
1815                            command_id, e
1816                        );
1817                    }
1818                }
1819            }
1820        }
1821
1822        // If params are still Storage (either too large or re-inline failed),
1823        // always mint a fresh GET request. A command may remain Pending longer
1824        // than the URL minted at upload completion, and leasing must never hand
1825        // a Worker an already-expired storage credential.
1826        if let BodySpec::Storage {
1827            size,
1828            storage_get_request: _,
1829            storage_put_used,
1830        } = &params
1831        {
1832            let get_request = self.generate_storage_get_request(command_id).await?;
1833            params = BodySpec::Storage {
1834                size: *size,
1835                storage_get_request: Some(get_request),
1836                storage_put_used: *storage_put_used,
1837            };
1838        }
1839
1840        Ok(Envelope::new(
1841            metadata.deployment_id.clone(),
1842            metadata.target.clone(),
1843            command_id.to_string(),
1844            metadata.attempt,
1845            metadata.deadline,
1846            metadata.command.clone(),
1847            params,
1848            response_handling,
1849        ))
1850    }
1851
1852    async fn create_response_handling(&self, command_id: &str) -> Result<ResponseHandling> {
1853        let upload_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1854        let expires_in = Duration::from_secs(Self::RESPONSE_CREDENTIAL_LIFETIME_SECS);
1855        let presigned = self
1856            .storage
1857            .presigned_put(&upload_path, expires_in)
1858            .await
1859            .context(ErrorData::StorageOperationFailed {
1860                message: "Failed to create response upload URL".to_string(),
1861                operation: Some("presigned_put".to_string()),
1862                path: Some(upload_path.to_string()),
1863            })?;
1864
1865        let (response_token, expires) = self.sign_response_url(command_id);
1866
1867        Ok(ResponseHandling {
1868            max_inline_bytes: self.inline_max_bytes as u64,
1869            submit_response_url: format!(
1870                "{}/commands/{}/response?response_token={}&expires={}",
1871                self.base_url.trim_end_matches('/'),
1872                command_id,
1873                response_token,
1874                expires,
1875            ),
1876            storage_upload_request: presigned,
1877        })
1878    }
1879
1880    async fn generate_params_upload(&self, command_id: &str) -> Result<StorageUpload> {
1881        let upload_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1882        let expires_in = Duration::from_secs(3600);
1883        let presigned = self
1884            .storage
1885            .presigned_put(&upload_path, expires_in)
1886            .await
1887            .into_alien_error()
1888            .context(ErrorData::StorageOperationFailed {
1889                message: "Failed to create presigned URL".to_string(),
1890                operation: Some("presigned_put".to_string()),
1891                path: Some(upload_path.to_string()),
1892            })?;
1893
1894        Ok(StorageUpload {
1895            put_request: presigned.clone(),
1896            expires_at: presigned.expiration,
1897        })
1898    }
1899
1900    async fn generate_storage_get_request(&self, command_id: &str) -> Result<PresignedRequest> {
1901        let path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1902        let expires_in = Duration::from_secs(3600);
1903        self.storage.presigned_get(&path, expires_in).await.context(
1904            ErrorData::StorageOperationFailed {
1905                message: "Failed to create storage get request".to_string(),
1906                operation: Some("presigned_get".to_string()),
1907                path: Some(path.to_string()),
1908            },
1909        )
1910    }
1911
1912    async fn generate_response_storage_get_request(
1913        &self,
1914        command_id: &str,
1915    ) -> Result<PresignedRequest> {
1916        let path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1917        let expires_in = Duration::from_secs(3600);
1918        self.storage.presigned_get(&path, expires_in).await.context(
1919            ErrorData::StorageOperationFailed {
1920                message: "Failed to create response storage get request".to_string(),
1921                operation: Some("presigned_get".to_string()),
1922                path: Some(path.to_string()),
1923            },
1924        )
1925    }
1926
1927    fn extract_command_id_from_index_key(&self, index_key: &str) -> Result<String> {
1928        index_key
1929            .split(':')
1930            .last()
1931            .ok_or_else(|| {
1932                AlienError::new(ErrorData::Other {
1933                    message: format!("Invalid index key format: {}", index_key),
1934                })
1935            })
1936            .map(|s| s.to_string())
1937    }
1938}
1939
1940#[cfg(test)]
1941mod relative_url_tests {
1942    use super::*;
1943    use alien_core::presigned::PresignedOperation;
1944    use std::collections::HashMap;
1945
1946    fn http_request(url: &str, operation: PresignedOperation) -> PresignedRequest {
1947        PresignedRequest::new_http(
1948            url.to_string(),
1949            match operation {
1950                PresignedOperation::Get => "GET",
1951                PresignedOperation::Put => "PUT",
1952                PresignedOperation::Delete => "DELETE",
1953            }
1954            .to_string(),
1955            HashMap::new(),
1956            operation,
1957            "commands/test".to_string(),
1958            Utc::now() + chrono::Duration::minutes(5),
1959        )
1960    }
1961
1962    #[test]
1963    fn leased_manager_urls_are_relative_to_lease_endpoint() {
1964        let mut envelope = Envelope::new(
1965            "deployment",
1966            CommandTarget::new("daemon", CommandTargetType::Daemon),
1967            "command",
1968            1,
1969            None,
1970            "run",
1971            BodySpec::Storage {
1972                size: Some(2048),
1973                storage_get_request: Some(http_request(
1974                    "http://manager.internal/storage/params?signature=params",
1975                    PresignedOperation::Get,
1976                )),
1977                storage_put_used: Some(true),
1978            },
1979            ResponseHandling {
1980                max_inline_bytes: 1024,
1981                submit_response_url:
1982                    "http://manager.internal/v1/commands/command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1983                        .to_string(),
1984                storage_upload_request: http_request(
1985                    "http://manager.internal/v1/storage/response?signature=DoNotCanonicalize%2FUpload",
1986                    PresignedOperation::Put,
1987                ),
1988            },
1989        );
1990
1991        CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
1992
1993        assert_eq!(
1994            envelope.response_handling.submit_response_url,
1995            "command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1996        );
1997        assert_eq!(
1998            envelope.response_handling.storage_upload_request.url(),
1999            "../storage/response?signature=DoNotCanonicalize%2FUpload"
2000        );
2001        let BodySpec::Storage {
2002            storage_get_request: Some(params),
2003            ..
2004        } = &envelope.params
2005        else {
2006            panic!("storage params request");
2007        };
2008        assert_eq!(params.url(), "../../storage/params?signature=params");
2009
2010        envelope.response_handling.storage_upload_request = http_request(
2011            "https://storage.example.com/result?signature=cloud",
2012            PresignedOperation::Put,
2013        );
2014        CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
2015        assert_eq!(
2016            envelope.response_handling.storage_upload_request.url(),
2017            "https://storage.example.com/result?signature=cloud",
2018            "cloud-presigned URLs must remain byte-for-byte absolute"
2019        );
2020    }
2021}
2022
2023#[cfg(test)]
2024mod idempotency_key_tests {
2025    use super::*;
2026    use crate::server::{validate_command_name, validate_command_target_id};
2027
2028    #[test]
2029    fn definite_dispatch_rejection_is_classified_by_typed_error_data() {
2030        let rejected = AlienError::new(ErrorData::TransportDispatchRejected {
2031            message: "not accepted".to_string(),
2032            transport_type: Some("http".to_string()),
2033            target: Some("command-id".to_string()),
2034        });
2035        assert!(is_definite_dispatch_rejection(&rejected));
2036
2037        let ambiguous = AlienError::new(ErrorData::TransportDispatchFailed {
2038            message: "acknowledgement lost".to_string(),
2039            transport_type: Some("http".to_string()),
2040            target: Some("command-id".to_string()),
2041        });
2042        assert!(!is_definite_dispatch_rejection(&ambiguous));
2043    }
2044
2045    /// Idempotency keys are `{dep}:{rid}:{command}:{key}`. If a target id could
2046    /// contain ':', the `rid` segment would be ambiguous: the two triples
2047    ///   (rid="svc",   command="a:b", key="k")
2048    ///   (rid="svc:a", command="b",   key="k")
2049    /// both compose to `dep:svc:a:b:k`. The shared guard forbids ':' in a target
2050    /// id, so the second can never be a resolved target — closing the collision
2051    /// at the rid boundary. Only the colon-free composition is exercised here;
2052    /// the colliding one is proven unreachable via the guard.
2053    #[test]
2054    fn target_id_colon_guard_prevents_idempotency_key_collision() {
2055        let colliding = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "k");
2056        assert_eq!(colliding, "dep:svc:a:b:k");
2057
2058        // The alternate triple that would collide needs target id "svc:a".
2059        assert!(
2060            validate_command_target_id("svc:a").is_err(),
2061            "a ':'-bearing target id must be rejected so it cannot forge the rid segment"
2062        );
2063        assert!(validate_command_target_id("svc").is_ok());
2064    }
2065
2066    /// The command name is the other forgeable segment. Client keys routinely
2067    /// contain ':', so without a command-name guard these two distinct inputs
2068    ///   (command="a:b", key="c")
2069    ///   (command="a",   key="b:c")
2070    /// both compose to `dep:svc:a:b:c` — a cross-command idempotency collision.
2071    /// The guard rejects the ':'-bearing command name, so only the second input
2072    /// is ever composed; the first can never reach key composition.
2073    #[test]
2074    fn command_name_colon_guard_prevents_idempotency_key_collision() {
2075        // Without the guard, both inputs compose to the identical string.
2076        let forged = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "c");
2077        let legitimate = CommandServer::compose_idempotency_key("dep", "svc", "a", "b:c");
2078        assert_eq!(
2079            forged, legitimate,
2080            "these inputs are exactly the colliding pair the guard must separate"
2081        );
2082        assert_eq!(legitimate, "dep:svc:a:b:c");
2083
2084        // The guard makes the colliding input unreachable: a ':'-bearing command
2085        // name is rejected, while the legitimate command name is accepted. With
2086        // the forger blocked, `a:b`+`c` can never compose the shared key — only
2087        // the distinct `a`+`b:c` command can.
2088        let err = validate_command_name("a:b").expect_err("':'-bearing command must be rejected");
2089        assert_eq!(err.code, "INVALID_COMMAND");
2090        assert!(validate_command_name("a").is_ok());
2091    }
2092}