fix: hook up repodata fetch reporter to gateway queries (#6112)

Author: baszalmstra

crates/pixi_command_dispatcher/src/command_dispatcher/builder.rs

@@ -13,7 +13,7 @@ use crate::injected_config::{
     BackendOverrideKey, ChannelConfigKey, EnabledProtocolsKey, ToolBuildEnvironmentKey,
 };
 use crate::reporter::{
-    BackendSourceBuildReporter, BuildBackendMetadataReporter, CondaSolveReporter,
+    BackendSourceBuildReporter, BuildBackendMetadataReporter, CondaSolveReporter, GatewayReporter,
     InstantiateBackendReporter, PixiInstallReporter, PixiSolveReporter, SourceMetadataReporter,
     SourceRecordReporter,
 };
@@ -78,6 +78,7 @@ pub struct CommandDispatcherBuilder {
     source_metadata_reporter: Option<Arc<dyn SourceMetadataReporter>>,
     source_record_reporter: Option<Arc<dyn SourceRecordReporter>>,
     backend_source_build_reporter: Option<Arc<dyn BackendSourceBuildReporter>>,
+    gateway_reporter: Option<Arc<dyn GatewayReporter>>,
 }
 
 impl CommandDispatcherBuilder {
@@ -209,6 +210,16 @@ impl CommandDispatcherBuilder {
         }
     }
 
+    /// Register the [`GatewayReporter`](crate::GatewayReporter) used by
+    /// every site that performs a repodata fetch through the
+    /// [`Gateway`].
+    pub fn with_gateway_reporter(self, reporter: Arc<dyn GatewayReporter>) -> Self {
+        Self {
+            gateway_reporter: Some(reporter),
+            ..self
+        }
+    }
+
     /// Sets the reqwest client to use for network fetches.
     pub fn with_download_client(self, client: LazyClient) -> Self {
         Self {
@@ -484,6 +495,9 @@ impl CommandDispatcherBuilder {
         if let Some(r) = self.backend_source_build_reporter.clone() {
             engine_builder = engine_builder.with_data(r);
         }
+        if let Some(r) = self.gateway_reporter.clone() {
+            engine_builder = engine_builder.with_data(r);
+        }
         if let Some(sem) = data.git_checkout_semaphore.clone() {
             engine_builder = engine_builder.with_data(GitCheckoutSemaphore(sem));
         }

crates/pixi_command_dispatcher/src/compute_data.rs

@@ -17,7 +17,7 @@ use tokio::sync::Semaphore;
 
 use crate::cache::BuildBackendMetadataCache;
 use crate::reporter::{
-    BackendSourceBuildReporter, BuildBackendMetadataReporter, CondaSolveReporter,
+    BackendSourceBuildReporter, BuildBackendMetadataReporter, CondaSolveReporter, GatewayReporter,
     InstantiateBackendReporter, PixiInstallReporter, PixiSolveReporter, SourceMetadataReporter,
     SourceRecordReporter,
 };
@@ -55,6 +55,18 @@ impl HasCondaSolveReporter for DataStore {
     }
 }
 
+/// Access the gateway-fetch reporter shared by every repodata-query
+/// site.
+pub trait HasGatewayReporter {
+    fn gateway_reporter(&self) -> Option<&Arc<dyn GatewayReporter>>;
+}
+
+impl HasGatewayReporter for DataStore {
+    fn gateway_reporter(&self) -> Option<&Arc<dyn GatewayReporter>> {
+        self.try_get::<Arc<dyn GatewayReporter>>()
+    }
+}
+
 /// Access the per-key pixi-solve reporter.
 pub trait HasPixiSolveReporter {
     fn pixi_solve_reporter(&self) -> Option<&Arc<dyn PixiSolveReporter>>;

crates/pixi_command_dispatcher/src/ephemeral_env/mod.rs

@@ -37,13 +37,14 @@ use xxhash_rust::xxh3::Xxh3;
 
 use crate::SolveCondaEnvironmentSpec;
 use crate::cache::markers::BuildBackendsDir;
-use crate::compute_data::{HasGateway, HasInstantiateBackendReporter};
+use crate::compute_data::{HasGateway, HasGatewayReporter, HasInstantiateBackendReporter};
 use crate::injected_config::{ChannelConfigKey, ToolBuildEnvironmentKey};
 use crate::install_binary::install_binary_records;
-use crate::reporter::InstantiateBackendReporter;
+use crate::reporter::{InstantiateBackendReporter, WrappingGatewayReporter};
 use crate::solve_binary::SolveCondaExt;
 use crate::solve_conda::SolveCondaEnvironmentError;
 use pixi_compute_cache_dirs::CacheDirsExt;
+use pixi_compute_reporters::OperationId;
 
 /// Specification for an ephemeral, binary-only conda environment.
 ///
@@ -414,6 +415,13 @@ async fn fetch_binary_repodata(
 
     let channel_config = ctx.compute(&ChannelConfigKey).await;
     let gateway = ctx.global_data().gateway().clone();
+    // `fetch_binary_repodata` is invoked from `EphemeralEnvKey::compute`,
+    // which runs inside the backend-instantiate op's `scope_active`.
+    let gateway_reporter = OperationId::current().and_then(|op_id| {
+        ctx.global_data()
+            .gateway_reporter()
+            .and_then(|r| r.create_gateway_reporter(op_id))
+    });
 
     let match_specs = binary_specs
         .clone()
@@ -425,13 +433,17 @@ async fn fetch_binary_repodata(
         .into_match_specs(&channel_config)
         .map_err(|e| EphemeralEnvError::SpecConversion(Arc::new(e)))?;
 
-    gateway
+    let mut query = gateway
         .query(
             spec.channels.iter().cloned().map(Channel::from_url),
             [build_env.host_platform, Platform::NoArch],
             match_specs.into_iter().chain(constraint_specs),
         )
-        .recursive(true)
+        .recursive(true);
+    if let Some(reporter) = gateway_reporter {
+        query = query.with_reporter(WrappingGatewayReporter(reporter));
+    }
+    query
         .await
         .map_err(|e| EphemeralEnvError::Gateway(Arc::new(e)))
 }

crates/pixi_command_dispatcher/src/keys/solve_conda.rs

@@ -29,9 +29,13 @@ use thiserror::Error;
 use tracing::instrument;
 
 use crate::{
-    SolveCondaEnvironmentSpec, SourceMetadata, compute_data::HasGateway,
-    solve_binary::SolveCondaExt, solve_conda::SolveCondaEnvironmentError,
+    SolveCondaEnvironmentSpec, SourceMetadata,
+    compute_data::{HasGateway, HasGatewayReporter},
+    reporter::WrappingGatewayReporter,
+    solve_binary::SolveCondaExt,
+    solve_conda::SolveCondaEnvironmentError,
 };
+use pixi_compute_reporters::OperationId;
 
 /// Input to [`SolveCondaKey`]. All fields participate in the Key's
 /// identity so two callers with equal specs share one compute.
@@ -234,7 +238,16 @@ impl Key for SolveCondaKey {
         // Clone the gateway handle so we don't hold an immutable
         // borrow on `ctx` across the subsequent mutable-borrow solve.
         let gateway = ctx.global_data().gateway().clone();
-        let binary_repodata = gateway
+        // `solve-conda` runs inside the parent operation's
+        // `scope_active` (e.g. the pixi-solve or backend-instantiate
+        // op that triggered it), so `OperationId::current()` gives
+        // the already-started op the gateway query belongs to.
+        let gateway_reporter = OperationId::current().and_then(|op_id| {
+            ctx.global_data()
+                .gateway_reporter()
+                .and_then(|r| r.create_gateway_reporter(op_id))
+        });
+        let mut query = gateway
             .query(
                 spec.channels.iter().cloned().map(Channel::from_url),
                 [spec.platform, Platform::NoArch],
@@ -244,8 +257,11 @@ impl Key for SolveCondaKey {
                     .chain(source_repodata_fetch_specs)
                     .chain(dev_source_fetch_specs),
             )
-            .recursive(true)
-            .await?;
+            .recursive(true);
+        if let Some(reporter) = gateway_reporter {
+            query = query.with_reporter(WrappingGatewayReporter(reporter));
+        }
+        let binary_repodata = query.await?;
 
         // Build the full solve spec and hand off to ctx.solve_conda
         // (semaphore + reporter lifecycle).

crates/pixi_command_dispatcher/src/lib.rs

@@ -129,7 +129,7 @@ pub use pixi_compute_sources::{
     SourceCheckoutExt, UrlCheckoutReporter, UrlDir,
 };
 pub use reporter::{
-    BackendSourceBuildReporter, BuildBackendMetadataReporter, CondaSolveReporter,
+    BackendSourceBuildReporter, BuildBackendMetadataReporter, CondaSolveReporter, GatewayReporter,
     InstantiateBackendReporter, PixiInstallReporter, PixiSolveEnvironmentSpec, PixiSolveReporter,
     SourceMetadataReporter, SourceMetadataReporterSpec, SourceRecordReporter,
     SourceRecordReporterSpec,

crates/pixi_command_dispatcher/src/reporter.rs

@@ -79,6 +79,37 @@ pub trait CondaSolveReporter: Send + Sync {
     fn on_finished(&self, solve_id: OperationId);
 }
 
+/// Reporter for repodata fetches performed by the
+/// [`Gateway`](rattler_repodata_gateway::Gateway). Stored independently
+/// on the engine [`DataStore`](pixi_compute_engine::DataStore) and
+/// looked up at every gateway-query site.
+///
+/// `op_id` identifies the started operation under which the query
+/// runs (e.g. the surrounding pixi-solve or backend-instantiate
+/// op). Implementations can use it for parent lookups (e.g. to nest
+/// the repodata bar under the right solve), or ignore it.
+pub trait GatewayReporter: Send + Sync {
+    /// Build a per-call rattler gateway reporter for the gateway
+    /// query about to run inside the started operation `op_id`.
+    /// `None` skips repodata-fetch progress reporting.
+    fn create_gateway_reporter(
+        &self,
+        op_id: OperationId,
+    ) -> Option<Box<dyn rattler_repodata_gateway::Reporter>>;
+}
+
+/// Adapts a `Box<dyn rattler_repodata_gateway::Reporter>` returned by
+/// [`GatewayReporter::create_gateway_reporter`] into the
+/// `impl Reporter + 'static` value expected by
+/// `rattler_repodata_gateway::Query::with_reporter`.
+pub struct WrappingGatewayReporter(pub Box<dyn rattler_repodata_gateway::Reporter>);
+
+impl rattler_repodata_gateway::Reporter for WrappingGatewayReporter {
+    fn download_reporter(&self) -> Option<&dyn rattler_repodata_gateway::DownloadReporter> {
+        self.0.download_reporter()
+    }
+}
+
 /// Reporter for the compute-engine [`InstantiateBackendKey`](crate::InstantiateBackendKey).
 pub trait InstantiateBackendReporter: Send + Sync {
     fn on_queued(&self, spec: &JsonRpcBackendSpec) -> OperationId;

crates/pixi_command_dispatcher/tests/integration/event_reporter.rs

@@ -6,10 +6,10 @@ use std::{
 use futures::{Stream, StreamExt};
 use pixi_build_discovery::JsonRpcBackendSpec;
 use pixi_command_dispatcher::{
-    BackendSourceBuildSpec, BuildBackendMetadataInner, CondaSolveReporter, GitCheckoutReporter,
-    InstallPixiEnvironmentSpec, PixiInstallReporter, PixiSolveEnvironmentSpec, PixiSolveReporter,
-    SolveCondaEnvironmentSpec, SourceMetadataReporterSpec, SourceRecordReporterSpec,
-    UrlCheckoutReporter,
+    BackendSourceBuildSpec, BuildBackendMetadataInner, CondaSolveReporter, GatewayReporter,
+    GitCheckoutReporter, InstallPixiEnvironmentSpec, PixiInstallReporter, PixiSolveEnvironmentSpec,
+    PixiSolveReporter, SolveCondaEnvironmentSpec, SourceMetadataReporterSpec,
+    SourceRecordReporterSpec, UrlCheckoutReporter,
     reporter::{
         BackendSourceBuildReporter, BuildBackendMetadataReporter, InstantiateBackendReporter,
         SourceMetadataReporter, SourceRecordReporter,
@@ -141,6 +141,15 @@ pub enum Event {
     InstantiateBackendFinished {
         id: OperationId,
     },
+
+    /// Recorded each time a gateway-query site requests a per-call
+    /// repodata reporter. `id` is a fresh op id for the query itself;
+    /// its parent in the registry is the started operation that
+    /// triggered the fetch (carried in `op_id`).
+    GatewayQuery {
+        id: OperationId,
+        op_id: OperationId,
+    },
 }
 
 pub struct EventReporter {
@@ -413,6 +422,17 @@ impl InstantiateBackendReporter for EventReporter {
     }
 }
 
+impl GatewayReporter for EventReporter {
+    fn create_gateway_reporter(
+        &self,
+        op_id: OperationId,
+    ) -> Option<Box<dyn rattler_repodata_gateway::Reporter>> {
+        let id = self.alloc();
+        self.record(Event::GatewayQuery { id, op_id });
+        None
+    }
+}
+
 impl BackendSourceBuildReporter for EventReporter {
     fn on_queued(&self, spec: &BackendSourceBuildSpec) -> OperationId {
         let id = self.alloc();
@@ -462,6 +482,7 @@ impl EventReporter {
             .with_source_metadata_reporter(self.clone())
             .with_source_record_reporter(self.clone())
             .with_backend_source_build_reporter(self.clone())
+            .with_gateway_reporter(self.clone())
     }
 }
 

crates/pixi_command_dispatcher/tests/integration/event_tree.rs

@@ -113,6 +113,9 @@ impl EventTree {
                 }
                 Event::InstantiateBackendStarted { id } => builder.alloc_pending(*id),
                 Event::InstantiateBackendFinished { .. } => {}
+                Event::GatewayQuery { id, .. } => {
+                    builder.alloc_node(*id, "Repodata fetch".to_owned());
+                }
                 Event::UrlCheckoutQueued { .. }
                 | Event::UrlCheckoutStarted { .. }
                 | Event::UrlCheckoutFinished { .. } => {

crates/pixi_command_dispatcher/tests/integration/snapshots/integration__simple_test.snap

@@ -1,18 +1,21 @@
 ---
 source: crates/pixi_command_dispatcher/tests/integration/main.rs
 expression: output
+snapshot_kind: text
 ---
 Pixi solve (test@[PLATFORM])
 ├── Git Checkout (file://[LOCAL_GIT_REPO]#[GIT_HASH])
 ├── Source metadata (foobar-desktop @ git+file://[LOCAL_GIT_REPO]?subdirectory=recipe&rev=[GIT_REF]#[GIT_HASH])
 │   ├── Build backend metadata (git+file://[LOCAL_GIT_REPO]?subdirectory=recipe&rev=[GIT_REF]#[GIT_HASH])
 │   │   ├── Git Checkout (file://[LOCAL_GIT_REPO]#[GIT_HASH])
 │   │   └── Instantiate backend (pixi-build-rattler-build)
-│   │       └── Conda solve #6
+│   │       ├── Repodata fetch
+│   │       └── Conda solve #7
 │   └── Source record (foobar-desktop @ git+file://[LOCAL_GIT_REPO]?subdirectory=recipe&rev=[GIT_REF]#[GIT_HASH])
 ├── Source metadata (foobar @ git+file://[LOCAL_GIT_REPO]?subdirectory=recipe&rev=[GIT_REF]#[GIT_HASH])
 │   └── Source record (foobar @ git+file://[LOCAL_GIT_REPO]?subdirectory=recipe&rev=[GIT_REF]#[GIT_HASH])
-└── Conda solve #10
+├── Repodata fetch
+└── Conda solve #12
 Pixi install (test-env)
 ├── Backend source build (foobar-desktop)
 └── Backend source build (foobar)

crates/pixi_reporters/src/lib.rs

@@ -34,9 +34,6 @@ pub struct TopLevelProgress {
     /// `OperationId` → bar slot in `conda_solve_reporter`. Lets
     /// `on_started` / `on_finished` find the bar created at `on_queued`.
     solve_bars: Mutex<HashMap<OperationId, usize>>,
-    // Held so on_clear can wipe it; the gateway integration that would
-    // populate it currently goes through a different code path, so the
-    // bar is effectively a placeholder for future wiring.
     repodata_reporter: RepodataReporter,
     sync_reporter: SyncReporter,
 }
@@ -114,6 +111,7 @@ impl TopLevelProgress {
             .with_instantiate_backend_reporter(self.clone())
             .with_git_checkout_reporter(self.source_checkout_reporter.clone())
             .with_backend_source_build_reporter(backend_source_build_reporter)
+            .with_gateway_reporter(self.clone())
     }
 
     /// Clear the current progress bars without tearing down the reporter.
@@ -216,3 +214,12 @@ impl pixi_command_dispatcher::CondaSolveReporter for TopLevelProgress {
         }
     }
 }
+
+impl pixi_command_dispatcher::GatewayReporter for TopLevelProgress {
+    fn create_gateway_reporter(
+        &self,
+        _op_id: OperationId,
+    ) -> Option<Box<dyn rattler_repodata_gateway::Reporter>> {
+        Some(Box::new(self.repodata_reporter.clone()))
+    }
+}