fix(example): Fix auto-memory wrapper user ID lookup

[yczhang-nv]

Description

Updating the auto-memory wrapper to resolve the runtime user ID from Context.user_id before falling back to legacy/custom user_manager, X-User-ID, and default_user.

The previous implementation directly dereferenced Context.user_manager, which is not present in the current nat run context and caused the auto-memory wrapper examples to fail with:

AttributeError: 'Context' object has no attribute 'user_manager'

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

Release Notes

• Improvements
  • Enhanced user identification handling with more robust fallback mechanisms when retrieving user IDs from multiple sources, including context, user manager, and request headers.
  • Updated tests to verify user ID extraction behavior.

[coderabbitai]

Walkthrough

The PR refactors _get_user_id_from_context in the AutoMemoryWrapperGraph to establish a new user ID retrieval priority: first attempting Context.user_id with error handling and logging, then falling back to Context.user_manager.get_id(), then extracting from the X-User-ID header using case-insensitive lookup, and finally returning "default_user". Tests are updated to match this behavior.

Changes

User ID Retrieval Priority and Defensive Extraction

Layer / File(s)	Summary
User ID retrieval method refactor packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py	_get_user_id_from_context now prioritizes Context.user_id with try/except error handling and debug logging, falls back to Context.user_manager.get_id(), then extracts from the X-User-ID header using case-insensitive metadata.headers lookup, returning "default_user" if all sources fail. Docstring is updated to reflect the new priority order.
Test fixture and user ID retrieval test packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py	Mock context fixture is updated to initialize Context.user_id to None instead of user_manager. New test test_get_user_id_from_context patches Context.get, sets wrapper_graph._context to a provided mock, and verifies _get_user_id_from_context() returns the context user id.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main change: fixing the auto-memory wrapper's user ID lookup mechanism, which aligns with the PR's core objective of resolving a runtime AttributeError.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (2)

packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py (1)

84-91: ⚡ Quick win

Consider catching AttributeError instead of broad Exception.

The PR description indicates the original bug was an AttributeError when Context.user_manager was absent. Catching the specific exception type makes the intent clearer and avoids masking unexpected errors.

♻️ Proposed fix

         # Priority 1: Get user_id from the runtime context.
         try:
             user_id = self._context.user_id
             if user_id:
                 logger.debug(f"Using user_id from context: {user_id}")
                 return user_id
-        except Exception as e:
+        except AttributeError as e:
             logger.debug(f"Failed to get user_id from context: {e}")

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py`
around lines 84 - 91, Replace the broad Exception catch in the block that reads
self._context.user_id with a specific AttributeError handler: in the method
where user_id = self._context.user_id is accessed (inside the
auto_memory_wrapper/agent.py code), catch AttributeError instead of Exception so
missing Context.user_manager/user_id raises are handled explicitly and other
unexpected errors are not masked; keep the same logger.debug message for the
AttributeError case.

packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py (1)

148-154: 💤 Low value

Test covers the happy path for Priority 1 well.

Consider adding a test that verifies priority order by setting both user_id and user_manager, then asserting that user_id wins and user_manager.get_id() is never called. This would strengthen confidence in the fallback hierarchy.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py` around
lines 148 - 154, Add a new test that sets both mock_context.user_id and
mock_context.user_manager (with a mock get_id) and patches Context.get to return
mock_context, then call wrapper_graph._get_user_id_from_context() and assert it
returns mock_context.user_id, and also assert mock_context.user_manager.get_id
was never called; reference the existing test_get_user_id_from_context pattern
and the _get_user_id_from_context method to locate where to add this check.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py`:
- Around line 84-91: Replace the broad Exception catch in the block that reads
self._context.user_id with a specific AttributeError handler: in the method
where user_id = self._context.user_id is accessed (inside the
auto_memory_wrapper/agent.py code), catch AttributeError instead of Exception so
missing Context.user_manager/user_id raises are handled explicitly and other
unexpected errors are not masked; keep the same logger.debug message for the
AttributeError case.

In `@packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py`:
- Around line 148-154: Add a new test that sets both mock_context.user_id and
mock_context.user_manager (with a mock get_id) and patches Context.get to return
mock_context, then call wrapper_graph._get_user_id_from_context() and assert it
returns mock_context.user_id, and also assert mock_context.user_manager.get_id
was never called; reference the existing test_get_user_id_from_context pattern
and the _get_user_id_from_context method to locate where to add this check.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e29edca3-1a02-4738-aefe-52511588dbd6

📥 Commits

Reviewing files that changed from the base of the PR and between f3b7c5a and 0da7055.

📒 Files selected for processing (2)

• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py
• packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py