fix(py): race condition when starting dev server, implement reflection api parity and action cancellation

- Fixed race condition where the runtime metadata was written before the server was ready.
- Implemented action cancellation support in the Reflection API.
- Added endpoints for retrieving registered values and environment info for JS SDK parity.
- Improved signal handling and server startup/shutdown logic.
- Updated samples to use standard startup commands and dependencies.
- Fixed rate limit issues in prompt demo by updating model configuration.
- Improved error message formatting and performed general code cleanup.
- Fixed the xai-hello sample run.sh
- Fixed xai-hello dependency
- Updated readmes to include more information about obtaining keys.

Author: yesudeep

py/bin/run_sample

@@ -53,10 +53,39 @@ elif command -v go &> /dev/null; then
     fi
 fi
 
+# Check if fzf is available or install it
+FZF_CMD=""
+if command -v fzf &> /dev/null; then
+    FZF_CMD="fzf"
+else
+    read -p "fzf not found. Would you like to install it for a better experience (includes preview)? (y/N): " -n 1 -r
+    echo
+    if [[ $REPLY =~ ^[Yy]$ ]]; then
+        if command -v brew &> /dev/null; then
+           echo "Installing fzf via brew..."
+           brew install fzf
+           FZF_CMD="fzf"
+        elif command -v go &> /dev/null; then
+            echo "Installing fzf via go..."
+            if go install github.com/junegunn/fzf@latest; then
+                FZF_CMD="$HOME/go/bin/fzf"
+            fi
+        else
+            echo "Neither brew nor go found. Skipping fzf installation."
+        fi
+    fi
+fi
+
 if [[ $# -lt 1 ]]; then
     SAMPLES=$(list_samples)
 
-    if [[ -n "$GUM_CMD" ]]; then
+    if [[ -n "$FZF_CMD" ]]; then
+        # Use fzf with preview window if available
+        SAMPLE_NAME=$(echo "$SAMPLES" | $FZF_CMD --preview "if [ -f $SAMPLES_DIR/{}/README.md ]; then cat $SAMPLES_DIR/{}/README.md; else echo 'No README found.'; fi" --preview-window=right:60%:wrap --height=100% --reverse --header="Select a sample to run")
+        if [[ -z "$SAMPLE_NAME" ]]; then
+            exit 0
+        fi
+    elif [[ -n "$GUM_CMD" ]]; then
         SAMPLE_NAME=$($GUM_CMD filter --placeholder "Select a sample to run..." <<< "$SAMPLES")
         if [[ -z "$SAMPLE_NAME" ]]; then
             exit 0
@@ -85,6 +114,15 @@ fi
 
 cd "$SAMPLE_DIR"
 
+if [[ -f "./README.md" ]]; then
+    if [[ -n "$GUM_CMD" ]]; then
+        $GUM_CMD format -t markdown < ./README.md
+    else
+        cat ./README.md
+    fi
+    echo ""
+fi
+
 if [[ -f "./run.sh" ]]; then
     exec ./run.sh "$@"
 else

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

@@ -16,6 +16,9 @@
 
 """Asynchronous server gateway interface implementation for Genkit."""
 
+import signal
+import urllib.error
+import urllib.request
 from collections.abc import Coroutine
 from typing import Any, TypeVar
 
@@ -139,8 +142,23 @@ async def run_user_coro_wrapper():
 
             reflection_server = _make_reflection_server(self.registry, spec)
 
+            # Setup signal handlers for graceful shutdown (parity with JS)
+
+            # Actually, anyio.run handles Ctrl+C (SIGINT) by raising KeyboardInterrupt/CancelledError
+            # For SIGTERM, we might need to be explicit if we run in a container/process manager.
+            # JS uses: process.on('SIGTERM', shutdown); process.on('SIGINT', shutdown);
+
+            # Since anyio/asyncio handles SIGINT well, let's add a task to catch SIGTERM
+            async def handle_sigterm(tg_to_cancel):
+                with anyio.open_signal_receiver(signal.SIGTERM) as signals:
+                    async for signum in signals:
+                        logger.info('Received SIGTERM, cancelling tasks...')
+                        tg_to_cancel.cancel_scope.cancel()
+                        return
+
             try:
-                async with RuntimeManager(spec):
+                # Use lazy_write=True to prevent race condition where file exists before server is up
+                async with RuntimeManager(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
@@ -150,6 +168,46 @@ async def run_user_coro_wrapper():
                         tg.start_soon(reflection_server.serve, name='genkit-reflection-server')
                         await logger.ainfo(f'Started Genkit reflection server at {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: 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():
+                                    with urllib.request.urlopen(f'{spec.url}/api/__health', timeout=0.5) as response:
+                                        return response.status == 200
+
+                                is_healthy = await anyio.to_thread.run_sync(check_health)
+                                if is_healthy:
+                                    break
+                            except Exception:
+                                await anyio.sleep(0.1)
+                        else:
+                            logger.warning(f'Reflection server at {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')

py/packages/genkit/src/genkit/ai/_runtime.py

@@ -146,29 +146,38 @@ class RuntimeManager:
     that the context manager exits cleanly and allows exceptions to propagate.
     """
 
-    def __init__(self, spec: ServerSpec, runtime_dir: str | Path | None = None):
+    def __init__(
+        self,
+        spec: ServerSpec,
+        runtime_dir: str | Path | None = None,
+        lazy_write: bool = False,
+    ):
         """Initialize the RuntimeManager.
 
         Args:
             spec: The server specification for the reflection server.
             runtime_dir: The directory to store the runtime file in.
                          Defaults to .genkit/runtimes in the current directory.
+            lazy_write: If True, the runtime file will not be written immediately
+                        on context entry. It must be written manually by calling
+                        write_runtime_file().
         """
         self.spec = spec
         if runtime_dir is None:
             self._runtime_dir = Path(os.getcwd()) / DEFAULT_RUNTIME_DIR_NAME
         else:
             self._runtime_dir = Path(runtime_dir)
 
+        self.lazy_write = lazy_write
         self._runtime_file_path: Path | None = None
 
     async def __aenter__(self) -> RuntimeManager:
         """Create the runtime directory and file."""
         try:
             await logger.adebug(f'Ensuring runtime directory exists: {self._runtime_dir}')
             self._runtime_dir.mkdir(parents=True, exist_ok=True)
-            self._runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec)
-            _register_atexit_cleanup_handler(self._runtime_file_path)
+            if not self.lazy_write:
+                self.write_runtime_file()
 
         except Exception as e:
             logger.error(f'Failed to initialize runtime file: {e}', exc_info=True)
@@ -189,18 +198,19 @@ async def __aexit__(
             exc_tb: The traceback of the exception that occurred.
 
         Returns:
-            True if cleanup was successful, False if cleanup failed.
+            False to indicate exceptions should propagate.
         """
+        self.cleanup()
         await logger.adebug('RuntimeManager async context exited.')
-        return True
+        return False
 
     def __enter__(self) -> RuntimeManager:
         """Synchronous entry point: Create the runtime directory and file."""
         try:
             logger.debug(f'[sync] Ensuring runtime directory exists: {self._runtime_dir}')
             self._runtime_dir.mkdir(parents=True, exist_ok=True)
-            self._runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec)
-            _register_atexit_cleanup_handler(self._runtime_file_path)
+            if not self.lazy_write:
+                self.write_runtime_file()
 
         except Exception as e:
             logger.error(f'[sync] Failed to initialize runtime file: {e}', exc_info=True)
@@ -219,5 +229,25 @@ def __exit__(self, exc_type: Exception | None, exc_val: Exception | None, exc_tb
         Returns:
             False to indicate exceptions should propagate.
         """
+        self.cleanup()
         logger.debug('RuntimeManager sync context exited.')
         return False
+
+    def write_runtime_file(self) -> Path:
+        """Calculates metadata, creates filename, and writes the runtime file.
+
+        Returns:
+            The Path object of the created file.
+        """
+        if self._runtime_file_path:
+            return self._runtime_file_path
+
+        self._runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec)
+        _register_atexit_cleanup_handler(self._runtime_file_path)
+        return self._runtime_file_path
+
+    def cleanup(self) -> None:
+        """Explicitly cleanup the runtime file."""
+        if self._runtime_file_path:
+            logger.debug(f'Cleaning up runtime file: {self._runtime_file_path}')
+            _remove_file(self._runtime_file_path)

py/packages/genkit/src/genkit/core/action/_action.py

@@ -124,6 +124,7 @@ def __init__(
         self,
         on_chunk: StreamingCallback | None = None,
         context: dict[str, Any] | None = None,
+        on_trace_start: Callable[[str], None] | None = None,
     ):
         """Initializes an ActionRunContext instance.
 
@@ -136,9 +137,12 @@ def __init__(
             context: An optional dictionary containing context data to be made
                      available within the action execution. Defaults to an empty
                      dictionary.
+            on_trace_start: A callable to be invoked with the trace ID when
+                            the trace is started.
         """
         self._on_chunk = on_chunk if on_chunk is not None else noop_streaming_callback
         self._context = context if context is not None else {}
+        self._on_trace_start = on_trace_start if on_trace_start else lambda _: None
 
     @property
     def context(self) -> dict[str, Any]:
@@ -302,6 +306,7 @@ async def arun(
         input: Any = None,
         on_chunk: StreamingCallback | None = None,
         context: dict[str, Any] | None = None,
+        on_trace_start: Callable[[str], None] | None = None,
         telemetry_labels: dict[str, Any] | None = None,
     ) -> ActionResponse:
         """Executes the action asynchronously with the given input.
@@ -331,14 +336,15 @@ async def arun(
 
         return await self._afn(
             input,
-            ActionRunContext(on_chunk=on_chunk, context=_action_context.get(None)),
+            ActionRunContext(on_chunk=on_chunk, context=_action_context.get(None), on_trace_start=on_trace_start),
         )
 
     async def arun_raw(
         self,
         raw_input: Any,
         on_chunk: StreamingCallback | None = None,
         context: dict[str, Any] | None = None,
+        on_trace_start: Callable[[str], None] | None = None,
         telemetry_labels: dict[str, Any] | None = None,
     ):
         """Executes the action asynchronously with raw, unvalidated input.
@@ -367,6 +373,7 @@ async def arun_raw(
             input=input_action,
             on_chunk=on_chunk,
             context=context,
+            on_trace_start=on_trace_start,
             telemetry_labels=telemetry_labels,
         )
 
@@ -501,6 +508,7 @@ async def async_tracing_wrapper(input: Any | None, ctx: ActionRunContext) -> Act
         afn = ensure_async(fn)
         with tracer.start_as_current_span(name) as span:
             trace_id = str(span.get_span_context().trace_id)
+            ctx._on_trace_start(trace_id)
             record_input_metadata(
                 span=span,
                 kind=kind,
@@ -541,6 +549,7 @@ def sync_tracing_wrapper(input: Any | None, ctx: ActionRunContext) -> ActionResp
         """
         with tracer.start_as_current_span(name) as span:
             trace_id = str(span.get_span_context().trace_id)
+            ctx._on_trace_start(trace_id)
             record_input_metadata(
                 span=span,
                 kind=kind,

py/packages/genkit/src/genkit/core/error.py

@@ -79,17 +79,17 @@ def __init__(
             detail: Optional detail information.
             source: Optional source of the error.
         """
-        source_prefix = f'{source}: ' if source else ''
-        super().__init__(f'{source_prefix}{status}: {message}')
-        self.original_message = message
-
         self.status = status
         if not self.status and isinstance(cause, GenkitError):
             self.status = cause.status
 
         if not self.status:
             self.status = 'INTERNAL'
 
+        source_prefix = f'{source}: ' if source else ''
+        super().__init__(f'{source_prefix}{self.status}: {message}')
+        self.original_message = message
+
         self.http_code = http_status_code(self.status)
 
         if not details: