feat(py): implement Reflection API v2 with WebSocket and JSON-RPC 2.0

This implements the Reflection API v2 for Python following RFC #4211.

Architecture Change:
- V1 (HTTP server): Runtime hosts an HTTP server, CLI/DevUI connect to it
- V2 (WebSocket client): CLI hosts a WebSocket server, Runtimes connect outbound

The v2 API reverses the connection direction, allowing better support for
bidirectional actions and environments where binding a port is impractical.

Module Reorganization:
- reflection.py: New v2 WebSocket client (primary when enabled)
- reflection_v1.py: Existing HTTP server implementation (default)

V2 Implementation:
- ReflectionClientV2 connects to a runtime manager via WebSocket
- JSON-RPC 2.0 protocol for all communication
- Supports: listActions, runAction, cancelAction, listValues
- Streaming: runActionState and streamChunk notifications
- Auto-reconnection with exponential backoff (1s to 60s max)
- Proper task-based cancellation via asyncio.current_task().cancel()

Activation:
- V2 is activated when GENKIT_REFLECTION_V2_SERVER env var is set
- V1 remains the default when env var is not set

Dependencies:
- websockets>=15.0 as core dependency

Tests:
- 25 tests for v2 WebSocket client
- 7 tests for v1 HTTP server

See: RFC #4211

Author: yesudeep

py/packages/genkit/pyproject.toml

@@ -54,6 +54,7 @@ dependencies = [
   "uvloop>=0.21.0; sys_platform != 'win32'",
   "anyio>=4.9.0",
   "opentelemetry-instrumentation-logging>=0.60b1",
+  "websockets>=15.0",
 ]
 description = "Genkit AI Framework"
 license = "Apache-2.0"

py/packages/genkit/src/genkit/ai/_base_async.py

@@ -30,7 +30,11 @@
 from genkit.core.environment import is_dev_environment
 from genkit.core.logging import get_logger
 from genkit.core.plugin import Plugin
-from genkit.core.reflection import create_reflection_asgi_app
+from genkit.core.reflection import (
+    ReflectionClientV2,
+    get_reflection_v2_url,
+)
+from genkit.core.reflection_v1 import create_reflection_asgi_app
 from genkit.core.registry import Registry
 from genkit.web.manager._ports import find_free_port_sync
 
@@ -145,8 +149,6 @@ async def run_user_coro_wrapper() -> None:
                 finally:
                     user_task_finished_event.set()
 
-            reflection_server = _make_reflection_server(self.registry, server_spec)
-
             # Setup signal handlers for graceful shutdown (parity with JS)
 
             # Actually, anyio.run handles Ctrl+C (SIGINT) by raising KeyboardInterrupt/CancelledError
@@ -162,65 +164,94 @@ async def handle_sigterm(tg_to_cancel: anyio.abc.TaskGroup) -> None:  # type: ig
                         return
 
             try:
-                # Use lazy_write=True to prevent race condition where file exists before server is up
-                async with RuntimeManager(server_spec, lazy_write=True) as runtime_manager:
-                    # We use anyio.TaskGroup because it is compatible with
-                    # asyncio's event loop and works with Python 3.10
-                    # (asyncio.TaskGroup was added in 3.11, and we can switch to
-                    # that when we drop support for 3.10).
+                # Check if Reflection API v2 is enabled
+                v2_url = get_reflection_v2_url()
+
+                if v2_url:
+                    # Reflection API v2: Use WebSocket client connecting to runtime manager
+                    client = ReflectionClientV2(self.registry, v2_url)
+
                     async with anyio.create_task_group() as tg:
-                        # Start reflection server in the background.
-                        tg.start_soon(reflection_server.serve, name='genkit-reflection-server')
-                        await logger.ainfo(f'Started Genkit reflection server at {server_spec.url}')
+                        # Start v2 client in background (handles its own reconnection)
+                        tg.start_soon(client.run, name='genkit-reflection-v2-client')
+                        await logger.ainfo(f'Started Genkit Reflection v2 client connecting to {v2_url}')
 
                         # Start SIGTERM handler
                         tg.start_soon(handle_sigterm, tg, name='genkit-sigterm-handler')
 
-                        # Wait for server to be responsive
-                        # We need to loop and poll the health endpoint or wait for uvicorn to be ready
-                        # Since uvicorn run is blocking (but we are in a task), we can't easily hook into its startup
-                        # unless we use uvicorn's server object directly which we do.
-                        # reflection_server.started is set when uvicorn starts.
-
-                        # Simple polling loop
-
-                        max_retries = 20  # 2 seconds total roughly
-                        for _i in range(max_retries):
-                            try:
-                                # TODO(#4334): Use async http client if available to avoid blocking loop?
-                                # But we are in dev mode, so maybe okay.
-                                # Actually we should use anyio.to_thread to avoid blocking event loop
-                                # or assume standard lib urllib is fast enough for localhost.
-
-                                # Using sync urllib in async loop blocks the loop!
-                                # We must use anyio.to_thread or a non-blocking check.
-                                # But let's check if reflection_server object has a 'started' flag we can trust.
-                                # uvicorn.Server has 'started' attribute but it might be internal state.
-
-                                # Let's stick to simple polling with to_thread for safety
-                                def check_health() -> bool:
-                                    health_url = f'{server_spec.url}/api/__health'
-                                    with urllib.request.urlopen(health_url, timeout=0.5) as response:
-                                        return response.status == 200
-
-                                is_healthy = await anyio.to_thread.run_sync(check_health)  # type: ignore[attr-defined]
-                                if is_healthy:
-                                    break
-                            except Exception:
-                                await anyio.sleep(0.1)
-                        else:
-                            logger.warning(f'Reflection server at {server_spec.url} did not become healthy in time.')
-
-                        # Now write the file (or verify it persisted)
-                        _ = runtime_manager.write_runtime_file()
-
-                        # Start the (potentially short-lived) user coroutine wrapper
+                        # Start the user coroutine
                         tg.start_soon(run_user_coro_wrapper, name='genkit-user-coroutine')
                         await logger.ainfo('Started Genkit user coroutine')
 
                         # Block here until the task group is canceled (e.g. Ctrl+C)
-                        # or a task raises an unhandled exception. It should not
-                        # exit just because the user coroutine finishes.
+                        # or a task raises an unhandled exception
+
+                else:
+                    # Reflection API v1: Start HTTP server
+                    reflection_server = _make_reflection_server(self.registry, server_spec)
+
+                    # Use lazy_write=True to prevent race condition where file exists before server is up
+                    async with RuntimeManager(server_spec, lazy_write=True) as runtime_manager:
+                        # We use anyio.TaskGroup because it is compatible with
+                        # asyncio's event loop and works with Python 3.10
+                        # (asyncio.TaskGroup was added in 3.11, and we can switch to
+                        # that when we drop support for 3.10).
+                        async with anyio.create_task_group() as tg:
+                            # Start reflection server in the background.
+                            tg.start_soon(reflection_server.serve, name='genkit-reflection-server')
+                            await logger.ainfo(f'Started Genkit reflection server at {server_spec.url}')
+
+                            # Start SIGTERM handler
+                            tg.start_soon(handle_sigterm, tg, name='genkit-sigterm-handler')
+
+                            # Wait for server to be responsive
+                            # We need to loop and poll the health endpoint or wait for uvicorn to be ready
+                            # Since uvicorn run is blocking (but we are in a task), we can't
+                            # easily hook into its startup
+                            # unless we use uvicorn's server object directly which we do.
+                            # reflection_server.started is set when uvicorn starts.
+
+                            # Simple polling loop
+
+                            max_retries = 20  # 2 seconds total roughly
+                            for _i in range(max_retries):
+                                try:
+                                    # TODO(#4334): Use async http client if available to avoid blocking loop?
+                                    # But we are in dev mode, so maybe okay.
+                                    # Actually we should use anyio.to_thread to avoid blocking event loop
+                                    # or assume standard lib urllib is fast enough for localhost.
+
+                                    # Using sync urllib in async loop blocks the loop!
+                                    # We must use anyio.to_thread or a non-blocking check.
+                                    # But let's check if reflection_server object has a 'started' flag we can trust.
+                                    # uvicorn.Server has 'started' attribute but it might be internal state.
+
+                                    # Let's stick to simple polling with to_thread for safety
+                                    def check_health() -> bool:
+                                        health_url = f'{server_spec.url}/api/__health'
+                                        with urllib.request.urlopen(health_url, timeout=0.5) as response:
+                                            return response.status == 200
+
+                                    is_healthy = await anyio.to_thread.run_sync(check_health)  # type: ignore[attr-defined]
+                                    if is_healthy:
+                                        break
+                                except Exception:
+                                    await anyio.sleep(0.1)
+                            else:
+                                logger.warning(
+                                    f'Reflection server at {server_spec.url} did not become healthy in time.'
+                                )
+
+                            # Now write the file (or verify it persisted)
+                            _ = runtime_manager.write_runtime_file()
+
+                            # Start the (potentially short-lived) user coroutine wrapper
+                            tg.start_soon(run_user_coro_wrapper, name='genkit-user-coroutine')
+                            await logger.ainfo('Started Genkit user coroutine')
+
+                            # Block here until the task group is canceled (e.g. Ctrl+C)
+                            # or a task raises an unhandled exception. It should not
+                            # exit just because the user coroutine finishes.
 
             except anyio.get_cancelled_exc_class():
                 logger.info('Development server task group cancelled (e.g., Ctrl+C).')

py/packages/genkit/src/genkit/core/constants.py

@@ -23,3 +23,7 @@
 GENKIT_VERSION = DEFAULT_GENKIT_VERSION
 
 GENKIT_CLIENT_HEADER = f'genkit-python/{DEFAULT_GENKIT_VERSION}'
+
+# Reflection API specification version.
+# This should match the value in JS (genkit-tools).
+GENKIT_REFLECTION_API_SPEC_VERSION = 1