Forward-merge release/1.3 into develop

[willkill07]

Description

This is a manual forward merge of #895 due to merge conflicts

Closes #895

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

  • MCP: session-aware client, configurable token storage (incl. object stores), new client config; deprecated legacy mcp_tool.
  • Observability: workflow run/trace IDs, traceparent/header ingestion, enriched span attribution.
  • Agents/CLI: usage metrics in responses; improved workflow scaffolding with prerelease-aware installs and updated templates.
  • Docker: build args for Python/UV/NAT across images.

• Documentation

  • Overhauled notebooks (new Getting Started/BYO Agent/Tools & Agents/Observability; removed older notebooks).
  • Expanded CLI reference (MCP, Object Store, Optimize, GPU sizing).
  • Secure token storage guide; install extras; contributor ulimit note.

• Refactor

  • Chat models: typed roles, split choices (streaming vs non-streaming), updated defaults.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

Expands docs and CI to process notebooks; broadens path checks; parameterizes Dockerfiles; restructures Sphinx API tree; overhauls notebooks/examples; introduces MCP OAuth2 token storage and session-aware MCP client; changes public chat data models (roles enum, choices split, usage handling); adds workflow run/trace IDs and events; updates FastAPI/OpenAI endpoint handling; multiple minor doc/docstring edits.

Changes

Cohort / File(s)	Summary
CI: docs/notebook checks ci/scripts/documentation_checks.sh	Add temporary export dir, convert .ipynb to Markdown via nbconvert, include converted notebooks in Vale checks; optional preserve/cleanup.
CI: path checks for notebooks ci/scripts/path_checks.py	Broaden ignore patterns to any paths under examples/notebooks and any containing configs/data/src.
Vale vocabulary ci/vale/styles/config/vocabularies/nat/accept.txt	Add/adjust accepted terms and regex variants (e.g., Colab, Gantt, impactful, uncomment, reimplement).
Container build args docker/Dockerfile, examples/.../Dockerfile (agno_personal_finance, getting_started/simple_* , evaluation_and_profiling/email_phishing_analyzer)	Add ARGs PYTHON_VERSION, UV_VERSION, NAT_VERSION; use UV_VERSION in multi-stage COPY.
Sphinx API tree builder docs/source/conf.py	Replace static paths with _build_api_tree(); copy nat and NVIDIA plugins into api_tree; set autoapi_dirs to absolute path.
Docs content updates docs/source/quick-start/installing.md, docs/source/reference/cli.md, docs/source/resources/contributing.md, docs/source/workflows/mcp/*	Add extras (all, ingestion, mcp, profiling); expand CLI reference (mcp, object-store, optimize, sizing); add ulimit note; add MCP secure token storage doc and index link; update auth/client docs.
Examples: README and commands examples/README.md, examples/agents/rewoo/README.md, examples/.../README.md	Reorganize notebooks listing; tweak curl JSON escaping; fix repository wording; expand MinIO instructions.
Examples: framework demo imports examples/frameworks/adk_demo/src/nat_adk_demo/register.py	Import new tools; add lint-suppression comments.
Examples: tests path rename examples/frameworks/haystack_deep_research_agent/tests/test_...py	Rename aiq→nat module paths.
Notebooks overhaul (add/replace/remove) examples/notebooks/*	Replace/rename getting-started series; add BYO agent and tools/agents notebooks; add observability/profiling notebook; remove prior first_search_agent and retail_sales_agent packages, configs, data, and multiple older notebooks; minor content edits in launchable.
HITL example agent usage examples/HITL/.../retry_react_agent.py	Return ChatResponse with Usage in error paths.
pyproject and CI deps pyproject.toml	Remove two example entries; add nbconvert to dev group.
Agent usage accounting src/nat/agent/react_agent/register.py, src/nat/agent/rewoo_agent/register.py	Compute simple token usage and return ChatResponse.with Usage.
API models: chat/roles/choices src/nat/data_models/api_server.py	Add UserMessageContentRoleType; change roles to enum; split Choice into ChatResponseChoice and ChatResponseChunkChoice; make Usage fields optional; default model "unknown-model"; adjust helpers; remove some converters.
Span context ID generation/validation src/nat/data_models/span.py, packages/nvidia_nat_opentelemetry/src/nat/plugins/opentelemetry/otel_span.py	Generate non-zero 128-bit/64-bit IDs; add validators; adjust OtelSpan ID derivation.
Runtime workflow IDs and events src/nat/runtime/runner.py, src/nat/builder/context.py, src/nat/observability/exporter/span_exporter.py, src/nat/runtime/session.py	Add workflow_run_id and workflow_trace_id in context; emit WORKFLOW_START/END; trace-id propagation and HTTP header parsing (traceparent/workflow-trace-id/workflow-run-id).
FastAPI OpenAI endpoint and MCP tools src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py, src/nat/front_ends/fastapi/message_validator.py	Set JSON content-type; simplify non-streaming to direct single response; update MCPClientConfig import; adjust message conversion for new Choice types.
CLI: workflow templates and create flow src/nat/cli/commands/workflow/templates/*, src/nat/cli/commands/workflow/workflow_commands.py	New starter workflow function and config (echo with prefix); register via relative import; determine NAT version for deps and prerelease handling; use python_safe_workflow_name; add description param.
MCP: auth token storage packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/*	Add token storage subsystem (ObjectStore/InMemory); MCPOAuth2Provider accepts builder and resolves object-store storage; wire into OAuth2 code flow; add config field token_storage_object_store; adjust registration.
MCP: client config and session-aware client packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_config.py, client_base.py, client_impl.py, tool.py, utils.py, src/nat/cli/commands/mcp/mcp.py	Add typed client/server config models; add user_id threading; implement per-session MCP clients with limits/cleanup; deprecate mcp_tool; add truncate_session_id; update CLI to new config import.
Data Flywheel ES schema/adapter packages/nvidia_nat_data_flywheel/...	Rename Request→ESRequest and update usages; move span_to_dfw_record import; change constructor kwargs to dict in exporter; update tests and adapters.
OpenTelemetry exporters packages/nvidia_nat_opentelemetry/...	Change constructors to accept otlp_kwargs dict instead of **kwargs; docstring formatting updates.
Phoenix/RagaAI mixins packages/nvidia_nat_phoenix/..., packages/nvidia_nat_ragaai/...	Docstring formatting tweaks only.
Plugins docstrings packages/nvidia_nat_adk/..., packages/nvidia_nat_agno/..., packages/nvidia_nat_crewai/..., packages/nvidia_nat_mem0ai/..., packages/nvidia_nat_zep_cloud/...	Minor doc/docstring edits (format/params).
Test utilities usage tracking packages/nvidia_nat_test/src/nat/test/functions.py	Compute naive token counts and include Usage in ChatResponse.
Experimental functions src/nat/experimental/test_time_compute/functions/*	Await builder.get_function; widen type annotations for schemas.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant Client as MCP Client
  participant Auth as MCPOAuth2Provider
  participant Store as TokenStorage (InMemory/ObjectStore)
  participant Server as MCP Server

  Note over Client,Auth: Session-aware auth with per-user token storage
  User->>Client: call_tool(request, user_id)
  Client->>Auth: get_auth_headers(user_id)
  alt Token exists
    Auth->>Store: retrieve(user_id)
    Store-->>Auth: AuthResult
  else No token
    Auth->>Server: OAuth2 code flow
    Server-->>Auth: AuthResult (access/refresh)
    Auth->>Store: store(user_id, AuthResult)
  end
  Auth-->>Client: Authorization header
  Client->>Server: Tool invocation (authorized)
  Server-->>Client: Result
  Client-->>User: Response

Loading

sequenceDiagram
  autonumber
  actor Caller
  participant Runner as NAT Runtime
  participant Ctx as Context
  participant Obs as SpanExporter

  Note over Runner,Ctx: Workflow run/trace lifecycle
  Caller->>Runner: run(workflow, input)
  Runner->>Ctx: set workflow_run_id, workflow_trace_id
  Runner->>Obs: WORKFLOW_START (ids, metadata)
  Runner->>Runner: execute entry function
  Runner->>Obs: WORKFLOW_END (result/summary)
  Runner->>Ctx: reset workflow ids
  Runner-->>Caller: result/stream

Loading

sequenceDiagram
  autonumber
  actor User
  participant NAT as MCPFunctionGroup
  participant Sess as SessionManager
  participant MClient as Per-Session MCP Client
  participant Srv as MCP Server

  Note over NAT,Sess: Session-aware tool routing
  User->>NAT: invoke tool
  NAT->>Sess: get_session_id()
  NAT->>Sess: _get_session_client(session_id)
  Sess->>MClient: create if missing (limit/cleanup)
  NAT->>MClient: call_tool()
  MClient->>Srv: request
  Srv-->>MClient: response
  MClient-->>NAT: tool result
  NAT-->>User: result

Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120+ minutes

Possibly related PRs

• feat: Improve the developer journey for example notebooks #912: Also restructures example notebooks and related example packages/configs; overlaps with CI/path and pyproject edits.

Suggested labels

breaking

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The title clearly states the core action of merging the release branch into develop, uses an imperative verb form, remains concise under the recommended length, and accurately reflects the PR’s purpose.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

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

---

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.}

[dagardner-nv]

/ok to test fb2d93c

[coderabbitai]

Actionable comments posted: 21

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (15)

packages/nvidia_nat_agno/src/nat/plugins/agno/tool_wrapper.py (1)

142-156: Switch the docstring to Google style for compliance.
Per the repo’s coding guidelines, public functions must use Google-style docstrings. This block still follows the old NumPy-style “Parameters/Returns” format. Please reformat it into Google style so we stay aligned with the documented standards. As per coding guidelines

packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/exporter/dfw_elasticsearch_exporter.py (1)

34-34: Add type hint for **elasticsearch_kwargs parameter.

The **elasticsearch_kwargs parameter lacks a type hint, violating the coding guidelines that require "All public APIs must have Python 3.11+ type hints on parameters and return values."

Apply this diff to add a type hint:

+from typing import Any
+
 ...
-                 **elasticsearch_kwargs):
+                 **elasticsearch_kwargs: Any):

Alternatively, for more precise typing, consider defining a TypedDict that matches the documented structure (endpoint, index, elasticsearch_auth, headers) and use Unpack[ElasticsearchKwargsDict] (Python 3.11+).

As per coding guidelines.

src/nat/experimental/test_time_compute/functions/execute_score_select_function.py (2)

42-43: Add return type hint and docstring.

The function is missing:

1. A return type hint: Since this is an async generator, it should be annotated as AsyncGenerator[FunctionInfo, None]
2. A Google-style docstring explaining the function's purpose, parameters, and return value

Apply this diff to add the missing return type hint:

+from typing import AsyncGenerator
+
 @register_function(config_type=ExecuteScoreSelectFunctionConfig)
-async def execute_score_select_function(config: ExecuteScoreSelectFunctionConfig, builder: Builder):
+async def execute_score_select_function(
+    config: ExecuteScoreSelectFunctionConfig, 
+    builder: Builder
+) -> AsyncGenerator[FunctionInfo, None]:
+    """Execute a function multiple times and select the best output.
+
+    Executes the configured function num_executions times concurrently, optionally
+    scores the outputs using a scoring strategy, and selects the best result using
+    a selection strategy.
+
+    Args:
+        config: Configuration specifying the function to execute, scoring strategy,
+            selection strategy, and number of executions.
+        builder: Builder instance for retrieving functions and strategies.
+
+    Yields:
+        FunctionInfo: Metadata about the wrapped execution function.
+    """
     import asyncio

As per coding guidelines.

---

91-99: Type safety issue in fallback return path.

If the selected output is not found in its_items (resulting in selected_index == -1), line 99 returns selected_output which is a string (from line 92). However, the return type annotation on line 71 specifies executable_fn.single_output_type, which may not be compatible with str.

This type mismatch could cause runtime errors if the calling code expects a properly typed result rather than a raw string.

Consider one of these solutions:

1. Log an error and raise an exception if no match is found:

-        return results[selected_index] if selected_index != -1 else selected_output
+        if selected_index == -1:
+            logger.error("Selected output not found in results. This should not happen.")
+            raise RuntimeError("Failed to match selected output with original results")
+        return results[selected_index]

2. Parse selected_output back to the expected type (if feasible based on the type):

-        return results[selected_index] if selected_index != -1 else selected_output
+        if selected_index == -1:
+            logger.warning("Selected output not found in results, returning first result as fallback")
+            return results[0]
+        return results[selected_index]

packages/nvidia_nat_opentelemetry/src/nat/plugins/opentelemetry/otlp_span_adapter_exporter.py (1)

70-70: Wrap extra OTLP parameters in otlp_kwargs for all OTLPSpanAdapterExporter instantiations. Change calls like:

OTLPSpanAdapterExporter(endpoint=…, batch_size=1, flush_interval=0.1, headers=…, resource_attributes=…)

to:

OTLPSpanAdapterExporter(endpoint=…, otlp_kwargs={"batch_size": 1, "flush_interval": 0.1, "headers": …, "resource_attributes": …})

packages/nvidia_nat_opentelemetry/src/nat/plugins/opentelemetry/otlp_span_redaction_adapter_exporter.py (1)

99-99: Update all instantiations to pass OTLP-specific args via otlp_kwargs. Replace direct keyword args (batch_size, flush_interval, max_queue_size, timeout, compression, resource_attributes, etc.) with

OTLPSpanHeaderRedactionAdapterExporter(
    endpoint=…,
    otlp_kwargs={ "batch_size": …, "flush_interval": …, … }
)

in every call site (including all tests).

src/nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py (2)

86-90: Remove duplicate validation check.

Lines 86-87 and 89-90 contain identical validation logic that checks has_single_output and raises the same error. This is redundant code duplication.

Apply this diff to remove the duplicate:

     if not augmented_function.has_single_output:
         raise ValueError("TTCToolWrapperFunction only supports functions with a single output.")
 
-    if not augmented_function.has_single_output:
-        raise ValueError("TTCToolWrapperFunction only supports functions with a single output.")
-
     if augmented_function.description and augmented_function.description != "":

---

116-144: Return type annotation is inconsistent with actual return values.

The function signature declares -> fn_output_schema:, but the function returns:

• Line 139: "Not enough information" (a str)
• Line 144: result from augmented_function.acall_invoke (presumably a BaseModel instance)

Additionally, if fn_output_schema could be type[None] (per line 102), using it directly as a return type annotation is semantically unclear.

Consider these options:

Option 1: Use a proper union type annotation if the function can return different types:

+    ReturnType = fn_output_schema if fn_output_schema else str
+
-    async def single_inner(input_message: str) -> fn_output_schema:
+    async def single_inner(input_message: str) -> BaseModel | str:

Option 2: Change the error case to raise an exception instead of returning a string:

     async def single_inner(input_message: str) -> fn_output_schema:
         ...
         if not llm_parsed:
             logger.warning("TTCToolWrapperFunction: LLM parsing error")
-            return "Not enough information"
+            raise ValueError("Not enough information to generate structured input")

src/nat/data_models/api_server.py (1)

39-206: Restore support for tool/function roles.

OpenAI-compatible chat payloads legitimately use roles like "tool" (and legacy "function"). Constraining Message.role, ChoiceMessage.role, and ChoiceDelta.role to USER|ASSISTANT|SYSTEM will now reject those requests/responses, breaking function/tool call scenarios that previously worked. Expand the enum (or revert to plain str) so existing integrations still accept the full OpenAI role set.

src/nat/agent/rewoo_agent/register.py (1)

169-171: Preserve stack trace and avoid double-logging on errors

Use bare raise (or raise … from ex) and log with logger.error() when re-raising to keep the original traceback and avoid duplicate stack traces. As per coding guidelines.

Apply:

-        except Exception as ex:
-            logger.exception("ReWOO Agent failed with exception: %s", ex)
-            raise RuntimeError
+        except Exception as ex:
+            logger.error("ReWOO Agent failed with exception: %s", ex)
+            raise

As per coding guidelines

docs/source/quick-start/installing.md (1)

18-18: Align product name with style guide

Use “NVIDIA NeMo Agent toolkit” (lowercase “toolkit”) for the first mention.

-# Installing NVIDIA NeMo Agent Toolkit
+# Installing NVIDIA NeMo Agent toolkit

As per coding guidelines

src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py (1)

690-710: Add authentication callback to session and handle non-streaming in streaming-only workflows

• Pass user_authentication_callback=self._http_flow_handler.authenticate to session():

  - async with session_manager.session(http_connection=request):
  + async with session_manager.session(
  +     http_connection=request,
  +     user_authentication_callback=self._http_flow_handler.authenticate,
  + ):

• generate_single_response now raises if workflow.has_single_output is false. Either reintroduce a fallback aggregation via generate_streaming_response_as_str for non-streaming clients, or enforce/validate single-output workflows and return a clear error when stream=False on streaming-only endpoints.

src/nat/cli/commands/workflow/workflow_commands.py (1)

373-377: Non-interactive uninstall for uv.

uv uninstall likely prompts; add -y for parity with pip.

-            uninstall_cmd = ['uv', 'pip', 'uninstall', package_name]
+            uninstall_cmd = ['uv', 'pip', 'uninstall', '-y', package_name]

packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (2)

37-39: Incorrect imports; module doesn’t export these types.

Import AuthenticatedContext and AuthFlowType from nat.data_models.authentication, not nat.authentication.interfaces.

-from nat.authentication.interfaces import AuthenticatedContext
-from nat.authentication.interfaces import AuthFlowType
-from nat.authentication.interfaces import AuthProviderBase
+from nat.data_models.authentication import AuthenticatedContext, AuthFlowType
+from nat.authentication.interfaces import AuthProviderBase

---

394-398: Fix callback signature mismatch; wrap to match provider’s awaitable (config, flow) callable.

Current signature Callable[[AuthFlowType], AuthenticatedContext] conflicts with OAuth2 provider which awaits a (config, flow) callable. Wrap provided callback to adapt.

-    def set_user_auth_callback(self, auth_callback: Callable[[AuthFlowType], AuthenticatedContext]):
-        """Set the user authentication callback."""
-        if self._auth_provider and hasattr(self._auth_provider, "_set_custom_auth_callback"):
-            self._auth_provider._set_custom_auth_callback(auth_callback)
+    def set_user_auth_callback(self, auth_callback: Callable[[AuthFlowType], AuthenticatedContext | Awaitable[AuthenticatedContext]]):
+        """Set the user authentication callback."""
+        if self._auth_provider and hasattr(self._auth_provider, "_set_custom_auth_callback"):
+            async def _wrapped(config, flow: AuthFlowType) -> AuthenticatedContext:
+                # Support sync or async 1‑arg callbacks
+                result = auth_callback(flow)
+                if asyncio.iscoroutine(result):
+                    return await result  # type: ignore[no-any-return]
+                return result  # type: ignore[return-value]
+            self._auth_provider._set_custom_auth_callback(_wrapped)

🧹 Nitpick comments (34)

packages/nvidia_nat_mcp/src/nat/plugins/mcp/tool.py (1)

113-116: Enhance deprecation notice with replacement and removal_version parameters.

The deprecation decorator is functional but could be more informative. Consider adding:

• replacement: Specify the exact function or configuration to use instead (e.g., "mcp_client function group" or a specific function name)
• removal_version: Provide a target version when this will be removed to help users plan migration

Apply this diff to enhance the deprecation notice:

 @deprecated(
+    removal_version="2.0",
+    replacement="mcp_client function group",
     reason=
     "This function is being replaced with the new mcp_client function group that supports additional MCP features",
     feature_name="mcp_tool_wrapper")

packages/nvidia_nat_mcp/src/nat/plugins/mcp/utils.py (1)

24-38: LGTM! Well-documented utility function.

The implementation is clear and correctly documented. The function appropriately truncates session IDs for logging purposes.

Optional: Consider input validation for edge cases.

For additional robustness, you could validate that max_length is positive:

 def truncate_session_id(session_id: str, max_length: int = 10) -> str:
     """
     Truncate a session ID for logging purposes.
 
     Args:
         session_id: The session ID to truncate
         max_length: Maximum length before truncation (default: 10)
 
     Returns:
         Truncated session ID with "..." if longer than max_length, otherwise full ID
     """
+    if max_length <= 0:
+        raise ValueError("max_length must be positive")
     if len(session_id) > max_length:
         return session_id[:max_length] + "..."
     return session_id

packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)

139-152: Consider improving the docstring for clarity.

The class rename from Request to ESRequest is good—it's more descriptive and avoids potential naming conflicts. However, the docstring could be more specific about its purpose.

Consider this more descriptive docstring:

 class ESRequest(BaseModel):
-    """Request structure used in requests."""
+    """Elasticsearch-specific request structure for OpenAI chat completion requests."""

packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/exporter/dfw_elasticsearch_exporter.py (1)

46-50: Consider using standard Google-style docstring format for nested kwargs.

The nested bullet list format for documenting elasticsearch_kwargs contents is unconventional for Google-style docstrings. While clear, a more standard approach would improve consistency.

Consider one of these alternatives:

Option 1: Note section

-            elasticsearch_kwargs: Additional arguments for ElasticsearchMixin:
-            - endpoint: The elasticsearch endpoint.
-            - index: The elasticsearch index name.
-            - elasticsearch_auth: The elasticsearch authentication credentials.
-            - headers: The elasticsearch headers.
+            **elasticsearch_kwargs: Additional keyword arguments for ElasticsearchMixin.
+
+        Note:
+            Expected elasticsearch_kwargs keys include:
+            
+            - endpoint (str): The elasticsearch endpoint
+            - index (str): The elasticsearch index name  
+            - elasticsearch_auth: The elasticsearch authentication credentials
+            - headers (dict): The elasticsearch headers

Option 2: Reference the mixin documentation

-            elasticsearch_kwargs: Additional arguments for ElasticsearchMixin:
-            - endpoint: The elasticsearch endpoint.
-            - index: The elasticsearch index name.
-            - elasticsearch_auth: The elasticsearch authentication credentials.
-            - headers: The elasticsearch headers.
+            **elasticsearch_kwargs: Additional keyword arguments for ElasticsearchMixin.
+                See `ElasticsearchMixin.__init__` for supported parameters.

As per coding guidelines.

examples/frameworks/haystack_deep_research_agent/tests/test_haystack_deep_research_agent.py (1)

50-51: Consider simplifying the dynamic import pattern.

The dynamic import pattern using importlib.import_module and getattr can be simplified to a direct import for better readability and static analysis support.

Apply this diff to simplify the import:

-    loader_mod = importlib.import_module("nat.runtime.loader")
-    load_workflow = getattr(loader_mod, "load_workflow")
+    from nat.runtime.loader import load_workflow

Note: If the dynamic import is intentional (e.g., to handle optional dependencies or make the test more flexible), you can ignore this suggestion.

examples/frameworks/adk_demo/src/nat_adk_demo/register.py (1)

19-20: Consider adding __all__ to make the registration pattern explicit.

For a registration module using import side effects, defining __all__ can help document which tools are being registered and improve code clarity.

Apply this diff to add __all__:

 from . import nat_time_mcp_tool
 from . import weather_update_tool
+
+__all__ = [
+    "nat_time_mcp_tool",
+    "weather_update_tool",
+]

src/nat/data_models/span.py (4)

142-146: Fix typos in field descriptions ("OTel-style").

Both descriptions say "OTel-syle". Correct to "OTel-style".

-    trace_id: int = Field(default_factory=_generate_nonzero_trace_id,
-                          description="The OTel-syle 128-bit trace ID of the span.")
-    span_id: int = Field(default_factory=_generate_nonzero_span_id,
-                         description="The OTel-syle 64-bit span ID of the span.")
+    trace_id: int = Field(default_factory=_generate_nonzero_trace_id,
+                          description="The OTel-style 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=_generate_nonzero_span_id,
+                         description="The OTel-style 64-bit span ID of the span.")

---

131-139: Ensure generators actually return non-zero IDs (align with docstrings).

Current implementation can (extremely rarely) return 0. Loop until non-zero.

 def _generate_nonzero_trace_id() -> int:
-    """Generate a non-zero 128-bit trace ID."""
-    return uuid.uuid4().int
+    """Generate a non-zero 128-bit trace ID."""
+    tid = 0
+    while tid == 0:
+        tid = uuid.uuid4().int  # 128-bit
+    return tid

 def _generate_nonzero_span_id() -> int:
-    """Generate a non-zero 64-bit span ID."""
-    return uuid.uuid4().int >> 64
+    """Generate a non-zero 64-bit span ID."""
+    sid = 0
+    while sid == 0:
+        sid = uuid.uuid4().int >> 64  # upper 64 bits
+    return sid

---

147-158: Validator cleanups: idiomatic None check; guard UUID parsing; optional TRY003 fix.

• Prefer v is None over isinstance(v, type(None)).
• Wrap UUID parsing to raise a clearer error; use from None to satisfy Ruff B904.
• Optionally assign error messages to variables to appease TRY003.

 @field_validator("trace_id", mode="before")
 @classmethod
 def _validate_trace_id(cls, v: int | str | None) -> int:
     """Regenerate if trace_id is None; raise an exception if trace_id is invalid;"""
-    if isinstance(v, str):
-        v = uuid.UUID(v).int
-    if isinstance(v, type(None)):
-        v = _generate_nonzero_trace_id()
-    if v <= 0 or v >> 128:
-        raise ValueError(f"Invalid trace_id: must be a non-zero 128-bit integer, got {v}")
+    if isinstance(v, str):
+        try:
+            v = uuid.UUID(v).int
+        except ValueError as err:
+            msg = f"Invalid trace_id string: {v!r}"  # noqa: TRY003
+            raise ValueError(msg) from None
+    if v is None:
+        v = _generate_nonzero_trace_id()
+    if v <= 0 or (v >> 128):
+        msg = "Invalid trace_id: must be a non-zero 128-bit integer"  # noqa: TRY003
+        raise ValueError(msg)
     return v

As per coding guidelines.

---

159-172: Use raise ... from None on parse failure; idiomatic None check; optional TRY003.

• Satisfies Ruff B904, removes broad message inline (TRY003) by assigning to variable.

 @field_validator("span_id", mode="before")
 @classmethod
 def _validate_span_id(cls, v: int | str | None) -> int:
     """Regenerate if span_id is None; raise an exception if span_id is invalid;"""
     if isinstance(v, str):
-        try:
-            v = int(v, 16)
-        except ValueError:
-            raise ValueError(f"span_id unable to be parsed: {v}")
-    if isinstance(v, type(None)):
+        try:
+            v = int(v, 16)
+        except ValueError as err:
+            msg = f"Invalid span_id hex string: {v!r}"  # noqa: TRY003
+            raise ValueError(msg) from None
+    if v is None:
         v = _generate_nonzero_span_id()
-    if v <= 0 or v >> 64:
-        raise ValueError(f"Invalid span_id: must be a non-zero 64-bit integer, got {v}")
+    if v <= 0 or (v >> 64):
+        msg = "Invalid span_id: must be a non-zero 64-bit integer"  # noqa: TRY003
+        raise ValueError(msg)
     return v

Based on hints from static analysis.

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

165-177: Avoid bare except + pass; log and narrow exceptions for traceparent parsing.

Add a module logger and catch ValueError only; log the failure instead of silently swallowing.

@@
-import typing
+import typing
+import logging
@@
 from contextlib import nullcontext
@@
-from nat.data_models.interactive import InteractionPrompt
+from nat.data_models.interactive import InteractionPrompt
@@
-_T = typing.TypeVar("_T")
+_T = typing.TypeVar("_T")
+logger = logging.getLogger(__name__)
@@
-        traceparent = request.headers.get("traceparent")
+        traceparent = request.headers.get("traceparent")
         if traceparent:
             try:
                 parts = traceparent.split("-")
                 if len(parts) >= 4:
                     trace_id_hex = parts[1]
                     if len(trace_id_hex) == 32:
                         trace_id_int = uuid.UUID(trace_id_hex).int
                         self._context_state.workflow_trace_id.set(trace_id_int)
-            except Exception:
-                pass
+            except ValueError:
+                logger.exception("Invalid traceparent header; ignoring: %r", traceparent)

As per coding guidelines.

---

178-189: Same here: narrow exception and log when parsing workflow-trace-id.

Use ValueError and log instead of silent pass.

-        if not self._context_state.workflow_trace_id.get():
+        if not self._context_state.workflow_trace_id.get():
             workflow_trace_id = request.headers.get("workflow-trace-id")
             if workflow_trace_id:
                 try:
                     self._context_state.workflow_trace_id.set(uuid.UUID(workflow_trace_id).int)
-                except Exception:
-                    pass
+                except ValueError:
+                    logger.exception("Invalid workflow-trace-id header; ignoring: %r", workflow_trace_id)

As per coding guidelines.

src/nat/agent/react_agent/register.py (1)

153-160: Consider more accurate token counting.

The current token counting uses simple whitespace splitting, which may not align with actual LLM tokenization. While this provides approximate usage statistics, consider using a tokenizer library (e.g., tiktoken for OpenAI models) for more accurate token counts in the future.

For now, the implementation is acceptable for approximate tracking and maintains consistency with the pattern used in related files (rewoo_agent, retry_react_agent).

examples/notebooks/3_adding_tools_and_agents.ipynb (1)

1683-1717: Model selection for OpenAI client may break

llm.model_name may not exist on the LangChain wrapper. Prefer deriving model from config or fallback attributes.

-        response = client.responses.create(
-            model=llm.model_name,
+        model_name = getattr(llm, "model", getattr(llm, "model_name", None)) or "gpt-4o"
+        response = client.responses.create(
+            model=model_name,
             input=[{

Please confirm whether the LangChain LLM wrapper exposes model or model_name. If not, read from config.llm_name accordingly. Based on learnings

docs/source/workflows/mcp/mcp-auth.md (1)

18-19: Fix product name casing in docs heading.

Use “NVIDIA NeMo Agent toolkit” (lowercase “toolkit”) in docs. As per naming rules.

As per coding guidelines

examples/notebooks/README.md (2)

35-35: Improve Colab links.

Optional: add per‑notebook “Open in Colab” badges that point directly to each notebook path for 1‑click launch.

---

40-44: Clarify uv prerequisite or offer fallback.

Note that uv must be installed, or provide a pure-Python fallback.

-uv venv --seed .venv
-source .venv/bin/activate
-uv pip install jupyterlab
+# Option A (recommended): uv
+uv venv --seed .venv && source .venv/bin/activate && uv pip install jupyterlab
+# If uv is not installed:
+# python -m venv .venv && source .venv/bin/activate && pip install --upgrade pip && pip install jupyterlab

docs/source/workflows/mcp/mcp-auth-token-storage.md (2)

20-20: Fix first-use naming.

Per docs guidance, first mention should read “NVIDIA NeMo Agent toolkit”.

-The NeMo Agent toolkit provides a configurable, secure token storage mechanism
+The NVIDIA NeMo Agent toolkit provides a configurable, secure token storage mechanism

As per coding guidelines

---

118-121: Avoid key collisions across providers.

Hashing only user_id can collide across different auth providers/servers. Include provider/server into the key.

-- **Key format**: `tokens/{sha256_hash}` where the hash is computed from the `user_id` to ensure S3 compatibility
+- **Key format**: `tokens/{sha256_hash}` where the hash is computed from `provider_id + ":" + user_id` (or server URL + user_id) to avoid cross‑provider collisions

docs/source/reference/cli.md (4)

22-26: Fix first-use naming.

Use “NVIDIA NeMo Agent toolkit” on first mention in docs.

-While the NeMo Agent toolkit library provides the capability
+While the NVIDIA NeMo Agent toolkit library provides the capability

As per coding guidelines

---

1019-1022: Align wording with available commands.

The group offers upload and delete (no download). Update description.

-The `nat object-store` command group provides utilities to interact with object stores. This command group is used to
-upload and download files to and from object stores.
+The `nat object-store` command group provides utilities to interact with object stores. This command group is used to
+upload files to, and delete files from, object stores.

---

352-359: Example consistency.

Consider adding explicit transport in example or note default (streamable-http) to match preceding docs.

---

150-159: Fix typo in help text
Replace “submit the the workflow” with “submit to the workflow” in:

• src/nat/front_ends/console/console_front_end_config.py:30
• docs/source/reference/cli.md:155, 380

src/nat/cli/commands/workflow/workflow_commands.py (5)

63-72: Use robust prerelease detection.

Use packaging.version to detect prereleases instead of segment count.

-def _is_nat_version_prerelease() -> bool:
-    """
-    Check if the NAT version is a prerelease.
-    """
-    version = _get_nat_version()
-    if version is None:
-        return False
-
-    return len(version.split(".")) >= 3
+def _is_nat_version_prerelease() -> bool:
+    """
+    Check if the NAT version is a prerelease.
+    """
+    from packaging.version import Version
+    from nat.cli.entrypoint import get_version
+    v = get_version()
+    if not v or v == "unknown":
+        return False
+    try:
+        return Version(v).is_prerelease
+    except Exception:
+        return False

---

299-301: Avoid symlink issues on Windows; fallback to copy.

Symlinks often require elevation on Windows and can fail. Fallback to copying if symlink creation fails or on Windows.

-        os.symlink(config_dir_source, config_dir_link)
-        os.symlink(data_dir_source, data_dir_link)
+        def _safe_link(src: Path, dst: Path) -> None:
+            try:
+                os.symlink(src, dst)
+            except Exception:
+                if dst.exists():
+                    return
+                if src.is_dir():
+                    shutil.copytree(src, dst)
+                else:
+                    shutil.copy2(src, dst)
+        _safe_link(config_dir_source, config_dir_link)
+        _safe_link(data_dir_source, data_dir_link)

---

305-311: subprocess.run misuse: check=True makes returncode check dead code.

Either set check=False and test returncode, or keep check=True and rely on exceptions. Prefer consistent pattern.

-            result = subprocess.run(install_cmd, capture_output=True, text=True, check=True)
-
-            if result.returncode != 0:
-                click.echo(f"An error occurred during installation:\n{result.stderr}")
-                return
+            result = subprocess.run(install_cmd, capture_output=True, text=True, check=False)
+            if result.returncode != 0:
+                click.echo(f"An error occurred during installation:\n{result.stderr}")
+                return

Apply similarly in reinstall_command (Line 343) and delete_command (Line 381).

---

343-349: Align with subprocess pattern above.

Make reinstall consistent: set check=False and handle nonzero return code.

-        result = subprocess.run(reinstall_cmd, capture_output=True, text=True, check=True)
+        result = subprocess.run(reinstall_cmd, capture_output=True, text=True, check=False)
         if result.returncode != 0:
             click.echo(f"An error occurred during installation:\n{result.stderr}")
             return

---

208-213: Fix minor typo in docstring.

“pre-popluate” → “pre-populate”.

-        description (str): Description to pre-popluate the workflow docstring.
+        description (str): Description to pre-populate the workflow docstring.

src/nat/authentication/oauth2/oauth2_auth_code_flow_provider.py (2)

86-86: Silence Ruff ARG002: rename unused kwargs.

Change to avoid the unused-arg warning while keeping signature compatibility:

-    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
+    async def authenticate(self, user_id: str | None = None, **_kwargs) -> AuthResult:

---

129-135: Be explicit about expires_at type.

AuthenticatedContext.metadata["expires_at"] might be a string (ISO) from frontends. Consider normalizing to datetime for clarity:

-        auth_result = AuthResult(
+        expires_at = metadata.get("expires_at")
+        if isinstance(expires_at, str):
+            try:
+                from datetime import datetime
+                expires_at = datetime.fromisoformat(expires_at)
+            except Exception:
+                expires_at = None
+        auth_result = AuthResult(
             credentials=[BearerTokenCred(token=SecretStr(token))],
-            token_expires_at=metadata.get("expires_at"),
+            token_expires_at=expires_at,
             raw=metadata.get("raw_token") or {},
         )

src/nat/cli/commands/workflow/templates/workflow.py.j2 (1)

31-33: Docstring nit: use “Yields” not “Returns”.

The function yields a FunctionInfo. Update for accuracy:

-    Returns:
-        FunctionInfo: The function info object for the function.
+    Yields:
+        FunctionInfo: The function info object for the function.

packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (2)

292-316: Token cache check uses legacy in‑memory attribute; consider token_storage awareness.

_with_reconnect() chooses timeout based on _has_cached_auth_token(), but this checks provider._authenticated_tokens which no longer exists with token_storage. Improve to consult storage when available:

• If self._auth_provider has _auth_code_provider with _token_storage and an effective user_id is known, attempt retrieve(user_id) and check not expired.

This preserves extended timeout only when interactive auth is likely needed.

---

635-639: Avoid logging full tool arguments at info level.

Tool args may contain PII/secrets. Log at debug or redact.

-            logger.info("Calling tool %s with arguments %s", self._tool_name, tool_args)
+            logger.debug("Calling tool %s", self._tool_name)

packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_impl.py (1)

200-236: Session limit pre-check outside writer lock may race.

Minor: compute and enforce the max_sessions check while holding the writer lock (you already re-check later). Consider dropping the pre-check or keep as soft-trigger to run cleanup only.