test(runtime): stabilize early cancellation checks

[xianlubird]

Summary

Make the early-cancellation tests wait for stream closure and worker acknowledgement instead of assuming cancellation prevents every response from reaching the client.

I encountered this while running CI for #14846 (fix(responses): preserve text instead of inferring tool calls), at head f41e5dcd531c3d4116b7238535b5111658104484. That PR changes the Responses converter and its Rust coverage; the failing test uses the runtime client directly. This follow-up keeps the cancellation-test fix separate from the Responses change.

CI failure

The failed job was dynamo-runtime / test / sequential cuda13.0, amd64. It reported:

_____________ test_client_context_cancel_before_await_request[tcp] _____________
lib/bindings/python/tests/cancellation/test_cancellation.py:390:
E   AssertionError: Request should be cancelled before any responses are generated
_ = Annotated(data=0, event=None, comment=[], id=None)

1 failed, 3616 passed, 143 skipped, 1989 deselected

The worker log shows the ordering:

11:06:27.964617Z request received
Processing iteration 0
Sending iteration 0
Processing iteration 1
11:06:28.084912Z request cancellation received
Context stopped at iteration 2

The cancellation reached the worker about 120 ms after the request. The mock yields every 100 ms, so a response was already in flight. dynamo-status-check then failed because this job failed. Rust tests, Clippy, the arm64 sequential job, and the backend tests passed. The matching job passed in the three main runs inspected, so this is not a claim that the exact CI failure was also reproduced on main.

Why this changes the assertion

client.generate() schedules the Rust future immediately; awaiting the returned Python future does not start the request. stop_generating() signals cancellation asynchronously, and the response stream may still deliver buffered data. Requiring zero responses therefore depends on cancellation winning a scheduling race.

Both the already-cancelled and cancel-before-await tests now drain in-flight responses and require stream closure plus an explicit event from the mock worker when it observes cancellation, within a shared 5-second deadline. They retain the stopped/not-killed assertions. A 15-second pytest timeout bounds the overall test, and the event replaces the fixed 200 ms wait. Production runtime behavior is unchanged.

Validation

• Complete cancellation module before rebasing onto current main: 14 passed, covering NATS and TCP.
• After rebasing onto current main: all four changed early-cancellation cases passed again.
• Controlled scheduling check: inserting a 150 ms event-loop delay between request creation and cancellation, still before awaiting the request, reproduced the original Annotated(data=0) assertion failure. This demonstrates the ordering assumption rather than reproducing the exact CI scheduling delay.
• Patched stress batch: 63 passed before an unrelated NATS No workers available discovery error stopped the batch. A separate delayed TCP batch passed all 10 repetitions.
• Negative control: omitting the cancellation call made both NATS and TCP tests fail at the 5-second timeout, as expected.
• Applicable repository pre-commit hooks and git diff --check passed.

These checks used macOS arm64, Python 3.12.7, the existing local Dynamo extension, NATS 2.12.14, and etcd 3.5.33. The exact private ECR CI image was unavailable without registry credentials; the extension was not rebuilt from the PR.

Related Issues

• Confirmed — no related issue.
• Related PR: fix(responses): preserve text instead of inferring tool calls #14846.

Summary by CodeRabbit

• Tests
  • Improved cancellation test reliability by waiting for worker acknowledgement before verifying request termination.
  • Added coverage for already-cancelled requests and cancellation before awaiting.
  • Added safeguards to prevent cancellation tests from hanging indefinitely.

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

[dynamo-ops]

/ok to test 11ff4ef

[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: 1d79e67f-7841-4cbf-8f4c-f0f29424279a

📥 Commits

Reviewing files that changed from the base of the PR and between c52945c and 11ff4ef.

📒 Files selected for processing (1)

• lib/bindings/python/tests/cancellation/test_cancellation.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The cancellation tests now use an event to acknowledge worker cancellation. A shared helper drains responses, waits for acknowledgement, and verifies context state. Two cancellation tests use the helper and 15-second timeouts.

Changes

Cancellation test synchronization

Layer / File(s)	Summary
Cancellation acknowledgement lib/bindings/python/tests/cancellation/test_cancellation.py	MockServer clears cancellation_observed when a request starts and sets it when cancellation is detected.
Shared cancellation verification lib/bindings/python/tests/cancellation/test_cancellation.py	assert_request_cancelled drains responses, waits up to five seconds for acknowledgement, and checks that the context stopped without being killed. The affected tests use the helper and 15-second timeouts.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 11ff4

The cancellation tests synchronize with the worker acknowledgement without a supported regression or merge-blocking risk.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 7 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly explains the test-stabilization change, failure cause, implementation, validation, and related issue status. It does not use the template headings exactly and does not identify…
Title check	✅ Passed	The title clearly and concisely describes the main change: stabilizing runtime tests for early cancellation.
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.

• Fix all pre-merge checks with AI

---

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

[dynamo-ops]

/ok to test 4b5f2ea

[xianlubird]

This is covered by #14974, which has already landed on main. It drains responses that were already in flight and waits for the worker to observe cancellation, so the fix in this PR would duplicate it.

I'm closing this PR rather than carrying the duplicate change through the merge conflict.