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