Repository navigation
Conversation
…elled _create_session reserves a capacity slot (_sessions[sid] = None plus an executor) before building the environment outside the lock, so that a slow factory does not block other sessions. Both failure handlers catch Exception, which excludes CancelledError, so a task cancelled between the reservation and registration leaves the reservation behind. That leak is permanent for the life of the process: get_capacity_status counts len(_sessions), while the idle reaper only scans _session_info, which an interrupted creation never reached. At the default max_concurrent_envs=1 the server then refuses every later client with SessionCapacityError. Release the reservation on cancellation, close whatever the interrupted creation built, and re-raise unchanged. The release is shielded so it still completes when a closing loop delivers another cancellation, matching the shielded teardown the WebSocket handler already uses. Cancelling the await does not stop the thread running the factory, so freeing the slot alone would leave the environment that thread goes on to build unclosed, keeping a container alive while the server reports the slot as free. Keep the factory's thread future so the environment is still reachable, and close it from a done callback. add_done_callback gives no thread guarantee, so the callback submits close() back to the executor the environment was created on rather than running it inline, which is what _cleanup_session_resources already requires for thread-sensitive libraries such as Playwright. The executor is shut down once that close is queued. Signed-off-by: Karthik Suresh <7954591+k21993@users.noreply.github.com>
Signed-off-by: Karthik Suresh <7954591+k21993@users.noreply.github.com>
burtenshaw
left a comment
burtenshaw
left a comment
There was a problem hiding this comment.
Request changes for three lifecycle failures. PR line references below name src/openenv/core/env_server/http_server.py at 820e1736.
-
[P1] Protect MCP entry's cancellation unwind (
http_server.py:462,http_server.py:488). Intest_second_cancel_during_mcp_enter_unwind_finishes_mcp_cleanup, an asyncmcp_session()blocks before yielding and awaits resource release in itsfinally. Cancel creation, wait for thatfinallyto start, then cancel creation again. The second cancellation aborts MCP cleanup before the outer handler starts its shielded task. The test fails withsecond cancellation abandoned MCP cleanup.AsyncExitStackhas not registered the context's exit callback because entry never completed, so closing the stack cannot recover that cleanup. The existing test's synchronousenter_unwound.set()does not exercise this window. Smallest fix: keep a reference to an independently awaited MCP entry task, cancel it at most once, and drain its unwind under the repeated-cancellation shield before disposing of the environment. Add an asynchronous cleanup barrier to the regression test. -
[P2] Give failure and cancellation one cleanup owner (
http_server.py:472,http_server.py:492). Intest_cancel_during_failed_mcp_enter_closes_environment_once, make MCP entry raiseRuntimeError, block the resulting executor-threadenv.close(), cancel creation, then releaseclose(). The failure handler has already dispatched one close. Cancellation enters the new handler with the sameenv, and_discard_reserved_session()dispatches another. The assertion reportsenvironment closed 2 times, on both the original head and the merged tree. The same test passes on current main. Smallest fix: start one protected cleanup task for startup failure and cancellation, retain its reference, and await that existing task when cancellation interrupts failure cleanup. Do not resubmitenv.close(). -
[P2] Release the construction lock before MCP startup (
http_server.py:420,http_server.py:508). Intest_blocked_mcp_entry_does_not_starve_other_shared_factories, configure two sessions withREQUIRES_SINGLE_THREAD_EXECUTOR = True. Let the first constructor return, then hold its asynchronous MCP entry open. Start a second session. Its constructor never runs despite the shared executor being idle, because the lock remains held through MCP startup and registration. The test fails withsecond factory starved behind MCP entry with an idle executor; it passes on current main. One stalled MCP connection now prevents every later session from constructing. Smallest fix: release the lock after successful delivery of the factory result, before awaiting MCP entry. Retain it through late-factory disposal when construction itself is cancelled. This narrows the lock to the ordering requirement and follows CLAUDE.md's Simplicity rule.
Session creation during server shutdown is a separate race that predates this PR, and #1447 does not cover it. Its current-main shutdown callback (src/openenv/core/env_server/http_server.py:919) snapshots session IDs and destroys their current values. In test_shutdown_during_factory_does_not_resurrect_session, hold a per-session factory, finish application shutdown, then release the factory. Creation publishes a session after shutdown: active_sessions == 1, with zero close calls. This fails with both current main and the PR merged onto it. Track this separately: shutdown must mark the server as closing under the session lock, coordinate in-flight creators before executor teardown, and prevent late publication. #1465 is already included in the reviewed head; its sync-client loop cleanup does not resolve this server race.
The remaining probes passed. Repeated cancellation retained capacity until late factory completion and closed each successfully built environment exactly once on its creating thread, for both executor modes. Factories failing after cancellation released their mappings. Twenty cancelled waiters behind a blocked shared factory released their own reservations without starting another factory. Capacity stayed within bounds; it is derived from dictionary length, with no decrementing counter. A factory that never returns still occupies its slot and the shared worker indefinitely, as on main. These probes establish cancellable waiters and recovery after the factory returns, not recovery from an unkillable thread.
The reproduction tests are in the collapsed block below. Save them as tests/core/test_pr_1301_review.py and run:
PYTHONPATH=src:envs uv run pytest tests/core/test_pr_1301_review.py -qAt 820e1736, with current main merged in, four tests fail and five pass. Setting OPENENV_REVIEW_REF=origin/main loads the HTTP server from main instead: the MCP unwind and shutdown tests still fail there, and the double-close and MCP-entry starvation tests pass.
Validation used current main fe18aae4bf7b787dc107ab304cf8de8462c5cb99, merged locally without conflicts as 0d9b26101620fba0e26960284c3cffdde4e45320. On Python 3.13.14/macOS, tests/core/test_production_mode_routes.py reported 67 passed and 1 skipped. tests/core reported 923 passed, 9 skipped, and one failure: test_interrupted_readiness_cleans_up_process[False], with PermissionError at src/openenv/core/harness/process.py:188 in os.killpg. That source and test are identical to current main; a harness-only rerun also hit the process-group permission error. MCP mode-routing and generic-client tests reported another 151 passed and 5 skipped. Changed-file usort, Ruff format, and Ruff rules passed. Repository-wide checks hit existing import-sort and formatting failures outside the changed files; repository-wide Ruff rules passed. All 12 recorded GitHub checks passed. There are no review comments, author replies, or unresolved Bugbot threads.
Reproduction tests (tests/core/test_pr_1301_review.py)
"""Deterministic review probes for OpenEnv PR 1301.
Set OPENENV_REVIEW_REF=origin/main to compare with the unmodified base.
"""
import asyncio
import os
import subprocess
import sys
import threading
import types
from contextlib import asynccontextmanager
import pytest
from fastapi import FastAPI
from openenv.core.env_server.exceptions import SessionCapacityError
from openenv.core.env_server.http_server import HTTPEnvServer
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import Action, Observation, State
if review_ref := os.environ.get("OPENENV_REVIEW_REF"):
source = subprocess.run(
["git", "show", f"{review_ref}:src/openenv/core/env_server/http_server.py"],
check=True,
capture_output=True,
text=True,
).stdout
module_name = "openenv.core.env_server._review_http_server"
review_module = types.ModuleType(module_name)
review_module.__package__ = "openenv.core.env_server"
sys.modules[module_name] = review_module
exec(
compile(source, f"{review_ref}/http_server.py", "exec"), review_module.__dict__
)
HTTPEnvServer = review_module.HTTPEnvServer
class ProbeEnv(Environment):
SUPPORTS_CONCURRENT_SESSIONS = True
def reset(self, **kwargs):
return Observation()
def step(self, action, **kwargs):
return Observation()
@property
def state(self):
return State()
async def thread_event(event):
assert await asyncio.to_thread(event.wait, 3), "worker did not reach barrier"
async def cancelled(task):
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, 3)
@pytest.mark.asyncio
async def test_second_cancel_during_mcp_enter_unwind_finishes_mcp_cleanup():
entered = asyncio.Event()
cleanup_started = asyncio.Event()
release_cleanup = asyncio.Event()
cleanup_finished = asyncio.Event()
closed = []
class Env(ProbeEnv):
@asynccontextmanager
async def mcp_session(self):
entered.set()
try:
await asyncio.Event().wait()
yield self
finally:
cleanup_started.set()
await release_cleanup.wait()
cleanup_finished.set()
def close(self):
closed.append(threading.get_ident())
server = HTTPEnvServer(Env, Action, Observation)
creation = asyncio.create_task(server._create_session())
try:
await asyncio.wait_for(entered.wait(), 3)
creation.cancel()
await asyncio.wait_for(cleanup_started.wait(), 3)
creation.cancel()
# Let cancellation reach the outer shield before allowing MCP cleanup.
for _ in range(10):
await asyncio.sleep(0)
release_cleanup.set()
await cancelled(creation)
assert cleanup_finished.is_set(), "second cancellation abandoned MCP cleanup"
assert len(closed) == 1
assert server.active_sessions == 0
finally:
release_cleanup.set()
await asyncio.gather(creation, return_exceptions=True)
server._executor.shutdown(wait=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("shared", [False, True], ids=["per-session", "shared"])
@pytest.mark.parametrize(
"factory_fails", [False, True], ids=["late-env", "factory-error"]
)
async def test_cancelled_factory_retains_capacity_and_disposes_once(
shared, factory_fails
):
started = threading.Event()
release = threading.Event()
instances = []
class Env(ProbeEnv):
REQUIRES_SINGLE_THREAD_EXECUTOR = shared
def __init__(self):
self.created_thread = threading.get_ident()
self.closes = []
instances.append(self)
started.set()
assert release.wait(3)
if factory_fails:
raise RuntimeError("factory failed after cancellation")
def close(self):
self.closes.append(threading.get_ident())
server = HTTPEnvServer(Env, Action, Observation)
creation = asyncio.create_task(server._create_session())
try:
await thread_event(started)
for _ in range(5):
creation.cancel()
await asyncio.sleep(0)
assert not creation.done()
status = server.get_capacity_status()
assert status.active_sessions == 1
assert status.available_slots == 0
with pytest.raises(SessionCapacityError):
await server._create_session()
release.set()
await cancelled(creation)
assert not server._sessions
assert not server._session_executors
assert not server._session_stacks
assert not server._session_info
assert server.get_capacity_status().active_sessions == 0
if not factory_fails:
assert instances[0].closes == [instances[0].created_thread]
await server._destroy_session("absent")
await server._destroy_session("absent")
assert server.get_capacity_status().active_sessions == 0
finally:
release.set()
await asyncio.gather(creation, return_exceptions=True)
if server._shared_session_executor is not None:
server._shared_session_executor.shutdown(wait=True)
server._executor.shutdown(wait=True)
@pytest.mark.asyncio
async def test_waiters_cancel_promptly_while_shared_factory_is_stuck():
started = threading.Event()
release = threading.Event()
instances = []
closes = []
class Env(ProbeEnv):
REQUIRES_SINGLE_THREAD_EXECUTOR = True
def __init__(self):
instances.append(self)
started.set()
assert release.wait(3)
def close(self):
closes.append(self)
server = HTTPEnvServer(Env, Action, Observation, max_concurrent_envs=2)
first = asyncio.create_task(server._create_session())
try:
await thread_event(started)
first.cancel()
for _ in range(3):
await asyncio.sleep(0)
for _ in range(20):
waiter = asyncio.create_task(server._create_session())
await asyncio.sleep(0)
assert server.active_sessions == 2
waiter.cancel()
await cancelled(waiter)
assert server.active_sessions == 1
assert len(instances) == 1
assert not first.done()
release.set()
await cancelled(first)
assert closes == instances
assert server.active_sessions == 0
finally:
release.set()
await asyncio.gather(first, return_exceptions=True)
server._shared_session_executor.shutdown(wait=True)
server._executor.shutdown(wait=True)
@pytest.mark.asyncio
async def test_blocked_mcp_entry_does_not_starve_other_shared_factories():
loop = asyncio.get_running_loop()
first_entered = asyncio.Event()
release_first_enter = asyncio.Event()
second_constructed = asyncio.Event()
instances = []
class Env(ProbeEnv):
REQUIRES_SINGLE_THREAD_EXECUTOR = True
def __init__(self):
self.number = len(instances) + 1
instances.append(self)
if self.number == 2:
loop.call_soon_threadsafe(second_constructed.set)
@asynccontextmanager
async def mcp_session(self):
if self.number == 1:
first_entered.set()
await release_first_enter.wait()
yield self
server = HTTPEnvServer(Env, Action, Observation, max_concurrent_envs=2)
first = asyncio.create_task(server._create_session())
second = None
try:
await asyncio.wait_for(first_entered.wait(), 3)
second = asyncio.create_task(server._create_session())
try:
await asyncio.wait_for(second_constructed.wait(), 0.3)
except TimeoutError:
pytest.fail("second factory starved behind MCP entry with an idle executor")
finally:
release_first_enter.set()
tasks = [first] + ([second] if second is not None else [])
await asyncio.gather(*tasks, return_exceptions=True)
for sid in list(server._sessions):
await server._destroy_session(sid)
server._shared_session_executor.shutdown(wait=True)
server._executor.shutdown(wait=True)
@pytest.mark.asyncio
async def test_cancel_during_failed_mcp_enter_closes_environment_once():
close_started = threading.Event()
release_close = threading.Event()
closes = []
class Env(ProbeEnv):
@asynccontextmanager
async def mcp_session(self):
raise RuntimeError("MCP handshake failed")
yield self
def close(self):
closes.append(threading.get_ident())
close_started.set()
assert release_close.wait(3)
server = HTTPEnvServer(Env, Action, Observation)
creation = asyncio.create_task(server._create_session())
try:
await thread_event(close_started)
creation.cancel()
for _ in range(10):
await asyncio.sleep(0)
release_close.set()
await cancelled(creation)
assert len(closes) == 1, f"environment closed {len(closes)} times"
assert server.active_sessions == 0
finally:
release_close.set()
await asyncio.gather(creation, return_exceptions=True)
server._executor.shutdown(wait=True)
@pytest.mark.asyncio
async def test_shutdown_during_factory_does_not_resurrect_session():
factory_started = threading.Event()
release_factory = threading.Event()
closed = []
class Env(ProbeEnv):
def __init__(self):
self.created_thread = threading.get_ident()
factory_started.set()
assert release_factory.wait(3)
def close(self):
closed.append(threading.get_ident())
server = HTTPEnvServer(Env, Action, Observation)
app = FastAPI()
server.register_routes(app, mode="production")
lifespan = app.router.lifespan_context(app)
await lifespan.__aenter__()
creation = asyncio.create_task(server._create_session())
try:
await thread_event(factory_started)
await asyncio.wait_for(lifespan.__aexit__(None, None, None), 3)
release_factory.set()
session_id, env = await asyncio.wait_for(creation, 3)
assert server.active_sessions == 0, (
f"session {session_id} registered after shutdown; close calls: {closed}"
)
assert closed == [env.created_thread]
finally:
release_factory.set()
await asyncio.gather(creation, return_exceptions=True)
for sid in list(server._sessions):
await server._destroy_session(sid)
server._executor.shutdown(wait=True)Addresses the three lifecycle failures raised in review of openenv#1301. MCP session entry that blocks before yielding may release what it already acquired in its own finally. Awaiting it inline tied it to the creating task, so a second cancellation landed on that finally and abandoned the unwind part-way. AsyncExitStack cannot recover it either: it registers the exit callback only once __aenter__ returns, so an entry that never returned leaves the stack with nothing to close. Enter in its own task and await it shielded, then let the cleanup owner cancel that task exactly once and drain it under the repeated-cancellation shield before disposing of the environment. Startup failure and cancellation each ran their own cleanup. When a cancellation arrived while the failure handler was already awaiting env.close() in the executor, the cancellation handler reached _discard_reserved_session with the same environment and submitted a second close. Both paths now start at most one cleanup task and a later caller joins the task already running, so an environment is closed once. The shared creation lock was held through MCP startup and registration, not just the factory call. One stalled handshake then kept every later session's constructor off an idle executor thread. Release it once the factory result is delivered, which is the point the shared thread is free. Cancellation during construction keeps the lock, so disposing of a factory that returns late still precedes the next constructor. Session creation racing server shutdown is a separate pre-existing race that reproduces on main and is not addressed here. Signed-off-by: Karthik Suresh <7954591+k21993@users.noreply.github.com>
|
Thanks for fbbbbe6. I re-ran the reproduction tests from the last review against it, and all three findings are fixed:
The capacity, late-factory and waiter probes still pass. The only probe that still fails is |
|
@burtenshaw all three are addressed in 0282ce4, with current main merged in. 1. MCP entry unwind: The entry runs as its own task, awaited through 2. One cleanup owner: Failure and cancellation both go through a 3. Creation lock: Released once the factory result is delivered, before MCP startup.
Your probes give 8 passed, 1 failed. The remainder is the shutdown race, which reproduces on |
Summary
HTTPEnvServer._create_session()reserves a capacity slot and executor before it constructs the environment outside the session lock. If the request task is cancelled during that construction,asyncio.CancelledErrorbypasses the existingexcept Exceptioncleanup and leaves the reservation permanently counted againstmax_concurrent_envs.This change:
CancelledError.At the default capacity of one, this prevents one interrupted connection from making the server reject every future client until restart. It also prevents a late-built container-backed environment from surviving after the server reports the slot as free.
Test plan
tests/core/test_production_mode_routes.py::TestMCPSessionResourceLeaks: 7 passedtests/core/test_production_mode_routes.py: 61 passed, 1 skippedusort, Ruff format, and Ruff rules: passedtests/envs/test_harbor_capture_server.py::test_start_refuses_a_port_held_by_a_non_capture_listenerorigin/main(expectsalready in use, receives the existing 30-second startup timeout)git diff --check: passedAlignment / RFC
No public API changes. This is a lifecycle bug fix that strengthens the explicit resource-cleanup requirement in RFC 002; it does not alter the WebSocket/MCP boundary or session semantics, so no new RFC is required.
Note
Medium Risk
Touches core session lifecycle and concurrency under cancellation; incorrect cleanup could still leak capacity or double-close envs, but behavior is heavily test-covered and scoped to failed/cancelled creation paths.
Overview
Fixes
HTTPEnvServer._create_sessionso a cancelled client no longer leaves a reserved capacity slot, executor, or half-built environment behind.Session creation now keeps the factory’s thread
Future, runs MCP__aenter__under a shielded task, and onCancelledErroror failed startup funnels teardown through a single_discard_reserved_sessionpath that drains in-flight factory/MCP work, closes the env on its executor thread, and clears all session maps. For single-thread executors, a_shared_session_creation_lockserializes factory runs (released before MCP) so a replacement session cannot start before a cancelled env is disposed.Adds seven regression tests in
test_production_mode_routes.pycovering cancellation mid-factory, double cancellation during cleanup/MCP enter, shared-executor ordering, and stalled MCP not starving other factories.Reviewed by Cursor Bugbot for commit 18b398f. Bugbot is set up for automated code reviews on this repo. Configure here.