Add CLI flags for explicit OCI authentication credentials

Cargo.toml

@@ -45,7 +45,7 @@ wasmtime-wasi-config = "38.0.4"
 [dependencies]
 anyhow = { workspace = true }
 axum = "0.8"
-clap = { version = "4.5", features = ["derive"] }
+clap = { version = "4.5", features = ["derive", "env"] }
 clap_complete = "4.5"
 etcetera = { workspace = true }
 figment = { version = "0.10", features = ["env", "toml"] }

crates/mcp-server/src/components.rs

@@ -11,7 +11,7 @@ use rmcp::{Peer, RoleServer};
 use serde_json::{json, Value};
 use tracing::{debug, error, info, instrument};
 use wassette::schema::{canonicalize_output_schema, ensure_structured_result};
-use wassette::{ComponentLoadOutcome, LifecycleManager, LoadResult};
+use wassette::{ComponentLoadOutcome, LifecycleManager, LoadResult, OciCredentials};
 
 #[instrument(skip(lifecycle_manager))]
 pub(crate) async fn get_component_tools(lifecycle_manager: &LifecycleManager) -> Result<Vec<Tool>> {
@@ -388,6 +388,16 @@ async fn handle_tool_list_notification(
 pub async fn handle_load_component_cli(
     req: &CallToolRequestParam,
     lifecycle_manager: &LifecycleManager,
+) -> Result<CallToolResult> {
+    handle_load_component_cli_with_credentials(req, lifecycle_manager, None).await
+}
+
+/// CLI-specific version of handle_load_component with optional credentials
+#[instrument(skip(lifecycle_manager))]
+pub async fn handle_load_component_cli_with_credentials(
+    req: &CallToolRequestParam,
+    lifecycle_manager: &LifecycleManager,
+    credentials: Option<OciCredentials>,
 ) -> Result<CallToolResult> {
     let args = extract_args_from_request(req)?;
     let path = args
@@ -397,7 +407,10 @@ pub async fn handle_load_component_cli(
 
     info!(path, "Loading component (CLI mode)");
 
-    match lifecycle_manager.load_component(path).await {
+    match lifecycle_manager
+        .load_component_with_credentials(path, credentials)
+        .await
+    {
         Ok(outcome) => {
             handle_tool_list_notification(None, &outcome.component_id, "load").await;
             create_load_component_success_result(&outcome)

crates/wassette/src/lib.rs

@@ -42,6 +42,7 @@ use component_storage::ComponentStorage;
 pub use config::{LifecycleBuilder, LifecycleConfig};
 pub use http::WassetteWasiState;
 use loader::{ComponentResource, DownloadedResource};
+pub use oci_auth::OciCredentials;
 use policy_internal::PolicyManager;
 pub use policy_internal::{PermissionGrantRequest, PermissionRule, PolicyInfo};
 use runtime_context::RuntimeContext;
@@ -424,15 +425,26 @@ impl LifecycleManager {
         self.policy_manager.restore_from_disk(component_id).await
     }
 
+    #[allow(dead_code)]
     async fn resolve_component_resource(&self, uri: &str) -> Result<(String, DownloadedResource)> {
+        self.resolve_component_resource_with_credentials(uri, None)
+            .await
+    }
+
+    async fn resolve_component_resource_with_credentials(
+        &self,
+        uri: &str,
+        credentials: Option<OciCredentials>,
+    ) -> Result<(String, DownloadedResource)> {
         // Show progress when running in CLI mode (stderr is a TTY)
         let show_progress = std::io::stderr().is_terminal();
 
-        let resource = loader::load_resource_with_progress::<ComponentResource>(
+        let resource = loader::load_resource_with_progress_and_credentials::<ComponentResource>(
             uri,
             &self.oci_client,
             &self.http_client,
             show_progress,
+            credentials,
         )
         .await?;
         let id = resource.id()?;
@@ -527,8 +539,29 @@ impl LifecycleManager {
     /// component and whether it replaced an existing instance.
     #[instrument(skip(self))]
     pub async fn load_component(&self, uri: &str) -> Result<ComponentLoadOutcome> {
+        self.load_component_with_credentials(uri, None).await
+    }
+
+    /// Loads a WebAssembly component from the specified URI with optional explicit credentials
+    ///
+    /// If a component with the given id already exists, it will be updated with the new component.
+    /// Returns rich [`ComponentLoadOutcome`] information describing the loaded
+    /// component and whether it replaced an existing instance.
+    ///
+    /// # Arguments
+    ///
+    /// * `uri` - The URI to load the component from (file://, oci://, or https://)
+    /// * `credentials` - Optional explicit OCI registry credentials (takes priority over Docker config)
+    #[instrument(skip(self))]
+    pub async fn load_component_with_credentials(
+        &self,
+        uri: &str,
+        credentials: Option<OciCredentials>,
+    ) -> Result<ComponentLoadOutcome> {
         debug!(uri, "Loading component");
-        let (component_id, resource) = self.resolve_component_resource(uri).await?;
+        let (component_id, resource) = self
+            .resolve_component_resource_with_credentials(uri, credentials)
+            .await?;
         let staged_path = self
             .stage_component_artifact(&component_id, resource)
             .await?;

crates/wassette/src/loader.rs

@@ -153,11 +153,18 @@ pub trait Loadable: Sized {
     const RESOURCE_TYPE: &'static str;
 
     async fn from_local_file(path: &Path) -> Result<DownloadedResource>;
+    #[allow(dead_code)]
     async fn from_oci_reference_with_progress(
         reference: &str,
         oci_client: &oci_client::Client,
         show_progress: bool,
     ) -> Result<DownloadedResource>;
+    async fn from_oci_reference_with_progress_and_credentials(
+        reference: &str,
+        oci_client: &oci_client::Client,
+        show_progress: bool,
+        credentials: Option<crate::oci_auth::OciCredentials>,
+    ) -> Result<DownloadedResource>;
     async fn from_url(url: &str, http_client: &reqwest::Client) -> Result<DownloadedResource>;
 }
 
@@ -192,6 +199,21 @@ impl Loadable for ComponentResource {
         reference: &str,
         oci_client: &oci_client::Client,
         show_progress: bool,
+    ) -> Result<DownloadedResource> {
+        Self::from_oci_reference_with_progress_and_credentials(
+            reference,
+            oci_client,
+            show_progress,
+            None,
+        )
+        .await
+    }
+
+    async fn from_oci_reference_with_progress_and_credentials(
+        reference: &str,
+        oci_client: &oci_client::Client,
+        show_progress: bool,
+        credentials: Option<crate::oci_auth::OciCredentials>,
     ) -> Result<DownloadedResource> {
         let reference: oci_client::Reference =
             reference.parse().context("Failed to parse OCI reference")?;
@@ -201,7 +223,7 @@ impl Loadable for ComponentResource {
         }
 
         // Get authentication credentials for this registry
-        let auth = crate::oci_auth::get_registry_auth(&reference)
+        let auth = crate::oci_auth::get_registry_auth_with_credentials(&reference, credentials)
             .context("Failed to get registry authentication")?;
 
         // First try oci-wasm for backwards compatibility with single-layer artifacts
@@ -344,6 +366,15 @@ impl Loadable for PolicyResource {
         bail!("OCI references are not supported for policy resources. Use 'file://' or 'https://' schemes instead.")
     }
 
+    async fn from_oci_reference_with_progress_and_credentials(
+        _reference: &str,
+        _oci_client: &oci_client::Client,
+        _show_progress: bool,
+        _credentials: Option<crate::oci_auth::OciCredentials>,
+    ) -> Result<DownloadedResource> {
+        bail!("OCI references are not supported for policy resources. Use 'file://' or 'https://' schemes instead.")
+    }
+
     async fn from_url(url: &str, http_client: &reqwest::Client) -> Result<DownloadedResource> {
         let url_obj = reqwest::Url::parse(url)?;
         let filename = url_obj
@@ -392,6 +423,24 @@ pub(crate) async fn load_resource_with_progress<T: Loadable>(
     oci_client: &oci_wasm::WasmClient,
     http_client: &reqwest::Client,
     show_progress: bool,
+) -> Result<DownloadedResource> {
+    load_resource_with_progress_and_credentials::<T>(
+        uri,
+        oci_client,
+        http_client,
+        show_progress,
+        None,
+    )
+    .await
+}
+
+/// Generic resource loading function with optional progress reporting and credentials
+pub(crate) async fn load_resource_with_progress_and_credentials<T: Loadable>(
+    uri: &str,
+    oci_client: &oci_wasm::WasmClient,
+    http_client: &reqwest::Client,
+    show_progress: bool,
+    credentials: Option<crate::oci_auth::OciCredentials>,
 ) -> Result<DownloadedResource> {
     let uri = uri.trim();
     let error_message = format!(
@@ -402,7 +451,15 @@ pub(crate) async fn load_resource_with_progress<T: Loadable>(
 
     match scheme {
         "file" => T::from_local_file(Path::new(reference)).await,
-        "oci" => T::from_oci_reference_with_progress(reference, oci_client, show_progress).await,
+        "oci" => {
+            T::from_oci_reference_with_progress_and_credentials(
+                reference,
+                oci_client,
+                show_progress,
+                credentials,
+            )
+            .await
+        }
         "https" => T::from_url(uri, http_client).await,
         _ => bail!("Unsupported {} scheme: {}", T::RESOURCE_TYPE, scheme),
     }