agent_engine/runtime/
google_vertex.rs

1//! Provider-local Google Vertex public catalog and generateContent adapter.
2use serde::Deserialize;
3use serde_json::{json, Value};
4use url::Url;
5
6pub const MAX_PAGE_BYTES: usize = 1024 * 1024;
7pub const MAX_SSE_BUFFER: usize = 1024 * 1024;
8
9#[derive(Debug, Clone)]
10pub struct VertexRuntimeContext {
11    project: String,
12    location: String,
13}
14impl VertexRuntimeContext {
15    pub fn new(project: &str, location: &str) -> Result<Self, VertexRuntimeError> {
16        agent_core::auth::google_vertex::VertexContext::new(project, location)
17            .map_err(|_| VertexRuntimeError::InvalidContext)?;
18        Ok(Self {
19            project: project.into(),
20            location: location.into(),
21        })
22    }
23    fn host(&self) -> String {
24        format!("{}-aiplatform.googleapis.com", self.location)
25    }
26}
27#[derive(Debug, thiserror::Error)]
28pub enum VertexRuntimeError {
29    #[error("google-vertex: invalid context")]
30    InvalidContext,
31    #[error("google-vertex: invalid path or model")]
32    InvalidPath,
33    #[error("google-vertex: malformed or oversized response")]
34    Malformed,
35    #[error("google-vertex: provider error")]
36    Provider,
37}
38
39pub fn catalog_path(
40    c: &VertexRuntimeContext,
41    token: Option<&str>,
42) -> Result<String, VertexRuntimeError> {
43    let mut p = format!(
44        "/v1/projects/{}/locations/{}/publishers/google/models",
45        c.project, c.location
46    );
47    if let Some(t) = token {
48        if t.is_empty()
49            || t.len() > 512
50            || !t
51                .bytes()
52                .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
53        {
54            return Err(VertexRuntimeError::InvalidPath);
55        }
56        p.push_str("?pageToken=");
57        p.push_str(t);
58    }
59    Ok(p)
60}
61pub fn runtime_path(
62    c: &VertexRuntimeContext,
63    id: &str,
64    stream: bool,
65) -> Result<String, VertexRuntimeError> {
66    let model = id
67        .strip_prefix("google-vertex/publishers/google/models/")
68        .ok_or(VertexRuntimeError::InvalidPath)?;
69    if model.is_empty()
70        || !model
71            .bytes()
72            .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'.' || b == b'_')
73    {
74        return Err(VertexRuntimeError::InvalidPath);
75    }
76    Ok(format!(
77        "/v1/projects/{}/locations/{}/publishers/google/models/{}:{}{}",
78        c.project,
79        c.location,
80        model,
81        if stream {
82            "streamGenerateContent"
83        } else {
84            "generateContent"
85        },
86        if stream { "?alt=sse" } else { "" }
87    ))
88}
89pub fn validate_runtime_url(s: &str, c: &VertexRuntimeContext) -> Result<(), VertexRuntimeError> {
90    let u = Url::parse(s).map_err(|_| VertexRuntimeError::InvalidPath)?;
91    if u.scheme() != "https" || u.host_str() != Some(&c.host()) || u.port().is_some() {
92        return Err(VertexRuntimeError::InvalidPath);
93    }
94    let prefix = format!(
95        "/v1/projects/{}/locations/{}/publishers/google/models/",
96        c.project, c.location
97    );
98    if !u.path().starts_with(&prefix)
99        || (!u.path().ends_with(":streamGenerateContent")
100            && !u.path().ends_with(":generateContent"))
101    {
102        return Err(VertexRuntimeError::InvalidPath);
103    }
104    if u.query().is_some_and(|q| q != "alt=sse") {
105        return Err(VertexRuntimeError::InvalidPath);
106    }
107    Ok(())
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub struct CatalogEntry {
112    pub id: String,
113    pub display_name: String,
114    pub streaming: bool,
115    pub tools: bool,
116}
117pub struct CatalogPage {
118    pub entries: Vec<CatalogEntry>,
119    pub next_page_token: Option<String>,
120}
121#[derive(Deserialize)]
122struct Page {
123    #[serde(rename = "publisherModels", default)]
124    models: Vec<Model>,
125    #[serde(rename = "nextPageToken")]
126    next: Option<String>,
127}
128#[derive(Deserialize)]
129struct Model {
130    name: String,
131    #[serde(rename = "displayName", default)]
132    display: String,
133    #[serde(rename = "supportedActions", default)]
134    actions: Value,
135}
136pub fn parse_catalog_page(body: &[u8]) -> Result<CatalogPage, VertexRuntimeError> {
137    if body.len() > MAX_PAGE_BYTES {
138        return Err(VertexRuntimeError::Malformed);
139    }
140    let p: Page = serde_json::from_slice(body).map_err(|_| VertexRuntimeError::Malformed)?;
141    if p.models.len() > 1000 {
142        return Err(VertexRuntimeError::Malformed);
143    }
144    let mut entries = vec![];
145    for m in p.models {
146        if m.name.starts_with("publishers/google/models/")
147            && m.actions
148                .get("streamGenerateContent")
149                .and_then(Value::as_bool)
150                == Some(true)
151        {
152            entries.push(CatalogEntry {
153                id: format!("google-vertex/{}", m.name),
154                display_name: if m.display.is_empty() {
155                    m.name
156                } else {
157                    m.display
158                },
159                streaming: true,
160                tools: true,
161            })
162        }
163    }
164    Ok(CatalogPage {
165        entries,
166        next_page_token: p.next,
167    })
168}
169
170pub enum Message {
171    User(String),
172    Model(String),
173}
174pub struct Tool {
175    pub name: String,
176    pub description: String,
177    pub parameters: Value,
178}
179pub fn build_request(messages: &[Message], tools: &[Tool]) -> Value {
180    let contents: Vec<_> = messages
181        .iter()
182        .map(|m| match m {
183            Message::User(t) => json!({"role":"user","parts":[{"text":t}]}),
184            Message::Model(t) => json!({"role":"model","parts":[{"text":t}]}),
185        })
186        .collect();
187    let declarations: Vec<_> = tools
188        .iter()
189        .map(|t| json!({"name":t.name,"description":t.description,"parameters":t.parameters}))
190        .collect();
191    let mut v = json!({"contents":contents});
192    if !declarations.is_empty() {
193        v["tools"] = json!([{"functionDeclarations":declarations}]);
194    }
195    v
196}
197
198#[derive(Debug, Clone, PartialEq)]
199pub enum VertexEvent {
200    Text(String),
201    ToolCall {
202        name: String,
203        args: Value,
204    },
205    Usage {
206        prompt: u64,
207        completion: u64,
208        total: u64,
209    },
210    Finish(Option<String>),
211}
212#[derive(Default)]
213pub struct SseDecoder {
214    buffer: Vec<u8>,
215}
216impl SseDecoder {
217    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<VertexEvent>, VertexRuntimeError> {
218        if self.buffer.len() + bytes.len() > MAX_SSE_BUFFER {
219            return Err(VertexRuntimeError::Malformed);
220        }
221        self.buffer.extend_from_slice(bytes);
222        let mut out = vec![];
223        while let Some(pos) = self.buffer.windows(2).position(|w| w == b"\n\n") {
224            let frame: Vec<_> = self.buffer.drain(..pos + 2).collect();
225            let text = std::str::from_utf8(&frame).map_err(|_| VertexRuntimeError::Malformed)?;
226            for line in text.lines() {
227                if let Some(data) = line.strip_prefix("data: ") {
228                    if data == "[DONE]" {
229                        continue;
230                    }
231                    out.extend(parse_event(data.as_bytes())?);
232                }
233            }
234        }
235        Ok(out)
236    }
237    pub fn finish(&mut self) -> Result<Vec<VertexEvent>, VertexRuntimeError> {
238        if self.buffer.iter().all(u8::is_ascii_whitespace) {
239            self.buffer.clear();
240            return Ok(vec![]);
241        }
242        let tail = std::mem::take(&mut self.buffer);
243        let s = std::str::from_utf8(&tail).map_err(|_| VertexRuntimeError::Malformed)?;
244        if let Some(data) = s.trim().strip_prefix("data: ") {
245            parse_event(data.as_bytes())
246        } else {
247            Err(VertexRuntimeError::Malformed)
248        }
249    }
250}
251fn parse_event(data: &[u8]) -> Result<Vec<VertexEvent>, VertexRuntimeError> {
252    let v: Value = serde_json::from_slice(data).map_err(|_| VertexRuntimeError::Malformed)?;
253    if v.get("error").is_some() {
254        return Err(VertexRuntimeError::Provider);
255    }
256    let mut out = vec![];
257    if let Some(cs) = v["candidates"].as_array() {
258        for c in cs {
259            if let Some(parts) = c["content"]["parts"].as_array() {
260                for p in parts {
261                    if let Some(t) = p["text"].as_str() {
262                        out.push(VertexEvent::Text(t.into()))
263                    }
264                    if let Some(fc) = p.get("functionCall") {
265                        out.push(VertexEvent::ToolCall {
266                            name: fc["name"]
267                                .as_str()
268                                .ok_or(VertexRuntimeError::Malformed)?
269                                .into(),
270                            args: fc["args"].clone(),
271                        })
272                    }
273                }
274            }
275            if let Some(r) = c["finishReason"].as_str() {
276                out.push(VertexEvent::Finish(Some(r.into())))
277            }
278        }
279    }
280    if let Some(u) = v.get("usageMetadata") {
281        out.push(VertexEvent::Usage {
282            prompt: u["promptTokenCount"].as_u64().unwrap_or(0),
283            completion: u["candidatesTokenCount"].as_u64().unwrap_or(0),
284            total: u["totalTokenCount"].as_u64().unwrap_or(0),
285        })
286    }
287    Ok(out)
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    #[test]
294    fn catalog_path_and_filter_are_pinned() {
295        let c = VertexRuntimeContext::new("my-project-123", "us-central1").unwrap();
296        assert_eq!(
297            catalog_path(&c, None).unwrap(),
298            "/v1/projects/my-project-123/locations/us-central1/publishers/google/models"
299        );
300        let p=parse_catalog_page(br#"{"publisherModels":[{"name":"publishers/google/models/gemini-2.0-flash","displayName":"Flash","supportedActions":{"streamGenerateContent":true}},{"name":"publishers/acme/models/evil","supportedActions":{"streamGenerateContent":true}},{"name":"publishers/google/models/embed","supportedActions":{"predict":true}}],"nextPageToken":"next"}"#).unwrap();
301        assert_eq!(p.entries.len(), 1);
302        assert_eq!(
303            p.entries[0].id,
304            "google-vertex/publishers/google/models/gemini-2.0-flash"
305        );
306        assert_eq!(p.next_page_token.as_deref(), Some("next"));
307    }
308    #[test]
309    fn rejects_page_token_injection_and_wrong_hosts_models() {
310        let c = VertexRuntimeContext::new("my-project-123", "us-central1").unwrap();
311        assert!(catalog_path(&c, Some("a/b")).is_err());
312        assert!(validate_runtime_url("https://us-central1-aiplatform.googleapis.com/v1/projects/my-project-123/locations/us-central1/publishers/google/models/gemini:streamGenerateContent?alt=sse",&c).is_ok());
313        assert!(
314            validate_runtime_url("https://us-central1-aiplatform.googleapis.com.evil/x", &c)
315                .is_err()
316        );
317        assert!(runtime_path(&c, "google-vertex/publishers/acme/models/x", true).is_err());
318    }
319    #[test]
320    fn public_request_has_no_code_assist_envelope() {
321        let c = VertexRuntimeContext::new("my-project-123", "us-central1").unwrap();
322        assert!(runtime_path(
323            &c,
324            "google-vertex/publishers/google/models/gemini-2.0-flash",
325            true
326        )
327        .unwrap()
328        .ends_with(":streamGenerateContent?alt=sse"));
329        let b = build_request(
330            &[Message::User("hello".into())],
331            &[Tool {
332                name: "weather".into(),
333                description: "Weather".into(),
334                parameters: json!({"type":"object"}),
335            }],
336        );
337        assert_eq!(b["contents"][0]["parts"][0]["text"], "hello");
338        assert!(b.get("request").is_none());
339        assert!(b.get("project").is_none());
340        assert!(b.get("model").is_none());
341    }
342    #[test]
343    fn sse_handles_fragmentation_tools_usage_and_eof_without_done() {
344        let mut d = SseDecoder::default();
345        assert!(d
346            .push(b"data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hel")
347            .unwrap()
348            .is_empty());
349        let e=d.push(b"lo\"},{\"functionCall\":{\"name\":\"weather\",\"args\":{\"city\":\"x\"}}}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":2,\"candidatesTokenCount\":3,\"totalTokenCount\":5}}\n\n").unwrap();
350        assert!(e
351            .iter()
352            .any(|e| matches!(e,VertexEvent::Text(t)if t=="hello")));
353        assert!(e
354            .iter()
355            .any(|e| matches!(e,VertexEvent::ToolCall{name,..}if name=="weather")));
356        assert!(e
357            .iter()
358            .any(|e| matches!(e, VertexEvent::Usage { total: 5, .. })));
359        assert!(e
360            .iter()
361            .any(|e| matches!(e,VertexEvent::Finish(Some(r))if r=="STOP")));
362        assert!(d.finish().unwrap().is_empty());
363    }
364}
agent_engine/runtime/google_vertex.rs

agent_engine/runtime/
google_vertex.rs
