google / google/adk-python

Graceful WebSocket drain on SIGTERM for api_server/web commands

Đang mở
#4,876 1 bình luận 0 reaction 2 người được giao Được @Jacksunwei nhận Xem trên GitHub
needs review web
Ngôn ngữ chính
Python
Star
21.5k
Fork
4k
Merge trung bình
1 ngày 14 giờ
Pull request đã merge (30 ngày)
37

Mô tả

## 🔴 Required Information

### Is your feature request related to a specific problem?

When deploying `adk api_server` or `adk web` on Cloud Run (or any container orchestrator), the platform sends SIGTERM on deployment, scale-down, or preemption. Uvicorn's default SIGTERM handler immediately closes WebSocket connections with close code 1012 (Service Restart), which drops active `/run_live` streaming sessions mid-conversation.

For agents using Gemini Live models (`gemini-live-*-native-audio`), every user interaction is a long-lived WebSocket connection. A routine deployment kills all active conversations instantly with no opportunity for the client to reconnect or the model to finish its response.

### Describe the Solution You'd Like

The `adk api_server` and `adk web` commands should drain active WebSocket connections before shutting down:

1. On SIGTERM, stop accepting new connections (close listeners).
2. Wait for in-flight WebSocket connections to finish naturally, up to a configurable timeout.
3. Then let uvicorn shut down cleanly.

This ensures zero dropped conversations during deployments when the orchestrator routes new connections to the new instance while the old one drains.

### Impact on your work

We run live audio agents on Cloud Run. Every deployment drops all active calls. We currently maintain a wrapper around `adk api_server` solely to add graceful drain behavior. This is critical infrastructure for any production deployment of live agents. Without it, deployments during business hours are disruptive.

### Willingness to contribute

Yes.

---

## 🟡 Recommended Information

### Describe Alternatives You've Considered

We wrote a custom CLI wrapper that intercepts SIGTERM, closes the server listeners, polls `uvicorn.Server.server_state.connections` until all WebSocket connections finish (or a timeout expires), then sets `server.should_exit = True`. This works but requires maintaining a parallel entry point and duplicating upstream's server startup logic.

### Proposed API / Implementation

Replace the current `uvicorn.run()` / `server.serve()` call in the `api_server` and `web` commands with a drain-aware server loop. Rough sketch:

```python
async def _serve_with_graceful_drain(server: uvicorn.Server, drain_timeout: float) -> None:
loop = asyncio.get_running_loop()
serve_task = asyncio.create_task(server._serve())

shutdown_event = asyncio.Event()

def _on_signal(sig: int) -> None:
for srv in server.servers:
srv.close() # stop accepting new connections
shutdown_event.set()

for sig in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(sig, _on_signal, sig)

wait_task = asyncio.create_task(shutdown_event.wait())
done, _ = await asyncio.wait(
[serve_task, wait_task], return_when=asyncio.FIRST_COMPLETED
)
if serve_task in done:
wait_task.cancel()
return

# Drain: wait for active connections to finish
deadline = time.monotonic() + drain_timeout
while server.server_state.connections and time.monotonic() < deadline:
await asyncio.sleep(0.1)

server.should_exit = True
await serve_task
```

A `--drain_timeout` CLI flag (defaulting to something reasonable like 300s) would let users tune it. Cloud Run's request timeout is the natural upper bound.

### Additional Context

- Cloud Run keeps the instance alive for active requests after SIGTERM (up to the configured request timeout), so the drain window is already provided by the platform — the server just needs to use it instead of exiting immediately.
- This affects any WebSocket-based command (`api_server`, `web`), not just live audio agents. SSE streaming sessions would also benefit.
- Related: #2620 (agent-level interrupt handling) covers cancelling agent tasks, this is about server-level connection draining.

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.