Improve gateway transport and usage runtime

crates/aether-ai-serving/src/execution_path.rs

@@ -38,6 +38,17 @@ pub enum AiStreamExecutionStep {
     RemoteDecision,
 }
 
+pub const DEFAULT_STREAM_EXECUTION_STEPS: &[AiStreamExecutionStep] = &[
+    AiStreamExecutionStep::LocalVideoContent,
+    AiStreamExecutionStep::LocalImage,
+    AiStreamExecutionStep::LocalOpenAiChat,
+    AiStreamExecutionStep::LocalOpenAiResponses,
+    AiStreamExecutionStep::LocalStandardFamily,
+    AiStreamExecutionStep::LocalSameFormatProvider,
+    AiStreamExecutionStep::LocalGeminiFiles,
+    AiStreamExecutionStep::RemoteDecision,
+];
+
 #[async_trait]
 pub trait AiSyncExecutionPathPort: Send + Sync {
     type Response: Send;
@@ -65,6 +76,10 @@ pub trait AiStreamExecutionPathPort: Send + Sync {
 
     fn scheduler_decision_supported(&self) -> bool;
 
+    fn stream_execution_steps(&self) -> &'static [AiStreamExecutionStep] {
+        DEFAULT_STREAM_EXECUTION_STEPS
+    }
+
     async fn execute_stream_step(
         &self,
         step: AiStreamExecutionStep,
@@ -135,29 +150,13 @@ where
 {
     let mut exhausted = None;
 
-    if let Some(response) = absorb_stream_step(
-        port,
-        AiStreamExecutionStep::LocalVideoContent,
-        &mut exhausted,
-    )
-    .await?
-    {
-        return Ok(response);
-    }
-
-    if port.scheduler_decision_supported() {
-        for step in [
-            AiStreamExecutionStep::LocalImage,
-            AiStreamExecutionStep::LocalOpenAiChat,
-            AiStreamExecutionStep::LocalOpenAiResponses,
-            AiStreamExecutionStep::LocalStandardFamily,
-            AiStreamExecutionStep::LocalSameFormatProvider,
-            AiStreamExecutionStep::LocalGeminiFiles,
-            AiStreamExecutionStep::RemoteDecision,
-        ] {
-            if let Some(response) = absorb_stream_step(port, step, &mut exhausted).await? {
-                return Ok(response);
-            }
+    for step in port.stream_execution_steps() {
+        if *step != AiStreamExecutionStep::LocalVideoContent && !port.scheduler_decision_supported()
+        {
+            continue;
+        }
+        if let Some(response) = absorb_stream_step(port, *step, &mut exhausted).await? {
+            return Ok(response);
         }
     }
 
@@ -279,6 +278,7 @@ mod tests {
     #[derive(Default)]
     struct TestStreamPort {
         scheduler_supported: bool,
+        stream_steps: Option<&'static [AiStreamExecutionStep]>,
         outcomes: Mutex<VecDeque<AiServingExecutionOutcome<&'static str, &'static str>>>,
         calls: Mutex<Vec<String>>,
     }
@@ -293,6 +293,11 @@ mod tests {
             self.scheduler_supported
         }
 
+        fn stream_execution_steps(&self) -> &'static [AiStreamExecutionStep] {
+            self.stream_steps
+                .unwrap_or(super::DEFAULT_STREAM_EXECUTION_STEPS)
+        }
+
         async fn execute_stream_step(
             &self,
             step: AiStreamExecutionStep,
@@ -405,6 +410,7 @@ mod tests {
     async fn stream_path_stops_at_first_response() {
         let port = TestStreamPort {
             scheduler_supported: true,
+            stream_steps: None,
             outcomes: Mutex::new(VecDeque::from([
                 AiServingExecutionOutcome::NoPath,
                 AiServingExecutionOutcome::Responded("image_response"),
@@ -424,10 +430,32 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn stream_path_runs_preferred_steps_only() {
+        const CHAT_ONLY: &[AiStreamExecutionStep] = &[AiStreamExecutionStep::LocalOpenAiChat];
+        let port = TestStreamPort {
+            scheduler_supported: true,
+            stream_steps: Some(CHAT_ONLY),
+            outcomes: Mutex::new(VecDeque::from([AiServingExecutionOutcome::Responded(
+                "chat_response",
+            )])),
+            calls: Mutex::default(),
+        };
+
+        let outcome = run_ai_stream_execution_path(&port).await.unwrap();
+
+        assert!(matches!(
+            outcome,
+            AiServingExecutionOutcome::Responded("chat_response")
+        ));
+        assert_eq!(port.calls.lock().unwrap().as_slice(), ["LocalOpenAiChat"]);
+    }
+
     #[tokio::test]
     async fn stream_path_returns_last_exhaustion_without_plan_fallback() {
         let port = TestStreamPort {
             scheduler_supported: true,
+            stream_steps: None,
             outcomes: Mutex::new(VecDeque::from([
                 AiServingExecutionOutcome::NoPath,
                 AiServingExecutionOutcome::Exhausted("local_image_exhausted"),

crates/aether-ai-serving/src/lib.rs

@@ -113,7 +113,7 @@ pub use dto::{
 pub use execution_path::{
     run_ai_stream_execution_path, run_ai_sync_execution_path, AiPlanFallbackReason,
     AiServingExecutionOutcome, AiStreamExecutionPathPort, AiStreamExecutionStep,
-    AiSyncExecutionPathPort, AiSyncExecutionStep,
+    AiSyncExecutionPathPort, AiSyncExecutionStep, DEFAULT_STREAM_EXECUTION_STEPS,
 };
 pub use failure_diagnostic::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind};
 pub use plan_payload::{