Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
edf8ac8
Implement WebSocket support for InvocationAgentServerHost (#46805)
007DXR May 12, 2026
c6bb227
Merge branch 'main' into feature/azure-ai-agentserver-invocations_ws
yulin-li May 12, 2026
c88024d
Merge branch 'main' into feature/azure-ai-agentserver-invocations_ws
yulin-li May 13, 2026
8df0573
[agentserver-invocations] Address PR review comments 18-52 for invoca…
yulin-li May 14, 2026
f31c8c9
Merge branch 'main' into feature/azure-ai-agentserver-invocations_ws
yulin-li May 14, 2026
901f907
Address Copilot AI review on PR 46841
May 14, 2026
be57e7a
Merge remote-tracking branch 'origin/main' into feature/azure-ai-agen…
May 15, 2026
265e620
Reorder CHANGELOG: error-source bullet (from main) first, WS bullets …
May 15, 2026
ca606f7
Merge branch 'main' into feature/azure-ai-agentserver-invocations_ws
yulin-li May 16, 2026
ea90b70
Address review threads 2-6 on PR 46841: move ws_ping_interval to Agen…
May 18, 2026
c49b804
Revert per-turn ws_invocation helper: WS has only one connection-scop…
May 18, 2026
db1e596
Drop README/CHANGELOG note about missing WS invocation_id and HTTP-on…
May 18, 2026
cffbb07
Merge branch 'main' into feature/azure-ai-agentserver-invocations_ws
yulin-li May 18, 2026
dc1cfc6
Fix stale module docstring: ws keep-alive is env-only now (no constru…
May 18, 2026
c562e2d
Add smoke tests for ws_bidirectional_streaming_agent sample
May 18, 2026
fe51b11
Drop unused interval kwarg from resolve_ws_ping_interval
May 18, 2026
6f529e5
Drop specific idle-timeout duration and APIM/Load Balancer mentions f…
May 18, 2026
ae06e2e
Drop stale ws_ping_interval= constructor mention from bidirectional s…
May 18, 2026
c793675
Merge branch 'main' into feature/azure-ai-agentserver-invocations_ws
yulin-li May 18, 2026
0822abe
Drop empty CHANGELOG sections to satisfy release-build verifier
May 18, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
### Features Added

- Added `_platform_headers` module with cross-cutting protocol header name constants (`x-request-id`, `x-platform-server`, `x-agent-session-id`, `x-platform-error-source`, `x-platform-error-detail`, and others). Protocol packages now import shared header name strings from core instead of maintaining their own copies.
- `AgentConfig.ws_ping_interval` — new field resolved from the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers). `0` disables; negative/non-finite values raise `ValueError` at startup. `AgentServerHost._build_hypercorn_config` wires this into Hypercorn's `websocket_ping_interval` so any protocol package serving WebSocket routes inherits keep-alive without per-package wiring.

## 2.0.0b3 (2026-04-22)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -219,13 +219,14 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF
cfg = self.config
logger.info(
"Platform environment: is_hosted=%s, agent_name=%s, agent_version=%s, "
"port=%s, session_id=%s, sse_keepalive_interval=%s",
"port=%s, session_id=%s, sse_keepalive_interval=%s, ws_ping_interval=%s",
cfg.is_hosted,
cfg.agent_name or _NOT_SET,
cfg.agent_version or _NOT_SET,
cfg.port,
cfg.session_id or _NOT_SET,
cfg.sse_keepalive_interval if cfg.sse_keepalive_interval > 0 else "disabled",
f"{cfg.ws_ping_interval}s" if cfg.ws_ping_interval > 0 else "disabled",
)
logger.info(
"Connectivity: project_endpoint=%s, otlp_endpoint=%s, appinsights_configured=%s",
Expand Down Expand Up @@ -417,6 +418,11 @@ def _build_hypercorn_config(self, host: str, port: int) -> object:
config.graceful_timeout = float(self._graceful_shutdown_timeout)
# Spec requires HTTP/1.1 only — disable HTTP/2
config.h2_max_concurrent_streams = 0
# WebSocket Ping/Pong keep-alive (RFC 6455 opcodes 0x9/0xA).
# ``0`` (disabled) maps to Hypercorn's ``None`` sentinel; any
# positive value is sent verbatim to Hypercorn.
ws_ping = self.config.ws_ping_interval
config.websocket_ping_interval = ws_ping if ws_ping > 0 else None # type: ignore[attr-defined]
# Access logging
if self._access_log is not None:
config.accesslog = self._access_log # type: ignore[assignment]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,11 @@
_ENV_APPLICATIONINSIGHTS_CONNECTION_STRING = "APPLICATIONINSIGHTS_CONNECTION_STRING"
_ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT"
_ENV_SSE_KEEPALIVE_INTERVAL = "SSE_KEEPALIVE_INTERVAL"
_ENV_WS_KEEPALIVE_INTERVAL = "WS_KEEPALIVE_INTERVAL"

_DEFAULT_PORT = 8088
_DEFAULT_SSE_KEEPALIVE_INTERVAL = 0
_DEFAULT_WS_PING_INTERVAL = 0.0


# ======================================================================
Expand Down Expand Up @@ -64,6 +66,8 @@ class AgentConfig: # pylint: disable=too-many-instance-attributes
:param appinsights_connection_string: Application Insights connection string.
:param otlp_endpoint: OTLP exporter endpoint.
:param sse_keepalive_interval: SSE keep-alive interval in seconds (0 = disabled).
:param ws_ping_interval: WebSocket protocol Ping interval in seconds
(``0`` disables keep-alive).
"""

def __init__(
Expand All @@ -80,6 +84,7 @@ def __init__(
appinsights_connection_string: str,
otlp_endpoint: str,
sse_keepalive_interval: int,
ws_ping_interval: float = 0.0,
) -> None:
self.agent_name = agent_name
self.agent_version = agent_version
Expand All @@ -92,6 +97,7 @@ def __init__(
self.appinsights_connection_string = appinsights_connection_string
self.otlp_endpoint = otlp_endpoint
self.sse_keepalive_interval = sse_keepalive_interval
self.ws_ping_interval = ws_ping_interval

@classmethod
def from_env(cls) -> Self:
Expand Down Expand Up @@ -123,6 +129,7 @@ def from_env(cls) -> Self:
_ENV_APPLICATIONINSIGHTS_CONNECTION_STRING, ""),
otlp_endpoint=os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT, ""),
sse_keepalive_interval=resolve_sse_keepalive_interval(None),
ws_ping_interval=resolve_ws_ping_interval(),
)


Expand Down Expand Up @@ -322,3 +329,40 @@ def resolve_otlp_endpoint() -> Optional[str]:
"""
value = os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT, "")
return value if value else None


def resolve_ws_ping_interval() -> float:
"""Resolve the WebSocket Ping/Pong keep-alive interval from the env var.

Reads the ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected
by AgentService into hosted-agent containers) and returns the parsed
value in seconds. ``0`` (or unset) disables keep-alive.

The keep-alive interval is intentionally env-only — there is no
programmatic constructor argument to override it. Hosted agents
inherit the platform-injected value automatically; in tests, set the
env var (e.g. via ``monkeypatch.setenv``).

:return: The resolved interval in seconds (``0`` means disabled).
:rtype: float
:raises ValueError: If the env var is set but not parseable as a
non-negative finite number.
"""
import math # local import — only used here

env_raw = os.environ.get(_ENV_WS_KEEPALIVE_INTERVAL)
if env_raw is None or env_raw == "":
return _DEFAULT_WS_PING_INTERVAL
try:
resolved = float(env_raw)
except ValueError as exc:
raise ValueError(
f"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: "
f"{env_raw!r} (expected a non-negative number)"
) from exc
if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0:
raise ValueError(
f"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: "
f"{env_raw!r} (expected a non-negative finite number)"
)
return resolved
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,10 @@
### Features Added

- Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics.
- WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404.
- WebSocket Ping/Pong keep-alive — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers; surfaced on `app.config.ws_ping_interval` in `azure-ai-agentserver-core>=2.0.0b4`). `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval` by `AgentServerHost._build_hypercorn_config`.
- WebSocket telemetry — single connection-scoped `websocket_session` OpenTelemetry span per WS connection, plus a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation.
- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages).

### Other Changes

Expand Down
72 changes: 71 additions & 1 deletion sdk/agentserver/azure-ai-agentserver-invocations/README.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
# Azure AI Agent Server Invocations client library for Python

The `azure-ai-agentserver-invocations` package provides the invocation protocol endpoints for Azure AI Hosted Agent containers. It plugs into the [`azure-ai-agentserver-core`](https://pypi.org/project/azure-ai-agentserver-core/) host framework and adds the full invocation lifecycle: `POST /invocations`, `GET /invocations/{id}`, `POST /invocations/{id}/cancel`, and `GET /invocations/docs/openapi.json`.
The `azure-ai-agentserver-invocations` package provides the invocation protocol endpoints for Azure AI Hosted Agent containers. It plugs into the [`azure-ai-agentserver-core`](https://pypi.org/project/azure-ai-agentserver-core/) host framework and supports two transports on the same host:

- **HTTP** (`invocations` protocol) — `POST /invocations`, `GET /invocations/{id}`, `POST /invocations/{id}/cancel`, `GET /invocations/docs/openapi.json`.
- **WebSocket** (`invocations_ws` protocol) — full-duplex streaming at `/invocations_ws`, registered with `@app.ws_handler`.

## Getting started

Expand All @@ -25,6 +28,7 @@ This automatically installs `azure-ai-agentserver-core` as a dependency.
- `@app.invoke_handler` — **Required.** Handles `POST /invocations`.
- `@app.get_invocation_handler` — Optional. Handles `GET /invocations/{id}`.
- `@app.cancel_invocation_handler` — Optional. Handles `POST /invocations/{id}/cancel`.
- `@app.ws_handler` — Optional. Handles WebSocket connections at `/invocations_ws`.

### Protocol endpoints

Expand All @@ -34,6 +38,7 @@ This automatically installs `azure-ai-agentserver-core` as a dependency.
| `GET` | `/invocations/{invocation_id}` | No | Retrieve invocation status or result |
| `POST` | `/invocations/{invocation_id}/cancel` | No | Cancel a running invocation |
| `GET` | `/invocations/docs/openapi.json` | No | Serve the agent's OpenAPI 3.x spec |
| `WS` | `/invocations_ws` | No | Full-duplex WebSocket transport (`invocations_ws` protocol) |

### Request and response headers

Expand Down Expand Up @@ -182,6 +187,69 @@ app = InvocationAgentServerHost(openapi_spec={
})
```

## WebSocket protocol (`invocations_ws`)

The same `InvocationAgentServerHost` object also exposes a WebSocket transport at `/invocations_ws`. Container authors do not install or import a second package — registering an `@app.ws_handler` is the only step. A multi-protocol agent shares one host, one session, and one process.

### Quick start

```python
from azure.ai.agentserver.invocations import InvocationAgentServerHost
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.websockets import WebSocket

app = InvocationAgentServerHost()


@app.invoke_handler # POST /invocations (HTTP)
async def invoke(request: Request) -> Response:
payload = await request.json()
return JSONResponse({"echo": payload})


@app.ws_handler # /invocations_ws (WebSocket)
async def ws(websocket: WebSocket) -> None:
async for message in websocket.iter_text():
await websocket.send_text(message)


app.run()
```

### What the SDK does for `@app.ws_handler`

- Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`.
- Calls `await websocket.accept()` before invoking your handler.
- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers). Set the value to `0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across upstream proxy / load-balancer idle timeouts without any extra application traffic.
- Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`.
- Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end.
- Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`.

### Per-connection tracing

A WebSocket connection is wrapped by the SDK in a single connection-scoped `websocket_session` OpenTelemetry span. The span carries the GenAI semantic-convention attributes plus `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Any child spans your handler opens — e.g. via `opentelemetry.trace.get_tracer(...).start_as_current_span(...)` — are automatically parented to the connection span.

### Handler signature

The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. The full WebSocket API — `iter_text`, `iter_bytes`, `iter_json`, `send_text`, `send_bytes`, `send_json`, `close`, `headers`, `query_params`, `client`, `state` — is available, so application protocols on top of `invocations_ws` are entirely under your control.

[starlette-ws]: https://www.starlette.io/websockets/

### Reference: configuration

| Environment variable | Default | Description |
|---|---|---|
| `WS_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected WebSocket Ping interval, in seconds. `0` disables keep-alive. Surfaced on `app.config.ws_ping_interval` and wired into Hypercorn's `websocket_ping_interval` by `AgentServerHost`. |

### Reference: close codes

| Close code | Meaning |
|---|---|
| `1000` | Handler returned cleanly (normal close). |
| `1011` | Handler raised an unhandled exception (mapped by the SDK). |
| `4000`-`4999` | Application-defined codes (set by the handler via `await websocket.close(code=...)` — surfaced unchanged to the client). |

## Troubleshooting

### Reporting issues
Expand All @@ -196,6 +264,8 @@ Visit the [Samples](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/
|---|---|
| [simple_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/simple_invoke_agent/) | Minimal synchronous request-response |
| [async_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/async_invoke_agent/) | Long-running operations with polling and cancellation |
| [ws_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/) | Combined `POST /invocations` (HTTP) and `/invocations_ws` (WebSocket) host |
| [ws_bidirectional_streaming_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/) | Full-duplex `/invocations_ws` agent: concurrent token streams + mid-flight cancel (relies on the SDK's WS protocol Ping/Pong keep-alive, not application-level heartbeats) |

## Contributing

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,25 @@ class InvocationConstants:
ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations.session_id"
ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations.error.code"
ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations.error.message"


class InvocationsWSConstants:
"""invocations_ws (WebSocket) protocol constants.

Route, span attribute keys, and ping/pong defaults for the
WebSocket endpoint hosted alongside the HTTP invocations protocol.
"""

# Route
ROUTE_PATH = "/invocations_ws"

# Close codes (RFC 6455)
CLOSE_NORMAL = 1000 # handler returned cleanly
CLOSE_INTERNAL_ERROR = 1011 # handler raised an unhandled exception

# Span attribute keys
ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations_ws.session_id"
ATTR_SPAN_CLOSE_CODE = "azure.ai.agentserver.invocations_ws.close_code"
ATTR_SPAN_DURATION_MS = "azure.ai.agentserver.invocations_ws.duration_ms"
ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code"
ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations_ws.error.message"
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
)

from ._constants import InvocationConstants
from ._invocation_ws import _WSHandlerMixin

logger = logging.getLogger("azure.ai.agentserver")

Expand Down Expand Up @@ -147,13 +148,18 @@ def _sanitize_id(value: str, fallback: str) -> str:
return value


class InvocationAgentServerHost(AgentServerHost):
class InvocationAgentServerHost(_WSHandlerMixin, AgentServerHost):
"""Invocation protocol host for Azure AI Hosted Agents.

A :class:`~azure.ai.agentserver.core.AgentServerHost` subclass that adds
the invocation protocol endpoints. Use the decorator methods to wire
handler functions to the endpoints.

The same host object also exposes the ``invocations_ws`` (WebSocket)
transport at :data:`/invocations_ws` — register a handler with the
:meth:`ws_handler` decorator. Multi-protocol agents share a single
host, session, and process.

For multi-protocol agents, compose via cooperative inheritance::

class MyHost(InvocationAgentServerHost, ResponsesAgentServerHost):
Expand All @@ -162,13 +168,19 @@ class MyHost(InvocationAgentServerHost, ResponsesAgentServerHost):
Usage::

from azure.ai.agentserver.invocations import InvocationAgentServerHost
from starlette.websockets import WebSocket

app = InvocationAgentServerHost()

@app.invoke_handler
@app.invoke_handler # POST /invocations
async def handle(request):
return JSONResponse({"ok": True})

@app.ws_handler # /invocations_ws
async def ws(websocket: WebSocket) -> None:
async for message in websocket.iter_text():
await websocket.send_text(message)

app.run()

:param openapi_spec: Optional OpenAPI spec dict. When provided, the spec
Expand All @@ -189,8 +201,13 @@ def __init__(
self._cancel_invocation_fn: Optional[Callable] = None
self._openapi_spec = openapi_spec

# Initialise WS handler slots (no parameters — the keep-alive
# interval lives on ``AgentConfig`` and is wired into Hypercorn
# by ``AgentServerHost._build_hypercorn_config``).
self._init_ws_state()

# Build invocation routes and pass to parent via routes kwarg
invocation_routes = [
invocation_routes: list[Any] = [
Route(
"/invocations/docs/openapi.json",
self._get_openapi_spec_endpoint,
Expand Down
Loading