Skip to content

Commit e8ee4ab

Browse files
fix: skip non-callable items in cleanup instead of logging TypeError
The root cause was that _cleanup_fns is a writable Python list, so external code could bypass register_cleanup() and append None, dicts, or other non-callable items directly. Both PySession::cleanup() and PyCoordinator::cleanup() now guard with is_none()/is_callable() checks before calling, matching the existing guard in register_cleanup(). 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
1 parent d32d906 commit e8ee4ab

File tree

2 files changed

+110
-0
lines changed

2 files changed

+110
-0
lines changed

bindings/python/src/lib.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -677,6 +677,17 @@ impl PySession {
677677
// Step 1: Call all cleanup functions in reverse order
678678
// ----------------------------------------------------------
679679
for callable in cleanup_callables.iter().rev() {
680+
// Guard: skip None and non-callable items (defense-in-depth)
681+
let should_skip = Python::try_attach(|py| -> bool {
682+
let bound = callable.bind(py);
683+
bound.is_none() || !bound.is_callable()
684+
})
685+
.unwrap_or(true);
686+
687+
if should_skip {
688+
continue;
689+
}
690+
680691
// 1a: Call the function inside the GIL
681692
let call_outcome: Option<PyResult<(bool, Py<PyAny>)>> =
682693
Python::try_attach(|py| -> PyResult<(bool, Py<PyAny>)> {
@@ -1486,6 +1497,10 @@ impl PyCoordinator {
14861497
// Execute in reverse order
14871498
for i in (0..len).rev() {
14881499
let cleanup_fn = list.get_item(i)?;
1500+
// Guard: skip None and non-callable items (defense-in-depth)
1501+
if cleanup_fn.is_none() || !cleanup_fn.is_callable() {
1502+
continue;
1503+
}
14891504
// Try calling; catch and log errors
14901505
match cleanup_fn.call0() {
14911506
Ok(result) => {

bindings/python/tests/test_rust_session_lifecycle.py

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,3 +384,98 @@ def test_coordinator_hooks_returns_rust_registry():
384384
assert isinstance(hooks, RustHookRegistry), (
385385
f"Expected RustHookRegistry, got {type(hooks)}"
386386
)
387+
388+
389+
# ---------------------------------------------------------------------------
390+
# Cleanup defense-in-depth: skip None and non-callable items in _cleanup_fns
391+
# ---------------------------------------------------------------------------
392+
393+
394+
@pytest.mark.asyncio
395+
async def test_session_cleanup_skips_non_callable_items():
396+
"""PySession.cleanup() must silently skip non-callable items in
397+
_cleanup_fns — no 'Error during cleanup' log messages."""
398+
import logging
399+
import io
400+
401+
session = await _make_initialized_session()
402+
403+
# Register a legitimate cleanup via the proper API
404+
called = []
405+
406+
def good_cleanup():
407+
called.append("good")
408+
409+
session.coordinator.register_cleanup(good_cleanup)
410+
411+
# Directly append non-callable items to the list — simulates what
412+
# happens when external code bypasses register_cleanup()
413+
fns = session.coordinator._cleanup_fns
414+
fns.append(None)
415+
fns.append({"name": "not-callable"})
416+
fns.append(42)
417+
418+
# Capture log output from the session logger
419+
log_stream = io.StringIO()
420+
handler = logging.StreamHandler(log_stream)
421+
handler.setLevel(logging.DEBUG)
422+
logger = logging.getLogger("amplifier_core.session")
423+
logger.addHandler(handler)
424+
try:
425+
await session.cleanup()
426+
finally:
427+
logger.removeHandler(handler)
428+
429+
# The good cleanup function must still have been called
430+
assert "good" in called, "Good cleanup function should have been called"
431+
432+
# No "Error during cleanup" messages should appear for non-callable items
433+
log_output = log_stream.getvalue()
434+
assert "Error during cleanup" not in log_output, (
435+
f"Non-callable items should be silently skipped, but got: {log_output}"
436+
)
437+
438+
439+
@pytest.mark.asyncio
440+
async def test_coordinator_cleanup_skips_non_callable_items():
441+
"""PyCoordinator.cleanup() must silently skip non-callable items in
442+
_cleanup_fns — no 'Error during cleanup' log messages."""
443+
import logging
444+
import io
445+
446+
session = await _make_initialized_session()
447+
coordinator = session.coordinator
448+
449+
# Register a legitimate cleanup via the proper API
450+
called = []
451+
452+
def good_cleanup():
453+
called.append("good")
454+
455+
coordinator.register_cleanup(good_cleanup)
456+
457+
# Directly append non-callable items to the list
458+
fns = coordinator._cleanup_fns
459+
fns.append(None)
460+
fns.append({"name": "not-callable"})
461+
fns.append(42)
462+
463+
# Capture log output from the coordinator logger
464+
log_stream = io.StringIO()
465+
handler = logging.StreamHandler(log_stream)
466+
handler.setLevel(logging.DEBUG)
467+
logger = logging.getLogger("amplifier_core.coordinator")
468+
logger.addHandler(handler)
469+
try:
470+
await coordinator.cleanup()
471+
finally:
472+
logger.removeHandler(handler)
473+
474+
# The good cleanup function must still have been called
475+
assert "good" in called, "Good cleanup function should have been called"
476+
477+
# No "Error during cleanup" messages should appear for non-callable items
478+
log_output = log_stream.getvalue()
479+
assert "Error during cleanup" not in log_output, (
480+
f"Non-callable items should be silently skipped, but got: {log_output}"
481+
)

0 commit comments

Comments
 (0)