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    CommandEnvelopeData, CommandMetadata, CommandRegistry, CommandStatus, InMemoryCommandRegistry,
51    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    /// Resolve a lease_id to `(command_id, owner_deployment_id)` via the
1038    /// `lease:{lease_id}` reverse index, re-verifying against the live lease
1039    /// record (the reverse entry can outlive a released lease briefly, and a
1040    /// command can have been re-leased under a new lease_id since).
1041    ///
1042    /// Used by the manager's auth layer to check that the caller may act on
1043    /// the deployment that holds the lease.
1044    pub async fn get_lease_owner(&self, lease_id: &str) -> Result<Option<(String, String)>> {
1045        let reverse_key = format!("lease:{}", lease_id);
1046        let Some(command_id_bytes) =
1047            self.kv
1048                .get(&reverse_key)
1049                .await
1050                .context(ErrorData::KvOperationFailed {
1051                    operation: "get".to_string(),
1052                    key: reverse_key.clone(),
1053                    message: "Failed to look up lease reverse index".to_string(),
1054                })?
1055        else {
1056            return Ok(None);
1057        };
1058        let command_id = String::from_utf8(command_id_bytes).map_err(|_| {
1059            AlienError::new(ErrorData::Other {
1060                message: format!("Lease reverse index '{}' is not valid UTF-8", reverse_key),
1061            })
1062        })?;
1063
1064        let lease_key = format!("cmd:{}:lease", command_id);
1065        let Some(lease_data) =
1066            self.kv
1067                .get(&lease_key)
1068                .await
1069                .context(ErrorData::KvOperationFailed {
1070                    operation: "get".to_string(),
1071                    key: lease_key,
1072                    message: "Failed to look up lease".to_string(),
1073                })?
1074        else {
1075            return Ok(None);
1076        };
1077        let lease: LeaseData = serde_json::from_slice(&lease_data)
1078            .into_alien_error()
1079            .context(ErrorData::SerializationFailed {
1080                message: "Failed to deserialize lease data".to_string(),
1081                data_type: Some("LeaseData".to_string()),
1082            })?;
1083        if lease.lease_id != lease_id {
1084            return Ok(None);
1085        }
1086
1087        Ok(Some((command_id, lease.owner)))
1088    }
1089
1090    /// Release a lease by lease_id only (for the API).
1091    pub async fn release_lease_by_id(&self, lease_id: &str) -> Result<()> {
1092        match self.get_lease_owner(lease_id).await? {
1093            Some((command_id, _owner)) => self.release_lease(&command_id, lease_id).await,
1094            None => Err(AlienError::new(ErrorData::LeaseNotFound {
1095                lease_id: lease_id.to_string(),
1096            })),
1097        }
1098    }
1099
1100    // =========================================================================
1101    // Internal Helper Methods
1102    // =========================================================================
1103
1104    async fn validate_create_command(&self, request: &CreateCommandRequest) -> Result<()> {
1105        if request.command.is_empty() {
1106            return Err(AlienError::new(ErrorData::InvalidCommand {
1107                message: "Command name cannot be empty".to_string(),
1108            }));
1109        }
1110
1111        // The command name occupies one segment of the `:`-delimited
1112        // idempotency key (`{dep}:{rid}:{command}:{key}`). A ':' in the name
1113        // would let it bleed into the client-key segment — which routinely
1114        // contains ':' — so (command="a:b", key="c") and (command="a",
1115        // key="b:c") would forge the same key and be treated as the same
1116        // command. Reject ':' here, mirroring the target-id colon guard, so the
1117        // command segment is always unambiguous.
1118        validate_command_name(&request.command)?;
1119
1120        if request.deployment_id.is_empty() {
1121            return Err(AlienError::new(ErrorData::InvalidCommand {
1122                message: "Deployment ID cannot be empty".to_string(),
1123            }));
1124        }
1125
1126        if let Some(deadline) = request.deadline {
1127            if deadline <= Utc::now() {
1128                return Err(AlienError::new(ErrorData::InvalidCommand {
1129                    message: "Deadline must be in the future".to_string(),
1130                }));
1131            }
1132        }
1133
1134        Ok(())
1135    }
1136
1137    // --- Idempotency ---
1138
1139    /// Compose the target-scoped idempotency key:
1140    /// `{deploymentId}:{targetResourceId}:{commandName}:{key}`.
1141    ///
1142    /// Scoping by target means the same client key addressed to two different
1143    /// targets creates two distinct commands.
1144    fn compose_idempotency_key(
1145        deployment_id: &str,
1146        target_resource_id: &str,
1147        command_name: &str,
1148        idem_key: &str,
1149    ) -> String {
1150        // Invariant: at every real call site the target id and command name are
1151        // both `:`-free (the former enforced at resolution via
1152        // `validate_command_target_id`, the latter via `validate_command_name`
1153        // in `validate_create_command`), so each occupies exactly one segment of
1154        // `{dep}:{rid}:{command}:{key}` and only the trailing client key may
1155        // carry ':'. Those guards live upstream; this is a pure formatter, so it
1156        // is not asserted on `command_name` here (the collision tests below
1157        // deliberately format a ':'-bearing command to demonstrate what the
1158        // upstream guard prevents).
1159        debug_assert!(
1160            !target_resource_id.contains(':'),
1161            "target_resource_id must be ':'-free before key composition: {target_resource_id}"
1162        );
1163        format!(
1164            "{}:{}:{}:{}",
1165            deployment_id, target_resource_id, command_name, idem_key
1166        )
1167    }
1168
1169    async fn check_idempotency(&self, idem_key: &str) -> Result<Option<String>> {
1170        let key = format!("idem:{}", idem_key);
1171        if let Some(data) = self
1172            .kv
1173            .get(&key)
1174            .await
1175            .context(ErrorData::KvOperationFailed {
1176                operation: "get".to_string(),
1177                key: key.clone(),
1178                message: "Failed to check idempotency".to_string(),
1179            })?
1180        {
1181            let command_id = String::from_utf8(data).into_alien_error().context(
1182                ErrorData::SerializationFailed {
1183                    message: "Invalid idempotency data".to_string(),
1184                    data_type: Some("String".to_string()),
1185                },
1186            )?;
1187            return Ok(Some(command_id));
1188        }
1189        Ok(None)
1190    }
1191
1192    /// Claim the idempotency key for `command_id`.
1193    ///
1194    /// Returns `None` when this command won the key. Returns
1195    /// `Some(winner_id)` when a concurrent create with the same key won the
1196    /// conditional put first — both requests passed the pre-create
1197    /// `check_idempotency` before either stored, so the loser must be
1198    /// detected here, after its command was already created.
1199    async fn store_idempotency(&self, idem_key: &str, command_id: &str) -> Result<Option<String>> {
1200        let key = format!("idem:{}", idem_key);
1201        let ttl = Duration::from_secs(24 * 60 * 60); // 24 hours
1202        let won = self
1203            .kv
1204            .put(
1205                &key,
1206                command_id.as_bytes().to_vec(),
1207                Some(PutOptions {
1208                    ttl: Some(ttl),
1209                    if_not_exists: true,
1210                }),
1211            )
1212            .await
1213            .context(ErrorData::KvOperationFailed {
1214                operation: "put".to_string(),
1215                key: key.clone(),
1216                message: "Failed to store idempotency".to_string(),
1217            })?;
1218        if won {
1219            return Ok(None);
1220        }
1221        // Lost the conditional put: read back who won. The winner entry can
1222        // only be absent if it TTL-expired in the microseconds since — treat
1223        // that as an inconsistency rather than silently duplicating.
1224        match self.check_idempotency(idem_key).await? {
1225            Some(winner_id) => Ok(Some(winner_id)),
1226            None => Err(AlienError::new(ErrorData::Other {
1227                message: format!(
1228                    "Idempotency key '{}' was concurrently claimed but has no winner entry",
1229                    key
1230                ),
1231            })),
1232        }
1233    }
1234
1235    // --- Params ---
1236
1237    pub async fn store_params(&self, command_id: &str, params: &BodySpec) -> Result<()> {
1238        let key = format!("cmd:{}:params", command_id);
1239
1240        // Try serializing as-is first
1241        let data = CommandParamsData {
1242            params: params.clone(),
1243        };
1244        let value = serde_json::to_vec(&data).into_alien_error().context(
1245            ErrorData::SerializationFailed {
1246                message: "Failed to serialize params".to_string(),
1247                data_type: Some("CommandParamsData".to_string()),
1248            },
1249        )?;
1250
1251        // If it fits in KV, store directly (fast path)
1252        if value.len() <= KV_VALUE_THRESHOLD {
1253            self.kv
1254                .put(&key, value, None)
1255                .await
1256                .context(ErrorData::KvOperationFailed {
1257                    operation: "put".to_string(),
1258                    key: key.clone(),
1259                    message: "Failed to store params".to_string(),
1260                })?;
1261            return Ok(());
1262        }
1263
1264        // Auto-promote: inline data exceeds KV limit, store raw bytes in blob
1265        if let BodySpec::Inline { inline_base64 } = params {
1266            let raw_bytes = general_purpose::STANDARD
1267                .decode(inline_base64)
1268                .into_alien_error()
1269                .context(ErrorData::SerializationFailed {
1270                    message: "Failed to decode inline base64 params for auto-promotion".to_string(),
1271                    data_type: Some("base64".to_string()),
1272                })?;
1273
1274            let raw_len = raw_bytes.len() as u64;
1275            let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1276
1277            self.storage
1278                .put(&blob_path, Bytes::from(raw_bytes).into())
1279                .await
1280                .into_alien_error()
1281                .context(ErrorData::StorageOperationFailed {
1282                    message: "Failed to auto-promote params to blob storage".to_string(),
1283                    operation: Some("put".to_string()),
1284                    path: Some(blob_path.to_string()),
1285                })?;
1286
1287            debug!(
1288                "Auto-promoted params for command {} to blob ({} bytes raw)",
1289                command_id, raw_len
1290            );
1291
1292            // Store tiny reference in KV instead
1293            let promoted = CommandParamsData {
1294                params: BodySpec::Storage {
1295                    size: Some(raw_len),
1296                    storage_get_request: None,
1297                    storage_put_used: Some(true),
1298                },
1299            };
1300            let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1301                ErrorData::SerializationFailed {
1302                    message: "Failed to serialize promoted params reference".to_string(),
1303                    data_type: Some("CommandParamsData".to_string()),
1304                },
1305            )?;
1306            self.kv.put(&key, promoted_value, None).await.context(
1307                ErrorData::KvOperationFailed {
1308                    operation: "put".to_string(),
1309                    key: key.clone(),
1310                    message: "Failed to store promoted params reference".to_string(),
1311                },
1312            )?;
1313            return Ok(());
1314        }
1315
1316        // Storage references are always tiny, store as-is
1317        self.kv
1318            .put(&key, value, None)
1319            .await
1320            .context(ErrorData::KvOperationFailed {
1321                operation: "put".to_string(),
1322                key: key.clone(),
1323                message: "Failed to store params".to_string(),
1324            })?;
1325        Ok(())
1326    }
1327
1328    pub async fn get_params(&self, command_id: &str) -> Result<Option<BodySpec>> {
1329        let key = format!("cmd:{}:params", command_id);
1330        if let Some(value) = self
1331            .kv
1332            .get(&key)
1333            .await
1334            .context(ErrorData::KvOperationFailed {
1335                operation: "get".to_string(),
1336                key: key.clone(),
1337                message: "Failed to get params".to_string(),
1338            })?
1339        {
1340            let data: CommandParamsData = serde_json::from_slice(&value)
1341                .into_alien_error()
1342                .context(ErrorData::SerializationFailed {
1343                    message: "Failed to deserialize params".to_string(),
1344                    data_type: Some("CommandParamsData".to_string()),
1345                })?;
1346            return Ok(Some(data.params));
1347        }
1348        Ok(None)
1349    }
1350
1351    // --- Response ---
1352
1353    pub async fn store_response(&self, command_id: &str, response: &CommandResponse) -> Result<()> {
1354        let key = format!("cmd:{}:response", command_id);
1355        let data = CommandResponseData {
1356            response: response.clone(),
1357        };
1358        let value = serde_json::to_vec(&data).into_alien_error().context(
1359            ErrorData::SerializationFailed {
1360                message: "Failed to serialize response".to_string(),
1361                data_type: Some("CommandResponseData".to_string()),
1362            },
1363        )?;
1364
1365        // If it fits in KV, store directly (fast path)
1366        if value.len() <= KV_VALUE_THRESHOLD {
1367            self.kv
1368                .put(&key, value, None)
1369                .await
1370                .context(ErrorData::KvOperationFailed {
1371                    operation: "put".to_string(),
1372                    key: key.clone(),
1373                    message: "Failed to store response".to_string(),
1374                })?;
1375            return Ok(());
1376        }
1377
1378        // Auto-promote: inline response exceeds KV limit
1379        if let CommandResponse::Success {
1380            response: BodySpec::Inline { inline_base64 },
1381        } = response
1382        {
1383            let raw_bytes = general_purpose::STANDARD
1384                .decode(inline_base64)
1385                .into_alien_error()
1386                .context(ErrorData::SerializationFailed {
1387                    message: "Failed to decode inline base64 response for auto-promotion"
1388                        .to_string(),
1389                    data_type: Some("base64".to_string()),
1390                })?;
1391
1392            let raw_len = raw_bytes.len() as u64;
1393            let blob_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1394
1395            self.storage
1396                .put(&blob_path, Bytes::from(raw_bytes).into())
1397                .await
1398                .into_alien_error()
1399                .context(ErrorData::StorageOperationFailed {
1400                    message: "Failed to auto-promote response to blob storage".to_string(),
1401                    operation: Some("put".to_string()),
1402                    path: Some(blob_path.to_string()),
1403                })?;
1404
1405            // Generate presigned GET URL for the caller
1406            let get_request = self
1407                .generate_response_storage_get_request(command_id)
1408                .await?;
1409
1410            debug!(
1411                "Auto-promoted response for command {} to blob ({} bytes raw)",
1412                command_id, raw_len
1413            );
1414
1415            // Store tiny reference in KV
1416            let promoted = CommandResponseData {
1417                response: CommandResponse::Success {
1418                    response: BodySpec::Storage {
1419                        size: Some(raw_len),
1420                        storage_get_request: Some(get_request),
1421                        storage_put_used: Some(true),
1422                    },
1423                },
1424            };
1425            let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1426                ErrorData::SerializationFailed {
1427                    message: "Failed to serialize promoted response reference".to_string(),
1428                    data_type: Some("CommandResponseData".to_string()),
1429                },
1430            )?;
1431            self.kv.put(&key, promoted_value, None).await.context(
1432                ErrorData::KvOperationFailed {
1433                    operation: "put".to_string(),
1434                    key: key.clone(),
1435                    message: "Failed to store promoted response reference".to_string(),
1436                },
1437            )?;
1438            return Ok(());
1439        }
1440
1441        // Error responses or storage references are always small, store as-is
1442        self.kv
1443            .put(&key, value, None)
1444            .await
1445            .context(ErrorData::KvOperationFailed {
1446                operation: "put".to_string(),
1447                key: key.clone(),
1448                message: "Failed to store response".to_string(),
1449            })?;
1450        Ok(())
1451    }
1452
1453    pub async fn get_response(&self, command_id: &str) -> Result<Option<CommandResponse>> {
1454        let key = format!("cmd:{}:response", command_id);
1455        if let Some(value) = self
1456            .kv
1457            .get(&key)
1458            .await
1459            .context(ErrorData::KvOperationFailed {
1460                operation: "get".to_string(),
1461                key: key.clone(),
1462                message: "Failed to get response".to_string(),
1463            })?
1464        {
1465            let data: CommandResponseData = serde_json::from_slice(&value)
1466                .into_alien_error()
1467                .context(ErrorData::SerializationFailed {
1468                    message: "Failed to deserialize response".to_string(),
1469                    data_type: Some("CommandResponseData".to_string()),
1470                })?;
1471            return Ok(Some(data.response));
1472        }
1473        Ok(None)
1474    }
1475
1476    // --- Pending Index ---
1477
1478    async fn create_pending_index(
1479        &self,
1480        deployment_id: &str,
1481        target_resource_id: &str,
1482        command_id: &str,
1483    ) -> Result<()> {
1484        // Invariant: the target id is `:`-free (enforced at resolution via
1485        // `validate_command_target_id`), so its prefix `target:{dep}:{rid}:`
1486        // cannot overlap another target's pending keys.
1487        debug_assert!(
1488            !target_resource_id.contains(':'),
1489            "target_resource_id must be ':'-free in the pending index: {target_resource_id}"
1490        );
1491        let timestamp = Utc::now().timestamp_nanos_opt().unwrap_or(0);
1492        let key = format!(
1493            "target:{}:{}:pending:{}:{}",
1494            deployment_id, target_resource_id, timestamp, command_id
1495        );
1496
1497        // Store empty value - just for ordering
1498        self.kv
1499            .put(&key, vec![], None)
1500            .await
1501            .context(ErrorData::KvOperationFailed {
1502                operation: "put".to_string(),
1503                key: key.clone(),
1504                message: "Failed to create pending index".to_string(),
1505            })?;
1506        Ok(())
1507    }
1508
1509    async fn delete_pending_index(
1510        &self,
1511        deployment_id: &str,
1512        target_resource_id: &str,
1513        command_id: &str,
1514    ) -> Result<()> {
1515        // We need to scan to find the exact key since we don't know the timestamp
1516        let prefix = format!("target:{}:{}:pending:", deployment_id, target_resource_id);
1517        let scan_result = self
1518            .kv
1519            .scan_prefix(&prefix, Some(100), None)
1520            .await
1521            .into_alien_error()
1522            .context(ErrorData::KvOperationFailed {
1523                operation: "scan_prefix".to_string(),
1524                key: prefix.clone(),
1525                message: "Failed to scan pending index".to_string(),
1526            })?;
1527
1528        for (key, _) in scan_result.items {
1529            if key.ends_with(&format!(":{}", command_id)) {
1530                let _ = self.kv.delete(&key).await;
1531                break;
1532            }
1533        }
1534        Ok(())
1535    }
1536
1537    // --- Lease ---
1538
1539    async fn delete_lease(&self, command_id: &str) -> Result<()> {
1540        let key = format!("cmd:{}:lease", command_id);
1541        let _ = self.kv.delete(&key).await;
1542        Ok(())
1543    }
1544
1545    /// Expire every overdue non-terminal command recorded in the deadline
1546    /// index. Intended to run periodically from the hosting process.
1547    ///
1548    /// Deadlines are otherwise only enforced lazily (status polls and lease
1549    /// scans), which never reaches a command nobody polls — most notably a
1550    /// `PendingUpload` whose params upload never completed, which has no
1551    /// pending-index entry and would otherwise live forever.
1552    ///
1553    /// Uses the conditional terminal transition, so racing a concurrent
1554    /// submit is safe: whoever wins, the record stays consistent. Returns
1555    /// the number of commands expired.
1556    pub async fn reap_expired_commands(&self) -> Result<u32> {
1557        let now = Utc::now();
1558        let mut expired = 0u32;
1559        let mut cursor: Option<String> = None;
1560        // Paginate the whole index (bounded page count as a runaway guard):
1561        // keys are not numerically ordered — the timestamp segment isn't
1562        // zero-padded — so due entries can sit behind any number of
1563        // future-dated ones and a single capped scan would starve them.
1564        for _ in 0..64 {
1565            let scan = self
1566                .kv
1567                .scan_prefix("deadline:", Some(256), cursor.clone())
1568                .await
1569                .into_alien_error()
1570                .context(ErrorData::KvOperationFailed {
1571                    operation: "scan_prefix".to_string(),
1572                    key: "deadline:".to_string(),
1573                    message: "Failed to scan the deadline index".to_string(),
1574                })?;
1575            let next_cursor = scan.next_cursor.clone();
1576            for (key, value) in scan.items {
1577                let Ok(data) = serde_json::from_slice::<DeadlineIndexData>(&value) else {
1578                    warn!(key = %key, "Unparseable deadline index entry; deleting");
1579                    let _ = self.kv.delete(&key).await;
1580                    continue;
1581                };
1582                if data.deadline > now {
1583                    // Not due yet. (Keys are not sortable numerically — the
1584                    // timestamp segment isn't zero-padded — so keep scanning.)
1585                    continue;
1586                }
1587
1588                let status = self
1589                    .command_registry
1590                    .get_command_status(&data.command_id)
1591                    .await?;
1592                match status {
1593                    None => {
1594                        let _ = self.kv.delete(&key).await;
1595                    }
1596                    Some(status) if status.state.is_terminal() => {
1597                        let _ = self.kv.delete(&key).await;
1598                    }
1599                    Some(status) => {
1600                        let won = self
1601                        .command_registry
1602                        .complete_command(
1603                            &data.command_id,
1604                            CommandState::Expired,
1605                            now,
1606                            None,
1607                            Some(serde_json::json!({
1608                                "code": "COMMAND_EXPIRED",
1609                                "message": format!("Deadline {} elapsed", data.deadline.to_rfc3339()),
1610                            })),
1611                        )
1612                        .await?;
1613                        if won {
1614                            expired += 1;
1615                            info!(command_id = %data.command_id, "Expired overdue command");
1616                            let _ = self.delete_lease(&data.command_id).await;
1617                            let _ = self
1618                                .delete_pending_index(
1619                                    &status.deployment_id,
1620                                    &status.target.resource_id,
1621                                    &data.command_id,
1622                                )
1623                                .await;
1624                        }
1625                        let _ = self.kv.delete(&key).await;
1626                    }
1627                }
1628            }
1629            match next_cursor {
1630                Some(next) => cursor = Some(next),
1631                None => break,
1632            }
1633        }
1634        Ok(expired)
1635    }
1636
1637    // --- Deadline Index ---
1638
1639    async fn create_deadline_index(&self, command_id: &str, deadline: DateTime<Utc>) -> Result<()> {
1640        let key = format!(
1641            "deadline:{}:{}",
1642            deadline.timestamp_nanos_opt().unwrap_or(0),
1643            command_id
1644        );
1645
1646        let data = DeadlineIndexData {
1647            command_id: command_id.to_string(),
1648            deadline,
1649        };
1650        let value = serde_json::to_vec(&data).into_alien_error().context(
1651            ErrorData::SerializationFailed {
1652                message: "Failed to serialize deadline index".to_string(),
1653                data_type: Some("DeadlineIndexData".to_string()),
1654            },
1655        )?;
1656
1657        // The entry must remain VISIBLE well past the deadline: scan_prefix
1658        // treats logically-expired keys as absent on every provider, so a
1659        // TTL equal to the deadline would hide the entry at exactly the
1660        // moment the reaper needs it (the bug that made the reaper inert).
1661        // The reaper deletes entries as it processes them; the 7-day grace
1662        // is only self-cleaning for processes that never run a reaper.
1663        const DEADLINE_INDEX_GRACE: chrono::Duration = chrono::Duration::days(7);
1664        let ttl = deadline
1665            .signed_duration_since(Utc::now())
1666            .checked_add(&DEADLINE_INDEX_GRACE)
1667            .unwrap_or(DEADLINE_INDEX_GRACE);
1668        let options = (ttl.num_seconds() > 0).then(|| PutOptions {
1669            ttl: Some(Duration::from_secs(ttl.num_seconds() as u64)),
1670            if_not_exists: false,
1671        });
1672
1673        self.kv
1674            .put(&key, value, options)
1675            .await
1676            .context(ErrorData::KvOperationFailed {
1677                operation: "put".to_string(),
1678                key: key.clone(),
1679                message: "Failed to create deadline index".to_string(),
1680            })?;
1681        Ok(())
1682    }
1683
1684    // --- Dispatch ---
1685
1686    async fn dispatch_command_push(
1687        &self,
1688        command_id: &str,
1689        deployment_id: &str,
1690    ) -> Result<CommandState> {
1691        // Get metadata from registry
1692        let metadata = self
1693            .command_registry
1694            .get_command_metadata(command_id)
1695            .await?
1696            .ok_or_else(|| {
1697                AlienError::new(ErrorData::CommandNotFound {
1698                    command_id: command_id.to_string(),
1699                })
1700            })?;
1701
1702        // Get params from KV
1703        let params = self.get_params(command_id).await?.ok_or_else(|| {
1704            AlienError::new(ErrorData::CommandNotFound {
1705                command_id: command_id.to_string(),
1706            })
1707        })?;
1708
1709        // Build envelope
1710        let envelope = self.build_envelope(command_id, &metadata, params).await?;
1711
1712        // Mark Dispatched BEFORE invoking the transport: a fast worker can
1713        // execute and submit before this function resumes, and submit
1714        // validates state == Dispatched — the old dispatch-then-mark order
1715        // could reject that submit (losing the response) or stomp its
1716        // terminal state back to Dispatched.
1717        if !self
1718            .command_registry
1719            .mark_dispatched_if_not_terminal(command_id, Utc::now())
1720            .await?
1721        {
1722            return Ok(self
1723                .command_registry
1724                .get_command_status(command_id)
1725                .await?
1726                .map(|status| status.state)
1727                .unwrap_or(CommandState::Dispatched));
1728        }
1729
1730        // A definite pre-delivery rejection (connection refusal, request
1731        // builder failure, or any HTTP status other than the runtime's exact
1732        // 202 acceptance) is safe to record as terminal DELIVERY_FAILED.
1733        // Other transport errors are ambiguous: the target may have accepted
1734        // the envelope before the response was lost. Keep those Dispatched so
1735        // a late response remains valid and return the durable ID for polling;
1736        // reverting/retrying could execute the command twice.
1737        if let Err(error) = self.command_dispatcher.dispatch(&envelope).await {
1738            if is_definite_dispatch_rejection(&error) {
1739                let delivery_failure = CommandResponse::error(
1740                    "DELIVERY_FAILED",
1741                    "Worker runtime did not accept command delivery",
1742                );
1743                self.submit_command_response(command_id, delivery_failure)
1744                    .await?;
1745                warn!(
1746                    command_id,
1747                    deployment_id,
1748                    error = %error,
1749                    "Push dispatch was definitely rejected; command marked Failed"
1750                );
1751                return Ok(CommandState::Failed);
1752            }
1753
1754            let error = error.context(ErrorData::TransportDispatchFailed {
1755                message: "Failed to dispatch command".to_string(),
1756                transport_type: None,
1757                target: Some(deployment_id.to_string()),
1758            });
1759            warn!(
1760                command_id,
1761                deployment_id,
1762                error = %error,
1763                "Push dispatch acknowledgement failed; command remains Dispatched for a possible late response"
1764            );
1765            return Ok(CommandState::Dispatched);
1766        }
1767
1768        info!("Command {} dispatched via push", envelope.command_id);
1769        Ok(CommandState::Dispatched)
1770    }
1771
1772    async fn build_envelope(
1773        &self,
1774        command_id: &str,
1775        metadata: &CommandEnvelopeData,
1776        mut params: BodySpec,
1777    ) -> Result<Envelope> {
1778        let response_handling = self.create_response_handling(command_id).await?;
1779
1780        // Re-inline: if params are in blob but fit in transport limit, read and embed inline.
1781        // This avoids unnecessary storage downloads for medium-sized params (18KB–150KB).
1782        if let BodySpec::Storage { size, .. } = &params {
1783            let raw_size = size.unwrap_or(0) as usize;
1784            if raw_size > 0 && raw_size <= self.inline_max_bytes {
1785                let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1786                match self.storage.get(&blob_path).await {
1787                    Ok(get_result) => match get_result.bytes().await {
1788                        Ok(raw_bytes) => {
1789                            params = BodySpec::inline(&raw_bytes);
1790                            debug!(
1791                                "Re-inlined params for command {} ({} bytes) into envelope",
1792                                command_id, raw_size
1793                            );
1794                        }
1795                        Err(e) => {
1796                            debug!(
1797                                    "Failed to read blob bytes for re-inline (command {}), falling back to presigned URL: {}",
1798                                    command_id, e
1799                                );
1800                        }
1801                    },
1802                    Err(e) => {
1803                        debug!(
1804                            "Failed to read blob for re-inline (command {}), falling back to presigned URL: {}",
1805                            command_id, e
1806                        );
1807                    }
1808                }
1809            }
1810        }
1811
1812        // If params are still Storage (either too large or re-inline failed),
1813        // always mint a fresh GET request. A command may remain Pending longer
1814        // than the URL minted at upload completion, and leasing must never hand
1815        // a Worker an already-expired storage credential.
1816        if let BodySpec::Storage {
1817            size,
1818            storage_get_request: _,
1819            storage_put_used,
1820        } = &params
1821        {
1822            let get_request = self.generate_storage_get_request(command_id).await?;
1823            params = BodySpec::Storage {
1824                size: *size,
1825                storage_get_request: Some(get_request),
1826                storage_put_used: *storage_put_used,
1827            };
1828        }
1829
1830        Ok(Envelope::new(
1831            metadata.deployment_id.clone(),
1832            metadata.target.clone(),
1833            command_id.to_string(),
1834            metadata.attempt,
1835            metadata.deadline,
1836            metadata.command.clone(),
1837            params,
1838            response_handling,
1839        ))
1840    }
1841
1842    async fn create_response_handling(&self, command_id: &str) -> Result<ResponseHandling> {
1843        let upload_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1844        let expires_in = Duration::from_secs(Self::RESPONSE_CREDENTIAL_LIFETIME_SECS);
1845        let presigned = self
1846            .storage
1847            .presigned_put(&upload_path, expires_in)
1848            .await
1849            .context(ErrorData::StorageOperationFailed {
1850                message: "Failed to create response upload URL".to_string(),
1851                operation: Some("presigned_put".to_string()),
1852                path: Some(upload_path.to_string()),
1853            })?;
1854
1855        let (response_token, expires) = self.sign_response_url(command_id);
1856
1857        Ok(ResponseHandling {
1858            max_inline_bytes: self.inline_max_bytes as u64,
1859            submit_response_url: format!(
1860                "{}/commands/{}/response?response_token={}&expires={}",
1861                self.base_url.trim_end_matches('/'),
1862                command_id,
1863                response_token,
1864                expires,
1865            ),
1866            storage_upload_request: presigned,
1867        })
1868    }
1869
1870    async fn generate_params_upload(&self, command_id: &str) -> Result<StorageUpload> {
1871        let upload_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1872        let expires_in = Duration::from_secs(3600);
1873        let presigned = self
1874            .storage
1875            .presigned_put(&upload_path, expires_in)
1876            .await
1877            .into_alien_error()
1878            .context(ErrorData::StorageOperationFailed {
1879                message: "Failed to create presigned URL".to_string(),
1880                operation: Some("presigned_put".to_string()),
1881                path: Some(upload_path.to_string()),
1882            })?;
1883
1884        Ok(StorageUpload {
1885            put_request: presigned.clone(),
1886            expires_at: presigned.expiration,
1887        })
1888    }
1889
1890    async fn generate_storage_get_request(&self, command_id: &str) -> Result<PresignedRequest> {
1891        let path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1892        let expires_in = Duration::from_secs(3600);
1893        self.storage.presigned_get(&path, expires_in).await.context(
1894            ErrorData::StorageOperationFailed {
1895                message: "Failed to create storage get request".to_string(),
1896                operation: Some("presigned_get".to_string()),
1897                path: Some(path.to_string()),
1898            },
1899        )
1900    }
1901
1902    async fn generate_response_storage_get_request(
1903        &self,
1904        command_id: &str,
1905    ) -> Result<PresignedRequest> {
1906        let path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1907        let expires_in = Duration::from_secs(3600);
1908        self.storage.presigned_get(&path, expires_in).await.context(
1909            ErrorData::StorageOperationFailed {
1910                message: "Failed to create response storage get request".to_string(),
1911                operation: Some("presigned_get".to_string()),
1912                path: Some(path.to_string()),
1913            },
1914        )
1915    }
1916
1917    fn extract_command_id_from_index_key(&self, index_key: &str) -> Result<String> {
1918        index_key
1919            .split(':')
1920            .last()
1921            .ok_or_else(|| {
1922                AlienError::new(ErrorData::Other {
1923                    message: format!("Invalid index key format: {}", index_key),
1924                })
1925            })
1926            .map(|s| s.to_string())
1927    }
1928}
1929
1930#[cfg(test)]
1931mod relative_url_tests {
1932    use super::*;
1933    use alien_core::presigned::PresignedOperation;
1934    use std::collections::HashMap;
1935
1936    fn http_request(url: &str, operation: PresignedOperation) -> PresignedRequest {
1937        PresignedRequest::new_http(
1938            url.to_string(),
1939            match operation {
1940                PresignedOperation::Get => "GET",
1941                PresignedOperation::Put => "PUT",
1942                PresignedOperation::Delete => "DELETE",
1943            }
1944            .to_string(),
1945            HashMap::new(),
1946            operation,
1947            "commands/test".to_string(),
1948            Utc::now() + chrono::Duration::minutes(5),
1949        )
1950    }
1951
1952    #[test]
1953    fn leased_manager_urls_are_relative_to_lease_endpoint() {
1954        let mut envelope = Envelope::new(
1955            "deployment",
1956            CommandTarget::new("daemon", CommandTargetType::Daemon),
1957            "command",
1958            1,
1959            None,
1960            "run",
1961            BodySpec::Storage {
1962                size: Some(2048),
1963                storage_get_request: Some(http_request(
1964                    "http://manager.internal/storage/params?signature=params",
1965                    PresignedOperation::Get,
1966                )),
1967                storage_put_used: Some(true),
1968            },
1969            ResponseHandling {
1970                max_inline_bytes: 1024,
1971                submit_response_url:
1972                    "http://manager.internal/v1/commands/command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1973                        .to_string(),
1974                storage_upload_request: http_request(
1975                    "http://manager.internal/v1/storage/response?signature=DoNotCanonicalize%2FUpload",
1976                    PresignedOperation::Put,
1977                ),
1978            },
1979        );
1980
1981        CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
1982
1983        assert_eq!(
1984            envelope.response_handling.submit_response_url,
1985            "command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1986        );
1987        assert_eq!(
1988            envelope.response_handling.storage_upload_request.url(),
1989            "../storage/response?signature=DoNotCanonicalize%2FUpload"
1990        );
1991        let BodySpec::Storage {
1992            storage_get_request: Some(params),
1993            ..
1994        } = &envelope.params
1995        else {
1996            panic!("storage params request");
1997        };
1998        assert_eq!(params.url(), "../../storage/params?signature=params");
1999
2000        envelope.response_handling.storage_upload_request = http_request(
2001            "https://storage.example.com/result?signature=cloud",
2002            PresignedOperation::Put,
2003        );
2004        CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
2005        assert_eq!(
2006            envelope.response_handling.storage_upload_request.url(),
2007            "https://storage.example.com/result?signature=cloud",
2008            "cloud-presigned URLs must remain byte-for-byte absolute"
2009        );
2010    }
2011}
2012
2013#[cfg(test)]
2014mod idempotency_key_tests {
2015    use super::*;
2016    use crate::server::{validate_command_name, validate_command_target_id};
2017
2018    #[test]
2019    fn definite_dispatch_rejection_is_classified_by_typed_error_data() {
2020        let rejected = AlienError::new(ErrorData::TransportDispatchRejected {
2021            message: "not accepted".to_string(),
2022            transport_type: Some("http".to_string()),
2023            target: Some("command-id".to_string()),
2024        });
2025        assert!(is_definite_dispatch_rejection(&rejected));
2026
2027        let ambiguous = AlienError::new(ErrorData::TransportDispatchFailed {
2028            message: "acknowledgement lost".to_string(),
2029            transport_type: Some("http".to_string()),
2030            target: Some("command-id".to_string()),
2031        });
2032        assert!(!is_definite_dispatch_rejection(&ambiguous));
2033    }
2034
2035    /// Idempotency keys are `{dep}:{rid}:{command}:{key}`. If a target id could
2036    /// contain ':', the `rid` segment would be ambiguous: the two triples
2037    ///   (rid="svc",   command="a:b", key="k")
2038    ///   (rid="svc:a", command="b",   key="k")
2039    /// both compose to `dep:svc:a:b:k`. The shared guard forbids ':' in a target
2040    /// id, so the second can never be a resolved target — closing the collision
2041    /// at the rid boundary. Only the colon-free composition is exercised here;
2042    /// the colliding one is proven unreachable via the guard.
2043    #[test]
2044    fn target_id_colon_guard_prevents_idempotency_key_collision() {
2045        let colliding = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "k");
2046        assert_eq!(colliding, "dep:svc:a:b:k");
2047
2048        // The alternate triple that would collide needs target id "svc:a".
2049        assert!(
2050            validate_command_target_id("svc:a").is_err(),
2051            "a ':'-bearing target id must be rejected so it cannot forge the rid segment"
2052        );
2053        assert!(validate_command_target_id("svc").is_ok());
2054    }
2055
2056    /// The command name is the other forgeable segment. Client keys routinely
2057    /// contain ':', so without a command-name guard these two distinct inputs
2058    ///   (command="a:b", key="c")
2059    ///   (command="a",   key="b:c")
2060    /// both compose to `dep:svc:a:b:c` — a cross-command idempotency collision.
2061    /// The guard rejects the ':'-bearing command name, so only the second input
2062    /// is ever composed; the first can never reach key composition.
2063    #[test]
2064    fn command_name_colon_guard_prevents_idempotency_key_collision() {
2065        // Without the guard, both inputs compose to the identical string.
2066        let forged = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "c");
2067        let legitimate = CommandServer::compose_idempotency_key("dep", "svc", "a", "b:c");
2068        assert_eq!(
2069            forged, legitimate,
2070            "these inputs are exactly the colliding pair the guard must separate"
2071        );
2072        assert_eq!(legitimate, "dep:svc:a:b:c");
2073
2074        // The guard makes the colliding input unreachable: a ':'-bearing command
2075        // name is rejected, while the legitimate command name is accepted. With
2076        // the forger blocked, `a:b`+`c` can never compose the shared key — only
2077        // the distinct `a`+`b:c` command can.
2078        let err = validate_command_name("a:b").expect_err("':'-bearing command must be rejected");
2079        assert_eq!(err.code, "INVALID_COMMAND");
2080        assert!(validate_command_name("a").is_ok());
2081    }
2082}