fix(genkit-tools): Refactor Tools API to prevent proxy 503 errors

[MichaelDoyle]

Refactored the Genkit Tools server API endpoints (/api/runAction, /api/streamAction, /api/streamTrace) to improve reliability in proxied environments like Google Cloud Shell. Shifts the "early information" contract from HTTP headers to the response body, requiring the client to parse these endpoints as streams.

Rationale

The previous implementation used manual header flushing (res.flushHeaders()) to send the X-Genkit-Trace-Id header early. In Google Cloud Shell, which routes traffic through a web proxy, this "headers-only" response followed by a long silence (while waiting for AI models) triggered a 503 Service Unavailable error (due to perceived inactivity?).

Key Changes

• Body-First Trace ID: Removed the X-Genkit-Trace-Id HTTP header and its CORS exposure. The traceId is now sent as the very first line of the response body in JSON format.
• Protocol-Level Keep-Alive: Removed manual flushHeaders() and res.statusCode = 200 calls. The first res.write() of the Trace ID chunk now naturally
  triggers header transmission and provides immediate activity (TTFB), preventing proxy timeouts.
• Standardized NDJSON: All three endpoints now consistently use a newline-delimited JSON (NDJSON) pattern, where metadata (telemetry) precedes the actual result or data chunks.
• Improved Error Handling: Unified the catch blocks to ensure an HTTP 500 status is sent if a failure occurs before the first chunk. If a failure occurs during streaming, it is appended as a final JSON chunk.
• Refined Documentation: Updated internal comments to clarify the "keep-alive" mechanism and the dependency on telemetry states.

Fixes #4796

[gemini-code-assist]

Summary of Changes

Hello @MichaelDoyle, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the Genkit Tools server API endpoints to enhance reliability in environments with restrictive proxies, such as Google Cloud Shell. The core change involves shifting the early transmission of the trace ID from HTTP headers to the response body, which helps prevent 503 Service Unavailable errors caused by perceived inactivity. The update standardizes streaming responses to use NDJSON and improves error handling for both initial request processing and ongoing streams.

Highlights

• Body-First Trace ID: The X-Genkit-Trace-Id HTTP header and its CORS exposure have been removed. The traceId is now transmitted as the very first line of the response body in JSON format.
• Protocol-Level Keep-Alive: Manual flushHeaders() and res.statusCode = 200 calls were removed. The initial res.write() of the Trace ID chunk now naturally triggers header transmission and provides immediate activity, preventing proxy timeouts.
• Standardized NDJSON: All three affected endpoints (/api/runAction, /api/streamAction, /api/streamTrace) now consistently use a newline-delimited JSON (NDJSON) pattern, with metadata (telemetry) preceding actual data chunks.
• Improved Error Handling: Unified catch blocks ensure an HTTP 500 status is sent if a failure occurs before the first chunk. If an error happens during streaming, it is appended as a final JSON chunk.
• Refined Documentation: Internal comments were updated to clarify the new 'keep-alive' mechanism and its dependency on telemetry states.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• genkit-tools/common/src/server/server.ts
  • Removed X-Genkit-Trace-Id from the CORS exposedHeaders configuration.
  • Updated /api/runAction to send the traceId as the first JSON object in the response body, rather than in an HTTP header.
  • Modified /api/runAction to use text/plain as the Content-Type and eliminated explicit res.flushHeaders() and res.statusCode settings.
  • Improved error handling for /api/runAction to conditionally set a 500 status code only if response headers have not yet been sent.
  • Refactored /api/streamAction to transmit the traceId as the initial JSON chunk in the response body and removed related header flushing logic.
  • Enhanced error handling for /api/streamAction to ensure a 500 status code is set if an error occurs before any headers are sent.
  • Adjusted /api/streamTrace to send the traceId as the first body chunk, removing the previous header-based traceId transmission and explicit header flushing.
  • Standardized error handling in /api/streamTrace to set a 500 status code if an error occurs before the initial body chunk is written.

Activity

• No human activity has been recorded on this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request successfully refactors the Genkit Tools server API endpoints (/api/runAction, /api/streamAction, /api/streamTrace) to address proxy 503 errors. This is achieved by shifting the traceId from HTTP headers to the response body, sending it as the very first line in newline-delimited JSON (NDJSON) format. This approach correctly prevents proxy timeouts by providing immediate activity (TTFB) and standardizes the streaming protocol. Error handling is also improved to ensure a 500 status code is sent if a failure occurs before the first chunk.

However, this change introduces a critical breaking change for existing clients, specifically the RuntimeManager in manager.ts. The RuntimeManager currently expects the X-Genkit-Trace-Id in HTTP headers and is not updated in this pull request to parse the response body as an NDJSON stream to extract the trace ID. This client-side update is essential for the functionality of these endpoints and needs to be addressed to ensure the system works as intended with the new server-side protocol.

[MichaelDoyle]

Dev UI side: FirebasePrivate/genkit-ui/pull/1760

[pavelgj]

doesn't this require Transfer-Encoding: chunked to work properly?

[MichaelDoyle]

Yes, you are right. Essentially, though, we are letting Nodejs/express handle this for us automatically. It ends up looking like this:

[pavelgj]

It's one of those things that just need work in practice... but it might be related to your other comment (buffering by the proxy)... there may not be enough signals to the proxy that it should not be buffering things.
chunked transfer encoding should be supported by a properly implemented proxy...

[MichaelDoyle]

One thing I noticed, the data does not "stream" in Cloud Shell - trying to figure out a path for that, but it doesn't seem like we're introducing that problem here, so we might want to punt for the next PR. I can see all the chunks, but the proxy is buffering the data.

[MichaelDoyle]

I've been testing in Cloud Shell. Here is the current state of the world with this PR:

• runAction is working properly
• streamAction - the payload is delivered all at once, instead of in chunks.
• streamTrace is not yet working at all; the connection hangs until the server disconnects