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 requires_anthropic_adaptive_thinking(&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.requires_anthropic_adaptive_thinking()
276            && matches!(thinking.mode, ThinkingMode::Enabled { .. })
277        {
278            return Err(anyhow::anyhow!(
279                "budget_tokens thinking is deprecated for provider={} model={}; use ThinkingConfig::adaptive() 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        Box::pin(async_stream::stream! {
485            let thinking = match self.resolve_thinking_config(request.thinking.as_ref()) {
486                Ok(thinking) => thinking,
487                Err(error) => {
488                    yield Ok(StreamDelta::Error {
489                        message: error.to_string(),
490                        kind: StreamErrorKind::InvalidRequest,
491                    });
492                    return;
493                }
494            };
495            if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
496                yield Ok(StreamDelta::Error {
497                    message: error.to_string(),
498                    kind: StreamErrorKind::InvalidRequest,
499                });
500                return;
501            }
502
503            let contents = build_api_contents(&request.messages);
504            let tools = request
505                .tools
506                .as_ref()
507                .map(|t| convert_tools_to_config(t.clone()));
508            let tool_config = request
509                .tool_choice
510                .as_ref()
511                .map(ApiFunctionCallingConfig::from_tool_choice);
512            let system_instruction = build_gemini_system_instruction(&request.system);
513            let thinking_config = thinking.as_ref().map(map_thinking_config);
514            let (response_mime_type, response_schema) =
515                gemini_response_format(request.response_format.as_ref());
516
517            let max_tokens = self.effective_max_tokens(&request);
518            let api_request = ApiGenerateContentRequest {
519                contents: &contents,
520                system_instruction: system_instruction.as_ref(),
521                tools: tools.as_ref().map(std::slice::from_ref),
522                tool_config,
523                generation_config: Some(ApiGenerationConfig {
524                    max_output_tokens: Some(max_tokens),
525                    thinking_config,
526                    response_mime_type,
527                    response_schema,
528                }),
529                cached_content: request.cached_content.as_deref(),
530            };
531
532            log::debug!(
533                "Vertex AI streaming LLM request model={} max_tokens={}",
534                self.model,
535                max_tokens
536            );
537
538            let url = format!("{}:streamGenerateContent?alt=sse", self.base_url("google"));
539
540            let response = match self.send_gemini_stream_request(&url, &api_request).await {
541                Ok(response) => response,
542                Err(item) => {
543                    yield item;
544                    return;
545                }
546            };
547
548            let mut inner = stream_gemini_response(response);
549            while let Some(item) = futures::StreamExt::next(&mut inner).await {
550                yield item;
551            }
552        })
553    }
554
555    /// Issue the Vertex Gemini streaming request, returning the raw response on
556    /// success or a ready-to-`yield` stream item describing the failure.
557    ///
558    /// The `Err` payload is the exact `StreamBox` item to `yield`: an
559    /// classified [`StreamDelta::Error`] for a transport failure or a
560    /// non-success HTTP status. This keeps the
561    /// generator's failure handling to a single `yield`.
562    async fn send_gemini_stream_request(
563        &self,
564        url: &str,
565        api_request: &ApiGenerateContentRequest<'_>,
566    ) -> Result<reqwest::Response, anyhow::Result<StreamDelta>> {
567        let response = match self
568            .client
569            .post(url)
570            .header("Content-Type", "application/json")
571            .bearer_auth(&self.access_token)
572            .json(api_request)
573            .send()
574            .await
575        {
576            Ok(response) => response,
577            Err(error) => return Err(Ok(reqwest_error_delta("request failed", &error))),
578        };
579
580        let status = response.status();
581        if !status.is_success() {
582            // Headers are read before the body: `text()` consumes the response.
583            let header_hint = crate::http::retry_after_from_headers(response.headers());
584            let body = response.text().await.unwrap_or_default();
585            let kind = if status == StatusCode::TOO_MANY_REQUESTS {
586                StreamErrorKind::RateLimited(vertex_retry_delay(header_hint, &body))
587            } else if status.is_server_error() {
588                StreamErrorKind::ServerError
589            } else {
590                StreamErrorKind::InvalidRequest
591            };
592            log::warn!("Vertex AI error status={status} body={body}");
593            return Err(Ok(StreamDelta::Error {
594                message: body,
595                kind,
596            }));
597        }
598
599        Ok(response)
600    }
601}
602
603/// Resolve the retry delay a Vertex 429 carries: the `Retry-After` header
604/// first, then a `google.rpc.RetryInfo` detail in the error body.
605///
606/// Vertex quota rejections come from Google's front end for every model family,
607/// so a Claude 429 can carry the same `RetryInfo` detail a Gemini one does.
608fn vertex_retry_delay(
609    header_hint: Option<std::time::Duration>,
610    body: &str,
611) -> Option<std::time::Duration> {
612    header_hint.or_else(|| crate::retry_hints::google_retry_delay(body))
613}
614
615/// Classify a non-success status from the Vertex Claude endpoint into an early
616/// [`ChatOutcome`], or `None` when the body should be parsed as a response.
617fn classify_vertex_claude_status(
618    status: StatusCode,
619    bytes: &[u8],
620    header_hint: Option<std::time::Duration>,
621) -> Option<ChatOutcome> {
622    let body = String::from_utf8_lossy(bytes);
623    if status == StatusCode::TOO_MANY_REQUESTS {
624        return Some(ChatOutcome::RateLimited(vertex_retry_delay(
625            header_hint,
626            &body,
627        )));
628    }
629    if status.is_server_error() {
630        log::error!("Vertex AI (Claude) server error status={status} body={body}");
631        return Some(ChatOutcome::ServerError(body.into_owned()));
632    }
633    if status.is_client_error() {
634        log::warn!("Vertex AI (Claude) client error status={status} body={body}");
635        return Some(ChatOutcome::InvalidRequest(body.into_owned()));
636    }
637    None
638}
639
640/// Build the Gemini `system_instruction` content from the request system prompt,
641/// or `None` when the prompt is empty.
642fn build_gemini_system_instruction(system: &str) -> Option<ApiContent> {
643    if system.is_empty() {
644        None
645    } else {
646        Some(ApiContent {
647            role: None,
648            parts: vec![ApiPart::Text {
649                text: system.to_owned(),
650                thought_signature: None,
651            }],
652        })
653    }
654}
655
656/// Map an optional response format into Gemini's `(responseMimeType,
657/// responseSchema)` pair, sanitizing the schema to the subset Gemini accepts.
658fn gemini_response_format(
659    response_format: Option<&ResponseFormat>,
660) -> (Option<&'static str>, Option<serde_json::Value>) {
661    response_format.map_or((None, None), |rf| {
662        (
663            Some("application/json"),
664            Some(gemini_response_schema(&rf.schema)),
665        )
666    })
667}
668
669// ============================================================================
670// Claude path (publishers/anthropic)
671// ============================================================================
672
673impl VertexProvider {
674    async fn chat_claude(&self, request: ChatRequest) -> Result<ChatOutcome> {
675        let thinking_config = match self.resolve_thinking_config(request.thinking.as_ref()) {
676            // Forcing a specific tool is incompatible with extended thinking on
677            // Claude (the API 400s), so drop thinking at the wire boundary even
678            // when it was resurrected from the provider-configured default.
679            Ok(thinking) => thinking_for_forced_tool(thinking, request.tool_choice.as_ref()),
680            Err(error) => return Ok(ChatOutcome::InvalidRequest(error.to_string())),
681        };
682        if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
683            return Ok(ChatOutcome::InvalidRequest(error.to_string()));
684        }
685        let messages = Self::build_cached_vertex_claude_messages(&request);
686        let tools = build_vertex_claude_tools(&request);
687        let thinking = thinking_config
688            .as_ref()
689            .map(anthropic_data::ApiThinkingConfig::from_thinking_config);
690        let output_config = thinking_config
691            .as_ref()
692            .and_then(|t| t.effort)
693            .map(|effort| anthropic_data::ApiOutputConfig { effort });
694        let system = Self::build_vertex_claude_system_prompt(&request.system);
695        let tool_choice = request
696            .tool_choice
697            .as_ref()
698            .map(anthropic_data::ApiToolChoice::from_tool_choice);
699
700        let max_tokens = self.effective_max_tokens(&request);
701        let api_request = anthropic_data::ApiMessagesRequest {
702            model: None, // model is in the URL for Vertex
703            max_tokens,
704            system,
705            messages: &messages,
706            tools: tools.as_deref(),
707            tool_choice,
708            stream: false,
709            thinking,
710            output_config,
711            anthropic_version: Some(VERTEX_ANTHROPIC_VERSION),
712        };
713
714        log::debug!(
715            "Vertex AI (Claude) LLM request model={} max_tokens={}",
716            self.model,
717            max_tokens
718        );
719
720        if log::log_enabled!(log::Level::Debug) {
721            match serde_json::to_string_pretty(&api_request) {
722                Ok(json) => log::debug!("Vertex AI (Claude) request payload:\n{json}"),
723                Err(e) => log::debug!("Failed to serialize request for logging: {e}"),
724            }
725        }
726
727        let url = format!("{}:rawPredict", self.base_url("anthropic"));
728
729        let response = self
730            .client
731            .post(&url)
732            .header("Content-Type", "application/json")
733            .timeout(std::time::Duration::from_secs(CHAT_READ_TIMEOUT_SECS))
734            .bearer_auth(&self.access_token)
735            .json(&api_request)
736            .send()
737            .await
738            .map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
739
740        let status = response.status();
741        // Read `Retry-After` off the 429 response before the body is consumed.
742        let retry_after = if status == StatusCode::TOO_MANY_REQUESTS {
743            crate::http::retry_after_from_headers(response.headers())
744        } else {
745            None
746        };
747        let bytes = response
748            .bytes()
749            .await
750            .map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?;
751
752        log::debug!(
753            "Vertex AI (Claude) response status={} body_len={}",
754            status,
755            bytes.len()
756        );
757
758        if let Some(outcome) = classify_vertex_claude_status(status, &bytes, retry_after) {
759            return Ok(outcome);
760        }
761
762        let api_response: anthropic_data::ApiResponse = serde_json::from_slice(&bytes)
763            .map_err(|e| anyhow::anyhow!("failed to parse response: {e}"))?;
764
765        log::debug!(
766            "Vertex AI (Claude) response: id={} model={} stop_reason={:?} usage={{input_tokens={}, output_tokens={}}} content_blocks={}",
767            api_response.id,
768            api_response.model,
769            api_response.stop_reason,
770            api_response.usage.total_input_tokens(),
771            api_response.usage.output,
772            api_response.content.len()
773        );
774
775        Ok(ChatOutcome::Success(Self::map_claude_response(
776            api_response,
777        )))
778    }
779
780    #[allow(clippy::too_many_lines)]
781    fn chat_stream_claude(&self, request: ChatRequest) -> StreamBox<'_> {
782        Box::pin(async_stream::stream! {
783            let thinking_config = match self.resolve_thinking_config(request.thinking.as_ref()) {
784                // Forcing a specific tool is incompatible with extended thinking
785                // on Claude (the API 400s), so drop thinking at the wire boundary
786                // even when it was resurrected from the provider-configured
787                // default.
788                Ok(thinking) => thinking_for_forced_tool(thinking, request.tool_choice.as_ref()),
789                Err(error) => {
790                    yield Ok(StreamDelta::Error {
791                        message: error.to_string(),
792                        kind: StreamErrorKind::InvalidRequest,
793                    });
794                    return;
795                }
796            };
797            if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
798                yield Ok(StreamDelta::Error {
799                    message: error.to_string(),
800                    kind: StreamErrorKind::InvalidRequest,
801                });
802                return;
803            }
804            let messages = Self::build_cached_vertex_claude_messages(&request);
805            let tools = build_vertex_claude_tools(&request);
806            let thinking = thinking_config
807                .as_ref()
808                .map(anthropic_data::ApiThinkingConfig::from_thinking_config);
809            let output_config = thinking_config
810                .as_ref()
811                .and_then(|t| t.effort)
812                .map(|effort| anthropic_data::ApiOutputConfig { effort });
813            let system = Self::build_vertex_claude_system_prompt(&request.system);
814            let tool_choice = request
815                .tool_choice
816                .as_ref()
817                .map(anthropic_data::ApiToolChoice::from_tool_choice);
818
819            let max_tokens = self.effective_max_tokens(&request);
820            let api_request = anthropic_data::ApiMessagesRequest {
821                model: None, // model is in the URL for Vertex
822                max_tokens,
823                system,
824                messages: &messages,
825                tools: tools.as_deref(),
826                tool_choice,
827                stream: true,
828                thinking,
829                output_config,
830                anthropic_version: Some(VERTEX_ANTHROPIC_VERSION),
831            };
832
833            log::debug!(
834                "Vertex AI (Claude) streaming request model={} max_tokens={}",
835                self.model,
836                max_tokens
837            );
838
839            if log::log_enabled!(log::Level::Debug) {
840                match serde_json::to_string_pretty(&api_request) {
841                    Ok(json) => log::debug!("Vertex AI (Claude) streaming request payload:\n{json}"),
842                    Err(e) => log::debug!("Failed to serialize request for logging: {e}"),
843                }
844            }
845
846            let url = format!("{}:streamRawPredict", self.base_url("anthropic"));
847
848            let response = match self
849                .client
850                .post(&url)
851                .header("Content-Type", "application/json")
852                .bearer_auth(&self.access_token)
853                .json(&api_request)
854                .send()
855                .await
856            {
857                Ok(r) => r,
858                Err(error) => {
859                    yield Ok(reqwest_error_delta("request failed", &error));
860                    return;
861                }
862            };
863
864            let status = response.status();
865
866            if status == StatusCode::TOO_MANY_REQUESTS {
867                // Headers are read before the body: `text()` consumes the response.
868                let header_hint = crate::http::retry_after_from_headers(response.headers());
869                let body = response.text().await.unwrap_or_default();
870                log::warn!("Vertex AI (Claude) rate limited status={status} body={body}");
871                yield Ok(StreamDelta::Error {
872                    message: "Rate limited".to_string(),
873                    kind: StreamErrorKind::RateLimited(vertex_retry_delay(header_hint, &body)),
874                });
875                return;
876            }
877
878            if status.is_server_error() {
879                let body = response.text().await.unwrap_or_default();
880                log::error!("Vertex AI (Claude) server error status={status} body={body}");
881                yield Ok(StreamDelta::Error {
882                    message: body,
883                    kind: StreamErrorKind::ServerError,
884                });
885                return;
886            }
887
888            if status.is_client_error() {
889                let body = response.text().await.unwrap_or_default();
890                log::warn!("Vertex AI (Claude) client error status={status} body={body}");
891                yield Ok(StreamDelta::Error {
892                    message: body,
893                    kind: StreamErrorKind::InvalidRequest,
894                });
895                return;
896            }
897
898            // Process SSE stream using the Anthropic SSE parser
899            let mut stream = response.bytes_stream();
900            let mut buffer = String::new();
901            let mut input_tokens: u32 = 0;
902            let mut output_tokens: u32 = 0;
903            let mut cached_input_tokens: u32 = 0;
904            let mut cache_creation_input_tokens: u32 = 0;
905            let mut tool_ids: std::collections::HashMap<usize, String> =
906                std::collections::HashMap::new();
907            let mut received_message_stop = false;
908            let mut pending_stop_reason: Option<agent_sdk_foundation::llm::StopReason> = None;
909
910            while let Some(chunk_result) = stream.next().await {
911                let chunk = match chunk_result {
912                    Ok(c) => c,
913                    Err(error) => {
914                        yield Ok(reqwest_body_error_delta("stream error", &error));
915                        return;
916                    }
917                };
918
919                buffer.push_str(&String::from_utf8_lossy(&chunk));
920
921                // Process complete SSE events (terminated by a blank line)
922                while let Some(event_block) = anthropic_data::take_next_sse_event(&mut buffer) {
923                    if anthropic_data::is_message_stop_event(&event_block) {
924                        received_message_stop = true;
925                    }
926
927                    if let Some(delta) = anthropic_data::parse_sse_event(
928                        &event_block,
929                        &mut input_tokens,
930                        &mut output_tokens,
931                        &mut cached_input_tokens,
932                        &mut cache_creation_input_tokens,
933                        &mut tool_ids,
934                        &mut pending_stop_reason,
935                    ) {
936                        yield Ok(delta);
937                    }
938                    if anthropic_data::is_message_stop_event(&event_block) {
939                        yield Ok(StreamDelta::Done {
940                            stop_reason: pending_stop_reason.take(),
941                        });
942                    }
943                }
944            }
945
946            // Process remaining buffer
947            let remaining = buffer.trim();
948            if !remaining.is_empty() {
949                if anthropic_data::is_message_stop_event(remaining) {
950                    received_message_stop = true;
951                }
952
953                if let Some(delta) = anthropic_data::parse_sse_event(
954                    remaining,
955                    &mut input_tokens,
956                    &mut output_tokens,
957                    &mut cached_input_tokens,
958                    &mut cache_creation_input_tokens,
959                    &mut tool_ids,
960                    &mut pending_stop_reason,
961                ) {
962                    yield Ok(delta);
963                }
964                if anthropic_data::is_message_stop_event(remaining) {
965                    yield Ok(StreamDelta::Done {
966                        stop_reason: pending_stop_reason.take(),
967                    });
968                }
969            }
970
971            if !received_message_stop {
972                log::warn!(
973                    "Vertex AI (Claude) SSE stream ended without message_stop"
974                );
975                yield Ok(StreamDelta::Error {
976                    message: "Stream ended unexpectedly without completion".to_string(),
977                    kind: StreamErrorKind::ServerError,
978                });
979            }
980        })
981    }
982}
983
984#[cfg(test)]
985mod tests {
986    use super::*;
987
988    #[test]
989    fn test_new_creates_provider() {
990        let provider = VertexProvider::new(
991            "token".to_string(),
992            "my-project".to_string(),
993            "us-central1".to_string(),
994            "custom-model".to_string(),
995        );
996
997        assert_eq!(provider.model(), "custom-model");
998        assert_eq!(provider.provider(), "vertex");
999    }
1000
1001    #[test]
1002    fn test_flash_factory() {
1003        let provider = VertexProvider::flash(
1004            "token".to_string(),
1005            "my-project".to_string(),
1006            "us-central1".to_string(),
1007        );
1008
1009        assert_eq!(provider.model(), MODEL_GEMINI_3_FLASH);
1010        assert_eq!(provider.provider(), "vertex");
1011    }
1012
1013    #[test]
1014    fn test_pro_factory() {
1015        let provider = VertexProvider::pro(
1016            "token".to_string(),
1017            "my-project".to_string(),
1018            "us-central1".to_string(),
1019        );
1020
1021        assert_eq!(provider.model(), MODEL_GEMINI_31_PRO);
1022        assert_eq!(provider.provider(), "vertex");
1023    }
1024
1025    #[test]
1026    fn test_provider_is_cloneable() {
1027        let provider = VertexProvider::new(
1028            "token".to_string(),
1029            "my-project".to_string(),
1030            "us-central1".to_string(),
1031            "test-model".to_string(),
1032        );
1033        let cloned = provider.clone();
1034
1035        assert_eq!(provider.model(), cloned.model());
1036        assert_eq!(provider.provider(), cloned.provider());
1037    }
1038
1039    #[test]
1040    fn test_is_claude_model() {
1041        let claude_provider = VertexProvider::new(
1042            "token".to_string(),
1043            "project".to_string(),
1044            "us-central1".to_string(),
1045            "claude-sonnet-4-20250514".to_string(),
1046        );
1047        assert!(claude_provider.is_claude_model());
1048
1049        let gemini_provider = VertexProvider::new(
1050            "token".to_string(),
1051            "project".to_string(),
1052            "us-central1".to_string(),
1053            "gemini-3-flash-preview".to_string(),
1054        );
1055        assert!(!gemini_provider.is_claude_model());
1056    }
1057
1058    #[test]
1059    fn test_base_url_gemini() {
1060        let provider = VertexProvider::new(
1061            "token".to_string(),
1062            "my-project".to_string(),
1063            "us-central1".to_string(),
1064            "gemini-3-flash-preview".to_string(),
1065        );
1066
1067        let url = provider.base_url("google");
1068        assert_eq!(
1069            url,
1070            "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview"
1071        );
1072    }
1073
1074    #[test]
1075    fn test_base_url_claude() {
1076        let provider = VertexProvider::new(
1077            "token".to_string(),
1078            "my-project".to_string(),
1079            "us-central1".to_string(),
1080            "claude-sonnet-4-20250514".to_string(),
1081        );
1082
1083        let url = provider.base_url("anthropic");
1084        assert_eq!(
1085            url,
1086            "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/anthropic/models/claude-sonnet-4-20250514"
1087        );
1088    }
1089
1090    #[test]
1091    fn test_base_url_with_different_region() {
1092        let provider = VertexProvider::new(
1093            "token".to_string(),
1094            "other-project".to_string(),
1095            "europe-west4".to_string(),
1096            "gemini-3.1-pro-preview".to_string(),
1097        );
1098
1099        let url = provider.base_url("google");
1100        assert!(url.starts_with("https://europe-west4-aiplatform.googleapis.com/"));
1101        assert!(url.contains("/projects/other-project/"));
1102        assert!(url.contains("/locations/europe-west4/"));
1103        assert!(url.ends_with("/models/gemini-3.1-pro-preview"));
1104    }
1105
1106    #[test]
1107    fn test_base_url_global_region_has_no_prefix() {
1108        let provider = VertexProvider::new(
1109            "token".to_string(),
1110            "my-project".to_string(),
1111            "global".to_string(),
1112            "gemini-3.1-pro-preview".to_string(),
1113        );
1114
1115        let url = provider.base_url("google");
1116        assert_eq!(
1117            url,
1118            "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-3.1-pro-preview"
1119        );
1120    }
1121
1122    #[test]
1123    fn test_vertex_claude_46_rejects_budgeted_thinking() {
1124        let provider = VertexProvider::new(
1125            "token".to_string(),
1126            "project".to_string(),
1127            "global".to_string(),
1128            MODEL_SONNET_46.to_string(),
1129        );
1130
1131        let error = provider
1132            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1133            .unwrap_err();
1134        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1135    }
1136
1137    #[test]
1138    fn test_vertex_claude_opus_47_rejects_budgeted_thinking() {
1139        let provider = VertexProvider::new(
1140            "token".to_string(),
1141            "project".to_string(),
1142            "global".to_string(),
1143            MODEL_OPUS_47.to_string(),
1144        );
1145
1146        let error = provider
1147            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1148            .unwrap_err();
1149        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1150    }
1151
1152    #[test]
1153    fn test_vertex_claude_opus_48_rejects_budgeted_thinking() {
1154        let provider = VertexProvider::new(
1155            "token".to_string(),
1156            "project".to_string(),
1157            "global".to_string(),
1158            MODEL_OPUS_48.to_string(),
1159        );
1160
1161        let error = provider
1162            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1163            .unwrap_err();
1164        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1165    }
1166
1167    #[test]
1168    fn test_vertex_claude_fable_5_rejects_budgeted_thinking() {
1169        let provider = VertexProvider::new(
1170            "token".to_string(),
1171            "project".to_string(),
1172            "global".to_string(),
1173            MODEL_FABLE_5.to_string(),
1174        );
1175
1176        let error = provider
1177            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1178            .unwrap_err();
1179        assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1180    }
1181
1182    #[test]
1183    fn test_model_constants() {
1184        assert_eq!(MODEL_GEMINI_3_FLASH, "gemini-3-flash-preview");
1185        assert_eq!(MODEL_GEMINI_31_PRO, "gemini-3.1-pro-preview");
1186        assert_eq!(MODEL_GEMINI_3_PRO, "gemini-3.0-pro");
1187    }
1188
1189    fn request_with_max_tokens(max_tokens: u32, explicit: bool) -> ChatRequest {
1190        ChatRequest {
1191            system: String::new(),
1192            messages: vec![agent_sdk_foundation::llm::Message::user("hi")],
1193            tools: None,
1194            max_tokens,
1195            max_tokens_explicit: explicit,
1196            session_id: None,
1197            cached_content: None,
1198            thinking: None,
1199            tool_choice: None,
1200            response_format: None,
1201            cache: None,
1202        }
1203    }
1204
1205    #[test]
1206    fn test_effective_max_tokens_honors_explicit_budget() {
1207        let provider = VertexProvider::new(
1208            "token".to_string(),
1209            "project".to_string(),
1210            "global".to_string(),
1211            MODEL_SONNET_46.to_string(),
1212        );
1213        let request = request_with_max_tokens(1234, true);
1214        assert_eq!(provider.effective_max_tokens(&request), 1234);
1215    }
1216
1217    #[test]
1218    fn test_effective_max_tokens_uses_clamped_default_when_implicit() {
1219        // An implicit budget for a Claude-on-Vertex model must fall back to the
1220        // clamped default (<= 32k), not the unclamped capability ceiling and
1221        // not ChatRequest::DEFAULT_MAX_TOKENS.
1222        let provider = VertexProvider::new(
1223            "token".to_string(),
1224            "project".to_string(),
1225            "global".to_string(),
1226            MODEL_SONNET_46.to_string(),
1227        );
1228        let request = request_with_max_tokens(4096, false);
1229        let effective = provider.effective_max_tokens(&request);
1230        assert_eq!(effective, provider.default_max_tokens());
1231        assert!(effective <= DEFAULT_SAFE_MAX_OUTPUT_TOKENS);
1232    }
1233}