imagegen_bridge_runtime/
orchestrator.rs

1//! Shared request orchestration used by every transport and SDK facade.
2
3use std::{
4    collections::{BTreeMap, BTreeSet},
5    future::Future,
6    sync::{
7        Arc,
8        atomic::{AtomicUsize, Ordering},
9    },
10    time::{Duration, Instant},
11};
12
13use imagegen_bridge_core::{
14    BridgeError, ErrorCode, FallbackPolicy, ImageRequest, ImageResponse, Normalization,
15    ProviderAttempt, ProviderAttemptOutcome, ProviderContext, ProviderEvent, ProviderRoute,
16    RequestLimits, RevisedPromptPolicy, negotiate_request, validate_request,
17};
18use tokio::sync::{Mutex, Notify, OwnedSemaphorePermit, mpsc};
19use tokio_util::sync::CancellationToken;
20use uuid::Uuid;
21
22use crate::{
23    admission::AdmissionGate,
24    idempotency::{IdempotencyAction, IdempotencyConfig, IdempotencyCoordinator},
25    materialize::{MaterializationConfig, OutputMaterializer},
26    registry::ProviderRegistry,
27    transparency::prepare_transparency,
28};
29
30/// One concurrency pool and its bounded waiting room.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct ConcurrencyLimit {
33    /// Maximum operations executing concurrently.
34    pub max_concurrent: usize,
35    /// Maximum operations waiting for this pool.
36    pub max_queued: usize,
37}
38
39impl Default for ConcurrencyLimit {
40    fn default() -> Self {
41        Self {
42            max_concurrent: 4,
43            max_queued: 16,
44        }
45    }
46}
47
48/// Complete shared runtime policy.
49#[derive(Debug, Clone)]
50pub struct RuntimeConfig {
51    /// Intrinsic request limits applied before provider selection.
52    pub request_limits: RequestLimits,
53    /// Deadline used when a request does not provide one.
54    pub default_timeout: Duration,
55    /// Cooperative cleanup window after cancellation or deadline expiry.
56    pub cancellation_grace: Duration,
57    /// Maximum time graceful shutdown waits for active calls to unwind.
58    pub shutdown_grace: Duration,
59    /// Global execution and queue bound.
60    pub global_limit: ConcurrencyLimit,
61    /// Default bound for each registered provider.
62    pub default_provider_limit: ConcurrencyLimit,
63    /// Per-provider limit overrides.
64    pub provider_limits: BTreeMap<String, ConcurrencyLimit>,
65    /// Scoped replay coordination policy.
66    pub idempotency: IdempotencyConfig,
67    /// Output verification and delivery policy.
68    pub materialization: MaterializationConfig,
69}
70
71impl Default for RuntimeConfig {
72    fn default() -> Self {
73        Self {
74            request_limits: RequestLimits::default(),
75            default_timeout: Duration::from_secs(5 * 60),
76            cancellation_grace: Duration::from_secs(1),
77            shutdown_grace: Duration::from_secs(10),
78            global_limit: ConcurrencyLimit {
79                max_concurrent: 16,
80                max_queued: 64,
81            },
82            default_provider_limit: ConcurrencyLimit::default(),
83            provider_limits: BTreeMap::new(),
84            idempotency: IdempotencyConfig::default(),
85            materialization: MaterializationConfig::default(),
86        }
87    }
88}
89
90/// Per-call transport context that is not forwarded as generation content.
91#[derive(Clone)]
92pub struct ExecutionContext {
93    /// Optional caller-generated safe request ID.
94    pub request_id: Option<String>,
95    /// Tenant/user boundary for an idempotency key.
96    pub idempotency_scope: String,
97    /// Transport disconnect or caller cancellation signal.
98    pub cancellation: CancellationToken,
99}
100
101impl std::fmt::Debug for ExecutionContext {
102    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103        formatter
104            .debug_struct("ExecutionContext")
105            .field("request_id", &self.request_id)
106            .field("idempotency_scope", &"[REDACTED SCOPE]")
107            .field("cancelled", &self.cancellation.is_cancelled())
108            .finish()
109    }
110}
111
112impl Default for ExecutionContext {
113    fn default() -> Self {
114        Self {
115            request_id: None,
116            idempotency_scope: "default".to_owned(),
117            cancellation: CancellationToken::new(),
118        }
119    }
120}
121
122/// Queue counters suitable for bounded operational metrics.
123#[derive(Debug, Clone, PartialEq, Eq)]
124pub struct RuntimeQueueSnapshot {
125    /// Callers waiting for the global pool.
126    pub global_queued: usize,
127    /// Callers waiting per provider.
128    pub providers_queued: BTreeMap<String, usize>,
129}
130
131/// Provider-neutral execution engine shared by CLI, HTTP, and libraries.
132pub struct ImagegenRuntime {
133    registry: ProviderRegistry,
134    config: RuntimeConfig,
135    global_gate: Arc<AdmissionGate>,
136    provider_gates: BTreeMap<String, Arc<AdmissionGate>>,
137    idempotency: IdempotencyCoordinator,
138    materializer: OutputMaterializer,
139    shutdown: CancellationToken,
140    activity: Arc<ActivityTracker>,
141    shutdown_result: Mutex<Option<Result<(), BridgeError>>>,
142}
143
144struct RouteExecution {
145    response: ImageResponse,
146    effective_request: ImageRequest,
147}
148
149impl std::fmt::Debug for ImagegenRuntime {
150    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        formatter
152            .debug_struct("ImagegenRuntime")
153            .field("registry", &self.registry)
154            .field("config", &self.config)
155            .field("shutdown", &self.shutdown.is_cancelled())
156            .finish_non_exhaustive()
157    }
158}
159
160impl ImagegenRuntime {
161    /// Validates configuration and constructs all fixed-capacity runtime state.
162    pub fn new(registry: ProviderRegistry, config: RuntimeConfig) -> Result<Self, BridgeError> {
163        if config.default_timeout.is_zero()
164            || config.default_timeout > Duration::from_millis(config.request_limits.max_timeout_ms)
165        {
166            return Err(configuration_error(
167                "default timeout must be within the configured request limit",
168            ));
169        }
170        if config.cancellation_grace.is_zero()
171            || config.cancellation_grace > Duration::from_secs(30)
172        {
173            return Err(configuration_error(
174                "cancellation grace must be greater than zero and at most 30 seconds",
175            ));
176        }
177        if config.shutdown_grace.is_zero() || config.shutdown_grace > Duration::from_secs(5 * 60) {
178            return Err(configuration_error(
179                "shutdown grace must be greater than zero and at most five minutes",
180            ));
181        }
182        for name in config.provider_limits.keys() {
183            registry.resolve(Some(name))?;
184        }
185        let global_gate = Arc::new(AdmissionGate::new(
186            config.global_limit.max_concurrent,
187            config.global_limit.max_queued,
188            "global",
189        )?);
190        let mut provider_gates = BTreeMap::new();
191        for (name, _) in registry.entries() {
192            let limit = config
193                .provider_limits
194                .get(name)
195                .copied()
196                .unwrap_or(config.default_provider_limit);
197            provider_gates.insert(
198                name.to_owned(),
199                Arc::new(AdmissionGate::new(
200                    limit.max_concurrent,
201                    limit.max_queued,
202                    format!("provider:{name}"),
203                )?),
204            );
205        }
206        let idempotency = IdempotencyCoordinator::new(config.idempotency)?;
207        let materializer = OutputMaterializer::new(config.materialization.clone())?;
208        Ok(Self {
209            registry,
210            config,
211            global_gate,
212            provider_gates,
213            idempotency,
214            materializer,
215            shutdown: CancellationToken::new(),
216            activity: Arc::new(ActivityTracker::default()),
217            shutdown_result: Mutex::new(None),
218        })
219    }
220
221    /// Returns the immutable provider registry.
222    #[must_use]
223    pub const fn registry(&self) -> &ProviderRegistry {
224        &self.registry
225    }
226
227    /// Validates intrinsic request bounds without starting provider work.
228    pub fn validate_request(&self, request: &ImageRequest) -> Result<(), BridgeError> {
229        validate_request(request, self.config.request_limits)
230    }
231
232    /// Whether verified bridge-owned artifact delivery is available.
233    #[must_use]
234    pub const fn has_artifact_store(&self) -> bool {
235        self.materializer.has_artifact_store()
236    }
237
238    /// Reads one ownership-verified artifact for authenticated delivery.
239    pub fn read_artifact(
240        &self,
241        artifact_id: &str,
242    ) -> Result<imagegen_bridge_artifacts::StoredArtifactContent, BridgeError> {
243        self.materializer.read_artifact(artifact_id)
244    }
245
246    /// Creates a bounded PNG thumbnail from one ownership-verified artifact.
247    pub fn read_artifact_thumbnail(
248        &self,
249        artifact_id: &str,
250        maximum_edge: u32,
251    ) -> Result<Vec<u8>, BridgeError> {
252        self.materializer.read_thumbnail(artifact_id, maximum_edge)
253    }
254
255    /// Executes with generated request metadata and the default idempotency scope.
256    pub async fn execute(&self, request: ImageRequest) -> Result<ImageResponse, BridgeError> {
257        self.execute_with(request, ExecutionContext::default())
258            .await
259    }
260
261    /// Executes the complete validation, negotiation, admission, provider, and
262    /// materialization pipeline.
263    pub async fn execute_with(
264        &self,
265        request: ImageRequest,
266        context: ExecutionContext,
267    ) -> Result<ImageResponse, BridgeError> {
268        self.execute_with_optional_events(request, context, None)
269            .await
270    }
271
272    /// Executes the complete pipeline and forwards bounded provider events.
273    pub async fn execute_with_events(
274        &self,
275        request: ImageRequest,
276        context: ExecutionContext,
277        events: mpsc::Sender<ProviderEvent>,
278    ) -> Result<ImageResponse, BridgeError> {
279        self.execute_with_optional_events(request, context, Some(events))
280            .await
281    }
282
283    async fn execute_with_optional_events(
284        &self,
285        request: ImageRequest,
286        context: ExecutionContext,
287        events: Option<mpsc::Sender<ProviderEvent>>,
288    ) -> Result<ImageResponse, BridgeError> {
289        if self.shutdown.is_cancelled() {
290            return Err(cancelled_error("runtime is shutting down"));
291        }
292        let _activity = self.activity.enter();
293        if self.shutdown.is_cancelled() {
294            return Err(cancelled_error("runtime is shutting down"));
295        }
296        validate_request(&request, self.config.request_limits)?;
297        let request_id = validated_request_id(context.request_id)?;
298        let timeout = request
299            .timeout_ms
300            .map_or(self.config.default_timeout, Duration::from_millis);
301        let deadline = Instant::now()
302            .checked_add(timeout)
303            .ok_or_else(|| configuration_error("request deadline overflowed"))?;
304        let operation = self.shutdown.child_token();
305
306        let mut token = None;
307        if let Some(key) = request.idempotency_key.as_deref() {
308            match self
309                .idempotency
310                .begin(&context.idempotency_scope, key, &request)
311                .await?
312            {
313                IdempotencyAction::Leader(leader) => token = Some(leader),
314                IdempotencyAction::Cached(response) => return Ok(*response),
315                IdempotencyAction::Wait(receiver) => {
316                    return self
317                        .await_controlled(
318                            IdempotencyAction::wait(receiver, deadline, &operation),
319                            deadline,
320                            &context.cancellation,
321                            &operation,
322                            "request was cancelled while waiting for idempotent result",
323                        )
324                        .await;
325                }
326            }
327        }
328
329        let result = self
330            .execute_leader(
331                request,
332                request_id,
333                deadline,
334                &context.cancellation,
335                &operation,
336                events,
337            )
338            .await;
339        match (token, &result) {
340            (Some(token), Ok(response)) => {
341                self.idempotency.complete(token, response.clone()).await;
342            }
343            (Some(token), Err(error)) => {
344                self.idempotency.fail(token, error.clone()).await;
345            }
346            (None, _) => {}
347        }
348        result
349    }
350
351    /// Returns current bounded queue depth without request identifiers.
352    #[must_use]
353    pub fn queue_snapshot(&self) -> RuntimeQueueSnapshot {
354        RuntimeQueueSnapshot {
355            global_queued: self.global_gate.queued(),
356            providers_queued: self
357                .provider_gates
358                .iter()
359                .map(|(name, gate)| (name.clone(), gate.queued()))
360                .collect(),
361        }
362    }
363
364    /// Stops admissions, cancels active operations, and releases providers.
365    pub async fn shutdown(&self) -> Result<(), BridgeError> {
366        let mut stored_result = self.shutdown_result.lock().await;
367        if let Some(result) = stored_result.as_ref() {
368            return result.clone();
369        }
370        self.shutdown.cancel();
371        self.global_gate.close();
372        for gate in self.provider_gates.values() {
373            gate.close();
374        }
375        let mut first_error =
376            if tokio::time::timeout(self.config.shutdown_grace, self.activity.wait_until_idle())
377                .await
378                .is_err()
379            {
380                Some(
381                    BridgeError::new(
382                        ErrorCode::Timeout,
383                        "runtime shutdown grace elapsed with active operations",
384                    )
385                    .retryable(true),
386                )
387            } else {
388                None
389            };
390        for (_, provider) in self.registry.entries() {
391            if let Err(error) = provider.shutdown().await
392                && first_error.is_none()
393            {
394                first_error = Some(error);
395            }
396        }
397        let result = first_error.map_or(Ok(()), Err);
398        *stored_result = Some(result.clone());
399        result
400    }
401
402    async fn execute_leader(
403        &self,
404        request: ImageRequest,
405        request_id: String,
406        deadline: Instant,
407        external_cancellation: &CancellationToken,
408        operation: &CancellationToken,
409        events: Option<mpsc::Sender<ProviderEvent>>,
410    ) -> Result<ImageResponse, BridgeError> {
411        let total_started = Instant::now();
412        let queue_started = Instant::now();
413        let global_permit = self
414            .acquire_controlled(
415                &self.global_gate,
416                deadline,
417                external_cancellation,
418                operation,
419            )
420            .await?;
421        let global_queue_ms = elapsed_ms(queue_started);
422        let primary_provider = request
423            .routing
424            .provider
425            .clone()
426            .unwrap_or_else(|| self.registry.default_name().to_owned());
427        let routes = std::iter::once(ProviderRoute {
428            provider: primary_provider.clone(),
429            model: request.routing.model.clone(),
430        })
431        .chain(request.routing.fallbacks.iter().cloned())
432        .collect::<Vec<_>>();
433        let mut unique_routes = BTreeSet::new();
434        for (index, route) in routes.iter().enumerate() {
435            if !unique_routes.insert((route.provider.as_str(), route.model.as_deref())) {
436                return Err(BridgeError::new(
437                    ErrorCode::InvalidRequest,
438                    "provider routing contains a duplicate resolved route",
439                )
440                .with_detail("field", format!("routing.fallbacks[{}]", index - 1))
441                .with_detail("provider", &route.provider)
442                .with_detail("model", route.model.as_deref()));
443            }
444        }
445        let mut attempts = Vec::with_capacity(routes.len());
446        let mut final_error = None;
447
448        for (index, route) in routes.iter().enumerate() {
449            let attempt_started = Instant::now();
450            match self
451                .execute_route(
452                    &request,
453                    &request_id,
454                    route,
455                    deadline,
456                    external_cancellation,
457                    operation,
458                    events.clone(),
459                )
460                .await
461            {
462                Ok(execution) => {
463                    let mut response = execution.response;
464                    attempts.push(ProviderAttempt {
465                        provider: response.provider.clone(),
466                        model: Some(response.model.clone()),
467                        outcome: ProviderAttemptOutcome::Succeeded,
468                        error_code: None,
469                        duration_ms: elapsed_ms(attempt_started),
470                    });
471                    if index > 0 {
472                        response.normalizations.insert(
473                            0,
474                            Normalization {
475                                field: "routing.provider".to_owned(),
476                                requested: Some(serde_json::json!(primary_provider)),
477                                effective: Some(serde_json::json!(response.provider)),
478                                reason: "provider_fallback".to_owned(),
479                            },
480                        );
481                        response.warnings.push("provider_fallback_used".to_owned());
482                    }
483                    if !request.routing.fallbacks.is_empty() {
484                        response.attempts = attempts;
485                    }
486                    response.timings.queue_ms =
487                        response.timings.queue_ms.saturating_add(global_queue_ms);
488                    response.timings.total_ms = elapsed_ms(total_started);
489                    self.materializer.attach_metadata(
490                        &request,
491                        &execution.effective_request,
492                        &mut response,
493                    )?;
494                    drop(global_permit);
495                    return Ok(response);
496                }
497                Err(error) => {
498                    attempts.push(ProviderAttempt {
499                        provider: route.provider.clone(),
500                        model: route.model.clone(),
501                        outcome: ProviderAttemptOutcome::Failed,
502                        error_code: Some(error.code),
503                        duration_ms: elapsed_ms(attempt_started),
504                    });
505                    let has_next = index + 1 < routes.len();
506                    if !has_next || !should_fallback(&error, request.routing.fallback_policy) {
507                        final_error = Some(error);
508                        break;
509                    }
510                    final_error = Some(error);
511                }
512            }
513        }
514        drop(global_permit);
515        Err(final_error
516            .unwrap_or_else(|| configuration_error("provider routing produced no attempts"))
517            .with_detail("attempts", attempts))
518    }
519
520    #[allow(clippy::too_many_arguments)]
521    async fn execute_route(
522        &self,
523        request: &ImageRequest,
524        request_id: &str,
525        route: &ProviderRoute,
526        deadline: Instant,
527        external_cancellation: &CancellationToken,
528        operation: &CancellationToken,
529        events: Option<mpsc::Sender<ProviderEvent>>,
530    ) -> Result<RouteExecution, BridgeError> {
531        let provider = self.registry.resolve(Some(&route.provider))?;
532        let descriptor = provider.descriptor();
533        let provider_gate = self
534            .provider_gates
535            .get(&descriptor.name)
536            .ok_or_else(|| configuration_error("provider has no configured admission gate"))?;
537        let queue_started = Instant::now();
538        let _provider_permit = self
539            .acquire_controlled(provider_gate, deadline, external_cancellation, operation)
540            .await?;
541        let queue_ms = elapsed_ms(queue_started);
542        let capabilities = self
543            .await_controlled(
544                provider.capabilities(route.model.as_deref()),
545                deadline,
546                external_cancellation,
547                operation,
548                "request was cancelled during capability discovery",
549            )
550            .await
551            .map_err(|error| attach_provider(error, &descriptor.name))?;
552        if capabilities.provider != descriptor.name {
553            return Err(protocol_error(
554                "provider capability identity does not match its registry identity",
555            )
556            .with_provider(descriptor.name));
557        }
558        let mut routed_request = request.clone();
559        routed_request.routing.provider = Some(route.provider.clone());
560        routed_request.routing.model.clone_from(&route.model);
561        routed_request.routing.fallbacks.clear();
562        let prepared = prepare_transparency(&routed_request, &capabilities)?;
563        validate_request(&prepared.provider_request, self.config.request_limits)?;
564        let negotiated = negotiate_request(&prepared.provider_request, &capabilities)?;
565        let effective_request = negotiated.effective_request;
566
567        let provider_started = Instant::now();
568        let provider_result = self
569            .await_provider_execution(
570                provider.execute(
571                    effective_request.clone(),
572                    ProviderContext {
573                        request_id: request_id.to_owned(),
574                        deadline,
575                        cancellation: operation.clone(),
576                        events,
577                    },
578                ),
579                deadline,
580                external_cancellation,
581                operation,
582            )
583            .await;
584        let provider_ms = elapsed_ms(provider_started);
585        let mut response =
586            provider_result.map_err(|error| attach_provider(error, &descriptor.name))?;
587        if let Some(expected_model) = capabilities.model.as_deref()
588            && response.model != expected_model
589        {
590            return Err(protocol_error(
591                "provider response model does not match discovered capabilities",
592            )
593            .with_provider(descriptor.name)
594            .with_detail("expected_model", expected_model)
595            .with_detail("actual_model", response.model));
596        }
597
598        response.id = request_id.to_owned();
599        response.provider = descriptor.name;
600        response.requested = request.parameters.clone();
601        response.effective = effective_request.parameters.clone();
602        if prepared.plan.is_some() {
603            response.effective.background = imagegen_bridge_core::Background::Transparent;
604            response
605                .warnings
606                .push("transparent_background_postprocessed".to_owned());
607        }
608        response.normalizations.splice(
609            0..0,
610            prepared
611                .normalizations
612                .into_iter()
613                .chain(negotiated.normalizations),
614        );
615        if effective_request.policies.revised_prompt == RevisedPromptPolicy::Omit {
616            response.revised_prompt = None;
617        }
618        response.timings.queue_ms = queue_ms;
619        response.timings.provider_ms = provider_ms;
620        let artifact_started = Instant::now();
621        response = self
622            .await_controlled(
623                self.materializer
624                    .materialize(response, request, &effective_request, prepared.plan),
625                deadline,
626                external_cancellation,
627                operation,
628                "request was cancelled during output materialization",
629            )
630            .await?;
631        response.timings.artifact_ms = elapsed_ms(artifact_started);
632        Ok(RouteExecution {
633            response,
634            effective_request,
635        })
636    }
637
638    async fn acquire_controlled(
639        &self,
640        gate: &AdmissionGate,
641        deadline: Instant,
642        external: &CancellationToken,
643        operation: &CancellationToken,
644    ) -> Result<OwnedSemaphorePermit, BridgeError> {
645        self.await_controlled(
646            gate.acquire(deadline, operation),
647            deadline,
648            external,
649            operation,
650            "request was cancelled while waiting for capacity",
651        )
652        .await
653    }
654
655    async fn await_provider_execution<T, F>(
656        &self,
657        future: F,
658        deadline: Instant,
659        external: &CancellationToken,
660        operation: &CancellationToken,
661    ) -> Result<T, BridgeError>
662    where
663        F: Future<Output = Result<T, BridgeError>>,
664    {
665        if deadline <= Instant::now() {
666            operation.cancel();
667            return Err(timeout_error());
668        }
669        tokio::pin!(future);
670        let interrupted = tokio::select! {
671            result = &mut future => return classify_provider_execution_result(result),
672            () = external.cancelled() => {
673                operation.cancel();
674                (ErrorCode::Cancelled, "request was cancelled during provider execution")
675            }
676            () = self.shutdown.cancelled() => {
677                operation.cancel();
678                (ErrorCode::Cancelled, "runtime shut down during provider execution")
679            }
680            () = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {
681                operation.cancel();
682                (ErrorCode::Timeout, "request deadline elapsed during provider execution")
683            }
684        };
685        match tokio::time::timeout(self.config.cancellation_grace, &mut future).await {
686            Ok(Ok(value)) => Ok(value),
687            Ok(Err(error)) if !matches!(error.code, ErrorCode::Timeout | ErrorCode::Cancelled) => {
688                Err(error)
689            }
690            Ok(Err(_)) | Err(_) => Err(unknown_outcome_error(interrupted.0, interrupted.1)),
691        }
692    }
693
694    async fn await_controlled<T, F>(
695        &self,
696        future: F,
697        deadline: Instant,
698        external: &CancellationToken,
699        operation: &CancellationToken,
700        cancellation_message: &'static str,
701    ) -> Result<T, BridgeError>
702    where
703        F: Future<Output = Result<T, BridgeError>>,
704    {
705        if deadline <= Instant::now() {
706            operation.cancel();
707            return Err(timeout_error());
708        }
709        tokio::pin!(future);
710        tokio::select! {
711            result = &mut future => result,
712            () = external.cancelled() => {
713                operation.cancel();
714                let _ = tokio::time::timeout(self.config.cancellation_grace, &mut future).await;
715                Err(cancelled_error(cancellation_message))
716            }
717            () = self.shutdown.cancelled() => {
718                operation.cancel();
719                let _ = tokio::time::timeout(self.config.cancellation_grace, &mut future).await;
720                Err(cancelled_error("runtime is shutting down"))
721            }
722            () = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {
723                operation.cancel();
724                let _ = tokio::time::timeout(self.config.cancellation_grace, &mut future).await;
725                Err(timeout_error())
726            }
727        }
728    }
729}
730
731#[derive(Default)]
732struct ActivityTracker {
733    active: AtomicUsize,
734    idle: Notify,
735}
736
737impl ActivityTracker {
738    fn enter(self: &Arc<Self>) -> ActivityGuard {
739        self.active.fetch_add(1, Ordering::AcqRel);
740        ActivityGuard {
741            tracker: Arc::clone(self),
742        }
743    }
744
745    async fn wait_until_idle(&self) {
746        loop {
747            let notified = self.idle.notified();
748            if self.active.load(Ordering::Acquire) == 0 {
749                return;
750            }
751            notified.await;
752        }
753    }
754}
755
756struct ActivityGuard {
757    tracker: Arc<ActivityTracker>,
758}
759
760impl Drop for ActivityGuard {
761    fn drop(&mut self) {
762        if self.tracker.active.fetch_sub(1, Ordering::AcqRel) == 1 {
763            self.tracker.idle.notify_waiters();
764        }
765    }
766}
767
768fn validated_request_id(value: Option<String>) -> Result<String, BridgeError> {
769    let Some(value) = value else {
770        return Ok(Uuid::now_v7().to_string());
771    };
772    let valid = !value.is_empty()
773        && value.len() <= 128
774        && value
775            .bytes()
776            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b':'));
777    if valid {
778        Ok(value)
779    } else {
780        Err(BridgeError::new(
781            ErrorCode::InvalidRequest,
782            "request ID is invalid",
783        ))
784    }
785}
786
787fn attach_provider(mut error: BridgeError, provider: &str) -> BridgeError {
788    if error.provider.is_none() {
789        error.provider = Some(provider.to_owned());
790    }
791    error
792}
793
794fn should_fallback(error: &BridgeError, policy: FallbackPolicy) -> bool {
795    if error
796        .details
797        .get("outcome")
798        .and_then(serde_json::Value::as_str)
799        == Some("unknown")
800    {
801        return false;
802    }
803    if matches!(
804        error.code,
805        ErrorCode::InvalidRequest
806            | ErrorCode::PermissionDenied
807            | ErrorCode::SafetyRejected
808            | ErrorCode::Cancelled
809            | ErrorCode::Session
810            | ErrorCode::IdempotencyConflict
811    ) {
812        return false;
813    }
814    if matches!(
815        error.code,
816        ErrorCode::Configuration
817            | ErrorCode::Authentication
818            | ErrorCode::UnsupportedCapability
819            | ErrorCode::RateLimited
820            | ErrorCode::Overloaded
821    ) {
822        return true;
823    }
824    policy == FallbackPolicy::OnError
825        && matches!(
826            error.code,
827            ErrorCode::Upstream
828                | ErrorCode::Protocol
829                | ErrorCode::Artifact
830                | ErrorCode::Input
831                | ErrorCode::Internal
832        )
833}
834
835fn elapsed_ms(started: Instant) -> u64 {
836    u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX)
837}
838
839fn timeout_error() -> BridgeError {
840    BridgeError::new(ErrorCode::Timeout, "request deadline elapsed").retryable(true)
841}
842
843fn unknown_outcome_error(code: ErrorCode, message: impl Into<String>) -> BridgeError {
844    BridgeError::new(code, message)
845        .retryable(false)
846        .with_detail("outcome", "unknown")
847}
848
849fn classify_provider_execution_result<T>(result: Result<T, BridgeError>) -> Result<T, BridgeError> {
850    match result {
851        Err(error)
852            if matches!(error.code, ErrorCode::Timeout | ErrorCode::Cancelled)
853                && error
854                    .details
855                    .get("outcome")
856                    .and_then(serde_json::Value::as_str)
857                    != Some("unknown") =>
858        {
859            Err(error.retryable(false).with_detail("outcome", "unknown"))
860        }
861        result => result,
862    }
863}
864
865fn cancelled_error(message: impl Into<String>) -> BridgeError {
866    BridgeError::new(ErrorCode::Cancelled, message)
867}
868
869fn configuration_error(message: impl Into<String>) -> BridgeError {
870    BridgeError::new(ErrorCode::Configuration, message)
871}
872
873fn protocol_error(message: impl Into<String>) -> BridgeError {
874    BridgeError::new(ErrorCode::Protocol, message)
875}
876
877#[cfg(test)]
878mod tests {
879    #![allow(clippy::panic, clippy::unwrap_used)]
880
881    use std::{
882        collections::BTreeSet,
883        sync::atomic::{AtomicBool, AtomicUsize, Ordering},
884    };
885
886    use async_trait::async_trait;
887    use base64::{Engine as _, engine::general_purpose::STANDARD};
888    use imagegen_bridge_artifacts::{ArtifactStore, ImageLimits, inspect_image};
889    use imagegen_bridge_core::{
890        Background, BatchCapabilities, BatchMode, CompatibilityMode, GeneratedImage,
891        GenerationParameters, ImageAction, ImagePayload, ImageProvider, ImageSize,
892        InputCapabilities, InputFidelity, Moderation, OutputFormat, ProviderCapabilities,
893        ProviderDescriptor, Quality, ResponseFormat, SizeCapabilities, SupportLevel, Timings,
894        U8Range, Usage,
895    };
896
897    use super::*;
898
899    const ONE_PIXEL_PNG: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=";
900
901    struct FakeProvider {
902        calls: AtomicUsize,
903        delay: Duration,
904        saw_cancellation: AtomicBool,
905        corrupt_metadata: bool,
906    }
907
908    impl FakeProvider {
909        fn new(delay: Duration) -> Self {
910            Self {
911                calls: AtomicUsize::new(0),
912                delay,
913                saw_cancellation: AtomicBool::new(false),
914                corrupt_metadata: false,
915            }
916        }
917
918        fn corrupting() -> Self {
919            Self {
920                corrupt_metadata: true,
921                ..Self::new(Duration::ZERO)
922            }
923        }
924    }
925
926    #[async_trait]
927    impl ImageProvider for FakeProvider {
928        fn descriptor(&self) -> ProviderDescriptor {
929            ProviderDescriptor {
930                name: "fake".to_owned(),
931                display_name: "Fake".to_owned(),
932                version: "test".to_owned(),
933                experimental: false,
934                models: vec!["test-image".to_owned()],
935            }
936        }
937
938        async fn capabilities(
939            &self,
940            model: Option<&str>,
941        ) -> Result<ProviderCapabilities, BridgeError> {
942            Ok(fake_capabilities(model))
943        }
944
945        async fn execute(
946            &self,
947            request: ImageRequest,
948            context: ProviderContext,
949        ) -> Result<ImageResponse, BridgeError> {
950            self.calls.fetch_add(1, Ordering::AcqRel);
951            tokio::select! {
952                () = tokio::time::sleep(self.delay) => {}
953                () = context.cancellation.cancelled() => {
954                    self.saw_cancellation.store(true, Ordering::Release);
955                    return Err(BridgeError::new(ErrorCode::Cancelled, "fake cancelled"));
956                }
957            }
958            let bytes = STANDARD.decode(ONE_PIXEL_PNG).unwrap();
959            let metadata = inspect_image(&bytes, ImageLimits::default()).unwrap();
960            let image = GeneratedImage {
961                index: 0,
962                payload: ImagePayload::B64Json {
963                    b64_json: ONE_PIXEL_PNG.to_owned(),
964                },
965                format: metadata.format,
966                width: metadata.width,
967                height: metadata.height,
968                bytes: metadata.bytes,
969                sha256: if self.corrupt_metadata {
970                    "0".repeat(64)
971                } else {
972                    metadata.sha256
973                },
974                generation_ms: None,
975                metadata_name: None,
976            };
977            Ok(ImageResponse {
978                id: context.request_id,
979                created: 0,
980                provider: "fake".to_owned(),
981                model: request
982                    .routing
983                    .model
984                    .clone()
985                    .unwrap_or_else(|| "fake-image".to_owned()),
986                requested: request.parameters.clone(),
987                effective: request.parameters.clone(),
988                normalizations: Vec::new(),
989                attempts: Vec::new(),
990                data: (0..request.parameters.n)
991                    .map(|index| GeneratedImage {
992                        index,
993                        ..image.clone()
994                    })
995                    .collect(),
996                failures: Vec::new(),
997                revised_prompt: Some("safe revised prompt".to_owned()),
998                usage: Some(Usage::default()),
999                session: None,
1000                timings: Timings::default(),
1001                warnings: Vec::new(),
1002            })
1003        }
1004
1005        async fn check_ready(&self) -> Result<(), BridgeError> {
1006            Ok(())
1007        }
1008    }
1009
1010    #[derive(Clone, Copy)]
1011    enum RouteFailure {
1012        Authentication,
1013        KnownUpstream,
1014        Safety,
1015        UnknownOutcome,
1016    }
1017
1018    struct RouteProvider {
1019        name: &'static str,
1020        failure: Option<RouteFailure>,
1021        execute_calls: AtomicUsize,
1022    }
1023
1024    impl RouteProvider {
1025        fn new(name: &'static str, failure: Option<RouteFailure>) -> Self {
1026            Self {
1027                name,
1028                failure,
1029                execute_calls: AtomicUsize::new(0),
1030            }
1031        }
1032    }
1033
1034    #[async_trait]
1035    impl ImageProvider for RouteProvider {
1036        fn descriptor(&self) -> ProviderDescriptor {
1037            ProviderDescriptor {
1038                name: self.name.to_owned(),
1039                display_name: self.name.to_owned(),
1040                version: "test".to_owned(),
1041                experimental: false,
1042                models: vec!["route-image".to_owned()],
1043            }
1044        }
1045
1046        async fn capabilities(
1047            &self,
1048            model: Option<&str>,
1049        ) -> Result<ProviderCapabilities, BridgeError> {
1050            if matches!(self.failure, Some(RouteFailure::Authentication)) {
1051                return Err(BridgeError::new(
1052                    ErrorCode::Authentication,
1053                    "test provider is unavailable",
1054                ));
1055            }
1056            let mut capabilities = fake_capabilities(model);
1057            capabilities.provider = self.name.to_owned();
1058            Ok(capabilities)
1059        }
1060
1061        async fn execute(
1062            &self,
1063            request: ImageRequest,
1064            context: ProviderContext,
1065        ) -> Result<ImageResponse, BridgeError> {
1066            self.execute_calls.fetch_add(1, Ordering::AcqRel);
1067            match self.failure {
1068                Some(RouteFailure::Safety) => {
1069                    return Err(BridgeError::safety_rejected("test safety rejection"));
1070                }
1071                Some(RouteFailure::UnknownOutcome) => {
1072                    return Err(
1073                        BridgeError::new(ErrorCode::Upstream, "test ambiguous failure")
1074                            .with_detail("outcome", "unknown"),
1075                    );
1076                }
1077                Some(RouteFailure::KnownUpstream) => {
1078                    return Err(BridgeError::new(ErrorCode::Upstream, "test known failure")
1079                        .with_detail("outcome", "failed"));
1080                }
1081                Some(RouteFailure::Authentication) | None => {}
1082            }
1083            let bytes = STANDARD.decode(ONE_PIXEL_PNG).unwrap();
1084            let metadata = inspect_image(&bytes, ImageLimits::default()).unwrap();
1085            Ok(ImageResponse {
1086                id: context.request_id,
1087                created: 0,
1088                provider: self.name.to_owned(),
1089                model: request
1090                    .routing
1091                    .model
1092                    .clone()
1093                    .unwrap_or_else(|| "route-image".to_owned()),
1094                requested: request.parameters.clone(),
1095                effective: request.parameters.clone(),
1096                normalizations: Vec::new(),
1097                attempts: Vec::new(),
1098                data: vec![GeneratedImage {
1099                    index: 0,
1100                    payload: ImagePayload::B64Json {
1101                        b64_json: ONE_PIXEL_PNG.to_owned(),
1102                    },
1103                    format: metadata.format,
1104                    width: metadata.width,
1105                    height: metadata.height,
1106                    bytes: metadata.bytes,
1107                    sha256: metadata.sha256,
1108                    generation_ms: None,
1109                    metadata_name: None,
1110                }],
1111                failures: Vec::new(),
1112                revised_prompt: None,
1113                usage: None,
1114                session: None,
1115                timings: Timings::default(),
1116                warnings: Vec::new(),
1117            })
1118        }
1119
1120        async fn check_ready(&self) -> Result<(), BridgeError> {
1121            Ok(())
1122        }
1123    }
1124
1125    fn fake_capabilities(model: Option<&str>) -> ProviderCapabilities {
1126        let unsupported_inputs = InputCapabilities {
1127            support: SupportLevel::Unsupported,
1128            max_count: 0,
1129            max_bytes_each: 0,
1130            max_bytes_total: 0,
1131        };
1132        ProviderCapabilities {
1133            provider: "fake".to_owned(),
1134            implementation_version: "test".to_owned(),
1135            model: model.map(str::to_owned),
1136            experimental: false,
1137            generation: true,
1138            edits: false,
1139            count: U8Range { min: 1, max: 1 },
1140            batching: BatchCapabilities {
1141                mode: BatchMode::Native,
1142                native_count: U8Range { min: 1, max: 1 },
1143                max_parallel_outputs: 1,
1144            },
1145            sizes: SizeCapabilities {
1146                auto: true,
1147                allowed: BTreeSet::from([ImageSize::exact(2, 2).unwrap()]),
1148                arbitrary: false,
1149                min_edge: None,
1150                max_edge: None,
1151                edge_multiple: None,
1152                min_pixels: None,
1153                max_pixels: None,
1154                max_aspect_ratio: None,
1155            },
1156            aspect_ratio: SupportLevel::Unsupported,
1157            resolution: SupportLevel::Unsupported,
1158            qualities: BTreeSet::from([Quality::Auto]),
1159            output_formats: BTreeSet::from([OutputFormat::Png]),
1160            backgrounds: BTreeSet::from([Background::Auto]),
1161            transparent_background: SupportLevel::Emulated,
1162            moderation: BTreeSet::from([Moderation::Auto]),
1163            negative_prompt: SupportLevel::Emulated,
1164            revised_prompt: SupportLevel::Native,
1165            user_attribution: SupportLevel::Native,
1166            input_fidelities: BTreeSet::from([InputFidelity::Low, InputFidelity::High]),
1167            actions: BTreeSet::from([ImageAction::Auto, ImageAction::Generate, ImageAction::Edit]),
1168            reference_images: unsupported_inputs.clone(),
1169            edit_images: unsupported_inputs.clone(),
1170            masks: unsupported_inputs,
1171            partial_images: U8Range { min: 0, max: 0 },
1172            persistent_sessions: false,
1173            explicit_threads: false,
1174        }
1175    }
1176
1177    fn runtime(
1178        provider: &Arc<FakeProvider>,
1179        mutate: impl FnOnce(&mut RuntimeConfig),
1180    ) -> ImagegenRuntime {
1181        let registry =
1182            ProviderRegistry::new([Arc::clone(provider) as Arc<dyn ImageProvider>], "fake")
1183                .unwrap();
1184        let mut config = RuntimeConfig::default();
1185        mutate(&mut config);
1186        ImagegenRuntime::new(registry, config).unwrap()
1187    }
1188
1189    fn routing_runtime(
1190        primary: &Arc<RouteProvider>,
1191        fallback: &Arc<RouteProvider>,
1192    ) -> ImagegenRuntime {
1193        let registry = ProviderRegistry::new(
1194            [
1195                Arc::clone(primary) as Arc<dyn ImageProvider>,
1196                Arc::clone(fallback) as Arc<dyn ImageProvider>,
1197            ],
1198            primary.name,
1199        )
1200        .unwrap();
1201        ImagegenRuntime::new(registry, RuntimeConfig::default()).unwrap()
1202    }
1203
1204    #[tokio::test]
1205    async fn negotiates_then_independently_verifies_and_projects_metadata() {
1206        let provider = Arc::new(FakeProvider::new(Duration::ZERO));
1207        let runtime = runtime(&provider, |_| {});
1208        let mut request = ImageRequest::generate("test");
1209        request.parameters = GenerationParameters {
1210            n: 2,
1211            ..GenerationParameters::default()
1212        };
1213        request.policies.compatibility = CompatibilityMode::Normalize;
1214        request.policies.revised_prompt = RevisedPromptPolicy::Omit;
1215        request.output.response_format = ResponseFormat::Metadata;
1216        let response = runtime
1217            .execute_with(
1218                request,
1219                ExecutionContext {
1220                    request_id: Some("request-1".to_owned()),
1221                    ..ExecutionContext::default()
1222                },
1223            )
1224            .await
1225            .unwrap();
1226        assert_eq!(response.id, "request-1");
1227        assert_eq!(response.requested.n, 2);
1228        assert_eq!(response.effective.n, 1);
1229        assert_eq!(response.data.len(), 1);
1230        assert!(matches!(response.data[0].payload, ImagePayload::Metadata));
1231        assert!(response.revised_prompt.is_none());
1232        assert_eq!(provider.calls.load(Ordering::Acquire), 1);
1233    }
1234
1235    #[tokio::test]
1236    async fn strict_mode_rejects_verified_dimension_mismatch() {
1237        let provider = Arc::new(FakeProvider::new(Duration::ZERO));
1238        let runtime = runtime(&provider, |_| {});
1239        let mut request = ImageRequest::generate("test");
1240        request.parameters.size = ImageSize::exact(2, 2).unwrap();
1241        let error = runtime.execute(request).await.unwrap_err();
1242        assert_eq!(error.code, ErrorCode::Protocol);
1243        assert_eq!(error.details["expected"], "2x2");
1244        assert_eq!(error.details["actual"], "1x1");
1245    }
1246
1247    #[tokio::test]
1248    async fn normalize_mode_reports_verified_dimension_mismatch() {
1249        let provider = Arc::new(FakeProvider::new(Duration::ZERO));
1250        let runtime = runtime(&provider, |_| {});
1251        let mut request = ImageRequest::generate("test");
1252        request.parameters.size = ImageSize::exact(2, 2).unwrap();
1253        request.policies.compatibility = CompatibilityMode::Normalize;
1254        let response = runtime.execute(request).await.unwrap();
1255        assert_eq!(response.effective.size, ImageSize::exact(1, 1).unwrap());
1256        assert!(response.normalizations.iter().any(|entry| {
1257            entry.field == "parameters.size"
1258                && entry.reason == "provider_output_dimensions_differed"
1259        }));
1260        assert!(
1261            response
1262                .warnings
1263                .iter()
1264                .any(|warning| warning == "provider_output_dimensions_differed")
1265        );
1266    }
1267
1268    #[tokio::test]
1269    async fn replays_idempotent_responses_without_a_second_provider_call() {
1270        let provider = Arc::new(FakeProvider::new(Duration::ZERO));
1271        let runtime = runtime(&provider, |_| {});
1272        let mut request = ImageRequest::generate("test");
1273        request.idempotency_key = Some("stable-key".to_owned());
1274        let first = runtime.execute(request.clone()).await.unwrap();
1275        let second = runtime.execute(request).await.unwrap();
1276        assert_eq!(first, second);
1277        assert_eq!(provider.calls.load(Ordering::Acquire), 1);
1278    }
1279
1280    #[tokio::test]
1281    async fn publishes_bridge_owned_artifacts_only_after_verification() {
1282        let directory = tempfile::tempdir().unwrap();
1283        let store = Arc::new(ArtifactStore::new(directory.path(), ImageLimits::default()).unwrap());
1284        let provider = Arc::new(FakeProvider::new(Duration::ZERO));
1285        let runtime = runtime(&provider, |config| {
1286            config.materialization.artifact_store = Some(store);
1287        });
1288        let mut request = ImageRequest::generate("test");
1289        request.output.response_format = ResponseFormat::Artifact;
1290        request.output.filename_prefix = Some("Result".to_owned());
1291        let response = runtime.execute(request).await.unwrap();
1292        assert!(matches!(
1293            &response.data[0].payload,
1294            ImagePayload::Artifact { name: Some(name), .. } if name.starts_with("result-")
1295        ));
1296    }
1297
1298    #[tokio::test]
1299    async fn deadline_keeps_an_unknown_idempotency_tombstone() {
1300        let provider = Arc::new(FakeProvider::new(Duration::from_secs(5)));
1301        let runtime = runtime(&provider, |_| {});
1302        let mut request = ImageRequest::generate("test");
1303        request.timeout_ms = Some(100);
1304        request.idempotency_key = Some("retryable-key".to_owned());
1305        let first = runtime.execute(request.clone()).await.unwrap_err();
1306        assert_eq!(first.code, ErrorCode::Timeout);
1307        assert!(!first.retryable);
1308        assert_eq!(first.details["outcome"], "unknown");
1309        tokio::task::yield_now().await;
1310        assert!(provider.saw_cancellation.load(Ordering::Acquire));
1311        let second = runtime.execute(request).await.unwrap_err();
1312        assert_eq!(second.code, ErrorCode::Timeout);
1313        assert!(!second.retryable);
1314        assert_eq!(second.details["outcome"], "unknown");
1315        assert_eq!(provider.calls.load(Ordering::Acquire), 1);
1316    }
1317
1318    #[tokio::test]
1319    async fn rejects_provider_metadata_that_does_not_match_decoded_output() {
1320        let provider = Arc::new(FakeProvider::corrupting());
1321        let runtime = runtime(&provider, |_| {});
1322        let error = runtime
1323            .execute(ImageRequest::generate("test"))
1324            .await
1325            .unwrap_err();
1326        assert_eq!(error.code, ErrorCode::Protocol);
1327    }
1328
1329    #[tokio::test]
1330    async fn runtime_queue_rejects_work_beyond_its_explicit_bound() {
1331        let provider = Arc::new(FakeProvider::new(Duration::from_millis(40)));
1332        let runtime = Arc::new(runtime(&provider, |config| {
1333            config.global_limit = ConcurrencyLimit {
1334                max_concurrent: 1,
1335                max_queued: 1,
1336            };
1337            config.default_provider_limit = ConcurrencyLimit {
1338                max_concurrent: 1,
1339                max_queued: 1,
1340            };
1341        }));
1342        let first_runtime = Arc::clone(&runtime);
1343        let first =
1344            tokio::spawn(
1345                async move { first_runtime.execute(ImageRequest::generate("first")).await },
1346            );
1347        while provider.calls.load(Ordering::Acquire) == 0 {
1348            tokio::task::yield_now().await;
1349        }
1350        let second_runtime = Arc::clone(&runtime);
1351        let second = tokio::spawn(async move {
1352            second_runtime
1353                .execute(ImageRequest::generate("second"))
1354                .await
1355        });
1356        while runtime.queue_snapshot().global_queued == 0 {
1357            tokio::task::yield_now().await;
1358        }
1359        let error = runtime
1360            .execute(ImageRequest::generate("rejected"))
1361            .await
1362            .unwrap_err();
1363        assert_eq!(error.code, ErrorCode::Overloaded);
1364        assert!(first.await.unwrap().is_ok());
1365        assert!(second.await.unwrap().is_ok());
1366    }
1367
1368    #[tokio::test]
1369    async fn concurrent_idempotent_callers_share_one_provider_operation() {
1370        let provider = Arc::new(FakeProvider::new(Duration::from_millis(20)));
1371        let runtime = Arc::new(runtime(&provider, |_| {}));
1372        let mut request = ImageRequest::generate("same");
1373        request.idempotency_key = Some("shared-key".to_owned());
1374        let first_runtime = Arc::clone(&runtime);
1375        let first_request = request.clone();
1376        let first = tokio::spawn(async move { first_runtime.execute(first_request).await });
1377        while provider.calls.load(Ordering::Acquire) == 0 {
1378            tokio::task::yield_now().await;
1379        }
1380        let second = runtime.execute(request).await.unwrap();
1381        let first = first.await.unwrap().unwrap();
1382        assert_eq!(first, second);
1383        assert_eq!(provider.calls.load(Ordering::Acquire), 1);
1384    }
1385
1386    #[tokio::test]
1387    async fn unavailable_primary_falls_back_in_order_and_records_attempts() {
1388        let primary = Arc::new(RouteProvider::new(
1389            "primary",
1390            Some(RouteFailure::Authentication),
1391        ));
1392        let fallback = Arc::new(RouteProvider::new("fallback", None));
1393        let runtime = routing_runtime(&primary, &fallback);
1394        let mut request = ImageRequest::generate("test routing");
1395        request.routing.fallbacks.push(ProviderRoute {
1396            provider: "fallback".to_owned(),
1397            model: None,
1398        });
1399        let response = runtime.execute(request).await.unwrap();
1400        assert_eq!(response.provider, "fallback");
1401        assert_eq!(response.attempts.len(), 2);
1402        assert_eq!(
1403            response.attempts[0].error_code,
1404            Some(ErrorCode::Authentication)
1405        );
1406        assert_eq!(
1407            response.attempts[1].outcome,
1408            ProviderAttemptOutcome::Succeeded
1409        );
1410        assert!(
1411            response
1412                .warnings
1413                .contains(&"provider_fallback_used".to_owned())
1414        );
1415        assert_eq!(primary.execute_calls.load(Ordering::Acquire), 0);
1416        assert_eq!(fallback.execute_calls.load(Ordering::Acquire), 1);
1417    }
1418
1419    #[tokio::test]
1420    async fn fallback_never_bypasses_safety_or_unknown_outcomes() {
1421        for failure in [RouteFailure::Safety, RouteFailure::UnknownOutcome] {
1422            let primary = Arc::new(RouteProvider::new("primary", Some(failure)));
1423            let fallback = Arc::new(RouteProvider::new("fallback", None));
1424            let runtime = routing_runtime(&primary, &fallback);
1425            let mut request = ImageRequest::generate("test routing guard");
1426            request.routing.fallback_policy = FallbackPolicy::OnError;
1427            request.routing.fallbacks.push(ProviderRoute {
1428                provider: "fallback".to_owned(),
1429                model: None,
1430            });
1431            let error = runtime.execute(request).await.unwrap_err();
1432            assert_eq!(fallback.execute_calls.load(Ordering::Acquire), 0);
1433            assert_eq!(error.details["attempts"].as_array().unwrap().len(), 1);
1434        }
1435    }
1436
1437    #[tokio::test]
1438    async fn on_error_expands_but_does_not_replace_unavailable_policy() {
1439        for (policy, expected_calls) in [
1440            (FallbackPolicy::OnUnavailable, 0),
1441            (FallbackPolicy::OnError, 1),
1442        ] {
1443            let primary = Arc::new(RouteProvider::new(
1444                "primary",
1445                Some(RouteFailure::KnownUpstream),
1446            ));
1447            let fallback = Arc::new(RouteProvider::new("fallback", None));
1448            let runtime = routing_runtime(&primary, &fallback);
1449            let mut request = ImageRequest::generate("test known outcome policy");
1450            request.routing.fallback_policy = policy;
1451            request.routing.fallbacks.push(ProviderRoute {
1452                provider: "fallback".to_owned(),
1453                model: None,
1454            });
1455            let result = runtime.execute(request).await;
1456            assert_eq!(
1457                fallback.execute_calls.load(Ordering::Acquire),
1458                expected_calls
1459            );
1460            assert_eq!(result.is_ok(), policy == FallbackPolicy::OnError);
1461        }
1462    }
1463
1464    #[tokio::test]
1465    async fn implicit_primary_cannot_be_repeated_as_a_fallback() {
1466        let primary = Arc::new(RouteProvider::new("primary", None));
1467        let fallback = Arc::new(RouteProvider::new("fallback", None));
1468        let runtime = routing_runtime(&primary, &fallback);
1469        let mut request = ImageRequest::generate("test duplicate route");
1470        request.routing.fallbacks.push(ProviderRoute {
1471            provider: "primary".to_owned(),
1472            model: None,
1473        });
1474        let error = runtime.execute(request).await.unwrap_err();
1475        assert_eq!(error.code, ErrorCode::InvalidRequest);
1476        assert_eq!(primary.execute_calls.load(Ordering::Acquire), 0);
1477    }
1478
1479    #[tokio::test]
1480    async fn shutdown_cancels_and_drains_active_work_before_provider_release() {
1481        let provider = Arc::new(FakeProvider::new(Duration::from_secs(5)));
1482        let runtime = Arc::new(runtime(&provider, |_| {}));
1483        let executing_runtime = Arc::clone(&runtime);
1484        let executing = tokio::spawn(async move {
1485            executing_runtime
1486                .execute(ImageRequest::generate("active"))
1487                .await
1488        });
1489        while provider.calls.load(Ordering::Acquire) == 0 {
1490            tokio::task::yield_now().await;
1491        }
1492        runtime.shutdown().await.unwrap();
1493        runtime.shutdown().await.unwrap();
1494        let error = executing.await.unwrap().unwrap_err();
1495        assert_eq!(error.code, ErrorCode::Cancelled);
1496        assert!(provider.saw_cancellation.load(Ordering::Acquire));
1497        let error = runtime
1498            .execute(ImageRequest::generate("after shutdown"))
1499            .await
1500            .unwrap_err();
1501        assert_eq!(error.code, ErrorCode::Cancelled);
1502    }
1503}
imagegen_bridge_runtime/orchestrator.rs

imagegen_bridge_runtime/
orchestrator.rs
