Call tool shutdown hooks after generation

[tamohannes]

Summary

Adds shutdown hooks for tool/code execution wrappers and ensures GenerationTask.generate() calls them from a finally block. This covers normal completion, errors, and early exits.

Why

Tool sessions, sandbox clients, or subprocess-backed providers should be closed before process exit. Without a guaranteed shutdown path, cleanup can hang or leak resources after outputs are already produced.

Tests

• python -m pytest tests/test_mcp_clients.py::test_tool_calling_wrapper_shutdown_calls_registered_tools tests/test_generation.py::test_generation_task_shutdown_awaits_async_llm_shutdown -q
• python -m ruff check nemo_skills/inference/generate.py nemo_skills/inference/model/code_execution.py nemo_skills/inference/model/tool_call.py tests/test_generation.py tests/test_mcp_clients.py

Summary by CodeRabbit

• Bug Fixes

  • Improved resource cleanup and lifecycle management for generation tasks to ensure models and sandboxes are properly shut down even when errors occur or tasks exit early.

• Tests

  • Added test coverage for shutdown procedures to verify proper cleanup behavior.

[coderabbitai]

Actionable comments posted: 0

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c23e66f0-6bc0-426f-a6df-99aa84b52df4

📥 Commits

Reviewing files that changed from the base of the PR and between b620e79 and cfa569b.

📒 Files selected for processing (5)

• nemo_skills/inference/generate.py
• nemo_skills/inference/model/code_execution.py
• nemo_skills/inference/model/tool_call.py
• tests/test_generation.py
• tests/test_mcp_clients.py

---

📝 Walkthrough

Walkthrough

This PR adds explicit shutdown lifecycle management to three inference classes. GenerationTask.shutdown() conditionally calls and awaits LLM shutdown; GenerationTask.generate() is refactored to guarantee shutdown execution via try/finally. CodeExecutionWrapper and ToolCallingWrapper each gain async shutdown methods that clean up wrapped model and resource-specific state (sandbox, tool_manager).

Changes

Shutdown Lifecycle for Inference Components

Layer / File(s)	Summary
GenerationTask shutdown and generate refactor nemo_skills/inference/generate.py, tests/test_generation.py	GenerationTask.shutdown() conditionally calls and awaits LLM shutdown. GenerationTask.generate() is wrapped in try/finally to guarantee shutdown() is called on all exit paths (early returns, errors). Test verifies async LLM shutdown is awaited correctly.
CodeExecutionWrapper shutdown lifecycle nemo_skills/inference/model/code_execution.py	CodeExecutionWrapper.shutdown() async method shuts down the wrapped model (awaiting result if awaitable) and closes the sandbox resource in a finally block.
ToolCallingWrapper shutdown and tool testing nemo_skills/inference/model/tool_call.py, tests/test_mcp_clients.py	ToolCallingWrapper.shutdown() async method shuts down tool_manager and wrapped model. Test infrastructure (ShutdownTool helper class and test case) verifies that registered tool shutdown is called exactly once.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Call tool shutdown hooks after generation' accurately describes the main change: adding and invoking shutdown hooks for tool/code execution wrappers after generation completes.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch tamohannes/waste-tool-shutdown

---

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