Everything is streaming (#7247)

Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-02-17 19:53:38 +01:00
parent 65d3a1550f
commit 9a91fccfb5
42 changed files with 801 additions and 1718 deletions
@@ -417,85 +417,85 @@ impl Provider for GithubCopilotProvider {
        self.model.clone()
    }

-    fn supports_streaming(&self) -> bool {
-        GITHUB_COPILOT_STREAM_MODELS
-            .iter()
-            .any(|prefix| self.model.model_name.starts_with(prefix))
-    }
-
-    #[tracing::instrument(
-        skip(self, model_config, system, messages, tools),
-        fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
-    )]
-    async fn complete_with_model(
-        &self,
-        session_id: Option<&str>,
-        model_config: &ModelConfig,
-        system: &str,
-        messages: &[Message],
-        tools: &[Tool],
-    ) -> Result<(Message, ProviderUsage), ProviderError> {
-        let payload = create_request(
-            model_config,
-            system,
-            messages,
-            tools,
-            &ImageFormat::OpenAi,
-            false,
-        )?;
-        let mut log = RequestLog::start(model_config, &payload)?;
-
-        // Make request with retry
-        let response = self
-            .with_retry(|| async {
-                let mut payload_clone = payload.clone();
-                self.post(session_id, &mut payload_clone).await
-            })
-            .await?;
-        let response = handle_response_openai_compat(response).await?;
-
-        let response = promote_tool_choice(response);
-
-        // Parse response
-        let message = response_to_message(&response)?;
-        let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
-            tracing::debug!("Failed to get usage data");
-            Usage::default()
-        });
-        let response_model = get_model(&response);
-        log.write(&response, Some(&usage))?;
-        Ok((message, ProviderUsage::new(response_model, usage)))
-    }
-
    async fn stream(
        &self,
+        model_config: &ModelConfig,
        session_id: &str,
        system: &str,
        messages: &[Message],
        tools: &[Tool],
    ) -> Result<MessageStream, ProviderError> {
-        let payload = create_request(
-            &self.model,
-            system,
-            messages,
-            tools,
-            &ImageFormat::OpenAi,
-            true,
-        )?;
-        let mut log = RequestLog::start(&self.model, &payload)?;
+        // Check if this model supports streaming
+        let supports_streaming = GITHUB_COPILOT_STREAM_MODELS
+            .iter()
+            .any(|prefix| model_config.model_name.starts_with(prefix));

-        let response = self
-            .with_retry(|| async {
-                let mut payload_clone = payload.clone();
-                let resp = self.post(Some(session_id), &mut payload_clone).await?;
-                handle_status_openai_compat(resp).await
-            })
-            .await
-            .inspect_err(|e| {
-                let _ = log.error(e);
-            })?;
+        if supports_streaming {
+            // Use streaming API
+            let payload = create_request(
+                model_config,
+                system,
+                messages,
+                tools,
+                &ImageFormat::OpenAi,
+                true,
+            )?;
+            let mut log = RequestLog::start(model_config, &payload)?;

-        stream_openai_compat(response, log)
+            let response = self
+                .with_retry(|| async {
+                    let mut payload_clone = payload.clone();
+                    let resp = self.post(Some(session_id), &mut payload_clone).await?;
+                    handle_status_openai_compat(resp).await
+                })
+                .await
+                .inspect_err(|e| {
+                    let _ = log.error(e);
+                })?;
+
+            stream_openai_compat(response, log)
+        } else {
+            // Use non-streaming API and wrap result
+            let session_id_opt = if session_id.is_empty() {
+                None
+            } else {
+                Some(session_id)
+            };
+            let payload = create_request(
+                model_config,
+                system,
+                messages,
+                tools,
+                &ImageFormat::OpenAi,
+                false,
+            )?;
+            let mut log = RequestLog::start(model_config, &payload)?;
+
+            // Make request with retry
+            let response = self
+                .with_retry(|| async {
+                    let mut payload_clone = payload.clone();
+                    self.post(session_id_opt, &mut payload_clone).await
+                })
+                .await?;
+            let response = handle_response_openai_compat(response).await?;
+
+            let response = promote_tool_choice(response);
+
+            // Parse response
+            let message = response_to_message(&response)?;
+            let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
+                tracing::debug!("Failed to get usage data");
+                Usage::default()
+            });
+            let response_model = get_model(&response);
+            log.write(&response, Some(&usage))?;
+
+            Ok(super::base::stream_from_single_message(
+                message,
+                ProviderUsage::new(response_model, usage),
+            ))
+        }
    }

    async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {