fix veo3 adapter - #1794
Conversation
WalkthroughRefactors the Vertex TaskAdaptor to be stateful and RelayInfo-driven, updates all method signatures accordingly, unifies validation, rewrites URL/header/body construction, adjusts request/response handling, adds status polling via FetchTask, and standardizes result mapping to shared model status constants within a single adaptor file. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant Adaptor as TaskAdaptor (stateful)
participant Vertex as Vertex AI API
Note over Adaptor: Initialized with ChannelType, apiKey, baseURL
Client->>Adaptor: Init(RelayInfo)
Client->>Adaptor: ValidateRequestAndSetAction(c, RelayInfo)
Adaptor->>Adaptor: ValidateBasicTaskRequest(...)
Client->>Adaptor: BuildRequestURL(RelayInfo)
Adaptor->>Adaptor: Derive model & region from RelayInfo
Client->>Adaptor: BuildRequestHeader(c, req, RelayInfo)
Adaptor->>Adaptor: Decode credentials, acquire access token
Client->>Adaptor: BuildRequestBody(c, RelayInfo)
Adaptor->>Adaptor: Compose Vertex requestPayload
Client->>Adaptor: DoRequest(c, RelayInfo, body)
Adaptor->>Vertex: POST submit request
Vertex-->>Adaptor: submitResponse (operation name)
Adaptor-->>Client: DoResponse(...) → task_id (from operation)
sequenceDiagram
autonumber
participant Client
participant Adaptor as TaskAdaptor
participant Vertex as Vertex Operations API
Client->>Adaptor: FetchTask(baseURL, key, {operationName})
Adaptor->>Vertex: GET /operations/{name}
Vertex-->>Adaptor: operationResponse (state, metadata)
Adaptor->>Adaptor: ParseTaskResult → model.TaskStatus*
Adaptor-->>Client: status, mapped result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
relay/channel/task/vertex/adaptor.go (2)
167-185: Handle non-2xx upstream responses and surface error details.
Currently all responses are parsed as success; a 4xx/5xx becomes a 500 "invalid_response". Preserve upstream status and message.Apply:
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) { responseBody, err := io.ReadAll(resp.Body) if err != nil { return "", nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) } _ = resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + var em struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + _ = json.Unmarshal(responseBody, &em) + msg := strings.TrimSpace(em.Error.Message) + if msg == "" { + msg = string(responseBody) + } + return "", nil, service.TaskErrorWrapper(fmt.Errorf(msg), "upstream_error", resp.StatusCode) + } var s submitResponse if err := json.Unmarshal(responseBody, &s); err != nil { return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError) }
112-129: Pass channel proxy to Vertex token acquisition and cache tokens.vertexcore.AcquireAccessToken(creds, proxy string) is implemented at relay/channel/vertex/service_account.go:142 — the Task adaptor currently calls AcquireAccessToken(..., "") which ignores proxy. Replace "" with info.ChannelSetting.Proxy at the two call sites: relay/channel/task/vertex/adaptor.go:122 and relay/channel/task/vertex/adaptor.go:224. Also add short-lived token caching (refresh near-expiry) to avoid repeated auth calls.
🧹 Nitpick comments (8)
relay/channel/task/vertex/adaptor.go (8)
24-27: Section banners are fine; keep consistent with repo style.
If there’s a linter/formatter rule for comments, align to it.
63-67: Unused struct field and missed opportunity for credential/token caching.
- baseURL is set in Init but never used; remove or wire it into URL construction.
- Repeated JSON unmarshal of apiKey occurs in multiple methods; decode once and cache.
Apply if you decide to remove the unused field:
type TaskAdaptor struct { ChannelType int - apiKey string - baseURL string + apiKey string }And in Init (see next comment) drop the assignment to baseURL.
69-73: Decode service account JSON once during Init and store the typed creds.
This avoids repeated json.Unmarshal and centralizes credential validation.I can provide a follow-up diff to add
creds vertexcore.Credentialsto TaskAdaptor and unmarshal here if you want.
81-110: URL building mostly correct; align defaults and avoid per-call JSON decode.
- Region fallback here is "global" but FetchTask assumes "us-central1" on parse failure; pick one default to avoid split behavior.
- Consider pre-parsed creds from Init to avoid repeating json.Unmarshal here.
Would you like a patch to centralize region defaulting and creds usage?
190-214: baseUrl parameter is unused; proxy and token concerns repeat here.
- baseUrl is not referenced; either use it to override domain or remove from signature.
- AcquireAccessToken again passes ""; use the same proxy strategy as request flow and share token caching.
If keeping baseUrl, consider honoring it when constructing url, or drop the parameter to avoid confusion.
205-206: Operation-name parsing: guard rails look good.
Ensure logs include the offending name when extraction fails to aid debugging.
239-302: Optional: include MIME inference once and reuse.
Current branches repeat encoding→MIME conversion. Factor into a small helper to reduce duplication.I can draft a tiny helper like
inferVideoMIME(enc string) stringif helpful.
320-355: Regex extractors look safe for Vertex op names.
Patterns match expected shapes; consider unit tests for edge cases.I can add tests for encode/decode and extractors; want me to open a small PR with those?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
relay/channel/task/vertex/adaptor.go(10 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
relay/channel/task/vertex/adaptor.go (5)
relay/common/relay_info.go (1)
RelayInfo(74-120)relay/common/relay_utils.go (1)
ValidateBasicTaskRequest(67-84)relay/channel/vertex/service_account.go (2)
Credentials(22-28)AcquireAccessToken(142-148)relay/channel/api_request.go (1)
DoTaskApiRequest(279-301)model/task.go (3)
TaskStatusFailure(18-18)TaskStatusInProgress(17-17)TaskStatusSuccess(19-19)
🔇 Additional comments (7)
relay/channel/task/vertex/adaptor.go (7)
10-10: Import of model status constants is correct.
Used appropriately in ParseTaskResult.
162-165: Delegation to shared request helper looks good.
Keeps behavior uniform across adaptors.
211-214: Region-aware fetch URLs look correct.
Matches global vs regional endpoints with fetchPredictOperation suffix.
246-256: Using centralized status constants is the right move.
Aligns with model.TaskStatus* across the codebase.
304-307: Helper section separation is clear.
No action needed.
308-318: ID encode/decode helpers look correct.
Base64 Raw URL encoding avoids slashes.
131-160: Harden BuildRequestBody: safe type-assert + coerce sampleCount
- Use a safe type assertion for c.Get("task_request") to avoid panics. Apply the original diff below in relay/channel/task/vertex/adaptor.go (BuildRequestBody):
- v, ok := c.Get("task_request") - if !ok { + v, ok := c.Get("task_request") + if !ok { return nil, fmt.Errorf("request not found in context") - } - req := v.(relaycommon.TaskSubmitReq) + } + req, ok := v.(relaycommon.TaskSubmitReq) + if !ok { + return nil, fmt.Errorf("invalid task_request type") + }
- Coerce metadata["sampleCount"] to an int before adding to body.Parameters to avoid sending floats/strings to Vertex. Example minimal approach (add import strconv):
if v, ok := req.Metadata["sampleCount"]; ok { switch t := v.(type) { case int: body.Parameters["sampleCount"] = t case float64: body.Parameters["sampleCount"] = int(t) case string: s := strings.TrimSpace(t) if n, err := strconv.Atoi(s); err == nil { body.Parameters["sampleCount"] = n } else if f, err := strconv.ParseFloat(s, 64); err == nil { body.Parameters["sampleCount"] = int(f) } } }
- Confirm upstream parameter names/types: verify that Vertex expects "storageUri" and integer "sampleCount" in Parameters (adjust names/types if upstream differs).
File/loc: relay/channel/task/vertex/adaptor.go — BuildRequestBody (around lines ~131–160).
| // ValidateRequestAndSetAction parses body, validates fields and sets default action. | ||
| func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) { | ||
| // Use the standard validation method for TaskSubmitReq | ||
| return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionTextGenerate) | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Confirm the action type for video generation.
Using constant.TaskActionTextGenerate may be incorrect for Veo video tasks.
Run to verify the right constant (e.g., TaskActionVideoGenerate) exists and is used elsewhere:
If video-specific action exists, update:
- return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionTextGenerate)
+ return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionVideoGenerate)🏁 Script executed:
#!/bin/bash
rg -n -C2 'TaskAction.*(Video|Text).*Generate' -- types
rg -n -C2 'ValidateBasicTaskRequest\('Length of output: 1754
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== TaskAction symbols (all occurrences) ==="
rg -n -S '\bTaskAction\w*\b' || true
echo
echo "=== Usages of constant.TaskAction.* ==="
rg -n -S 'constant\.TaskAction\w*' || true
echo
echo "=== Files declaring package constant ==="
rg -n -S 'package constant' --files || true
echo
echo "=== Search for specific TaskAction constant names ==="
rg -n -S 'TaskActionTextGenerate|TaskActionVideoGenerate|TaskActionGenerate' || true
echo
echo "=== Show files that likely define constants (const blocks) ==="
rg -n -S '^\s*const\s+\(|^\s*const\s+\w+' || trueLength of output: 265682
Use TaskActionGenerate for video generation (TaskActionVideoGenerate doesn't exist). Change the call in relay/channel/task/vertex/adaptor.go (ValidateRequestAndSetAction — ~line 78) from constant.TaskActionTextGenerate to constant.TaskActionGenerate if this adaptor should accept video-generation requests; if the adaptor is text-only, keep as-is.
🤖 Prompt for AI Agents
In relay/channel/task/vertex/adaptor.go around lines 75-79, the call to
relaycommon.ValidateBasicTaskRequest currently passes
constant.TaskActionTextGenerate but this adaptor should accept video-generation
requests; update the argument to constant.TaskActionGenerate so the validator
allows video generation actions, keeping the rest of the call identical.
fix veo3 adapter
2 participants
Summary by CodeRabbit
New Features
Improvements