fix: graceful shutdown tracks and cancels in-flight tasks (#356, #357)

[Luffy2208]

Summary

Implemented graceful shutdown task lifecycle handling for Python SDK agent server:

• Added in-flight task registry on AgentServer (_in_flight_tasks) and _track_task(...) helper so server-spawned tasks are retained and auto-removed when done.
• Wired reasoner request task creation through server task tracking in Agent request handler (both callback and direct execution paths).
• Updated _graceful_shutdown(...) to:
  • snapshot tracked tasks,
  • wait up to timeout_seconds,
  • cancel pending tasks after timeout,
  • gather(..., return_exceptions=True) to flush cancellation,
  • clear in-flight registry before exit.
• Kept scope focused on graceful shutdown tracking/drain/cancel behavior.

Testing

• ./scripts/test-all.sh
• Additional verification (please describe):
  • pytest sdk/python/tests/test_agent_graceful_shutdown.py::test_graceful_shutdown_cancels_in_flight_tasks_within_deadline -v
  • pytest sdk/python/tests/test_agent_graceful_shutdown.py::test_graceful_shutdown_force_cancels_tasks_after_timeout -v

Checklist

• I updated documentation where applicable.
• I added or updated tests (or none were needed).
• I updated CHANGELOG.md (or this change does not warrant a changelog entry).

Screenshots (if UI-related)

N/A (backend/test-only change)

Related issues

• Fixes [Python SDK] graceful shutdown does not track or cancel in-flight tasks #356
• Fixes [Python SDK] graceful shutdown does not enforce timeout-based task cancellation #357
• Related: [Python SDK] Agent.stop() is not implemented #355

[CLAassistant]

All committers have signed the CLA.

[github-actions]

Performance

SDK	Memory	Δ	Latency	Δ	Tests	Status
Python	7.9 KB	-13%	0.34 µs	-3%	✓	✓

✓ No regressions detected

[github-actions]

📊 Coverage gate

Thresholds from .coverage-gate.toml: per-surface ≥ 86%, aggregate ≥ 88%, max per-surface regression ≤ 1.0 pp, max aggregate regression ≤ 0.50 pp.

Surface	Current	Baseline	Δ
control-plane	87.30%	87.30%	→ +0.00 pp	🟡
sdk-go	90.70%	90.70%	→ +0.00 pp	🟢
sdk-python	93.63%	93.63%	↑ +0.00 pp	🟢
sdk-typescript	92.56%	92.56%	→ +0.00 pp	🟢
web-ui	90.02%	90.01%	↑ +0.01 pp	🟢
aggregate	89.02%	89.01%	↑ +0.01 pp	🟡

✅ Gate passed

No surface regressed past the allowed threshold and the aggregate stayed above the floor.

[github-actions]

📐 Patch coverage gate

Threshold: 80% on lines this PR touches vs origin/main (from .coverage-gate.toml:thresholds.min_patch).

Surface	Touched lines	Patch coverage	Status
control-plane	0	—	➖ no changes
sdk-go	0	—	➖ no changes
sdk-python	0	—	➖ no changes
sdk-typescript	0	—	➖ no changes
web-ui	0	—	➖ no changes

✅ Patch gate passed

Every surface whose lines were touched by this PR has patch coverage at or above the threshold.

superseded

[santoshkumarradha]

Implementation looks good, few things to fix before merge:

1. test_graceful_shutdown_allows_fast_tasks_to_finish — listed in PR description as tested but doesn't exist in the file. Need the happy-path test.

2. Dead code — comment out the old asyncio.create_task block in agent.py (~L1723-1728), the return await run_reasoner(), and the pytest.skip fallbacks in both shutdown tests. Just delete them, git has history.

3. Stale TODOs — top 3 lines of the test file still flag these as source bugs. This PR fixes them, remove the TODOs.

4. _track_task has no direct test — tests manually stuff _in_flight_tasks which skips the actual tracking path. Quick unit test that _track_task adds/removes via done-callback would close that gap.

5. _track_in_flight_task closure in agent.py untested — if server_handler is None it silently skips tracking. Worth a test to verify the wiring.

Core approach is solid, just cleanup + test gaps.

[Luffy2208]

About test_graceful_shutdown_allows_fast_tasks_to_finish , it was extra testcase created by me only for more clear testing purpose with one more testcase test_graceful_shutdown_records_task_outcomes_and_timing to test the shutdown process by creating a fast task and very slow task and then asserting the test on the basis of the timings recorded.
So its not part of the original testcases. So i should not add those , right?

[AbirAbbas]

About test_graceful_shutdown_allows_fast_tasks_to_finish , it was extra testcase created by me only for more clear testing purpose with one more testcase test_graceful_shutdown_records_task_outcomes_and_timing to test the shutdown process by creating a fast task and very slow task and then asserting the test on the basis of the timings recorded. So its not part of the original testcases. So i should not add those , right?

The tests aren't in the PR, but the description lists them as tested — just update the description to match what's actually included and we're good on that point

[Luffy2208]

@AbirAbbas @santoshkumarradha Changes requested till point 4 are done and I have commited them you can check on them , but about 5th , can you elaborate what do we have to test ?

5. _track_in_flight_task closure in agent.py untested — if server_handler is None it silently skips tracking. Worth a test to verify the wiring.

a test to check if its skipping the tracking when server_handler is None ?

[AbirAbbas]

its okay we can skip number 5, its not necessary