1use 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
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_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 fn effective_max_tokens(&self, request: &ChatRequest) -> u32 {
201 if request.max_tokens_explicit {
202 request.max_tokens
203 } else {
204 self.default_max_tokens()
205 }
206 }
207
208 fn map_claude_response(api_response: anthropic_data::ApiResponse) -> ChatResponse {
209 let content = anthropic_data::map_content_blocks(api_response.content);
210 let stop_reason = api_response
211 .stop_reason
212 .as_ref()
213 .map(anthropic_data::map_stop_reason);
214
215 ChatResponse {
216 id: api_response.id,
217 content,
218 model: api_response.model,
219 stop_reason,
220 usage: Usage {
221 input_tokens: api_response.usage.total_input_tokens(),
222 output_tokens: api_response.usage.output,
223 cached_input_tokens: api_response.usage.cached_input_tokens(),
224 cache_creation_input_tokens: api_response.usage.cache_creation_input_tokens(),
225 },
226 }
227 }
228}
229
230#[async_trait]
231impl LlmProvider for VertexProvider {
232 async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome> {
233 if self.is_claude_model() {
234 return self.chat_claude(request).await;
235 }
236 self.chat_gemini(request).await
237 }
238
239 fn chat_stream(&self, request: ChatRequest) -> StreamBox<'_> {
240 if self.is_claude_model() {
241 return self.chat_stream_claude(request);
242 }
243 self.chat_stream_gemini(request)
244 }
245
246 fn validate_thinking_config(&self, thinking: Option<&ThinkingConfig>) -> Result<()> {
247 let Some(thinking) = thinking else {
248 return Ok(());
249 };
250
251 if self
252 .capabilities()
253 .is_some_and(|caps| !caps.supports_thinking)
254 {
255 return Err(anyhow::anyhow!(
256 "thinking is not supported for provider={} model={}",
257 self.provider(),
258 self.model()
259 ));
260 }
261
262 if matches!(thinking.mode, ThinkingMode::Adaptive)
263 && !self
264 .capabilities()
265 .is_some_and(|caps| caps.supports_adaptive_thinking)
266 {
267 return Err(anyhow::anyhow!(
268 "adaptive thinking is not supported for provider={} model={}",
269 self.provider(),
270 self.model()
271 ));
272 }
273
274 if self.is_claude_model()
275 && self.is_anthropic_adaptive_thinking_model()
276 && matches!(thinking.mode, ThinkingMode::Enabled { .. })
277 {
278 return Err(anyhow::anyhow!(
279 "budget_tokens thinking is rejected for provider={} model={}; use ThinkingConfig::adaptive() or ThinkingConfig::default_with_effort(_) instead",
280 self.provider(),
281 self.model()
282 ));
283 }
284
285 Ok(())
286 }
287
288 async fn probe_connectivity(&self) -> bool {
289 let url = format!("https://{}/", self.endpoint_domain());
290 crate::provider::probe_http_reachability(&self.client, &url).await
291 }
292
293 fn model(&self) -> &str {
294 &self.model
295 }
296
297 fn provider(&self) -> &'static str {
298 "vertex"
299 }
300
301 fn configured_thinking(&self) -> Option<&ThinkingConfig> {
302 self.thinking.as_ref()
303 }
304
305 fn default_max_tokens(&self) -> u32 {
306 let provider = if self.is_claude_model() {
307 "anthropic"
308 } else {
309 "gemini"
310 };
311 let model_max = self
312 .capabilities()
313 .and_then(|caps| caps.max_output_tokens)
314 .or_else(|| {
315 crate::model_capabilities::default_max_output_tokens(provider, self.model())
316 })
317 .unwrap_or(4096);
318 model_max.clamp(4096, DEFAULT_SAFE_MAX_OUTPUT_TOKENS)
319 }
320}
321
322impl 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 let retry_after = if status == StatusCode::TOO_MANY_REQUESTS {
403 crate::http::retry_after_from_headers(response.headers())
404 } else {
405 None
406 };
407 let bytes = response
408 .bytes()
409 .await
410 .map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?;
411
412 log::debug!(
413 "Vertex AI LLM response status={} body_len={}",
414 status,
415 bytes.len()
416 );
417
418 if status == StatusCode::TOO_MANY_REQUESTS {
419 let body = String::from_utf8_lossy(&bytes);
420 return Ok(ChatOutcome::RateLimited(vertex_retry_delay(
421 retry_after,
422 &body,
423 )));
424 }
425
426 if status.is_server_error() {
427 let body = String::from_utf8_lossy(&bytes);
428 log::error!("Vertex AI server error status={status} body={body}");
429 return Ok(ChatOutcome::ServerError(body.into_owned()));
430 }
431
432 if status.is_client_error() {
433 let body = String::from_utf8_lossy(&bytes);
434 log::warn!("Vertex AI client error status={status} body={body}");
435 return Ok(ChatOutcome::InvalidRequest(body.into_owned()));
436 }
437
438 let api_response: ApiGenerateContentResponse = serde_json::from_slice(&bytes)
439 .map_err(|e| anyhow::anyhow!("failed to parse response: {e}"))?;
440
441 let candidate = api_response
442 .candidates
443 .into_iter()
444 .next()
445 .ok_or_else(|| anyhow::anyhow!("no candidates in response"))?;
446
447 let content = build_content_blocks(&candidate.content);
448
449 if content.is_empty() && !candidate.content.parts.is_empty() {
450 log::warn!(
451 "Vertex AI parts not converted to content blocks raw_parts={:?}",
452 candidate.content.parts
453 );
454 }
455
456 let has_tool_calls = content
457 .iter()
458 .any(|b| matches!(b, agent_sdk_foundation::llm::ContentBlock::ToolUse { .. }));
459
460 let stop_reason = candidate
461 .finish_reason
462 .as_ref()
463 .map(|r| map_finish_reason(r, has_tool_calls));
464
465 let usage = api_response
466 .usage_metadata
467 .unwrap_or(ApiUsageMetadata {
468 prompt: 0,
469 candidates: 0,
470 cached_content: 0,
471 })
472 .into_usage();
473
474 Ok(ChatOutcome::Success(ChatResponse {
475 id: String::new(),
476 content,
477 model: self.model.clone(),
478 stop_reason,
479 usage,
480 }))
481 }
482
483 fn chat_stream_gemini(&self, request: ChatRequest) -> StreamBox<'_> {
484 let served_route = self.route().to_owned();
485 Box::pin(async_stream::stream! {
486 let thinking = match self.resolve_thinking_config(request.thinking.as_ref()) {
487 Ok(thinking) => thinking,
488 Err(error) => {
489 yield Ok(StreamDelta::Error {
490 message: error.to_string(),
491 kind: StreamErrorKind::InvalidRequest,
492 });
493 return;
494 }
495 };
496 if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
497 yield Ok(StreamDelta::Error {
498 message: error.to_string(),
499 kind: StreamErrorKind::InvalidRequest,
500 });
501 return;
502 }
503
504 let contents = build_api_contents(&request.messages);
505 let tools = request
506 .tools
507 .as_ref()
508 .map(|t| convert_tools_to_config(t.clone()));
509 let tool_config = request
510 .tool_choice
511 .as_ref()
512 .map(ApiFunctionCallingConfig::from_tool_choice);
513 let system_instruction = build_gemini_system_instruction(&request.system);
514 let thinking_config = thinking.as_ref().map(map_thinking_config);
515 let (response_mime_type, response_schema) =
516 gemini_response_format(request.response_format.as_ref());
517
518 let max_tokens = self.effective_max_tokens(&request);
519 let api_request = ApiGenerateContentRequest {
520 contents: &contents,
521 system_instruction: system_instruction.as_ref(),
522 tools: tools.as_ref().map(std::slice::from_ref),
523 tool_config,
524 generation_config: Some(ApiGenerationConfig {
525 max_output_tokens: Some(max_tokens),
526 thinking_config,
527 response_mime_type,
528 response_schema,
529 }),
530 cached_content: request.cached_content.as_deref(),
531 };
532
533 log::debug!(
534 "Vertex AI streaming LLM request model={} max_tokens={}",
535 self.model,
536 max_tokens
537 );
538
539 let url = format!("{}:streamGenerateContent?alt=sse", self.base_url("google"));
540
541 let response = match self.send_gemini_stream_request(&url, &api_request).await {
542 Ok(response) => response,
543 Err(item) => {
544 yield item;
545 return;
546 }
547 };
548
549 let mut inner = stream_gemini_response(response);
550 while let Some(item) = futures::StreamExt::next(&mut inner).await {
551 yield match item {
552 Ok(StreamDelta::Done { stop_reason, .. }) => Ok(StreamDelta::Done {
553 stop_reason,
554 served_route: Some(served_route.clone()),
555 }),
556 other => other,
557 };
558 }
559 })
560 }
561
562 async fn send_gemini_stream_request(
570 &self,
571 url: &str,
572 api_request: &ApiGenerateContentRequest<'_>,
573 ) -> Result<reqwest::Response, anyhow::Result<StreamDelta>> {
574 let response = match self
575 .client
576 .post(url)
577 .header("Content-Type", "application/json")
578 .bearer_auth(&self.access_token)
579 .json(api_request)
580 .send()
581 .await
582 {
583 Ok(response) => response,
584 Err(error) => return Err(Ok(reqwest_error_delta("request failed", &error))),
585 };
586
587 let status = response.status();
588 if !status.is_success() {
589 let header_hint = crate::http::retry_after_from_headers(response.headers());
591 let body = response.text().await.unwrap_or_default();
592 let kind = if status == StatusCode::TOO_MANY_REQUESTS {
593 StreamErrorKind::RateLimited(vertex_retry_delay(header_hint, &body))
594 } else if status.is_server_error() {
595 StreamErrorKind::ServerError
596 } else {
597 StreamErrorKind::InvalidRequest
598 };
599 log::warn!("Vertex AI error status={status} body={body}");
600 return Err(Ok(StreamDelta::Error {
601 message: body,
602 kind,
603 }));
604 }
605
606 Ok(response)
607 }
608}
609
610fn vertex_retry_delay(
616 header_hint: Option<std::time::Duration>,
617 body: &str,
618) -> Option<std::time::Duration> {
619 header_hint.or_else(|| crate::retry_hints::google_retry_delay(body))
620}
621
622fn classify_vertex_claude_status(
625 status: StatusCode,
626 bytes: &[u8],
627 header_hint: Option<std::time::Duration>,
628) -> Option<ChatOutcome> {
629 let body = String::from_utf8_lossy(bytes);
630 if status == StatusCode::TOO_MANY_REQUESTS {
631 return Some(ChatOutcome::RateLimited(vertex_retry_delay(
632 header_hint,
633 &body,
634 )));
635 }
636 if status.is_server_error() {
637 log::error!("Vertex AI (Claude) server error status={status} body={body}");
638 return Some(ChatOutcome::ServerError(body.into_owned()));
639 }
640 if status.is_client_error() {
641 log::warn!("Vertex AI (Claude) client error status={status} body={body}");
642 return Some(ChatOutcome::InvalidRequest(body.into_owned()));
643 }
644 None
645}
646
647fn build_gemini_system_instruction(system: &str) -> Option<ApiContent> {
650 if system.is_empty() {
651 None
652 } else {
653 Some(ApiContent {
654 role: None,
655 parts: vec![ApiPart::Text {
656 text: system.to_owned(),
657 thought_signature: None,
658 }],
659 })
660 }
661}
662
663fn gemini_response_format(
666 response_format: Option<&ResponseFormat>,
667) -> (Option<&'static str>, Option<serde_json::Value>) {
668 response_format.map_or((None, None), |rf| {
669 (
670 Some("application/json"),
671 Some(gemini_response_schema(&rf.schema)),
672 )
673 })
674}
675
676impl VertexProvider {
681 async fn chat_claude(&self, request: ChatRequest) -> Result<ChatOutcome> {
682 let thinking_config = match self.resolve_thinking_config(request.thinking.as_ref()) {
683 Ok(thinking) => thinking_for_forced_tool(thinking, request.tool_choice.as_ref()),
687 Err(error) => return Ok(ChatOutcome::InvalidRequest(error.to_string())),
688 };
689 if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
690 return Ok(ChatOutcome::InvalidRequest(error.to_string()));
691 }
692 let messages = Self::build_cached_vertex_claude_messages(&request);
693 let tools = build_vertex_claude_tools(&request);
694 let thinking = thinking_config
695 .as_ref()
696 .and_then(anthropic_data::ApiThinkingConfig::from_thinking_config);
697 let output_config = thinking_config
698 .as_ref()
699 .and_then(|t| t.effort)
700 .map(|effort| anthropic_data::ApiOutputConfig { effort });
701 let system = Self::build_vertex_claude_system_prompt(&request.system);
702 let tool_choice = request
703 .tool_choice
704 .as_ref()
705 .map(anthropic_data::ApiToolChoice::from_tool_choice);
706
707 let max_tokens = self.effective_max_tokens(&request);
708 let api_request = anthropic_data::ApiMessagesRequest {
709 model: None, max_tokens,
711 system,
712 messages: &messages,
713 tools: tools.as_deref(),
714 tool_choice,
715 stream: false,
716 thinking,
717 output_config,
718 anthropic_version: Some(VERTEX_ANTHROPIC_VERSION),
719 };
720
721 log::debug!(
722 "Vertex AI (Claude) LLM request model={} max_tokens={}",
723 self.model,
724 max_tokens
725 );
726
727 if log::log_enabled!(log::Level::Debug) {
728 match serde_json::to_string_pretty(&api_request) {
729 Ok(json) => log::debug!("Vertex AI (Claude) request payload:\n{json}"),
730 Err(e) => log::debug!("Failed to serialize request for logging: {e}"),
731 }
732 }
733
734 let url = format!("{}:rawPredict", self.base_url("anthropic"));
735
736 let response = self
737 .client
738 .post(&url)
739 .header("Content-Type", "application/json")
740 .timeout(std::time::Duration::from_secs(CHAT_READ_TIMEOUT_SECS))
741 .bearer_auth(&self.access_token)
742 .json(&api_request)
743 .send()
744 .await
745 .map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
746
747 let status = response.status();
748 let retry_after = if status == StatusCode::TOO_MANY_REQUESTS {
750 crate::http::retry_after_from_headers(response.headers())
751 } else {
752 None
753 };
754 let bytes = response
755 .bytes()
756 .await
757 .map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?;
758
759 log::debug!(
760 "Vertex AI (Claude) response status={} body_len={}",
761 status,
762 bytes.len()
763 );
764
765 if let Some(outcome) = classify_vertex_claude_status(status, &bytes, retry_after) {
766 return Ok(outcome);
767 }
768
769 let api_response: anthropic_data::ApiResponse = serde_json::from_slice(&bytes)
770 .map_err(|e| anyhow::anyhow!("failed to parse response: {e}"))?;
771
772 log::debug!(
773 "Vertex AI (Claude) response: id={} model={} stop_reason={:?} usage={{input_tokens={}, output_tokens={}}} content_blocks={}",
774 api_response.id,
775 api_response.model,
776 api_response.stop_reason,
777 api_response.usage.total_input_tokens(),
778 api_response.usage.output,
779 api_response.content.len()
780 );
781
782 Ok(ChatOutcome::Success(Self::map_claude_response(
783 api_response,
784 )))
785 }
786
787 #[allow(clippy::too_many_lines)]
788 fn chat_stream_claude(&self, request: ChatRequest) -> StreamBox<'_> {
789 let served_route = self.route().to_owned();
790 Box::pin(async_stream::stream! {
791 let thinking_config = match self.resolve_thinking_config(request.thinking.as_ref()) {
792 Ok(thinking) => thinking_for_forced_tool(thinking, request.tool_choice.as_ref()),
797 Err(error) => {
798 yield Ok(StreamDelta::Error {
799 message: error.to_string(),
800 kind: StreamErrorKind::InvalidRequest,
801 });
802 return;
803 }
804 };
805 if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
806 yield Ok(StreamDelta::Error {
807 message: error.to_string(),
808 kind: StreamErrorKind::InvalidRequest,
809 });
810 return;
811 }
812 let messages = Self::build_cached_vertex_claude_messages(&request);
813 let tools = build_vertex_claude_tools(&request);
814 let thinking = thinking_config
815 .as_ref()
816 .and_then(anthropic_data::ApiThinkingConfig::from_thinking_config);
817 let output_config = thinking_config
818 .as_ref()
819 .and_then(|t| t.effort)
820 .map(|effort| anthropic_data::ApiOutputConfig { effort });
821 let system = Self::build_vertex_claude_system_prompt(&request.system);
822 let tool_choice = request
823 .tool_choice
824 .as_ref()
825 .map(anthropic_data::ApiToolChoice::from_tool_choice);
826
827 let max_tokens = self.effective_max_tokens(&request);
828 let api_request = anthropic_data::ApiMessagesRequest {
829 model: None, max_tokens,
831 system,
832 messages: &messages,
833 tools: tools.as_deref(),
834 tool_choice,
835 stream: true,
836 thinking,
837 output_config,
838 anthropic_version: Some(VERTEX_ANTHROPIC_VERSION),
839 };
840
841 log::debug!(
842 "Vertex AI (Claude) streaming request model={} max_tokens={}",
843 self.model,
844 max_tokens
845 );
846
847 if log::log_enabled!(log::Level::Debug) {
848 match serde_json::to_string_pretty(&api_request) {
849 Ok(json) => log::debug!("Vertex AI (Claude) streaming request payload:\n{json}"),
850 Err(e) => log::debug!("Failed to serialize request for logging: {e}"),
851 }
852 }
853
854 let url = format!("{}:streamRawPredict", self.base_url("anthropic"));
855
856 let response = match self
857 .client
858 .post(&url)
859 .header("Content-Type", "application/json")
860 .bearer_auth(&self.access_token)
861 .json(&api_request)
862 .send()
863 .await
864 {
865 Ok(r) => r,
866 Err(error) => {
867 yield Ok(reqwest_error_delta("request failed", &error));
868 return;
869 }
870 };
871
872 let status = response.status();
873
874 if status == StatusCode::TOO_MANY_REQUESTS {
875 let header_hint = crate::http::retry_after_from_headers(response.headers());
877 let body = response.text().await.unwrap_or_default();
878 log::warn!("Vertex AI (Claude) rate limited status={status} body={body}");
879 yield Ok(StreamDelta::Error {
880 message: "Rate limited".to_string(),
881 kind: StreamErrorKind::RateLimited(vertex_retry_delay(header_hint, &body)),
882 });
883 return;
884 }
885
886 if status.is_server_error() {
887 let body = response.text().await.unwrap_or_default();
888 log::error!("Vertex AI (Claude) server error status={status} body={body}");
889 yield Ok(StreamDelta::Error {
890 message: body,
891 kind: StreamErrorKind::ServerError,
892 });
893 return;
894 }
895
896 if status.is_client_error() {
897 let body = response.text().await.unwrap_or_default();
898 log::warn!("Vertex AI (Claude) client error status={status} body={body}");
899 yield Ok(StreamDelta::Error {
900 message: body,
901 kind: StreamErrorKind::InvalidRequest,
902 });
903 return;
904 }
905
906 let mut stream = response.bytes_stream();
908 let mut buffer = String::new();
909 let mut input_tokens: u32 = 0;
910 let mut output_tokens: u32 = 0;
911 let mut cached_input_tokens: u32 = 0;
912 let mut cache_creation_input_tokens: u32 = 0;
913 let mut tool_ids: std::collections::HashMap<usize, String> =
914 std::collections::HashMap::new();
915 let mut received_message_stop = false;
916 let mut pending_stop_reason: Option<agent_sdk_foundation::llm::StopReason> = None;
917
918 while let Some(chunk_result) = stream.next().await {
919 let chunk = match chunk_result {
920 Ok(c) => c,
921 Err(error) => {
922 yield Ok(reqwest_body_error_delta("stream error", &error));
923 return;
924 }
925 };
926
927 buffer.push_str(&String::from_utf8_lossy(&chunk));
928
929 while let Some(event_block) = anthropic_data::take_next_sse_event(&mut buffer) {
931 if anthropic_data::is_message_stop_event(&event_block) {
932 received_message_stop = true;
933 }
934
935 if let Some(delta) = anthropic_data::parse_sse_event(
936 &event_block,
937 &mut input_tokens,
938 &mut output_tokens,
939 &mut cached_input_tokens,
940 &mut cache_creation_input_tokens,
941 &mut tool_ids,
942 &mut pending_stop_reason,
943 ) {
944 yield Ok(delta);
945 }
946 if anthropic_data::is_message_stop_event(&event_block) {
947 yield Ok(StreamDelta::Done {
948 stop_reason: pending_stop_reason.take(),
949 served_route: Some(served_route.clone()),
950 });
951 }
952 }
953 }
954
955 let remaining = buffer.trim();
957 if !remaining.is_empty() {
958 if anthropic_data::is_message_stop_event(remaining) {
959 received_message_stop = true;
960 }
961
962 if let Some(delta) = anthropic_data::parse_sse_event(
963 remaining,
964 &mut input_tokens,
965 &mut output_tokens,
966 &mut cached_input_tokens,
967 &mut cache_creation_input_tokens,
968 &mut tool_ids,
969 &mut pending_stop_reason,
970 ) {
971 yield Ok(delta);
972 }
973 if anthropic_data::is_message_stop_event(remaining) {
974 yield Ok(StreamDelta::Done {
975 stop_reason: pending_stop_reason.take(),
976 served_route: Some(served_route.clone()),
977 });
978 }
979 }
980
981 if !received_message_stop {
982 log::warn!(
983 "Vertex AI (Claude) SSE stream ended without message_stop"
984 );
985 yield Ok(StreamDelta::Error {
986 message: "Stream ended unexpectedly without completion".to_string(),
987 kind: StreamErrorKind::ServerError,
988 });
989 }
990 })
991 }
992}
993
994#[cfg(test)]
995mod tests {
996 use super::*;
997
998 #[test]
999 fn test_new_creates_provider() {
1000 let provider = VertexProvider::new(
1001 "token".to_string(),
1002 "my-project".to_string(),
1003 "us-central1".to_string(),
1004 "custom-model".to_string(),
1005 );
1006
1007 assert_eq!(provider.model(), "custom-model");
1008 assert_eq!(provider.provider(), "vertex");
1009 }
1010
1011 #[test]
1012 fn test_flash_factory() {
1013 let provider = VertexProvider::flash(
1014 "token".to_string(),
1015 "my-project".to_string(),
1016 "us-central1".to_string(),
1017 );
1018
1019 assert_eq!(provider.model(), MODEL_GEMINI_3_FLASH);
1020 assert_eq!(provider.provider(), "vertex");
1021 }
1022
1023 #[test]
1024 fn test_pro_factory() {
1025 let provider = VertexProvider::pro(
1026 "token".to_string(),
1027 "my-project".to_string(),
1028 "us-central1".to_string(),
1029 );
1030
1031 assert_eq!(provider.model(), MODEL_GEMINI_31_PRO);
1032 assert_eq!(provider.provider(), "vertex");
1033 }
1034
1035 #[test]
1036 fn test_provider_is_cloneable() {
1037 let provider = VertexProvider::new(
1038 "token".to_string(),
1039 "my-project".to_string(),
1040 "us-central1".to_string(),
1041 "test-model".to_string(),
1042 );
1043 let cloned = provider.clone();
1044
1045 assert_eq!(provider.model(), cloned.model());
1046 assert_eq!(provider.provider(), cloned.provider());
1047 }
1048
1049 #[test]
1050 fn test_is_claude_model() {
1051 let claude_provider = VertexProvider::new(
1052 "token".to_string(),
1053 "project".to_string(),
1054 "us-central1".to_string(),
1055 "claude-sonnet-4-20250514".to_string(),
1056 );
1057 assert!(claude_provider.is_claude_model());
1058
1059 let gemini_provider = VertexProvider::new(
1060 "token".to_string(),
1061 "project".to_string(),
1062 "us-central1".to_string(),
1063 "gemini-3-flash-preview".to_string(),
1064 );
1065 assert!(!gemini_provider.is_claude_model());
1066 }
1067
1068 #[test]
1069 fn test_base_url_gemini() {
1070 let provider = VertexProvider::new(
1071 "token".to_string(),
1072 "my-project".to_string(),
1073 "us-central1".to_string(),
1074 "gemini-3-flash-preview".to_string(),
1075 );
1076
1077 let url = provider.base_url("google");
1078 assert_eq!(
1079 url,
1080 "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-3-flash-preview"
1081 );
1082 }
1083
1084 #[test]
1085 fn test_base_url_claude() {
1086 let provider = VertexProvider::new(
1087 "token".to_string(),
1088 "my-project".to_string(),
1089 "us-central1".to_string(),
1090 "claude-sonnet-4-20250514".to_string(),
1091 );
1092
1093 let url = provider.base_url("anthropic");
1094 assert_eq!(
1095 url,
1096 "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/anthropic/models/claude-sonnet-4-20250514"
1097 );
1098 }
1099
1100 #[test]
1101 fn test_base_url_with_different_region() {
1102 let provider = VertexProvider::new(
1103 "token".to_string(),
1104 "other-project".to_string(),
1105 "europe-west4".to_string(),
1106 "gemini-3.1-pro-preview".to_string(),
1107 );
1108
1109 let url = provider.base_url("google");
1110 assert!(url.starts_with("https://europe-west4-aiplatform.googleapis.com/"));
1111 assert!(url.contains("/projects/other-project/"));
1112 assert!(url.contains("/locations/europe-west4/"));
1113 assert!(url.ends_with("/models/gemini-3.1-pro-preview"));
1114 }
1115
1116 #[test]
1117 fn test_base_url_global_region_has_no_prefix() {
1118 let provider = VertexProvider::new(
1119 "token".to_string(),
1120 "my-project".to_string(),
1121 "global".to_string(),
1122 "gemini-3.1-pro-preview".to_string(),
1123 );
1124
1125 let url = provider.base_url("google");
1126 assert_eq!(
1127 url,
1128 "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-3.1-pro-preview"
1129 );
1130 }
1131
1132 #[test]
1133 fn test_vertex_claude_46_rejects_budgeted_thinking() {
1134 let provider = VertexProvider::new(
1135 "token".to_string(),
1136 "project".to_string(),
1137 "global".to_string(),
1138 MODEL_SONNET_46.to_string(),
1139 );
1140
1141 let error = provider
1142 .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1143 .unwrap_err();
1144 assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1145 }
1146
1147 #[test]
1148 fn test_vertex_claude_opus_47_rejects_budgeted_thinking() {
1149 let provider = VertexProvider::new(
1150 "token".to_string(),
1151 "project".to_string(),
1152 "global".to_string(),
1153 MODEL_OPUS_47.to_string(),
1154 );
1155
1156 let error = provider
1157 .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1158 .unwrap_err();
1159 assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1160 }
1161
1162 #[test]
1163 fn test_vertex_claude_opus_48_rejects_budgeted_thinking() {
1164 let provider = VertexProvider::new(
1165 "token".to_string(),
1166 "project".to_string(),
1167 "global".to_string(),
1168 MODEL_OPUS_48.to_string(),
1169 );
1170
1171 let error = provider
1172 .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1173 .unwrap_err();
1174 assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1175 }
1176
1177 #[test]
1178 fn test_vertex_claude_fable_5_rejects_budgeted_thinking() {
1179 let provider = VertexProvider::new(
1180 "token".to_string(),
1181 "project".to_string(),
1182 "global".to_string(),
1183 MODEL_FABLE_5.to_string(),
1184 );
1185
1186 let error = provider
1187 .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
1188 .unwrap_err();
1189 assert!(error.to_string().contains("ThinkingConfig::adaptive()"));
1190 }
1191
1192 #[test]
1193 fn test_model_constants() {
1194 assert_eq!(MODEL_GEMINI_3_FLASH, "gemini-3-flash-preview");
1195 assert_eq!(MODEL_GEMINI_31_PRO, "gemini-3.1-pro-preview");
1196 assert_eq!(MODEL_GEMINI_3_PRO, "gemini-3.0-pro");
1197 }
1198
1199 fn request_with_max_tokens(max_tokens: u32, explicit: bool) -> ChatRequest {
1200 ChatRequest {
1201 system: String::new(),
1202 messages: vec![agent_sdk_foundation::llm::Message::user("hi")],
1203 tools: None,
1204 max_tokens,
1205 max_tokens_explicit: explicit,
1206 session_id: None,
1207 cached_content: None,
1208 thinking: None,
1209 tool_choice: None,
1210 response_format: None,
1211 cache: None,
1212 }
1213 }
1214
1215 #[test]
1216 fn test_effective_max_tokens_honors_explicit_budget() {
1217 let provider = VertexProvider::new(
1218 "token".to_string(),
1219 "project".to_string(),
1220 "global".to_string(),
1221 MODEL_SONNET_46.to_string(),
1222 );
1223 let request = request_with_max_tokens(1234, true);
1224 assert_eq!(provider.effective_max_tokens(&request), 1234);
1225 }
1226
1227 #[test]
1228 fn test_effective_max_tokens_uses_clamped_default_when_implicit() {
1229 let provider = VertexProvider::new(
1233 "token".to_string(),
1234 "project".to_string(),
1235 "global".to_string(),
1236 MODEL_SONNET_46.to_string(),
1237 );
1238 let request = request_with_max_tokens(4096, false);
1239 let effective = provider.effective_max_tokens(&request);
1240 assert_eq!(effective, provider.default_max_tokens());
1241 assert!(effective <= DEFAULT_SAFE_MAX_OUTPUT_TOKENS);
1242 }
1243}