1use 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
37pub const MODEL_GEMINI_3_PRO: &str = "gemini-3.0-pro";
39
40const VERTEX_ANTHROPIC_VERSION: &str = "vertex-2023-10-16";
42const DEFAULT_SAFE_MAX_OUTPUT_TOKENS: u32 = 32_000;
43
44const CONNECT_TIMEOUT_SECS: u64 = 30;
46const TCP_KEEPALIVE_SECS: u64 = 30;
48const 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
57fn 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#[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 #[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 #[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 #[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 fn is_claude_model(&self) -> bool {
130 self.model.starts_with("claude-")
131 }
132
133 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 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 #[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 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
324impl 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 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 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 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
612fn 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
624fn 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
649fn 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
665fn 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
678impl 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 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, max_tokens,
713 system,
714 messages: &messages,
715 tools: tools.as_deref(),
716 tool_choice,
717 stream: false,
718 thinking,
719 output_config,
720 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 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 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, max_tokens,
837 system,
838 messages: &messages,
839 tools: tools.as_deref(),
840 tool_choice,
841 stream: true,
842 thinking,
843 output_config,
844 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 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 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 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 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 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 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}