Skip to main content

agent_sdk_providers/impls/
vertex.rs

1//! Google Vertex AI provider implementation.
2//!
3//! This module provides an implementation of `LlmProvider` for the Google Vertex AI
4//! platform. It supports both Gemini models (using the Gemini API format) and
5//! Claude models (using the Anthropic Messages API format via `rawPredict`).
6//!
7//! Publisher detection is automatic based on the model name:
8//! - `claude-*` models route to `publishers/anthropic` using `rawPredict`
9//! - All other models route to `publishers/google` using `generateContent`
10
11use crate::attachments::validate_request_attachments;
12use crate::impls::anthropic::{
13    MODEL_FABLE_5, MODEL_OPUS_46, MODEL_OPUS_47, MODEL_OPUS_48, MODEL_SONNET_5, MODEL_SONNET_46,
14    data as anthropic_data,
15};
16use crate::impls::gemini::data::{
17    ApiContent, ApiFunctionCallingConfig, ApiGenerateContentRequest, ApiGenerateContentResponse,
18    ApiGenerationConfig, ApiPart, ApiUsageMetadata, build_api_contents, build_content_blocks,
19    convert_tools_to_config, gemini_response_schema, map_finish_reason, map_thinking_config,
20    stream_gemini_response,
21};
22use crate::provider::{LlmProvider, thinking_for_forced_tool};
23use crate::streaming::{
24    StreamBox, StreamDelta, StreamErrorKind, reqwest_body_error_delta, reqwest_error_delta,
25};
26use agent_sdk_foundation::llm::{
27    ChatOutcome, ChatRequest, ChatResponse, ResponseFormat, ThinkingConfig, ThinkingMode, Usage,
28};
29use anyhow::Result;
30use async_trait::async_trait;
31use futures::StreamExt;
32use reqwest::StatusCode;
33
34pub const MODEL_GEMINI_3_FLASH: &str = "gemini-3-flash-preview";
35pub const MODEL_GEMINI_31_PRO: &str = "gemini-3.1-pro-preview";
36
37// Legacy Gemini 3.0 Pro model kept for explicit opt-in.
38pub const MODEL_GEMINI_3_PRO: &str = "gemini-3.0-pro";
39
40/// The Anthropic API version used for Claude models on Vertex AI.
41const VERTEX_ANTHROPIC_VERSION: &str = "vertex-2023-10-16";
42const DEFAULT_SAFE_MAX_OUTPUT_TOKENS: u32 = 32_000;
43
44/// Connect timeout for the HTTP client (matches Anthropic).
45const CONNECT_TIMEOUT_SECS: u64 = 30;
46/// TCP keepalive interval to keep long streaming connections from dropping.
47const TCP_KEEPALIVE_SECS: u64 = 30;
48/// Per-request read timeout for the **non-streaming** chat paths. Bounds a
49/// black-holed endpoint so a single turn cannot hang the agent loop forever.
50/// Streaming requests intentionally have no overall timeout.
51const CHAT_READ_TIMEOUT_SECS: u64 = 300;
52
53const fn vertex_cache_control() -> anthropic_data::ApiCacheControl {
54    anthropic_data::ApiCacheControl::ephemeral_with_ttl(None)
55}
56
57/// Build the Claude-on-Vertex tool list, caching the tool prefix with the
58/// Vertex ephemeral breakpoint.
59fn build_vertex_claude_tools(request: &ChatRequest) -> Option<Vec<anthropic_data::ApiTool>> {
60    anthropic_data::build_api_tools_with_cache(request, Some(vertex_cache_control()))
61}
62
63/// Google Vertex AI LLM provider.
64///
65/// Uses the same Gemini request/response format as `GeminiProvider` but
66/// authenticates via `OAuth2` Bearer tokens and routes through the Vertex AI
67/// regional endpoint.
68///
69/// Claude models are also supported — the provider detects the publisher from
70/// the model name and uses the appropriate API format automatically.
71#[derive(Clone)]
72pub struct VertexProvider {
73    client: reqwest::Client,
74    access_token: String,
75    project_id: String,
76    region: String,
77    model: String,
78    thinking: Option<ThinkingConfig>,
79}
80
81impl VertexProvider {
82    /// Create a new Vertex AI provider with full control over all parameters.
83    #[must_use]
84    pub fn new(access_token: String, project_id: String, region: String, model: String) -> Self {
85        let client = reqwest::Client::builder()
86            .connect_timeout(std::time::Duration::from_secs(CONNECT_TIMEOUT_SECS))
87            .tcp_keepalive(std::time::Duration::from_secs(TCP_KEEPALIVE_SECS))
88            .build()
89            .unwrap_or_else(|error| {
90                log::warn!(
91                    "failed to build Vertex HTTP client with timeouts ({error}); using default client"
92                );
93                reqwest::Client::new()
94            });
95
96        Self {
97            client,
98            access_token,
99            project_id,
100            region,
101            model,
102            thinking: None,
103        }
104    }
105
106    /// Create a provider using Gemini 3 Flash Preview on Vertex AI.
107    #[must_use]
108    pub fn flash(access_token: String, project_id: String, region: String) -> Self {
109        Self::new(
110            access_token,
111            project_id,
112            region,
113            MODEL_GEMINI_3_FLASH.to_owned(),
114        )
115    }
116
117    /// Create a provider using Gemini 3.1 Pro Preview on Vertex AI.
118    #[must_use]
119    pub fn pro(access_token: String, project_id: String, region: String) -> Self {
120        Self::new(
121            access_token,
122            project_id,
123            region,
124            MODEL_GEMINI_31_PRO.to_owned(),
125        )
126    }
127
128    /// Detect whether the model is a Claude model (Anthropic publisher).
129    fn is_claude_model(&self) -> bool {
130        self.model.starts_with("claude-")
131    }
132
133    /// The API domain for this provider's location.
134    ///
135    /// For the `global` location the domain is `aiplatform.googleapis.com`
136    /// (no region prefix). Regional locations use `{region}-aiplatform.googleapis.com`.
137    fn endpoint_domain(&self) -> String {
138        if self.region == "global" {
139            "aiplatform.googleapis.com".to_owned()
140        } else {
141            format!("{}-aiplatform.googleapis.com", self.region)
142        }
143    }
144
145    /// Build the base URL for the given publisher and model.
146    fn base_url(&self, publisher: &str) -> String {
147        let domain = self.endpoint_domain();
148        format!(
149            "https://{domain}/v1/projects/{project}/locations/{region}/publishers/{publisher}/models/{model}",
150            domain = domain,
151            region = self.region,
152            project = self.project_id,
153            publisher = publisher,
154            model = self.model,
155        )
156    }
157
158    /// Set the provider-owned thinking configuration for this model.
159    #[must_use]
160    pub const fn with_thinking(mut self, thinking: ThinkingConfig) -> Self {
161        self.thinking = Some(thinking);
162        self
163    }
164
165    fn is_anthropic_adaptive_thinking_model(&self) -> bool {
166        matches!(
167            self.model.as_str(),
168            MODEL_SONNET_46
169                | MODEL_SONNET_5
170                | MODEL_OPUS_46
171                | MODEL_OPUS_47
172                | MODEL_OPUS_48
173                | MODEL_FABLE_5
174        )
175    }
176
177    fn build_cached_vertex_claude_messages(
178        request: &ChatRequest,
179    ) -> Vec<anthropic_data::ApiMessage> {
180        let mut messages = anthropic_data::build_api_messages(request);
181        anthropic_data::apply_cache_control_to_last_user_message(
182            &mut messages,
183            vertex_cache_control(),
184        );
185        messages
186    }
187
188    fn build_vertex_claude_system_prompt(
189        system: &str,
190    ) -> Option<anthropic_data::ApiSystemPrompt<'_>> {
191        anthropic_data::build_api_system_prompt(system, Some(vertex_cache_control()))
192    }
193
194    /// Effective output-token budget for a request.
195    ///
196    /// Mirrors the Anthropic provider: an implicit budget falls back to the
197    /// provider/model default ([`default_max_tokens`](LlmProvider::default_max_tokens),
198    /// which clamps to Vertex's safe ceiling) instead of silently capping at
199    /// `ChatRequest::DEFAULT_MAX_TOKENS`.
200    fn effective_max_tokens(&self, request: &ChatRequest) -> u32 {
201        if request.max_tokens_explicit {
202            request.max_tokens
203        } else {
204            self.default_max_tokens()
205        }
206    }
207
208    fn map_claude_response(api_response: anthropic_data::ApiResponse) -> ChatResponse {
209        let content = anthropic_data::map_content_blocks(api_response.content);
210        let stop_reason = api_response
211            .stop_reason
212            .as_ref()
213            .map(anthropic_data::map_stop_reason);
214
215        ChatResponse {
216            id: api_response.id,
217            content,
218            model: api_response.model,
219            stop_reason,
220            usage: Usage {
221                input_tokens: api_response.usage.total_input_tokens(),
222                output_tokens: api_response.usage.output,
223                cached_input_tokens: api_response.usage.cached_input_tokens(),
224                cache_creation_input_tokens: api_response.usage.cache_creation_input_tokens(),
225            },
226        }
227    }
228}
229
230#[async_trait]
231impl LlmProvider for VertexProvider {
232    async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome> {
233        if self.is_claude_model() {
234            return self.chat_claude(request).await;
235        }
236        self.chat_gemini(request).await
237    }
238
239    fn chat_stream(&self, request: ChatRequest) -> StreamBox<'_> {
240        if self.is_claude_model() {
241            return self.chat_stream_claude(request);
242        }
243        self.chat_stream_gemini(request)
244    }
245
246    fn validate_thinking_config(&self, thinking: Option<&ThinkingConfig>) -> Result<()> {
247        let Some(thinking) = thinking else {
248            return Ok(());
249        };
250
251        if self
252            .capabilities()
253            .is_some_and(|caps| !caps.supports_thinking)
254        {
255            return Err(anyhow::anyhow!(
256                "thinking is not supported for provider={} model={}",
257                self.provider(),
258                self.model()
259            ));
260        }
261
262        if matches!(thinking.mode, ThinkingMode::Adaptive)
263            && !self
264                .capabilities()
265                .is_some_and(|caps| caps.supports_adaptive_thinking)
266        {
267            return Err(anyhow::anyhow!(
268                "adaptive thinking is not supported for provider={} model={}",
269                self.provider(),
270                self.model()
271            ));
272        }
273
274        if self.is_claude_model()
275            && self.is_anthropic_adaptive_thinking_model()
276            && matches!(thinking.mode, ThinkingMode::Enabled { .. })
277        {
278            return Err(anyhow::anyhow!(
279                "budget_tokens thinking is rejected for provider={} model={}; use ThinkingConfig::adaptive() or ThinkingConfig::default_with_effort(_) instead",
280                self.provider(),
281                self.model()
282            ));
283        }
284
285        Ok(())
286    }
287
288    async fn probe_connectivity(&self) -> bool {
289        let url = format!("https://{}/", self.endpoint_domain());
290        crate::provider::probe_http_reachability(&self.client, &url).await
291    }
292
293    fn model(&self) -> &str {
294        &self.model
295    }
296
297    fn provider(&self) -> &'static str {
298        "vertex"
299    }
300
301    fn configured_thinking(&self) -> Option<&ThinkingConfig> {
302        self.thinking.as_ref()
303    }
304
305    fn default_max_tokens(&self) -> u32 {
306        let provider = if self.is_claude_model() {
307            "anthropic"
308        } else {
309            "gemini"
310        };
311        let model_max = self
312            .capabilities()
313            .and_then(|caps| caps.max_output_tokens)
314            .or_else(|| {
315                crate::model_capabilities::default_max_output_tokens(provider, self.model())
316            })
317            .unwrap_or(4096);
318        model_max.clamp(4096, DEFAULT_SAFE_MAX_OUTPUT_TOKENS)
319    }
320}
321
322// ============================================================================
323// Gemini path (publishers/google)
324// ============================================================================
325
326impl VertexProvider {
327    #[allow(clippy::too_many_lines)]
328    async fn chat_gemini(&self, request: ChatRequest) -> Result<ChatOutcome> {
329        let thinking = match self.resolve_thinking_config(request.thinking.as_ref()) {
330            Ok(thinking) => thinking,
331            Err(error) => return Ok(ChatOutcome::InvalidRequest(error.to_string())),
332        };
333        if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
334            return Ok(ChatOutcome::InvalidRequest(error.to_string()));
335        }
336        let contents = build_api_contents(&request.messages);
337        let tools = request
338            .tools
339            .as_ref()
340            .map(|t| convert_tools_to_config(t.clone()));
341        let tool_config = request
342            .tool_choice
343            .as_ref()
344            .map(ApiFunctionCallingConfig::from_tool_choice);
345        let system_instruction = if request.system.is_empty() {
346            None
347        } else {
348            Some(ApiContent {
349                role: None,
350                parts: vec![ApiPart::Text {
351                    text: request.system.clone(),
352                    thought_signature: None,
353                }],
354            })
355        };
356
357        let thinking_config = thinking.as_ref().map(map_thinking_config);
358        let (response_mime_type, response_schema) =
359            request.response_format.as_ref().map_or((None, None), |rf| {
360                (
361                    Some("application/json"),
362                    Some(gemini_response_schema(&rf.schema)),
363                )
364            });
365
366        let max_tokens = self.effective_max_tokens(&request);
367        let api_request = ApiGenerateContentRequest {
368            contents: &contents,
369            system_instruction: system_instruction.as_ref(),
370            tools: tools.as_ref().map(std::slice::from_ref),
371            tool_config,
372            generation_config: Some(ApiGenerationConfig {
373                max_output_tokens: Some(max_tokens),
374                thinking_config,
375                response_mime_type,
376                response_schema,
377            }),
378            cached_content: request.cached_content.as_deref(),
379        };
380
381        log::debug!(
382            "Vertex AI LLM request model={} max_tokens={}",
383            self.model,
384            max_tokens
385        );
386
387        let url = format!("{}:generateContent", self.base_url("google"));
388
389        let response = self
390            .client
391            .post(&url)
392            .header("Content-Type", "application/json")
393            .timeout(std::time::Duration::from_secs(CHAT_READ_TIMEOUT_SECS))
394            .bearer_auth(&self.access_token)
395            .json(&api_request)
396            .send()
397            .await
398            .map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
399
400        let status = response.status();
401        // Read `Retry-After` off the 429 response before the body is consumed.
402        let retry_after = if status == StatusCode::TOO_MANY_REQUESTS {
403            crate::http::retry_after_from_headers(response.headers())
404        } else {
405            None
406        };
407        let bytes = response
408            .bytes()
409            .await
410            .map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?;
411
412        log::debug!(
413            "Vertex AI LLM response status={} body_len={}",
414            status,
415            bytes.len()
416        );
417
418        if status == StatusCode::TOO_MANY_REQUESTS {
419            let body = String::from_utf8_lossy(&bytes);
420            return Ok(ChatOutcome::RateLimited(vertex_retry_delay(
421                retry_after,
422                &body,
423            )));
424        }
425
426        if status.is_server_error() {
427            let body = String::from_utf8_lossy(&bytes);
428            log::error!("Vertex AI server error status={status} body={body}");
429            return Ok(ChatOutcome::ServerError(body.into_owned()));
430        }
431
432        if status.is_client_error() {
433            let body = String::from_utf8_lossy(&bytes);
434            log::warn!("Vertex AI client error status={status} body={body}");
435            return Ok(ChatOutcome::InvalidRequest(body.into_owned()));
436        }
437
438        let api_response: ApiGenerateContentResponse = serde_json::from_slice(&bytes)
439            .map_err(|e| anyhow::anyhow!("failed to parse response: {e}"))?;
440
441        let candidate = api_response
442            .candidates
443            .into_iter()
444            .next()
445            .ok_or_else(|| anyhow::anyhow!("no candidates in response"))?;
446
447        let content = build_content_blocks(&candidate.content);
448
449        if content.is_empty() && !candidate.content.parts.is_empty() {
450            log::warn!(
451                "Vertex AI parts not converted to content blocks raw_parts={:?}",
452                candidate.content.parts
453            );
454        }
455
456        let has_tool_calls = content
457            .iter()
458            .any(|b| matches!(b, agent_sdk_foundation::llm::ContentBlock::ToolUse { .. }));
459
460        let stop_reason = candidate
461            .finish_reason
462            .as_ref()
463            .map(|r| map_finish_reason(r, has_tool_calls));
464
465        let usage = api_response
466            .usage_metadata
467            .unwrap_or(ApiUsageMetadata {
468                prompt: 0,
469                candidates: 0,
470                cached_content: 0,
471            })
472            .into_usage();
473
474        Ok(ChatOutcome::Success(ChatResponse {
475            id: String::new(),
476            content,
477            model: self.model.clone(),
478            stop_reason,
479            usage,
480        }))
481    }
482
483    fn chat_stream_gemini(&self, request: ChatRequest) -> StreamBox<'_> {
484        let served_route = self.route().to_owned();
485        Box::pin(async_stream::stream! {
486            let thinking = match self.resolve_thinking_config(request.thinking.as_ref()) {
487                Ok(thinking) => thinking,
488                Err(error) => {
489                    yield Ok(StreamDelta::Error {
490                        message: error.to_string(),
491                        kind: StreamErrorKind::InvalidRequest,
492                    });
493                    return;
494                }
495            };
496            if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
497                yield Ok(StreamDelta::Error {
498                    message: error.to_string(),
499                    kind: StreamErrorKind::InvalidRequest,
500                });
501                return;
502            }
503
504            let contents = build_api_contents(&request.messages);
505            let tools = request
506                .tools
507                .as_ref()
508                .map(|t| convert_tools_to_config(t.clone()));
509            let tool_config = request
510                .tool_choice
511                .as_ref()
512                .map(ApiFunctionCallingConfig::from_tool_choice);
513            let system_instruction = build_gemini_system_instruction(&request.system);
514            let thinking_config = thinking.as_ref().map(map_thinking_config);
515            let (response_mime_type, response_schema) =
516                gemini_response_format(request.response_format.as_ref());
517
518            let max_tokens = self.effective_max_tokens(&request);
519            let api_request = ApiGenerateContentRequest {
520                contents: &contents,
521                system_instruction: system_instruction.as_ref(),
522                tools: tools.as_ref().map(std::slice::from_ref),
523                tool_config,
524                generation_config: Some(ApiGenerationConfig {
525                    max_output_tokens: Some(max_tokens),
526                    thinking_config,
527                    response_mime_type,
528                    response_schema,
529                }),
530                cached_content: request.cached_content.as_deref(),
531            };
532
533            log::debug!(
534                "Vertex AI streaming LLM request model={} max_tokens={}",
535                self.model,
536                max_tokens
537            );
538
539            let url = format!("{}:streamGenerateContent?alt=sse", self.base_url("google"));
540
541            let response = match self.send_gemini_stream_request(&url, &api_request).await {
542                Ok(response) => response,
543                Err(item) => {
544                    yield item;
545                    return;
546                }
547            };
548
549            let mut inner = stream_gemini_response(response);
550            while let Some(item) = futures::StreamExt::next(&mut inner).await {
551                yield match item {
552                    Ok(StreamDelta::Done { stop_reason, .. }) => Ok(StreamDelta::Done {
553                        stop_reason,
554                        served_route: Some(served_route.clone()),
555                    }),
556                    other => other,
557                };
558            }
559        })
560    }
561
562    /// Issue the Vertex Gemini streaming request, returning the raw response on
563    /// success or a ready-to-`yield` stream item describing the failure.
564    ///
565    /// The `Err` payload is the exact `StreamBox` item to `yield`: an
566    /// classified [`StreamDelta::Error`] for a transport failure or a
567    /// non-success HTTP status. This keeps the
568    /// generator's failure handling to a single `yield`.
569    async fn send_gemini_stream_request(
570        &self,
571        url: &str,
572        api_request: &ApiGenerateContentRequest<'_>,
573    ) -> Result<reqwest::Response, anyhow::Result<StreamDelta>> {
574        let response = match self
575            .client
576            .post(url)
577            .header("Content-Type", "application/json")
578            .bearer_auth(&self.access_token)
579            .json(api_request)
580            .send()
581            .await
582        {
583            Ok(response) => response,
584            Err(error) => return Err(Ok(reqwest_error_delta("request failed", &error))),
585        };
586
587        let status = response.status();
588        if !status.is_success() {
589            // Headers are read before the body: `text()` consumes the response.
590            let header_hint = crate::http::retry_after_from_headers(response.headers());
591            let body = response.text().await.unwrap_or_default();
592            let kind = if status == StatusCode::TOO_MANY_REQUESTS {
593                StreamErrorKind::RateLimited(vertex_retry_delay(header_hint, &body))
594            } else if status.is_server_error() {
595                StreamErrorKind::ServerError
596            } else {
597                StreamErrorKind::InvalidRequest
598            };
599            log::warn!("Vertex AI error status={status} body={body}");
600            return Err(Ok(StreamDelta::Error {
601                message: body,
602                kind,
603            }));
604        }
605
606        Ok(response)
607    }
608}
609
610/// Resolve the retry delay a Vertex 429 carries: the `Retry-After` header
611/// first, then a `google.rpc.RetryInfo` detail in the error body.
612///
613/// Vertex quota rejections come from Google's front end for every model family,
614/// so a Claude 429 can carry the same `RetryInfo` detail a Gemini one does.
615fn vertex_retry_delay(
616    header_hint: Option<std::time::Duration>,
617    body: &str,
618) -> Option<std::time::Duration> {
619    header_hint.or_else(|| crate::retry_hints::google_retry_delay(body))
620}
621
622/// Classify a non-success status from the Vertex Claude endpoint into an early
623/// [`ChatOutcome`], or `None` when the body should be parsed as a response.
624fn classify_vertex_claude_status(
625    status: StatusCode,
626    bytes: &[u8],
627    header_hint: Option<std::time::Duration>,
628) -> Option<ChatOutcome> {
629    let body = String::from_utf8_lossy(bytes);
630    if status == StatusCode::TOO_MANY_REQUESTS {
631        return Some(ChatOutcome::RateLimited(vertex_retry_delay(
632            header_hint,
633            &body,
634        )));
635    }
636    if status.is_server_error() {
637        log::error!("Vertex AI (Claude) server error status={status} body={body}");
638        return Some(ChatOutcome::ServerError(body.into_owned()));
639    }
640    if status.is_client_error() {
641        log::warn!("Vertex AI (Claude) client error status={status} body={body}");
642        return Some(ChatOutcome::InvalidRequest(body.into_owned()));
643    }
644    None
645}
646
647/// Build the Gemini `system_instruction` content from the request system prompt,
648/// or `None` when the prompt is empty.
649fn build_gemini_system_instruction(system: &str) -> Option<ApiContent> {
650    if system.is_empty() {
651        None
652    } else {
653        Some(ApiContent {
654            role: None,
655            parts: vec![ApiPart::Text {
656                text: system.to_owned(),
657                thought_signature: None,
658            }],
659        })
660    }
661}
662
663/// Map an optional response format into Gemini's `(responseMimeType,
664/// responseSchema)` pair, sanitizing the schema to the subset Gemini accepts.
665fn gemini_response_format(
666    response_format: Option<&ResponseFormat>,
667) -> (Option<&'static str>, Option<serde_json::Value>) {
668    response_format.map_or((None, None), |rf| {
669        (
670            Some("application/json"),
671            Some(gemini_response_schema(&rf.schema)),
672        )
673    })
674}
675
676// ============================================================================
677// Claude path (publishers/anthropic)
678// ============================================================================
679
680impl VertexProvider {
681    async fn chat_claude(&self, request: ChatRequest) -> Result<ChatOutcome> {
682        let thinking_config = match self.resolve_thinking_config(request.thinking.as_ref()) {
683            // Forcing a specific tool is incompatible with extended thinking on
684            // Claude (the API 400s), so drop thinking at the wire boundary even
685            // when it was resurrected from the provider-configured default.
686            Ok(thinking) => thinking_for_forced_tool(thinking, request.tool_choice.as_ref()),
687            Err(error) => return Ok(ChatOutcome::InvalidRequest(error.to_string())),
688        };
689        if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
690            return Ok(ChatOutcome::InvalidRequest(error.to_string()));
691        }
692        let messages = Self::build_cached_vertex_claude_messages(&request);
693        let tools = build_vertex_claude_tools(&request);
694        let thinking = thinking_config
695            .as_ref()
696            .and_then(anthropic_data::ApiThinkingConfig::from_thinking_config);
697        let output_config = thinking_config
698            .as_ref()
699            .and_then(|t| t.effort)
700            .map(|effort| anthropic_data::ApiOutputConfig { effort });
701        let system = Self::build_vertex_claude_system_prompt(&request.system);
702        let tool_choice = request
703            .tool_choice
704            .as_ref()
705            .map(anthropic_data::ApiToolChoice::from_tool_choice);
706
707        let max_tokens = self.effective_max_tokens(&request);
708        let api_request = anthropic_data::ApiMessagesRequest {
709            model: None, // model is in the URL for Vertex
710            max_tokens,
711            system,
712            messages: &messages,
713            tools: tools.as_deref(),
714            tool_choice,
715            stream: false,
716            thinking,
717            output_config,
718            anthropic_version: Some(VERTEX_ANTHROPIC_VERSION),
719        };
720
721        log::debug!(
722            "Vertex AI (Claude) LLM request model={} max_tokens={}",
723            self.model,
724            max_tokens
725        );
726
727        if log::log_enabled!(log::Level::Debug) {
728            match serde_json::to_string_pretty(&api_request) {
729                Ok(json) => log::debug!("Vertex AI (Claude) request payload:\n{json}"),
730                Err(e) => log::debug!("Failed to serialize request for logging: {e}"),
731            }
732        }
733
734        let url = format!("{}:rawPredict", self.base_url("anthropic"));
735
736        let response = self
737            .client
738            .post(&url)
739            .header("Content-Type", "application/json")
740            .timeout(std::time::Duration::from_secs(CHAT_READ_TIMEOUT_SECS))
741            .bearer_auth(&self.access_token)
742            .json(&api_request)
743            .send()
744            .await
745            .map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
746
747        let status = response.status();
748        // Read `Retry-After` off the 429 response before the body is consumed.
749        let retry_after = if status == StatusCode::TOO_MANY_REQUESTS {
750            crate::http::retry_after_from_headers(response.headers())
751        } else {
752            None
753        };
754        let bytes = response
755            .bytes()
756            .await
757            .map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?;
758
759        log::debug!(
760            "Vertex AI (Claude) response status={} body_len={}",
761            status,
762            bytes.len()
763        );
764
765        if let Some(outcome) = classify_vertex_claude_status(status, &bytes, retry_after) {
766            return Ok(outcome);
767        }
768
769        let api_response: anthropic_data::ApiResponse = serde_json::from_slice(&bytes)
770            .map_err(|e| anyhow::anyhow!("failed to parse response: {e}"))?;
771
772        log::debug!(
773            "Vertex AI (Claude) response: id={} model={} stop_reason={:?} usage={{input_tokens={}, output_tokens={}}} content_blocks={}",
774            api_response.id,
775            api_response.model,
776            api_response.stop_reason,
777            api_response.usage.total_input_tokens(),
778            api_response.usage.output,
779            api_response.content.len()
780        );
781
782        Ok(ChatOutcome::Success(Self::map_claude_response(
783            api_response,
784        )))
785    }
786
787    #[allow(clippy::too_many_lines)]
788    fn chat_stream_claude(&self, request: ChatRequest) -> StreamBox<'_> {
789        let served_route = self.route().to_owned();
790        Box::pin(async_stream::stream! {
791            let thinking_config = match self.resolve_thinking_config(request.thinking.as_ref()) {
792                // Forcing a specific tool is incompatible with extended thinking
793                // on Claude (the API 400s), so drop thinking at the wire boundary
794                // even when it was resurrected from the provider-configured
795                // default.
796                Ok(thinking) => thinking_for_forced_tool(thinking, request.tool_choice.as_ref()),
797                Err(error) => {
798                    yield Ok(StreamDelta::Error {
799                        message: error.to_string(),
800                        kind: StreamErrorKind::InvalidRequest,
801                    });
802                    return;
803                }
804            };
805            if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
806                yield Ok(StreamDelta::Error {
807                    message: error.to_string(),
808                    kind: StreamErrorKind::InvalidRequest,
809                });
810                return;
811            }
812            let messages = Self::build_cached_vertex_claude_messages(&request);
813            let tools = build_vertex_claude_tools(&request);
814            let thinking = thinking_config
815                .as_ref()
816                .and_then(anthropic_data::ApiThinkingConfig::from_thinking_config);
817            let output_config = thinking_config
818                .as_ref()
819                .and_then(|t| t.effort)
820                .map(|effort| anthropic_data::ApiOutputConfig { effort });
821            let system = Self::build_vertex_claude_system_prompt(&request.system);
822            let tool_choice = request
823                .tool_choice
824                .as_ref()
825                .map(anthropic_data::ApiToolChoice::from_tool_choice);
826
827            let max_tokens = self.effective_max_tokens(&request);
828            let api_request = anthropic_data::ApiMessagesRequest {
829                model: None, // model is in the URL for Vertex
830                max_tokens,
831                system,
832                messages: &messages,
833                tools: tools.as_deref(),
834                tool_choice,
835                stream: true,
836                thinking,
837                output_config,
838                anthropic_version: Some(VERTEX_ANTHROPIC_VERSION),
839            };
840
841            log::debug!(
842                "Vertex AI (Claude) streaming request model={} max_tokens={}",
843                self.model,
844                max_tokens
845            );
846
847            if log::log_enabled!(log::Level::Debug) {
848                match serde_json::to_string_pretty(&api_request) {
849                    Ok(json) => log::debug!("Vertex AI (Claude) streaming request payload:\n{json}"),
850                    Err(e) => log::debug!("Failed to serialize request for logging: {e}"),
851                }
852            }
853
854            let url = format!("{}:streamRawPredict", self.base_url("anthropic"));
855
856            let response = match self
857                .client
858                .post(&url)
859                .header("Content-Type", "application/json")
860                .bearer_auth(&self.access_token)
861                .json(&api_request)
862                .send()
863                .await
864            {
865                Ok(r) => r,
866                Err(error) => {
867                    yield Ok(reqwest_error_delta("request failed", &error));
868                    return;
869                }
870            };
871
872            let status = response.status();
873
874            if status == StatusCode::TOO_MANY_REQUESTS {
875                // Headers are read before the body: `text()` consumes the response.
876                let header_hint = crate::http::retry_after_from_headers(response.headers());
877                let body = response.text().await.unwrap_or_default();
878                log::warn!("Vertex AI (Claude) rate limited status={status} body={body}");
879                yield Ok(StreamDelta::Error {
880                    message: "Rate limited".to_string(),
881                    kind: StreamErrorKind::RateLimited(vertex_retry_delay(header_hint, &body)),
882                });
883                return;
884            }
885
886            if status.is_server_error() {
887                let body = response.text().await.unwrap_or_default();
888                log::error!("Vertex AI (Claude) server error status={status} body={body}");
889                yield Ok(StreamDelta::Error {
890                    message: body,
891                    kind: StreamErrorKind::ServerError,
892                });
893                return;
894            }
895
896            if status.is_client_error() {
897                let body = response.text().await.unwrap_or_default();
898                log::warn!("Vertex AI (Claude) client error status={status} body={body}");
899                yield Ok(StreamDelta::Error {
900                    message: body,
901                    kind: StreamErrorKind::InvalidRequest,
902                });
903                return;
904            }
905
906            // Process SSE stream using the Anthropic SSE parser
907            let mut stream = response.bytes_stream();
908            let mut buffer = String::new();
909            let mut input_tokens: u32 = 0;
910            let mut output_tokens: u32 = 0;
911            let mut cached_input_tokens: u32 = 0;
912            let mut cache_creation_input_tokens: u32 = 0;
913            let mut tool_ids: std::collections::HashMap<usize, String> =
914                std::collections::HashMap::new();
915            let mut received_message_stop = false;
916            let mut pending_stop_reason: Option<agent_sdk_foundation::llm::StopReason> = None;
917
918            while let Some(chunk_result) = stream.next().await {
919                let chunk = match chunk_result {
920                    Ok(c) => c,
921                    Err(error) => {
922                        yield Ok(reqwest_body_error_delta("stream error", &error));
923                        return;
924                    }
925                };
926
927                buffer.push_str(&String::from_utf8_lossy(&chunk));
928
929                // Process complete SSE events (terminated by a blank line)
930                while let Some(event_block) = anthropic_data::take_next_sse_event(&mut buffer) {
931                    if anthropic_data::is_message_stop_event(&event_block) {
932                        received_message_stop = true;
933                    }
934
935                    if let Some(delta) = anthropic_data::parse_sse_event(
936                        &event_block,
937                        &mut input_tokens,
938                        &mut output_tokens,
939                        &mut cached_input_tokens,
940                        &mut cache_creation_input_tokens,
941                        &mut tool_ids,
942                        &mut pending_stop_reason,
943                    ) {
944                        yield Ok(delta);
945                    }
946                    if anthropic_data::is_message_stop_event(&event_block) {
947                        yield Ok(StreamDelta::Done {
948                            stop_reason: pending_stop_reason.take(),
949                            served_route: Some(served_route.clone()),
950                        });
951                    }
952                }
953            }
954
955            // Process remaining buffer
956            let remaining = buffer.trim();
957            if !remaining.is_empty() {
958                if anthropic_data::is_message_stop_event(remaining) {
959                    received_message_stop = true;
960                }
961
962                if let Some(delta) = anthropic_data::parse_sse_event(
963                    remaining,
964                    &mut input_tokens,
965                    &mut output_tokens,
966                    &mut cached_input_tokens,
967                    &mut cache_creation_input_tokens,
968                    &mut tool_ids,
969                    &mut pending_stop_reason,
970                ) {
971                    yield Ok(delta);
972                }
973                if anthropic_data::is_message_stop_event(remaining) {
974                    yield Ok(StreamDelta::Done {
975                        stop_reason: pending_stop_reason.take(),
976                        served_route: Some(served_route.clone()),
977                    });
978                }
979            }
980
981            if !received_message_stop {
982                log::warn!(
983                    "Vertex AI (Claude) SSE stream ended without message_stop"
984                );
985                yield Ok(StreamDelta::Error {
986                    message: "Stream ended unexpectedly without completion".to_string(),
987                    kind: StreamErrorKind::ServerError,
988                });
989            }
990        })
991    }
992}
993
994#[cfg(test)]
995mod tests {
996    use super::*;
997
998    #[test]
999    fn test_new_creates_provider() {
1000        let provider = VertexProvider::new(
1001            "token".to_string(),
1002            "my-project".to_string(),
1003            "us-central1".to_string(),
1004            "custom-model".to_string(),
1005        );
1006
1007        assert_eq!(provider.model(), "custom-model");
1008        assert_eq!(provider.provider(), "vertex");
1009    }
1010
1011    #[test]
1012    fn test_flash_factory() {
1013        let provider = VertexProvider::flash(
1014            "token".to_string(),
1015            "my-project".to_string(),
1016            "us-central1".to_string(),
1017        );
1018
1019        assert_eq!(provider.model(), MODEL_GEMINI_3_FLASH);
1020        assert_eq!(provider.provider(), "vertex");
1021    }
1022
1023    #[test]
1024    fn test_pro_factory() {
1025        let provider = VertexProvider::pro(
1026            "token".to_string(),
1027            "my-project".to_string(),
1028            "us-central1".to_string(),
1029        );
1030
1031        assert_eq!(provider.model(), MODEL_GEMINI_31_PRO);
1032        assert_eq!(provider.provider(), "vertex");
1033    }
1034
1035    #[test]
1036    fn test_provider_is_cloneable() {
1037        let provider = VertexProvider::new(
1038            "token".to_string(),
1039            "my-project".to_string(),
1040            "us-central1".to_string(),
1041            "test-model".to_string(),
1042        );
1043        let cloned = provider.clone();
1044
1045        assert_eq!(provider.model(), cloned.model());
1046        assert_eq!(provider.provider(), cloned.provider());
1047    }
1048
1049    #[test]
1050    fn test_is_claude_model() {
1051        let claude_provider = VertexProvider::new(
1052            "token".to_string(),
1053            "project".to_string(),
1054            "us-central1".to_string(),
1055            "claude-sonnet-4-20250514".to_string(),
1056        );
1057        assert!(claude_provider.is_claude_model());
1058
1059        let gemini_provider = VertexProvider::new(
1060            "token".to_string(),
1061            "project".to_string(),
1062            "us-central1".to_string(),
1063            "gemini-3-flash-preview".to_string(),
1064        );
1065        assert!(!gemini_provider.is_claude_model());
1066    }
1067
1068    #[test]
1069    fn test_base_url_gemini() {
1070        let provider = VertexProvider::new(
1071            "token".to_string(),
1072            "my-project".to_string(),
1073            "us-central1".to_string(),
1074            "gemini-3-flash-preview".to_string(),
1075        );
1076
1077        let url = provider.base_url("google");
1078        assert_eq!(
1079            url,
1080            "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview"
1081        );
1082    }
1083
1084    #[test]
1085    fn test_base_url_claude() {
1086        let provider = VertexProvider::new(
1087            "token".to_string(),
1088            "my-project".to_string(),
1089            "us-central1".to_string(),
1090            "claude-sonnet-4-20250514".to_string(),
1091        );
1092
1093        let url = provider.base_url("anthropic");
1094        assert_eq!(
1095            url,
1096            "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/anthropic/models/claude-sonnet-4-20250514"
1097        );
1098    }
1099
1100    #[test]
1101    fn test_base_url_with_different_region() {
1102        let provider = VertexProvider::new(
1103            "token".to_string(),
1104            "other-project".to_string(),
1105            "europe-west4".to_string(),
1106            "gemini-3.1-pro-preview".to_string(),
1107        );
1108
1109        let url = provider.base_url("google");
1110        assert!(url.starts_with("https://europe-west4-aiplatform.googleapis.com/"));
1111        assert!(url.contains("/projects/other-project/"));
1112        assert!(url.contains("/locations/europe-west4/"));
1113        assert!(url.ends_with("/models/gemini-3.1-pro-preview"));
1114    }
1115
1116    #[test]
1117    fn test_base_url_global_region_has_no_prefix() {
1118        let provider = VertexProvider::new(
1119            "token".to_string(),
1120            "my-project".to_string(),
1121            "global".to_string(),
1122            "gemini-3.1-pro-preview".to_string(),
1123        );
1124
1125        let url = provider.base_url("google");
1126        assert_eq!(
1127            url,
1128            "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-3.1-pro-preview"
1129        );
1130    }
1131
1132    #[test]
1133    fn test_vertex_claude_46_rejects_budgeted_thinking() {
1134        let provider = VertexProvider::new(
1135            "token".to_string(),
1136            "project".to_string(),
1137            "global".to_string(),
1138            MODEL_SONNET_46.to_string(),
1139        );
1140
1141        let error = provider
1142            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1143            .unwrap_err();
1144        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1145    }
1146
1147    #[test]
1148    fn test_vertex_claude_opus_47_rejects_budgeted_thinking() {
1149        let provider = VertexProvider::new(
1150            "token".to_string(),
1151            "project".to_string(),
1152            "global".to_string(),
1153            MODEL_OPUS_47.to_string(),
1154        );
1155
1156        let error = provider
1157            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1158            .unwrap_err();
1159        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1160    }
1161
1162    #[test]
1163    fn test_vertex_claude_opus_48_rejects_budgeted_thinking() {
1164        let provider = VertexProvider::new(
1165            "token".to_string(),
1166            "project".to_string(),
1167            "global".to_string(),
1168            MODEL_OPUS_48.to_string(),
1169        );
1170
1171        let error = provider
1172            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1173            .unwrap_err();
1174        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1175    }
1176
1177    #[test]
1178    fn test_vertex_claude_fable_5_rejects_budgeted_thinking() {
1179        let provider = VertexProvider::new(
1180            "token".to_string(),
1181            "project".to_string(),
1182            "global".to_string(),
1183            MODEL_FABLE_5.to_string(),
1184        );
1185
1186        let error = provider
1187            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1188            .unwrap_err();
1189        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1190    }
1191
1192    #[test]
1193    fn test_model_constants() {
1194        assert_eq!(MODEL_GEMINI_3_FLASH, "gemini-3-flash-preview");
1195        assert_eq!(MODEL_GEMINI_31_PRO, "gemini-3.1-pro-preview");
1196        assert_eq!(MODEL_GEMINI_3_PRO, "gemini-3.0-pro");
1197    }
1198
1199    fn request_with_max_tokens(max_tokens: u32, explicit: bool) -> ChatRequest {
1200        ChatRequest {
1201            system: String::new(),
1202            messages: vec![agent_sdk_foundation::llm::Message::user("hi")],
1203            tools: None,
1204            max_tokens,
1205            max_tokens_explicit: explicit,
1206            session_id: None,
1207            cached_content: None,
1208            thinking: None,
1209            tool_choice: None,
1210            response_format: None,
1211            cache: None,
1212        }
1213    }
1214
1215    #[test]
1216    fn test_effective_max_tokens_honors_explicit_budget() {
1217        let provider = VertexProvider::new(
1218            "token".to_string(),
1219            "project".to_string(),
1220            "global".to_string(),
1221            MODEL_SONNET_46.to_string(),
1222        );
1223        let request = request_with_max_tokens(1234, true);
1224        assert_eq!(provider.effective_max_tokens(&request), 1234);
1225    }
1226
1227    #[test]
1228    fn test_effective_max_tokens_uses_clamped_default_when_implicit() {
1229        // An implicit budget for a Claude-on-Vertex model must fall back to the
1230        // clamped default (<= 32k), not the unclamped capability ceiling and
1231        // not ChatRequest::DEFAULT_MAX_TOKENS.
1232        let provider = VertexProvider::new(
1233            "token".to_string(),
1234            "project".to_string(),
1235            "global".to_string(),
1236            MODEL_SONNET_46.to_string(),
1237        );
1238        let request = request_with_max_tokens(4096, false);
1239        let effective = provider.effective_max_tokens(&request);
1240        assert_eq!(effective, provider.default_max_tokens());
1241        assert!(effective <= DEFAULT_SAFE_MAX_OUTPUT_TOKENS);
1242    }
1243}