stash: refactoring using yoke grind

crates/mvp-system/src/chat/node_image.rs

@@ -49,16 +49,6 @@ pub(super) struct NodeImageRequest {
     pub(super) enabled: bool,
 }
 
-#[allow(dead_code)]
-#[derive(Clone, Debug)]
-pub(super) struct PreparedNodeImage {
-    pub(super) image_ref: String,
-    pub(super) tag: String,
-    pub(super) already_available: bool,
-    pub(super) built: bool,
-    pub(super) pushed: bool,
-}
-
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(super) enum NodeImageProgressEventKind {
     ImageReference {
@@ -140,7 +130,7 @@ struct RealImageCommandRunner;
 pub(super) fn prepare_node_image_with_progress(
     request: NodeImageRequest,
     progress: Option<&mut dyn NodeImageProgressSink>,
-) -> Result<PreparedNodeImage, String> {
+) -> Result<String, String> {
     let mut progress = progress;
     let mut runner = RealImageCommandRunner;
     prepare_node_image_inner(request, &mut progress, &mut runner)
@@ -150,16 +140,10 @@ fn prepare_node_image_inner(
     request: NodeImageRequest,
     progress: &mut Option<&mut dyn NodeImageProgressSink>,
     runner: &mut dyn ImageCommandRunner,
-) -> Result<PreparedNodeImage, String> {
+) -> Result<String, String> {
     emit_image_reference(progress, "requested", &request.requested_image);
     if !request.enabled {
-        return Ok(PreparedNodeImage {
+        return Ok(request.requested_image);
-            image_ref: request.requested_image,
-            tag: String::new(),
-            already_available: false,
-            built: false,
-            pushed: false,
-        });
     }
 
     emit_image_reference(progress, "base", &request.base_image);
@@ -208,16 +192,9 @@ fn prepare_node_image_inner(
         docker_image_labels_match(runner, &root, &image_ref, &expected_node_labels)?;
     let remote_available = remote_required && runner.docker_manifest_exists(&root, &image_ref);
     if !request.force_refresh && remote_required && remote_available {
-        let pushed =
         ensure_aliases_for_remote(runner, progress, &root, &image_ref, &image, &alias_tags)?;
         prune_old_dirty_images(runner, &root, &image, &tag);
-        return Ok(PreparedNodeImage {
+        return Ok(image_ref);
-            image_ref,
-            tag,
-            already_available: true,
-            built: false,
-            pushed,
-        });
     }
     if !request.force_refresh && remote_required && local_image_matches {
         ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?;
@@ -226,24 +203,12 @@ fn prepare_node_image_inner(
             push_image(runner, progress, &root, &alias)?;
         }
         prune_old_dirty_images(runner, &root, &image, &tag);
-        return Ok(PreparedNodeImage {
+        return Ok(image_ref);
-            image_ref,
-            tag,
-            already_available: true,
-            built: false,
-            pushed: true,
-        });
     }
     if !request.force_refresh && !remote_required && local_image_matches {
         ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?;
         prune_old_dirty_images(runner, &root, &image, &tag);
-        return Ok(PreparedNodeImage {
+        return Ok(image_ref);
-            image_ref,
-            tag,
-            already_available: true,
-            built: false,
-            pushed: false,
-        });
     }
     let base_image_matches =
         docker_image_labels_match(runner, &root, &request.base_image, &expected_base_labels)?;
@@ -294,23 +259,15 @@ fn prepare_node_image_inner(
     )?;
     ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?;
 
-    let mut pushed = false;
     if remote_required {
         push_image(runner, progress, &root, &image_ref)?;
-        pushed = true;
         for alias in alias_refs(&image, &alias_tags) {
             push_image(runner, progress, &root, &alias)?;
         }
     }
 
     prune_old_dirty_images(runner, &root, &image, &tag);
-    Ok(PreparedNodeImage {
+    Ok(image_ref)
-        image_ref,
-        tag,
-        already_available: false,
-        built: true,
-        pushed,
-    })
 }
 
 fn workspace_root() -> Result<PathBuf, String> {

crates/mvp-system/src/chat/runtime.rs

@@ -25,7 +25,7 @@ use signal_hook::iterator::Signals;
 use crate::chat::config as chat_config;
 use crate::chat::node_image::{
     NodeImageProgressEvent, NodeImageProgressEventKind, NodeImageProgressSink, NodeImageProvider,
-    NodeImageRequest, PreparedNodeImage, prepare_node_image_with_progress,
+    NodeImageRequest, prepare_node_image_with_progress,
 };
 use crate::node_provisioning::{ProviderKind, provider_kind};
 use crate::observability::{benchmark, frame_archive::FrameArchive};
@@ -775,73 +775,21 @@ impl Config {
         let args = ParsedArgs::parse(provided_args)?;
         let loaded = load_chat_config(args.config_path.as_deref())?;
         let toml = loaded.overlay;
-        let provider = provider_from_sources(args.provider, toml.provider.kind.as_deref())?;
+        let provider = provider_from_sources(args.provider.clone(), toml.provider.kind.as_deref())?;
         let node_image = first_non_empty([toml.image.node.clone()]).unwrap_or_default();
         if provider != provider_kind::process() && node_image.is_empty() {
             return Err("node image is required for docker or vastai provider".to_owned());
         }
-        let pipeline_stages = args
+        let pipeline_stages = Self::pipeline_stages(&args, &toml)?;
-            .pipeline_stages
+        let max_tokens = Self::max_tokens(&toml)?;
-            .or(toml.runtime.pipeline_stages)
-            .unwrap_or(1);
-        if pipeline_stages == 0 {
-            return Err("--pipeline-stages must be greater than 0".to_owned());
-        }
-        let max_tokens = toml.runtime.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS);
-        if max_tokens == 0 {
-            return Err("[runtime].max_tokens must be greater than 0".to_owned());
-        }
         let gpu_run = args.gpu || env_flag(MVP_CHAT_GPU_RUN_ENV, false);
-        let endpoint_addr_mask = match first_non_empty([
+        let endpoint_addr_mask = Self::endpoint_addr_mask(&args, &toml)?;
-            args.endpoint_addr_mask.clone(),
+        let (relay_mode, relay_url) = Self::relay_settings(&args, &toml, endpoint_addr_mask)?;
-            toml.relay.endpoint_addr_mask.clone(),
+        let cached_model = Self::cached_model_source(&args, gpu_run, &provider)
-        ]) {
-            Some(mask) => EndpointAddrMask::parse(&mask)?,
-            None => EndpointAddrMask::Full,
-        };
-        let relay_mode = first_non_empty([args.relay_mode.clone(), toml.relay.mode.clone()]);
-        let mut relay_url = first_non_empty([args.relay_url.clone(), toml.relay.url.clone()]);
-        if endpoint_addr_mask.requires_relay() && relay_url.is_none() {
-            relay_url = first_non_empty([toml.vastai.relay_url.clone()]);
-        }
-        if endpoint_addr_mask.requires_relay() && relay_url.is_none() {
-            return Err("relay-only endpoint address mask requires [relay].url, --relay-url, or [vastai].relay_url".to_owned());
-        }
-        let relay_mode = relay_mode.or_else(|| relay_url.as_ref().map(|_| "default".to_owned()));
-        let cached_model_source = match args.cached_model {
-            Some(source) => Some(source),
-            None if gpu_run && provider == provider_kind::process() => {
-                Some(CachedModelSource::Discover)
-            }
-            None => None,
-        };
-        let cached_model = cached_model_source
             .map(CachedModelConfig::from_source)
             .transpose()?;
-        let model = if provider == provider_kind::vastai() {
+        let model = Self::model_config(&provider, &toml, cached_model.as_ref())?;
-            match &cached_model {
+        let datastream_frame_log = Self::datastream_frame_log(&args, &toml);
-                Some(cached_model) => {
-                    vastai_model_config_for_cached_model(toml.model.clone(), cached_model)?
-                }
-                None => toml.model.clone(),
-            }
-        } else {
-            toml.model.clone()
-        };
-        let datastream_frame_log = if args.dump_logs {
-            Some(
-                args.dump_log_path
-                    .unwrap_or_else(|| PathBuf::from("mvp-chat.log")),
-            )
-        } else if toml.observability.dump_logs.unwrap_or(false) {
-            Some(
-                first_non_empty([toml.observability.dump_log_path.clone()])
-                    .map(PathBuf::from)
-                    .unwrap_or_else(|| PathBuf::from("mvp-chat.log")),
-            )
-        } else {
-            None
-        };
         let vastai = if provider == provider_kind::vastai() {
             Some(resolve_vastai_config(&toml.vastai, &node_image)?)
         } else {
@@ -871,9 +819,114 @@ impl Config {
         })
     }
 
+    fn pipeline_stages(args: &ParsedArgs, toml: &ChatTomlConfig) -> Result<u32, String> {
+        let pipeline_stages = args
+            .pipeline_stages
+            .or(toml.runtime.pipeline_stages)
+            .unwrap_or(1);
+        if pipeline_stages == 0 {
+            return Err("--pipeline-stages must be greater than 0".to_owned());
+        }
+        Ok(pipeline_stages)
+    }
+
+    fn max_tokens(toml: &ChatTomlConfig) -> Result<u32, String> {
+        let max_tokens = toml.runtime.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS);
+        if max_tokens == 0 {
+            return Err("[runtime].max_tokens must be greater than 0".to_owned());
+        }
+        Ok(max_tokens)
+    }
+
+    fn endpoint_addr_mask(
+        args: &ParsedArgs,
+        toml: &ChatTomlConfig,
+    ) -> Result<EndpointAddrMask, String> {
+        match first_non_empty([
+            args.endpoint_addr_mask.clone(),
+            toml.relay.endpoint_addr_mask.clone(),
+        ]) {
+            Some(mask) => EndpointAddrMask::parse(&mask),
+            None => Ok(EndpointAddrMask::Full),
+        }
+    }
+
+    fn relay_settings(
+        args: &ParsedArgs,
+        toml: &ChatTomlConfig,
+        endpoint_addr_mask: EndpointAddrMask,
+    ) -> Result<(Option<String>, Option<String>), String> {
+        let relay_mode = first_non_empty([args.relay_mode.clone(), toml.relay.mode.clone()]);
+        let mut relay_url = first_non_empty([args.relay_url.clone(), toml.relay.url.clone()]);
+        if endpoint_addr_mask.requires_relay() && relay_url.is_none() {
+            relay_url = first_non_empty([toml.vastai.relay_url.clone()]);
+        }
+        if endpoint_addr_mask.requires_relay() && relay_url.is_none() {
+            return Err("relay-only endpoint address mask requires [relay].url, --relay-url, or [vastai].relay_url".to_owned());
+        }
+        let relay_mode = relay_mode.or_else(|| relay_url.as_ref().map(|_| "default".to_owned()));
+        Ok((relay_mode, relay_url))
+    }
+
+    fn cached_model_source(
+        args: &ParsedArgs,
+        gpu_run: bool,
+        provider: &ProviderKind,
+    ) -> Option<CachedModelSource> {
+        match &args.cached_model {
+            Some(source) => Some(source.clone()),
+            None if gpu_run && provider == &provider_kind::process() => {
+                Some(CachedModelSource::Discover)
+            }
+            None => None,
+        }
+    }
+
+    fn model_config(
+        provider: &ProviderKind,
+        toml: &ChatTomlConfig,
+        cached_model: Option<&CachedModelConfig>,
+    ) -> Result<ChatModelConfig, String> {
+        if provider != &provider_kind::vastai() {
+            return Ok(toml.model.clone());
+        }
+        match cached_model {
+            Some(cached_model) => {
+                vastai_model_config_for_cached_model(toml.model.clone(), cached_model)
+            }
+            None => Ok(toml.model.clone()),
+        }
+    }
+
+    fn datastream_frame_log(args: &ParsedArgs, toml: &ChatTomlConfig) -> Option<PathBuf> {
+        if args.dump_logs {
+            return Some(
+                args.dump_log_path
+                    .clone()
+                    .unwrap_or_else(|| PathBuf::from("mvp-chat.log")),
+            );
+        }
+        if !toml.observability.dump_logs.unwrap_or(false) {
+            return None;
+        }
+        Some(
+            first_non_empty([toml.observability.dump_log_path.clone()])
+                .map(PathBuf::from)
+                .unwrap_or_else(|| PathBuf::from("mvp-chat.log")),
+        )
+    }
+
     // The orchestrator launch spec is still pending. These flags are the current adapter;
     // adjust this mapping when the approved orchestrator launch contract is finalized.
     fn orchestrator_cli_args(&self, image_ref: &str) -> Vec<String> {
+        macro_rules! push_opt {
+            ($args:ident, $option:expr, $flag:expr, |$value:ident| $arg:expr) => {
+                if let Some($value) = $option {
+                    $args.extend([$flag.to_owned(), $arg]);
+                }
+            };
+        }
+
         let mut args = vec![
             "--provider".to_owned(),
             self.provider.as_str().to_owned(),
@@ -889,27 +942,36 @@ impl Config {
             self.pipeline_stages.to_string(),
             "--dashboard".to_owned(),
         ];
-        if let Some(model_id) = &self.model.id {
+        push_opt!(args, &self.model.id, "--model-id", |model_id| model_id
-            args.extend(["--model-id".to_owned(), model_id.clone()]);
+            .clone());
-        }
+        push_opt!(
-        if let Some(path) = &self.model.gguf_local_path {
+            args,
-            args.extend(["--gguf-local-path".to_owned(), path.clone()]);
+            &self.model.gguf_local_path,
-        }
+            "--gguf-local-path",
-        if let Some(repo) = &self.model.gguf_repo {
+            |path| path.clone()
-            args.extend(["--gguf-repo".to_owned(), repo.clone()]);
+        );
-        }
+        push_opt!(args, &self.model.gguf_repo, "--gguf-repo", |repo| repo
-        if let Some(file) = &self.model.gguf_file {
+            .clone());
-            args.extend(["--gguf-file".to_owned(), file.clone()]);
+        push_opt!(args, &self.model.gguf_file, "--gguf-file", |file| file
-        }
+            .clone());
-        if let Some(revision) = &self.model.gguf_revision {
+        push_opt!(
-            args.extend(["--gguf-revision".to_owned(), revision.clone()]);
+            args,
-        }
+            &self.model.gguf_revision,
-        if let Some(path) = &self.model.tokenizer_local_path {
+            "--gguf-revision",
-            args.extend(["--tokenizer-local-path".to_owned(), path.clone()]);
+            |revision| { revision.clone() }
-        }
+        );
-        if let Some(max_context) = self.model.max_context {
+        push_opt!(
-            args.extend(["--max-context".to_owned(), max_context.to_string()]);
+            args,
-        }
+            &self.model.tokenizer_local_path,
+            "--tokenizer-local-path",
+            |path| path.clone()
+        );
+        push_opt!(
+            args,
+            self.model.max_context,
+            "--max-context",
+            |max_context| { max_context.to_string() }
+        );
         if self.provider == provider_kind::process() {
             args.extend([
                 "--worker-bin".to_owned(),
@@ -922,18 +984,14 @@ impl Config {
                 cached_model.host_path.to_string_lossy().to_string(),
             ]);
         }
-        if let Some(path) = &self.datastream_frame_log {
+        push_opt!(
-            args.extend([
+            args,
-                "--datastream-frame-log".to_owned(),
+            &self.datastream_frame_log,
-                path.to_string_lossy().to_string(),
+            "--datastream-frame-log",
-            ]);
+            |path| { path.to_string_lossy().to_string() }
-        }
+        );
-        if let Some(mode) = &self.relay_mode {
+        push_opt!(args, &self.relay_mode, "--relay-mode", |mode| mode.clone());
-            args.extend(["--relay-mode".to_owned(), mode.clone()]);
+        push_opt!(args, &self.relay_url, "--relay-url", |url| url.clone());
-        }
-        if let Some(url) = &self.relay_url {
-            args.extend(["--relay-url".to_owned(), url.clone()]);
-        }
         if self.endpoint_addr_mask != EndpointAddrMask::Full {
             args.extend([
                 "--endpoint-addr-mask".to_owned(),
@@ -946,39 +1004,41 @@ impl Config {
                 vastai.bootstrap_command.clone(),
                 "--no-vastai-confirm-lease".to_owned(),
             ]);
-            if let Some(disk_gb) = vastai.disk_gb {
+            push_opt!(args, vastai.disk_gb, "--vastai-disk-gb", |disk_gb| disk_gb
-                args.extend(["--vastai-disk-gb".to_owned(), disk_gb.to_string()]);
+                .to_string());
-            }
+            push_opt!(args, &vastai.gpu_name, "--vastai-gpu-name", |gpu_name| {
-            if let Some(gpu_name) = &vastai.gpu_name {
+                gpu_name.clone()
-                args.extend(["--vastai-gpu-name".to_owned(), gpu_name.clone()]);
+            });
-            }
+            push_opt!(
-            if let Some(min_gpu_ram_mb) = vastai.min_gpu_ram_mb {
+                args,
-                args.extend([
+                vastai.min_gpu_ram_mb,
-                    "--vastai-min-gpu-ram-mb".to_owned(),
+                "--vastai-min-gpu-ram-mb",
-                    min_gpu_ram_mb.to_string(),
+                |min_gpu_ram_mb| min_gpu_ram_mb.to_string()
-                ]);
+            );
-            }
+            push_opt!(
-            if let Some(min_down_mbps) = vastai.min_down_mbps {
+                args,
-                args.extend([
+                vastai.min_down_mbps,
-                    "--vastai-min-down-mbps".to_owned(),
+                "--vastai-min-down-mbps",
-                    min_down_mbps.to_string(),
+                |min_down_mbps| min_down_mbps.to_string()
-                ]);
+            );
-            }
+            push_opt!(
-            if let Some(min_up_mbps) = vastai.min_up_mbps {
+                args,
-                args.extend(["--vastai-min-up-mbps".to_owned(), min_up_mbps.to_string()]);
+                vastai.min_up_mbps,
-            }
+                "--vastai-min-up-mbps",
-            if let Some(max_dph_total) = vastai.max_dph_total {
+                |min_up_mbps| { min_up_mbps.to_string() }
-                args.extend([
+            );
-                    "--vastai-max-dph-total".to_owned(),
+            push_opt!(
-                    max_dph_total.to_string(),
+                args,
-                ]);
+                vastai.max_dph_total,
-            }
+                "--vastai-max-dph-total",
-            if let Some(min_reliability) = vastai.min_reliability {
+                |max_dph_total| { max_dph_total.to_string() }
-                args.extend([
+            );
-                    "--vastai-min-reliability".to_owned(),
+            push_opt!(
-                    min_reliability.to_string(),
+                args,
-                ]);
+                vastai.min_reliability,
-            }
+                "--vastai-min-reliability",
+                |min_reliability| min_reliability.to_string()
+            );
             if let Some(require_verified) = vastai.require_verified {
                 args.push(if require_verified {
                     "--vastai-require-verified".to_owned()
@@ -989,12 +1049,14 @@ impl Config {
             for host_id in &vastai.blacklist_hosts {
                 args.extend(["--vastai-blacklist-host".to_owned(), host_id.to_string()]);
             }
-            if let Some(onstart) = &vastai.onstart {
+            push_opt!(args, &vastai.onstart, "--vastai-onstart", |onstart| onstart
-                args.extend(["--vastai-onstart".to_owned(), onstart.clone()]);
+                .clone());
-            }
+            push_opt!(
-            if let Some(ssh_identity) = &vastai.ssh_identity {
+                args,
-                args.extend(["--vastai-ssh-identity".to_owned(), ssh_identity.clone()]);
+                &vastai.ssh_identity,
-            }
+                "--vastai-ssh-identity",
+                |ssh_identity| { ssh_identity.clone() }
+            );
         }
         args
     }
@@ -1044,6 +1106,70 @@ impl ParsedArgs {
         Ok(())
     }
 
+    fn apply_provider_arg(&mut self, arg: &str) -> Result<bool, String> {
+        match arg {
+            "--help" | "-h" | "help" => self.help = true,
+            "--gpu" => self.gpu = true,
+            "--vastai" => self.set_provider_selector(provider_kind::vastai())?,
+            "--process" => self.set_provider_selector(provider_kind::process())?,
+            "--docker" => self.set_provider_selector(provider_kind::docker())?,
+            "--yes" | "-y" => self.vastai_yes = true,
+            "--dump-logs" => self.dump_logs = true,
+            "--cached-model" => self.cached_model = Some(CachedModelSource::Discover),
+            "--skip-rebuild" => self.skip_rebuild = true,
+            _ => return Ok(false),
+        }
+        Ok(true)
+    }
+
+    fn apply_config_arg<I>(&mut self, arg: &str, args: &mut I) -> Result<bool, String>
+    where
+        I: Iterator<Item = String>,
+    {
+        match arg {
+            "--config" => self.config_path = Some(PathBuf::from(next_arg(args, "--config")?)),
+            "--pipeline-stages" | "--pipeline-parallel" => {
+                if self.pipeline_stages.is_some() {
+                    return Err("pipeline stage count was provided more than once".to_owned());
+                }
+                self.pipeline_stages = Some(parse_pipeline_stages_value(args, arg)?);
+            }
+            "--relay-mode" => self.relay_mode = Some(next_arg(args, "--relay-mode")?),
+            "--relay-url" => self.relay_url = Some(next_arg(args, "--relay-url")?),
+            "--endpoint-addr-mask" => {
+                self.endpoint_addr_mask = Some(next_arg(args, "--endpoint-addr-mask")?)
+            }
+            "--run-id" => {
+                let run_id: u64 = parse_next(args, "--run-id")?;
+                if run_id == 0 {
+                    return Err("--run-id must be greater than 0".to_owned());
+                }
+                self.run_id = Some(run_id);
+            }
+            _ => return Ok(false),
+        }
+        Ok(true)
+    }
+
+    fn apply_assignment_arg(&mut self, arg: &str) -> Result<bool, String> {
+        if let Some(path) = arg.strip_prefix("--dump-logs=") {
+            if path.is_empty() {
+                return Err("--dump-logs path must not be empty".to_owned());
+            }
+            self.dump_logs = true;
+            self.dump_log_path = Some(PathBuf::from(path));
+            return Ok(true);
+        }
+        if let Some(path) = arg.strip_prefix("--cached-model=") {
+            if path.is_empty() {
+                return Err("--cached-model path must not be empty".to_owned());
+            }
+            self.cached_model = Some(CachedModelSource::Path(PathBuf::from(path)));
+            return Ok(true);
+        }
+        Ok(false)
+    }
+
     fn parse<I>(provided_args: I) -> Result<Self, String>
     where
         I: IntoIterator<Item = String>,
@@ -1051,61 +1177,13 @@ impl ParsedArgs {
         let mut parsed = Self::default();
         let mut args = provided_args.into_iter().peekable();
         while let Some(arg) = args.next() {
-            match arg.as_str() {
+            if parsed.apply_provider_arg(&arg)?
-                "--help" | "-h" | "help" => parsed.help = true,
+                || parsed.apply_config_arg(&arg, &mut args)?
-                "--gpu" => parsed.gpu = true,
+                || parsed.apply_assignment_arg(&arg)?
-                "--vastai" => parsed.set_provider_selector(provider_kind::vastai())?,
+            {
-                "--process" => parsed.set_provider_selector(provider_kind::process())?,
+                continue;
-                "--docker" => parsed.set_provider_selector(provider_kind::docker())?,
-                "--yes" | "-y" => parsed.vastai_yes = true,
-                "--config" => {
-                    parsed.config_path = Some(PathBuf::from(next_arg(&mut args, "--config")?))
-                }
-                "--pipeline-stages" | "--pipeline-parallel" => {
-                    if parsed.pipeline_stages.is_some() {
-                        return Err("pipeline stage count was provided more than once".to_owned());
-                    }
-                    parsed.pipeline_stages =
-                        Some(parse_pipeline_stages_value(&mut args, arg.as_str())?)
-                }
-                "--relay-mode" => parsed.relay_mode = Some(next_arg(&mut args, "--relay-mode")?),
-                "--relay-url" => parsed.relay_url = Some(next_arg(&mut args, "--relay-url")?),
-                "--endpoint-addr-mask" => {
-                    parsed.endpoint_addr_mask = Some(next_arg(&mut args, "--endpoint-addr-mask")?)
-                }
-                "--run-id" => {
-                    let run_id: u64 = parse_next(&mut args, "--run-id")?;
-                    if run_id == 0 {
-                        return Err("--run-id must be greater than 0".to_owned());
-                    }
-                    parsed.run_id = Some(run_id);
-                }
-                "--dump-logs" => {
-                    parsed.dump_logs = true;
-                }
-                value if value.starts_with("--dump-logs=") => {
-                    let path = value.strip_prefix("--dump-logs=").expect("prefix checked");
-                    if path.is_empty() {
-                        return Err("--dump-logs path must not be empty".to_owned());
-                    }
-                    parsed.dump_logs = true;
-                    parsed.dump_log_path = Some(PathBuf::from(path));
-                }
-                "--cached-model" => {
-                    parsed.cached_model = Some(CachedModelSource::Discover);
-                }
-                value if value.starts_with("--cached-model=") => {
-                    let path = value
-                        .strip_prefix("--cached-model=")
-                        .expect("prefix checked");
-                    if path.is_empty() {
-                        return Err("--cached-model path must not be empty".to_owned());
-                    }
-                    parsed.cached_model = Some(CachedModelSource::Path(PathBuf::from(path)));
-                }
-                "--skip-rebuild" => parsed.skip_rebuild = true,
-                other => return Err(format!("unsupported mvp-chat argument {other:?}")),
             }
+            return Err(format!("unsupported mvp-chat argument {arg:?}"));
         }
         Ok(parsed)
     }
@@ -1511,7 +1589,7 @@ fn signal_orch_process_group(child: &Child, signal: libc::c_int) -> io::Result<(
 fn prepare_node_image_progress_adapter(
     request: NodeImageRequest,
     progress: Option<&mut dyn NodeImageProgressSink>,
-) -> Result<PreparedNodeImage, String> {
+) -> Result<String, String> {
     prepare_node_image_with_progress(request, progress)
 }
 
@@ -1525,7 +1603,7 @@ fn prepare_runtime(config: &Config) -> Result<String, String> {
 #[allow(dead_code)]
 fn prepare_runtime_with<F>(config: &Config, prepare_node_image_fn: F) -> Result<String, String>
 where
-    F: FnMut(NodeImageRequest) -> Result<PreparedNodeImage, String>,
+    F: FnMut(NodeImageRequest) -> Result<String, String>,
 {
     let mut prepare_node_image_fn = prepare_node_image_fn;
     prepare_runtime_with_progress(
@@ -1541,10 +1619,7 @@ fn prepare_runtime_with_progress<F>(
     progress: Option<&mut ChatDatastream>,
 ) -> Result<String, String>
 where
-    F: FnMut(
+    F: FnMut(NodeImageRequest, Option<&mut dyn NodeImageProgressSink>) -> Result<String, String>,
-        NodeImageRequest,
-        Option<&mut dyn NodeImageProgressSink>,
-    ) -> Result<PreparedNodeImage, String>,
 {
     let mut progress = progress;
     let binary_mode = if config.skip_rebuild {
@@ -1758,9 +1833,9 @@ where
         CHAT_RUNTIME_CHANNEL,
         "prepare_node_image",
         "ready",
-        json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_ref": prepared.image_ref, "elapsed_ms": prepare_node_image_started.elapsed().as_millis()}),
+        json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_ref": &prepared, "elapsed_ms": prepare_node_image_started.elapsed().as_millis()}),
     );
-    Ok(prepared.image_ref)
+    Ok(prepared)
 }
 
 fn stdin_prompt_events() -> mpsc::Receiver<PromptInput> {
@@ -1852,32 +1927,23 @@ fn run_chat_loop_with_input_and_progress(
 }
 
 #[cfg(test)]
-fn run_chat_session_with_output<R, W, O>(
+fn run_chat_session_with_output(
-    writer: &mut W,
+    writer: &mut impl Write,
-    reader: R,
+    reader: impl BufRead,
     input_rx: mpsc::Receiver<PromptInput>,
     max_tokens: u32,
-    output: &mut O,
+    output: &mut impl Write,
-) -> Result<(), String>
+) -> Result<(), String> {
-where
-    R: BufRead,
-    W: Write,
-    O: Write,
-{
     run_chat_session_with_output_and_progress(writer, reader, input_rx, max_tokens, output, None)
 }
 
-fn run_chat_session_with_progress<R, W>(
+fn run_chat_session_with_progress(
-    writer: &mut W,
+    writer: &mut impl Write,
-    reader: R,
+    reader: impl BufRead,
     input_rx: mpsc::Receiver<PromptInput>,
     max_tokens: u32,
     progress: Option<&mut ChatDatastream>,
-) -> Result<(), String>
+) -> Result<(), String> {
-where
-    R: BufRead,
-    W: Write,
-{
     let mut output = io::stdout();
     run_chat_session_with_output_and_progress(
         writer,
@@ -1905,19 +1971,14 @@ fn prompt_hash_hex(prompt: &str) -> String {
     blake3::hash(prompt.as_bytes()).to_hex().to_string()
 }
 
-fn run_chat_session_with_output_and_progress<R, W, O>(
+fn run_chat_session_with_output_and_progress(
-    writer: &mut W,
+    writer: &mut impl Write,
-    mut reader: R,
+    mut reader: impl BufRead,
     input_rx: mpsc::Receiver<PromptInput>,
     max_tokens: u32,
-    output: &mut O,
+    output: &mut impl Write,
     progress: Option<&mut ChatDatastream>,
-) -> Result<(), String>
+) -> Result<(), String> {
-where
-    R: BufRead,
-    W: Write,
-    O: Write,
-{
     let mut progress = progress;
     let mut next_request_id = 1_u64;
     let mut next_prompt_index = 1_u64;
@@ -2375,10 +2436,8 @@ fn parse_pipeline_stages_value(
 mod tests {
     use super::*;
 
-    use std::ffi::{OsStr, OsString};
+    use std::ffi::OsString;
     use std::io::{Cursor, Read};
-    #[cfg(target_os = "linux")]
-    use std::os::unix::process::CommandExt;
     use std::path::Path;
     use std::sync::Mutex;
     use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
@@ -2487,50 +2546,6 @@ mod tests {
         path
     }
 
-    fn base_config(provider: ProviderKind) -> Config {
-        Config {
-            orch_bin: PathBuf::from("/tmp/mvp-orchestrator"),
-            worker_bin: PathBuf::from("/tmp/mvp-worker-node"),
-            rpc_addr: DEFAULT_RPC_ADDR.to_owned(),
-            node_image: "docker.io/acme/node:latest".to_owned(),
-            provider,
-            image_tag: None,
-            cached_model: None,
-            datastream_frame_log: None,
-            run_id: 1,
-            vastai_yes: false,
-            vastai: None,
-            model: ChatModelConfig::default(),
-            pipeline_stages: 1,
-            max_tokens: DEFAULT_MAX_TOKENS,
-            skip_rebuild: true,
-            gpu_run: false,
-            relay_mode: None,
-            relay_url: None,
-            endpoint_addr_mask: EndpointAddrMask::Full,
-        }
-    }
-
-    fn valid_vastai() -> ResolvedVastAiConfig {
-        ResolvedVastAiConfig {
-            api_key: "secret".to_owned(),
-            relay_url: "https://relay.example".to_owned(),
-            image: "docker.io/acme/node:latest".to_owned(),
-            bootstrap_command: "boot".to_owned(),
-            disk_gb: None,
-            gpu_name: None,
-            min_gpu_ram_mb: None,
-            min_down_mbps: None,
-            min_up_mbps: None,
-            max_dph_total: None,
-            min_reliability: None,
-            require_verified: None,
-            blacklist_hosts: Vec::new(),
-            onstart: None,
-            ssh_identity: None,
-        }
-    }
-
     fn channel_lines(lines: &[&str]) -> mpsc::Receiver<PromptInput> {
         let (tx, rx) = mpsc::channel();
         for line in lines {
@@ -2802,131 +2817,6 @@ kind = "mock"
         });
     }
 
-    struct MockApproval {
-        terminal: bool,
-        answer: Result<bool, String>,
-    }
-
-    impl VastAiApproval for MockApproval {
-        fn stdin_is_terminal(&self) -> bool {
-            self.terminal
-        }
-
-        fn ask(&mut self) -> Result<bool, String> {
-            self.answer.clone()
-        }
-    }
-
-    fn panic_prepare_node_image(_: NodeImageRequest) -> Result<PreparedNodeImage, String> {
-        panic!("image preparer must not be called when --skip-rebuild is set")
-    }
-
-    fn panic_prepare_node_image_with_progress(
-        _request: NodeImageRequest,
-        _progress: Option<&mut dyn NodeImageProgressSink>,
-    ) -> Result<PreparedNodeImage, String> {
-        panic!("image preparer must not be called when --skip-rebuild is set")
-    }
-
-    fn runtime_events(path: &Path) -> Vec<serde_json::Value> {
-        fs::read_to_string(path)
-            .expect("read progress archive")
-            .lines()
-            .filter_map(|line| {
-                let outer: serde_json::Value = serde_json::from_str(line).ok()?;
-                if outer.get("channel").and_then(serde_json::Value::as_str)
-                    != Some(CHAT_RUNTIME_CHANNEL)
-                {
-                    return None;
-                }
-                outer
-                    .get("payload")?
-                    .get("value")?
-                    .as_str()
-                    .and_then(|text| serde_json::from_str::<serde_json::Value>(text).ok())
-            })
-            .collect()
-    }
-
-    fn emit_fake_node_image_progress(
-        progress: Option<&mut dyn NodeImageProgressSink>,
-        success: bool,
-    ) {
-        let Some(sink) = progress else {
-            return;
-        };
-        sink.emit(NodeImageProgressEvent {
-            command_label: None,
-            image_ref: Some("docker.io/acme/node:prepared".to_owned()),
-            elapsed_ms: None,
-            kind: NodeImageProgressEventKind::ImageReference {
-                role: "resolved".to_owned(),
-                image_ref: "docker.io/acme/node:prepared".to_owned(),
-            },
-        });
-        sink.emit(NodeImageProgressEvent {
-            command_label: Some("build mvp node image".to_owned()),
-            image_ref: Some("docker.io/acme/node:prepared".to_owned()),
-            elapsed_ms: Some(0),
-            kind: NodeImageProgressEventKind::CommandStarted {
-                program: "fake-docker".to_owned(),
-                args: vec!["build".to_owned()],
-            },
-        });
-        sink.emit(NodeImageProgressEvent {
-            command_label: Some("build mvp node image".to_owned()),
-            image_ref: Some("docker.io/acme/node:prepared".to_owned()),
-            elapsed_ms: Some(1),
-            kind: NodeImageProgressEventKind::CommandStdout {
-                line: "building layer".to_owned(),
-            },
-        });
-        sink.emit(NodeImageProgressEvent {
-            command_label: Some("build mvp node image".to_owned()),
-            image_ref: Some("docker.io/acme/node:prepared".to_owned()),
-            elapsed_ms: Some(2),
-            kind: NodeImageProgressEventKind::CommandStderr {
-                line: "pushing metadata".to_owned(),
-            },
-        });
-        sink.emit(NodeImageProgressEvent {
-            command_label: Some("build mvp node image".to_owned()),
-            image_ref: Some("docker.io/acme/node:prepared".to_owned()),
-            elapsed_ms: Some(3),
-            kind: NodeImageProgressEventKind::CommandExited {
-                status: if success {
-                    "exit status: 0".to_owned()
-                } else {
-                    "exit status: 42".to_owned()
-                },
-                code: Some(if success { 0 } else { 42 }),
-                success,
-            },
-        });
-    }
-
-    fn fake_prepare_node_image_with_progress(
-        _request: NodeImageRequest,
-        progress: Option<&mut dyn NodeImageProgressSink>,
-    ) -> Result<PreparedNodeImage, String> {
-        emit_fake_node_image_progress(progress, true);
-        Ok(PreparedNodeImage {
-            image_ref: "docker.io/acme/node:prepared".to_owned(),
-            tag: "prepared".to_owned(),
-            already_available: false,
-            built: true,
-            pushed: false,
-        })
-    }
-
-    fn failing_prepare_node_image_with_progress(
-        _request: NodeImageRequest,
-        progress: Option<&mut dyn NodeImageProgressSink>,
-    ) -> Result<PreparedNodeImage, String> {
-        emit_fake_node_image_progress(progress, false);
-        Err("build mvp node image failed with exit status: 42".to_owned())
-    }
-
     #[test]
     fn prompt_loop_exits_cleanly_and_ignores_empty_prompts() {
         let mut rpc_writer = Vec::new();

crates/mvp-system/src/node/actor.rs

@@ -347,7 +347,89 @@ impl NodeAgentActor {
         }
     }
 
+    fn forward_prompt_or_snapshot(&mut self, ctx: &Ctx, msg: NodeAgentMsg) -> Option<NodeAgentMsg> {
+        match msg {
+            NodeAgentMsg::InferPrompt {
+                request_id,
+                prompt,
+                max_tokens,
+                reply_to,
+            } => {
+                if let Some(report_to) = self.report_to {
+                    let _ = ctx.send(
+                        report_to,
+                        NodeAgentReport::PromptRequested {
+                            request_id,
+                            prompt,
+                            max_tokens,
+                            reply_to,
+                        },
+                    );
+                }
+                None
+            }
+            NodeAgentMsg::EncodePrompt {
+                request_id,
+                prompt,
+                reply_to,
+            } => {
+                if let Some(report_to) = self.report_to {
+                    let _ = ctx.send(
+                        report_to,
+                        NodeAgentReport::EncodePromptRequested {
+                            request_id,
+                            prompt,
+                            reply_to,
+                        },
+                    );
+                }
+                None
+            }
+            NodeAgentMsg::DecodeTokens {
+                request_id,
+                tokens,
+                reply_to,
+            } => {
+                if let Some(report_to) = self.report_to {
+                    let _ = ctx.send(
+                        report_to,
+                        NodeAgentReport::DecodeTokensRequested {
+                            request_id,
+                            tokens,
+                            reply_to,
+                        },
+                    );
+                }
+                None
+            }
+            NodeAgentMsg::Snapshot { reply_to } => {
+                let _ = ctx.send(
+                    reply_to,
+                    NodeAgentReport::Snapshot {
+                        commands: self
+                            .core
+                            .commands()
+                            .iter()
+                            .map(StageCommandWire::from)
+                            .collect(),
+                        events: self
+                            .core
+                            .events()
+                            .iter()
+                            .map(StageLifecycleWire::from)
+                            .collect(),
+                    },
+                );
+                None
+            }
+            other => Some(other),
+        }
+    }
+
     fn observe(&mut self, ctx: &Ctx, msg: NodeAgentMsg) {
+        let Some(msg) = self.forward_prompt_or_snapshot(ctx, msg) else {
+            return;
+        };
         match msg {
             NodeAgentMsg::ProvisionStage(provision) => {
                 self.core.observe(stage::StageEvent::ProvisionStage {
@@ -480,79 +562,10 @@ impl NodeAgentActor {
                     run_id: stage::RunId(run_id),
                 })
             }
-            NodeAgentMsg::InferPrompt {
+            NodeAgentMsg::InferPrompt { .. }
-                request_id,
+            | NodeAgentMsg::EncodePrompt { .. }
-                prompt,
+            | NodeAgentMsg::DecodeTokens { .. }
-                max_tokens,
+            | NodeAgentMsg::Snapshot { .. } => unreachable!("prompt messages returned early"),
-                reply_to,
-            } => {
-                if let Some(report_to) = self.report_to {
-                    let _ = ctx.send(
-                        report_to,
-                        NodeAgentReport::PromptRequested {
-                            request_id,
-                            prompt,
-                            max_tokens,
-                            reply_to,
-                        },
-                    );
-                }
-                return;
-            }
-            NodeAgentMsg::EncodePrompt {
-                request_id,
-                prompt,
-                reply_to,
-            } => {
-                if let Some(report_to) = self.report_to {
-                    let _ = ctx.send(
-                        report_to,
-                        NodeAgentReport::EncodePromptRequested {
-                            request_id,
-                            prompt,
-                            reply_to,
-                        },
-                    );
-                }
-                return;
-            }
-            NodeAgentMsg::DecodeTokens {
-                request_id,
-                tokens,
-                reply_to,
-            } => {
-                if let Some(report_to) = self.report_to {
-                    let _ = ctx.send(
-                        report_to,
-                        NodeAgentReport::DecodeTokensRequested {
-                            request_id,
-                            tokens,
-                            reply_to,
-                        },
-                    );
-                }
-                return;
-            }
-            NodeAgentMsg::Snapshot { reply_to } => {
-                let _ = ctx.send(
-                    reply_to,
-                    NodeAgentReport::Snapshot {
-                        commands: self
-                            .core
-                            .commands()
-                            .iter()
-                            .map(StageCommandWire::from)
-                            .collect(),
-                        events: self
-                            .core
-                            .events()
-                            .iter()
-                            .map(StageLifecycleWire::from)
-                            .collect(),
-                    },
-                );
-                return;
-            }
         }
         self.drain_outputs(ctx);
     }

crates/mvp-system/src/node/worker_node_runtime.rs

@@ -1252,6 +1252,25 @@ impl WorkerEdgeRuntime {
         driver: &mut IrohDriver,
     ) -> Result<(), String> {
         loop {
+            let progressed =
+                self.drain_edge_commands(worker, arena_manager, config, datastream, driver)?
+                    || self.drain_driver_events()
+                    || self.drain_edge_events(stack, node_actor)?;
+            if !progressed {
+                break;
+            }
+        }
+        Ok(())
+    }
+
+    fn drain_edge_commands(
+        &mut self,
+        worker: &mut TinygradWorker,
+        arena_manager: &Arc<Mutex<arena::ArenaManager>>,
+        config: &DeploymentConfig,
+        datastream: &mut NodeDatastream,
+        driver: &mut IrohDriver,
+    ) -> Result<bool, String> {
         let mut progressed = false;
         while self.edge_command_cursor < self.establisher.commands().len() {
             let command = self.establisher.commands()[self.edge_command_cursor].clone();
@@ -1457,7 +1476,11 @@ impl WorkerEdgeRuntime {
                 }
             }
         }
+        Ok(progressed)
+    }
 
+    fn drain_driver_events(&mut self) -> bool {
+        let mut progressed = false;
         while self.driver_event_cursor < self.driver_model.events().len() {
             let event = self.driver_model.events()[self.driver_event_cursor].clone();
             self.driver_event_cursor += 1;
@@ -1494,7 +1517,15 @@ impl WorkerEdgeRuntime {
                 driver_model::DriverEventOut::StreamClosed { .. } => {}
             }
         }
+        progressed
+    }
 
+    fn drain_edge_events(
+        &mut self,
+        stack: &DistributionRuntimeStack,
+        node_actor: ActorAddress,
+    ) -> Result<bool, String> {
+        let mut progressed = false;
         while self.edge_event_cursor < self.establisher.events().len() {
             let event = self.establisher.events()[self.edge_event_cursor].clone();
             self.edge_event_cursor += 1;
@@ -1543,11 +1574,7 @@ impl WorkerEdgeRuntime {
                 edge::EdgeLifecycleEvent::EdgeStopped { .. } => {}
             }
         }
-            if !progressed {
+        Ok(progressed)
-                break;
-            }
-        }
-        Ok(())
     }
 }
 

crates/mvp-system/src/orchestration/engine_builder/mod.rs

@@ -20,8 +20,7 @@ pub mod runtime_stack;
 pub mod workload;
 
 pub use crate::run_plan::NodeId;
-pub use engine::{ClusterBuilder, ClusterHandle};
+pub use engine::ClusterBuilder;
-pub use error::{EngineBuildError, PlanningError};
 pub use events::EngineEvent;
 pub use launcher::StaticNodeLauncher;
 pub use model::{DTypeFamily, ModelArtifact, ModelSpec};
@@ -29,5 +28,3 @@ pub use node_image::{NodeImageSpec, WorkerRuntimeSpec};
 pub use planner::FixedLinearPipelinePlanner;
 pub use pool::{NodeCapability, NodeLease, ResourceFacts, StaticPoolProvider};
 pub use roles::RoleKind;
-pub use runtime_stack::RuntimeNode;
-pub use workload::WorkloadAdapter;

crates/mvp-system/src/orchestration/run_fsm.rs

@@ -230,14 +230,7 @@ impl OrchestratorRun {
             RunEvent::StageReady {
                 run_id,
                 stage_index,
-            } => {
+            } => self.stage_ready(run_id, stage_index),
-                if run_id != self.config.run_id || !self.plan_has_stage(stage_index) {
-                    self.fault(RunFaultReason::UnknownStageReady { stage_index });
-                    return;
-                }
-                self.ready_stages.insert(stage_index);
-                self.maybe_inject_initial();
-            }
             RunEvent::TokenInEndpointReady => {
                 self.token_in_ready = true;
                 self.maybe_inject_initial();
@@ -250,11 +243,52 @@ impl OrchestratorRun {
                 sequence,
                 token_id,
                 eos,
+            } => self.token_received(sequence, token_id, eos),
+            RunEvent::StageFault { run_id, .. }
+            | RunEvent::EndpointFault { run_id, .. }
+            | RunEvent::OperatorStop { run_id }
+            | RunEvent::MembershipLost { run_id, .. }
+            | RunEvent::StageStopped { run_id, .. }
+                if run_id != self.config.run_id => {}
+            RunEvent::StageFault {
+                stage_index,
+                reason,
+                ..
             } => {
-                if self.terminal {
+                self.fault(RunFaultReason::StageFault {
+                    stage_index,
+                    reason,
+                });
+            }
+            RunEvent::EndpointFault { endpoint, .. } => {
+                self.fault(RunFaultReason::EndpointFault { endpoint });
+            }
+            RunEvent::OperatorStop { .. } => self.operator_stop(),
+            RunEvent::MembershipLost { node_id, .. } => {
+                self.fault(RunFaultReason::MembershipLost { node_id });
+            }
+            RunEvent::StageStopped { stage_index, .. } => {
+                self.stopped_stages.insert(stage_index);
+                self.maybe_torn_down();
+            }
+            RunEvent::TokenEndpointsStopped => {
+                self.token_endpoints_stopped = true;
+                self.maybe_torn_down();
+            }
+        }
+    }
+
+    fn stage_ready(&mut self, run_id: RunId, stage_index: u32) {
+        if run_id != self.config.run_id || !self.plan_has_stage(stage_index) {
+            self.fault(RunFaultReason::UnknownStageReady { stage_index });
             return;
         }
-                if sequence != self.expected_token_sequence {
+        self.ready_stages.insert(stage_index);
+        self.maybe_inject_initial();
+    }
+
+    fn token_received(&mut self, sequence: u64, token_id: u32, eos: bool) {
+        if self.terminal || sequence != self.expected_token_sequence {
             return;
         }
         self.expected_token_sequence += 1;
@@ -266,43 +300,6 @@ impl OrchestratorRun {
             self.complete();
         }
     }
-            RunEvent::StageFault {
-                run_id,
-                stage_index,
-                reason,
-            } if run_id == self.config.run_id => {
-                self.fault(RunFaultReason::StageFault {
-                    stage_index,
-                    reason,
-                });
-            }
-            RunEvent::EndpointFault { run_id, endpoint } if run_id == self.config.run_id => {
-                self.fault(RunFaultReason::EndpointFault { endpoint });
-            }
-            RunEvent::OperatorStop { run_id } if run_id == self.config.run_id => {
-                self.operator_stop();
-            }
-            RunEvent::MembershipLost { run_id, node_id } if run_id == self.config.run_id => {
-                self.fault(RunFaultReason::MembershipLost { node_id });
-            }
-            RunEvent::StageStopped {
-                run_id,
-                stage_index,
-            } if run_id == self.config.run_id => {
-                self.stopped_stages.insert(stage_index);
-                self.maybe_torn_down();
-            }
-            RunEvent::TokenEndpointsStopped => {
-                self.token_endpoints_stopped = true;
-                self.maybe_torn_down();
-            }
-            RunEvent::StageFault { .. }
-            | RunEvent::EndpointFault { .. }
-            | RunEvent::OperatorStop { .. }
-            | RunEvent::MembershipLost { .. }
-            | RunEvent::StageStopped { .. } => {}
-        }
-    }
 
     pub fn advance_time_ms(&mut self, _delta: u64) {}
 

crates/mvp-system/src/staging/control.rs

@@ -302,22 +302,10 @@ impl StageController {
             StageEvent::WorkerReady => self.worker_ready = true,
             StageEvent::WeightsReady => self.weights_ready = true,
             StageEvent::InboundEdgeReady { edge_id } => {
-                if self
+                self.edge_ready(edge_id, EdgeDirection::Inbound)
-                    .provision
-                    .as_ref()
-                    .is_some_and(|p| p.inbound.edge_id == edge_id)
-                {
-                    self.inbound_ready = true;
-                }
             }
             StageEvent::OutboundEdgeReady { edge_id } => {
-                if self
+                self.edge_ready(edge_id, EdgeDirection::Outbound)
-                    .provision
-                    .as_ref()
-                    .is_some_and(|p| p.outbound.edge_id == edge_id)
-                {
-                    self.outbound_ready = true;
-                }
             }
             StageEvent::ObjectLoaded {
                 edge_id,
@@ -340,6 +328,21 @@ impl StageController {
         self.maybe_stage_ready();
     }
 
+    fn edge_ready(&mut self, edge_id: EdgeId, direction: EdgeDirection) {
+        let Some(provision) = &self.provision else {
+            return;
+        };
+        match direction {
+            EdgeDirection::Inbound if provision.inbound.edge_id == edge_id => {
+                self.inbound_ready = true
+            }
+            EdgeDirection::Outbound if provision.outbound.edge_id == edge_id => {
+                self.outbound_ready = true
+            }
+            _ => {}
+        }
+    }
+
     pub fn commands(&self) -> &[StageCommand] {
         &self.commands
     }
@@ -380,16 +383,18 @@ impl StageController {
         if self.stage_ready_emitted || self.faulted || self.stopped || self.stopping_run.is_some() {
             return;
         }
-        if self.worker_ready && self.weights_ready && self.inbound_ready && self.outbound_ready {
+        if !(self.worker_ready && self.weights_ready && self.inbound_ready && self.outbound_ready) {
-            if let Some(provision) = &self.provision {
+            return;
+        }
+        let Some(provision) = &self.provision else {
+            return;
+        };
         self.stage_ready_emitted = true;
         self.events.push(StageLifecycleEvent::StageReady {
             run_id: provision.run_id,
             stage_index: provision.stage_index,
         });
     }
-        }
-    }
 
     fn object_loaded(
         &mut self,

crates/mvp-system/src/staging/gguf_metadata.rs

@@ -174,23 +174,26 @@ enum GgufValueType {
 
 impl GgufValueType {
     fn read<R: Read>(reader: &mut R) -> Result<Self, String> {
+        const VALUE_TYPES: [GgufValueType; 13] = [
+            GgufValueType::Uint8,
+            GgufValueType::Int8,
+            GgufValueType::Uint16,
+            GgufValueType::Int16,
+            GgufValueType::Uint32,
+            GgufValueType::Int32,
+            GgufValueType::Float32,
+            GgufValueType::Bool,
+            GgufValueType::String,
+            GgufValueType::Array,
+            GgufValueType::Uint64,
+            GgufValueType::Int64,
+            GgufValueType::Float64,
+        ];
         let raw = read_u32(reader)?;
-        match raw {
+        VALUE_TYPES
-            0 => Ok(Self::Uint8),
+            .get(raw as usize)
-            1 => Ok(Self::Int8),
+            .copied()
-            2 => Ok(Self::Uint16),
+            .ok_or_else(|| format!("unsupported GGUF metadata value type {raw}"))
-            3 => Ok(Self::Int16),
-            4 => Ok(Self::Uint32),
-            5 => Ok(Self::Int32),
-            6 => Ok(Self::Float32),
-            7 => Ok(Self::Bool),
-            8 => Ok(Self::String),
-            9 => Ok(Self::Array),
-            10 => Ok(Self::Uint64),
-            11 => Ok(Self::Int64),
-            12 => Ok(Self::Float64),
-            other => Err(format!("unsupported GGUF metadata value type {other}")),
-        }
     }
 
     fn is_integer(self) -> bool {

crates/mvp-system/src/staging/gguf_shard.rs

@@ -1,6 +1,6 @@
 use std::fs::File;
 use std::io::{Read, Seek, SeekFrom, Write};
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 use serde::{Deserialize, Serialize};
 
@@ -537,54 +537,16 @@ where
     }
     let metadata_body = &metadata_prefix[24..];
     let alignment = u64::from(plan.alignment.max(1));
-    let mut data_offsets = Vec::with_capacity(plan.tensors.len());
+    let data_offsets = stage_shard_data_offsets(plan, alignment)?;
-    let mut data_cursor = 0_u64;
-    for tensor in &plan.tensors {
-        data_cursor = align_to(data_cursor, alignment)?;
-        data_offsets.push(data_cursor);
-        data_cursor = data_cursor
-            .checked_add(tensor.byte_len)
-            .ok_or_else(|| format!("stage shard data size overflow at tensor {}", tensor.name))?;
-    }
 
-    let partial_path = output_path.with_extension(format!(
+    let partial_path = partial_stage_shard_path(output_path);
-        "{}partial",
-        output_path
-            .extension()
-            .and_then(|value| value.to_str())
-            .map(|ext| format!("{ext}."))
-            .unwrap_or_default()
-    ));
     if let Some(parent) = output_path.parent() {
         std::fs::create_dir_all(parent)
             .map_err(|e| format!("create stage shard cache dir {}: {e}", parent.display()))?;
     }
     let mut out = File::create(&partial_path)
         .map_err(|e| format!("create stage shard {}: {e}", partial_path.display()))?;
-    out.write_all(GGUF_MAGIC)
+    write_stage_shard_header(&mut out, plan, metadata_body, &data_offsets, alignment)?;
-        .map_err(|e| format!("write stage shard magic: {e}"))?;
-    out.write_all(&SUPPORTED_GGUF_VERSION.to_le_bytes())
-        .map_err(|e| format!("write stage shard version: {e}"))?;
-    out.write_all(&(plan.tensors.len() as u64).to_le_bytes())
-        .map_err(|e| format!("write stage shard tensor count: {e}"))?;
-    out.write_all(&plan.metadata_count.to_le_bytes())
-        .map_err(|e| format!("write stage shard metadata count: {e}"))?;
-    out.write_all(metadata_body)
-        .map_err(|e| format!("write stage shard metadata: {e}"))?;
-    for (tensor, data_offset) in plan.tensors.iter().zip(data_offsets.iter().copied()) {
-        write_gguf_string(&mut out, &tensor.name)?;
-        out.write_all(&(tensor.dims.len() as u32).to_le_bytes())
-            .map_err(|e| format!("write tensor dim count for {}: {e}", tensor.name))?;
-        for dim in &tensor.dims {
-            out.write_all(&dim.to_le_bytes())
-                .map_err(|e| format!("write tensor dim for {}: {e}", tensor.name))?;
-        }
-        out.write_all(&tensor.ggml_type.to_le_bytes())
-            .map_err(|e| format!("write tensor type for {}: {e}", tensor.name))?;
-        out.write_all(&data_offset.to_le_bytes())
-            .map_err(|e| format!("write tensor offset for {}: {e}", tensor.name))?;
-    }
-    pad_writer_to_alignment(&mut out, alignment)?;
     let mut written_data = 0_u64;
     let mut tensor_index = 0_usize;
     for range in &plan.merged_tensor_ranges {
@@ -712,6 +674,63 @@ where
     Ok(())
 }
 
+fn stage_shard_data_offsets(plan: &StageShardPlan, alignment: u64) -> Result<Vec<u64>, String> {
+    let mut offsets = Vec::with_capacity(plan.tensors.len());
+    let mut cursor = 0_u64;
+    for tensor in &plan.tensors {
+        cursor = align_to(cursor, alignment)?;
+        offsets.push(cursor);
+        cursor = cursor
+            .checked_add(tensor.byte_len)
+            .ok_or_else(|| format!("stage shard data size overflow at tensor {}", tensor.name))?;
+    }
+    Ok(offsets)
+}
+
+fn partial_stage_shard_path(output_path: &Path) -> PathBuf {
+    output_path.with_extension(format!(
+        "{}partial",
+        output_path
+            .extension()
+            .and_then(|value| value.to_str())
+            .map(|ext| format!("{ext}."))
+            .unwrap_or_default()
+    ))
+}
+
+fn write_stage_shard_header(
+    out: &mut File,
+    plan: &StageShardPlan,
+    metadata_body: &[u8],
+    data_offsets: &[u64],
+    alignment: u64,
+) -> Result<(), String> {
+    out.write_all(GGUF_MAGIC)
+        .map_err(|e| format!("write stage shard magic: {e}"))?;
+    out.write_all(&SUPPORTED_GGUF_VERSION.to_le_bytes())
+        .map_err(|e| format!("write stage shard version: {e}"))?;
+    out.write_all(&(plan.tensors.len() as u64).to_le_bytes())
+        .map_err(|e| format!("write stage shard tensor count: {e}"))?;
+    out.write_all(&plan.metadata_count.to_le_bytes())
+        .map_err(|e| format!("write stage shard metadata count: {e}"))?;
+    out.write_all(metadata_body)
+        .map_err(|e| format!("write stage shard metadata: {e}"))?;
+    for (tensor, data_offset) in plan.tensors.iter().zip(data_offsets.iter().copied()) {
+        write_gguf_string(out, &tensor.name)?;
+        out.write_all(&(tensor.dims.len() as u32).to_le_bytes())
+            .map_err(|e| format!("write tensor dim count for {}: {e}", tensor.name))?;
+        for dim in &tensor.dims {
+            out.write_all(&dim.to_le_bytes())
+                .map_err(|e| format!("write tensor dim for {}: {e}", tensor.name))?;
+        }
+        out.write_all(&tensor.ggml_type.to_le_bytes())
+            .map_err(|e| format!("write tensor type for {}: {e}", tensor.name))?;
+        out.write_all(&data_offset.to_le_bytes())
+            .map_err(|e| format!("write tensor offset for {}: {e}", tensor.name))?;
+    }
+    pad_writer_to_alignment(out, alignment)
+}
+
 fn fetch_http_range(url: &str, start: u64, len: u64) -> Result<Vec<u8>, String> {
     if len == 0 {
         return Ok(Vec::new());
@@ -793,22 +812,26 @@ enum GgufValueType {
 
 impl GgufValueType {
     fn read<R: Read>(reader: &mut R) -> Result<Self, String> {
-        match read_u32(reader)? {
+        const VALUE_TYPES: [GgufValueType; 13] = [
-            0 => Ok(Self::Uint8),
+            GgufValueType::Uint8,
-            1 => Ok(Self::Int8),
+            GgufValueType::Int8,
-            2 => Ok(Self::Uint16),
+            GgufValueType::Uint16,
-            3 => Ok(Self::Int16),
+            GgufValueType::Int16,
-            4 => Ok(Self::Uint32),
+            GgufValueType::Uint32,
-            5 => Ok(Self::Int32),
+            GgufValueType::Int32,
-            6 => Ok(Self::Float32),
+            GgufValueType::Float32,
-            7 => Ok(Self::Bool),
+            GgufValueType::Bool,
-            8 => Ok(Self::String),
+            GgufValueType::String,
-            9 => Ok(Self::Array),
+            GgufValueType::Array,
-            10 => Ok(Self::Uint64),
+            GgufValueType::Uint64,
-            11 => Ok(Self::Int64),
+            GgufValueType::Int64,
-            12 => Ok(Self::Float64),
+            GgufValueType::Float64,
-            other => Err(format!("unsupported GGUF value type {other}")),
+        ];
-        }
+        let raw = read_u32(reader)?;
+        VALUE_TYPES
+            .get(raw as usize)
+            .copied()
+            .ok_or_else(|| format!("unsupported GGUF value type {raw}"))
     }
 
     fn fixed_width(self) -> Option<u64> {