fix: idle timeout wiring in vertex path

[TejasGhatte]

Summary

The SetStreamIdleTimeoutIfEmpty call was being made before invoking the shared HandleGemini*Stream functions, meaning the timeout was set in the provider-level methods rather than inside the shared handlers. This caused the Vertex provider — which delegates directly to those shared handlers — to never set the stream idle timeout at all. Moving the call into HandleGeminiChatCompletionStream and HandleGeminiResponsesStream, and passing streamIdleTimeoutInSeconds as an explicit parameter, ensures the timeout is consistently applied regardless of which provider invokes the handler.

Changes

• Removed SetStreamIdleTimeoutIfEmpty calls from GeminiProvider.ChatCompletionStream and GeminiProvider.ResponsesStream.
• Added streamIdleTimeoutInSeconds int parameter to HandleGeminiChatCompletionStream and HandleGeminiResponsesStream.
• Moved SetStreamIdleTimeoutIfEmpty to the top of each shared handler so it is always executed.
• Updated VertexProvider.ChatCompletionStream and VertexProvider.ResponsesStream to pass provider.networkConfig.StreamIdleTimeoutInSeconds to the shared handlers, fixing the missing timeout for Vertex streaming calls.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

Verify that streaming requests through the Vertex provider respect the configured StreamIdleTimeoutInSeconds by setting a short timeout and confirming the stream is terminated after the idle period elapses.

go test ./...

Breaking changes

• Yes
• No

Related issues

Security considerations

No security implications.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 5fc8f2ab-2b87-4000-b35f-6ae8a7806ec2

📥 Commits

Reviewing files that changed from the base of the PR and between e2f9f2d and 7d0f706.

📒 Files selected for processing (7)

• core/providers/cohere/cohere.go
• core/providers/gemini/gemini.go
• core/providers/huggingface/huggingface.go
• core/providers/mistral/mistral.go
• core/providers/replicate/replicate.go
• core/providers/vertex/vertex.go
• core/providers/vllm/vllm.go

🚧 Files skipped from review as they are similar to previous changes (2)

• core/providers/gemini/gemini.go
• core/providers/vertex/vertex.go

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved stream timeout handling across several AI provider streaming flows for more consistent behavior.
  • Ensured timeout settings are applied earlier or only after successful responses, reducing cases where streaming could stall or behave unexpectedly.
  • Updated Gemini, Cohere, Hugging Face, Mistral, Replicate, and vLLM streaming paths for better reliability during chat, transcription, and image-related streaming.

Walkthrough

This change centralizes or repositions stream idle-timeout initialization across Gemini, Vertex, Cohere, Hugging Face, Mistral, Replicate, and vLLM streaming paths. Gemini handlers now accept the timeout value directly, while other providers set it earlier in successful streaming setup.

Changes

Stream idle-timeout updates

Layer / File(s)	Summary
Gemini handlers take timeout input core/providers/gemini/gemini.go, core/providers/vertex/vertex.go	HandleGeminiChatCompletionStream and HandleGeminiResponsesStream now receive streamIdleTimeoutInSeconds and initialize it internally; Gemini call sites in ChatCompletionStream and ResponsesStream pass provider.networkConfig.StreamIdleTimeoutInSeconds. SpeechStream and TranscriptionStream also move timeout setup earlier in the streaming flow.
Timeout set after successful response setup core/providers/cohere/cohere.go, core/providers/huggingface/huggingface.go, core/providers/mistral/mistral.go, core/providers/vllm/vllm.go	These streaming paths move SetStreamIdleTimeoutIfEmpty to run after successful upstream response handling and before passthrough or stream-channel creation, removing later duplicate or earlier placements.
Timeout set after headers are stored core/providers/replicate/replicate.go	Replicate streaming methods move SetStreamIdleTimeoutIfEmpty to immediately follow response-header storage in context and before passthrough or buffered channel setup.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• maximhq/bifrost#3495: Threads StreamIdleTimeoutInSeconds into shared streaming helpers and moves timeout initialization into those helpers.
• maximhq/bifrost#3522: Touches the same Gemini streaming handlers and related streaming control flow.
• No additional related PRs identified with a strong code-level connection.

Suggested reviewers: akshaydeo, danpiths

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main fix: wiring stream idle timeout in the Vertex streaming path.
Description check	✅ Passed	The description covers the summary, changes, type, affected areas, testing, breaking changes, security, and related issues.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 07-06-fix_idle_timeout_wiring_in_vertex_path

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

_{Comment @coderabbitai help to get the list of available commands.}

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

tejas ghatte seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[TejasGhatte]

• fix: idle timeout wiring in vertex path #4937 👈 (View in Graphite)
• fix: set content type header consistently for responses api requests #4935
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 5/5

Safe to merge — the changes are narrowly scoped, mechanically consistent across all affected providers, and match the pattern already established by the OpenAI and Anthropic shared handlers.

The fix correctly wires the configured idle timeout into both the Gemini/Vertex shared handlers and the large-payload passthrough path. The SetStreamIdleTimeoutIfEmpty guard means the call is idempotent and transport-layer overrides are still respected. No concurrency, interface, or pooling concerns are introduced.

No files require special attention; the changes are uniform and self-contained within each provider's streaming methods.

Important Files Changed

Filename	Overview
core/providers/vertex/vertex.go	Passes StreamIdleTimeoutInSeconds to HandleGeminiChatCompletionStream and HandleGeminiResponsesStream — the primary fix for Vertex idle-timeout wiring.
core/providers/gemini/gemini.go	Adds streamIdleTimeoutInSeconds parameter to HandleGeminiChatCompletionStream and HandleGeminiResponsesStream; moves SetStreamIdleTimeoutIfEmpty before passthrough check in SpeechStream and TranscriptionStream.
core/providers/cohere/cohere.go	Moves SetStreamIdleTimeoutIfEmpty before SetupStreamingPassthrough in ChatCompletionStream and ResponsesStream, fixing the large-payload passthrough timeout path.
core/providers/replicate/replicate.go	Same passthrough-ordering fix applied consistently across TextCompletionStream, ChatCompletionStream, ResponsesStream, ImageGenerationStream, and ImageEditStream.
core/providers/huggingface/huggingface.go	Moves SetStreamIdleTimeoutIfEmpty before passthrough check in ImageGenerationStream and ImageEditStream.
core/providers/mistral/mistral.go	Same passthrough-ordering fix in TranscriptionStream.
core/providers/vllm/vllm.go	Same passthrough-ordering fix in TranscriptionStream.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant V as VertexProvider
    participant G as HandleGeminiChatCompletionStream
    participant U as SetStreamIdleTimeoutIfEmpty
    participant P as SetupStreamingPassthrough

    Note over V,P: After fix
    V->>G: call with streamIdleTimeoutInSeconds
    G->>U: SetStreamIdleTimeoutIfEmpty(ctx, seconds)
    U-->>G: timeout set in context
    G->>P: SetupStreamingPassthrough(ctx, resp)
    Note right of P: GetStreamIdleTimeout(ctx) now returns configured value
    P-->>G: passthrough configured or skipped
    G-->>V: "chan *BifrostStreamChunk"

    Note over V,P: Before fix (Vertex path)
    V->>G: call (no timeout param)
    Note over U: SetStreamIdleTimeoutIfEmpty never called for Vertex
    G->>P: SetupStreamingPassthrough(ctx, resp)
    Note right of P: GetStreamIdleTimeout(ctx) falls back to 120s default
    P-->>G: passthrough with wrong timeout
    G-->>V: "chan *BifrostStreamChunk"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant V as VertexProvider
    participant G as HandleGeminiChatCompletionStream
    participant U as SetStreamIdleTimeoutIfEmpty
    participant P as SetupStreamingPassthrough

    Note over V,P: After fix
    V->>G: call with streamIdleTimeoutInSeconds
    G->>U: SetStreamIdleTimeoutIfEmpty(ctx, seconds)
    U-->>G: timeout set in context
    G->>P: SetupStreamingPassthrough(ctx, resp)
    Note right of P: GetStreamIdleTimeout(ctx) now returns configured value
    P-->>G: passthrough configured or skipped
    G-->>V: "chan *BifrostStreamChunk"

    Note over V,P: Before fix (Vertex path)
    V->>G: call (no timeout param)
    Note over U: SetStreamIdleTimeoutIfEmpty never called for Vertex
    G->>P: SetupStreamingPassthrough(ctx, resp)
    Note right of P: GetStreamIdleTimeout(ctx) falls back to 120s default
    P-->>G: passthrough with wrong timeout
    G-->>V: "chan *BifrostStreamChunk"

Loading

_{Reviews (2): Last reviewed commit: "fix: idle timeout wiring in vertex path" | Re-trigger Greptile}

[coderabbitai]

🧹 Nitpick comments (1)

core/providers/vertex/vertex.go (1)

977-977: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Consider a regression test for the idle-timeout wiring fix.

This PR fixes a real gap where Vertex-Gemini streaming never had an idle timeout applied. A table-driven or unit test asserting providerUtils.GetStreamIdleTimeout(ctx) is non-zero after ChatCompletionStream/ResponsesStream dispatch (for the Gemini-family branch) would guard against regressing this fix.

Also applies to: 1505-1505

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@core/providers/vertex/vertex.go` at line 977, Add a regression test for the
Gemini streaming path in vertex.go to verify the idle-timeout wiring stays
intact. Focus on the ChatCompletionStream and ResponsesStream dispatch flow in
the Gemini-family branch, and assert that
providerUtils.GetStreamIdleTimeout(ctx) is non-zero after the call. Use a
table-driven or unit test that exercises the
provider.networkConfig.StreamIdleTimeoutInSeconds path so the fix cannot
regress.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@core/providers/vertex/vertex.go`:
- Line 977: Add a regression test for the Gemini streaming path in vertex.go to
verify the idle-timeout wiring stays intact. Focus on the ChatCompletionStream
and ResponsesStream dispatch flow in the Gemini-family branch, and assert that
providerUtils.GetStreamIdleTimeout(ctx) is non-zero after the call. Use a
table-driven or unit test that exercises the
provider.networkConfig.StreamIdleTimeoutInSeconds path so the fix cannot
regress.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: dd474357-70e7-4374-92cf-0973c82ff884

📥 Commits

Reviewing files that changed from the base of the PR and between e0c541d and e2f9f2d.

📒 Files selected for processing (2)

• core/providers/gemini/gemini.go
• core/providers/vertex/vertex.go

[Pratham-Mishra04]

Merge activity

• Jul 6, 9:25 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jul 6, 9:26 AM UTC: @Pratham-Mishra04 merged this pull request with Graphite.