Populate full connection attributes and payload for HTTP and WebSocket sessions

[ericevans-nv]

Description

Closes Issue #1578

Populate all connection attributes (URL, scheme, port, headers, query params, path params, cookies, client info) in set_metadata_from_websocket to achieve parity with HTTP metadata extraction. Add a payload field to RequestAttributes that stores the parsed HTTP request body and the full WebSocket user message dict (excluding the security field). Make set_metadata_from_http_request async to support await request.json().

Changes:

• api_server.py — Add payload: dict[str, Any] | None field to the Request model.
• user_metadata.py — Add payload property to RequestAttributes.
• session.py — Populate all WebSocket connection attributes (url, scheme, headers, query params, path params, client host/port) in set_metadata_from_websocket. Make set_metadata_from_http_request async and read the parsed JSON body into payload.
• message_handler.py — Store message.model_dump(exclude={"security"}) as the WebSocket payload and set it on the session metadata in _run_workflow.
• Tests — Replace the former test_user_attributes_from_http_request with comprehensive test_metadata_from_http_request_populates_all_request_attributes (integration) and test_metadata_from_websocket_populates_all_request_attributes (unit) tests that assert every request attribute field for both transport types. Update test_session_traceparent.py to await the now-async set_metadata_from_http_request. Add payload assertion to test_request_attributes_defaults.

By Submitting this PR I confirm:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

Summary by CodeRabbit

• New Features

  • Request payloads from HTTP and WebSocket connections are now captured and exposed in request metadata for downstream workflows.

• Chores

  • Metadata collection for HTTP requests was updated to support extracting the request payload.

• Tests

  • Tests expanded and updated to validate that metadata (including payload, headers, query params, cookies, and client info) is populated for HTTP and WebSocket flows.

[coderabbitai]

Walkthrough

Adds request payload capture and propagation across HTTP and WebSocket handling: a new payload field on the Request model, async extraction of HTTP JSON payload in session setup, websocket metadata expansion, exposure via RequestAttributes, and tests updated to verify the payload and other request attributes.

Changes

Cohort / File(s)	Summary
Core Data Models packages/nvidia_nat_core/src/nat/data_models/api_server.py	Added `payload: dict[str, typing.Any]
Front-end / Message Handling packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py	Message handler now captures the incoming user message payload into an internal attribute and propagates it into the session context before workflow execution.
Session Runtime packages/nvidia_nat_core/src/nat/runtime/session.py	set_metadata_from_http_request converted to async and now await request.json() to populate payload (with graceful fallback); websocket metadata population expanded to include url_path, url_port, url_scheme, headers, query_params, path_params, client_host, and client_port.
Public Metadata API packages/nvidia_nat_core/src/nat/runtime/user_metadata.py	Added payload property on RequestAttributes exposing self._request.payload as `dict[str, typing.Any]
Tests packages/nvidia_nat_core/tests/nat/runtime/test_session_traceparent.py, packages/nvidia_nat_core/tests/nat/runtime/test_user_metadata.py, packages/nvidia_nat_core/tests/nat/server/test_unified_api_server.py	Removed @pytest.mark.asyncio decorators where awaiting the new async method; added assertions for default payload and expanded HTTP/WebSocket tests to assert payload and additional metadata fields.

Sequence Diagram

sequenceDiagram
    participant Client
    participant FastAPI_Handler as FastAPI Handler
    participant SessionManager
    participant Context
    participant Workflow

    Client->>FastAPI_Handler: Send HTTP or WebSocket request with payload
    FastAPI_Handler->>FastAPI_Handler: Capture payload into internal attribute
    FastAPI_Handler->>SessionManager: Invoke set metadata (HTTP/WebSocket)
    activate SessionManager
    SessionManager->>SessionManager: await request.json() (HTTP) / read websocket fields
    SessionManager->>Context: Store metadata including payload
    deactivate SessionManager
    FastAPI_Handler->>Context: Propagate captured payload into session context
    FastAPI_Handler->>Workflow: Start workflow using populated context
    Workflow->>Context: Access payload via RequestAttributes.payload
    Context-->>Workflow: Provide payload and metadata

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 56.25% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: populating full connection attributes and payload for both HTTP and WebSocket sessions, which aligns with the substantial modifications across multiple files to achieve metadata extraction parity between transports.
Merge Conflict Detection	✅ Passed	✅ No merge conflicts detected when merging into develop

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

No actionable comments were generated in the recent review. 🎉

🧹 Recent nitpick comments

packages/nvidia_nat_core/src/nat/runtime/session.py (2)

600-603: Consider adding a debug log when payload parsing fails.

Silently swallowing the exception is fine for the happy path (non-JSON bodies, GET requests), but a logger.debug would help during troubleshooting when a JSON body was expected but failed to parse. The broad except Exception is pragmatic here since request.json() can fail in various ways (disconnect, decode error, empty body).

💡 Optional: add debug-level logging

         try:
             self._context.metadata._request.payload = await request.json()
-        except Exception:
+        except Exception:  # noqa: BLE001
+            logger.debug("Could not parse request body as JSON; payload will be None")
             self._context.metadata._request.payload = None

---

648-657: WebSocket metadata population looks correct.

The new fields mirror the HTTP metadata extraction and the websocket.client guard correctly handles the None case. Minor style note: the HTTP path uses request.client.host/.port (named attributes) while this uses index access ([0]/[1]). Both work on Starlette's Address namedtuple, but named attributes would be more readable and consistent.

💡 Optional: use named attributes for consistency

-        host = websocket.client[0] if websocket.client else None
-        port = websocket.client[1] if websocket.client else None
+        host = websocket.client.host if websocket.client else None
+        port = websocket.client.port if websocket.client else None

packages/nvidia_nat_core/tests/nat/server/test_unified_api_server.py (1)

528-534: Redundant local imports.

MagicMock (line 530) is already imported at line 22, and SessionManager (line 533) at line 76. The local re-imports are unnecessary.

💡 Remove redundant local imports

 def test_metadata_from_websocket_populates_all_request_attributes() -> None:
     """Unit test: set_metadata_from_websocket populates context metadata from a mock websocket."""
-    from unittest.mock import MagicMock
-
     from nat.builder.context import ContextState
-    from nat.runtime.session import SessionManager
     from nat.runtime.user_metadata import RequestAttributes

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py`:
- Line 122: The payload currently includes PII because
self._user_message_payload is set from message.model_dump(exclude={"security"})
but does not exclude the WebSocketUserMessage.user field; update the call in the
constructor where self._user_message_payload is assigned to exclude "user" as
well (e.g., message.model_dump(exclude={"security", "user"})) or explicitly
remove/redact the user entry before assigning to self._user_message_payload so
RequestAttributes.payload no longer contains name/email; if you intend to keep
user, add clear documentation and ensure downstream consumers treat it as
sensitive.

🧹 Nitpick comments (5)

packages/nvidia_nat_core/src/nat/runtime/user_metadata.py (1)

134-137: Docstring style inconsistency with sibling properties.

All other properties in this class use a multi-line Google-style docstring with a Returns: section. Consider matching that style for consistency.

Suggested docstring

     `@property`
     def payload(self) -> dict[str, typing.Any] | None:
-        """Request payload parsed as a dictionary."""
+        """
+        This property retrieves the parsed payload from the request.
+
+        Returns:
+            dict[str, typing.Any] | None
+        """
         return self._request.payload

packages/nvidia_nat_core/src/nat/runtime/session.py (2)

654-657: Prefer named attributes over tuple indexing for consistency.

Lines 597-598 use request.client.host and request.client.port for HTTP. Here, tuple indexing [0]/[1] is used for WebSocket. Starlette's Address namedtuple supports both, but using .host/.port would be more readable and consistent.

Suggested change

-        host = websocket.client[0] if websocket.client else None
-        port = websocket.client[1] if websocket.client else None
+        host = websocket.client.host if websocket.client else None
+        port = websocket.client.port if websocket.client else None

---

600-603: Silent fallback on payload parsing is appropriate but consider logging.

The broad except Exception (flagged by ruff BLE001) is reasonable here since request.json() can fail for many valid reasons (empty body, non-JSON content type, already-consumed stream). However, a debug-level log would aid troubleshooting without being noisy.

Optional: add debug logging

         try:
             self._context.metadata._request.payload = await request.json()
-        except Exception:
+        except Exception:  # noqa: BLE001
+            logger.debug("Could not parse request body as JSON payload")
             self._context.metadata._request.payload = None

packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py (1)

410-410: Deeply nested private attribute access bypasses the public API.

Accessing self._session_manager._context.metadata._request.payload reaches through three layers of private attributes. The same pattern exists in set_metadata_from_websocket, but consider adding a setter method on SessionManager (or the metadata object) to encapsulate this, especially since this is now called from an external class.

packages/nvidia_nat_core/tests/nat/server/test_unified_api_server.py (1)

528-575: Redundant import of MagicMock on Line 530.

MagicMock is already imported at line 22. The local re-import on Line 530 is unnecessary.

Remove redundant import

 def test_metadata_from_websocket_populates_all_request_attributes() -> None:
     """Unit test: set_metadata_from_websocket populates context metadata from a mock websocket."""
-    from unittest.mock import MagicMock
-
     from nat.builder.context import ContextState
     from nat.runtime.session import SessionManager

[ericevans-nv]

/merge