Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions infrastructure/filestorage/contracts.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,5 +55,13 @@ def put_object(self, key: str, data: bytes) -> None:
def describe(self) -> str:
"""Short human-readable destination, for logs and CLI output."""

def close(self) -> None:
"""Release owned network resources (for example an ``httpx.Client``).

Providers that do not own a client may no-op. Callers that build a
fresh store per operation should invoke this in a ``finally`` so
connection pools are not left to cyclic GC.
"""


__all__ = ["ObjectStore", "RemoteObject"]
25 changes: 14 additions & 11 deletions infrastructure/filestorage/operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,18 +173,21 @@ def run_remote_sync(
return None
roots = syncable_roots()
store = build_object_store(config)
return _owned_report(
run_sync(
store,
direction=resolved,
roots=roots,
exclusions=config.exclude,
dry_run=dry_run,
on_progress=on_progress,
# The backend, not the engine, knows how hard it can be pushed.
max_parallel_uploads=max_parallel_uploads_for_provider(config.provider),
try:
return _owned_report(
run_sync(
store,
direction=resolved,
roots=roots,
exclusions=config.exclude,
dry_run=dry_run,
on_progress=on_progress,
# The backend, not the engine, knows how hard it can be pushed.
max_parallel_uploads=max_parallel_uploads_for_provider(config.provider),
)
)
)
finally:
store.close()


__all__ = [
Expand Down
3 changes: 3 additions & 0 deletions infrastructure/filestorage/providers/_s3_shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,9 @@ def describe(self) -> str:
"""Human-facing ``scheme://bucket/prefix`` target for error messages."""
raise NotImplementedError

def close(self) -> None:
"""S3-family stores do not own an ``httpx.Client``; nothing to release."""

def list_objects(self, prefix: str) -> list[RemoteObject]:
full_prefix = (
self._config.key_for(prefix) if prefix else f"{self._config.prefix.rstrip('/')}/"
Expand Down
7 changes: 7 additions & 0 deletions infrastructure/filestorage/providers/azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,13 @@ def __init__(
)
self._token = token if token is not None else _get_azure_access_token()
self._client = client if client is not None else httpx.Client(timeout=30.0)
self._owns_client = client is None

def close(self) -> None:
"""Close the owned ``httpx.Client``; no-op when the caller injected one."""
if self._owns_client:
self._client.close()
self._owns_client = False

def describe(self) -> str:
return f"azure-blob://{self._account_name}/{self._config.bucket}/{self._config.prefix}"
Expand Down
3 changes: 3 additions & 0 deletions infrastructure/filestorage/providers/gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,9 @@ def __init__(self, config: RemoteSyncConfig, *, session: Any | None = None) -> N
self._config = config
self._session = session if session is not None else _build_session()

def close(self) -> None:
"""GCS uses a requests session managed elsewhere; nothing required here."""

def describe(self) -> str:
return f"gs://{self._config.bucket}/{self._config.prefix}"

Expand Down
6 changes: 6 additions & 0 deletions infrastructure/filestorage/providers/vercel.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,12 @@ def __init__(
self._client = client if client is not None else httpx.Client(timeout=_TIMEOUT)
self._owns_client = client is None

def close(self) -> None:
"""Close the owned ``httpx.Client``; no-op when the caller injected one."""
if self._owns_client:
self._client.close()
self._owns_client = False

def describe(self) -> str:
return f"vercel-blob://{self._config.bucket}/{self._config.prefix}"

Expand Down
24 changes: 24 additions & 0 deletions tests/filestorage/test_azure_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,3 +181,27 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
with pytest.raises(RemoteSyncUnavailableError) as caught_put:
store.put_object("test.json", b"data")
assert "403" in str(caught_put.value)


def test_close_releases_owned_httpx_client(monkeypatch: pytest.MonkeyPatch) -> None:
"""Owned clients must close; injected clients must stay open (#6989)."""
closed: list[httpx.Client] = []
real_close = httpx.Client.close

def _tracking_close(self: httpx.Client) -> None:
closed.append(self)
real_close(self)

monkeypatch.setattr(httpx.Client, "close", _tracking_close)
owned = AzureBlobObjectStore(_config(), token="fake")
owned.close()
assert closed == [owned._client]
# Idempotent: a second close must not touch the already-closed client again.
owned.close()
assert closed == [owned._client]

injected = httpx.Client(transport=_FakeTransport())
store = AzureBlobObjectStore(_config(), token="fake", client=injected)
store.close()
assert injected not in closed
injected.close()
3 changes: 3 additions & 0 deletions tests/filestorage/test_community_provider_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,9 @@ def put_object(self, key: str, data: bytes) -> None:
def describe(self) -> str:
return f"{_PROVIDER_NAME}://template/{self._config.prefix}"

def close(self) -> None:
return


@pytest.fixture
def fake_provider() -> Iterator[dict[str, tuple[bytes, datetime]]]:
Expand Down
3 changes: 3 additions & 0 deletions tests/filestorage/test_multi_provider_support.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ def put_object(self, key: str, data: bytes) -> None:
def describe(self) -> str:
return f"{self.label}://test"

def close(self) -> None:
return


@pytest.fixture
def roots(tmp_path: Path) -> tuple[SyncRoot, ...]:
Expand Down
3 changes: 3 additions & 0 deletions tests/filestorage/test_push_concurrency.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@ def put_object(self, key: str, data: bytes) -> None:
def describe(self) -> str:
return "recording://bucket"

def close(self) -> None:
return


class FailingStore(RecordingStore):
"""Fails the named keys, to pin which error a caller sees."""
Expand Down
3 changes: 3 additions & 0 deletions tests/filestorage/test_remote_sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,9 @@ def put_object(self, key: str, data: bytes) -> None:
def describe(self) -> str:
return "fake://bucket"

def close(self) -> None:
return


@pytest.fixture
def home(tmp_path: Path) -> Path:
Expand Down
44 changes: 44 additions & 0 deletions tests/filestorage/test_service_and_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ def put_object(self, key: str, data: bytes) -> None:
def describe(self) -> str:
return "mem://t"

def close(self) -> None:
return


@pytest.fixture(autouse=True)
def _restore_registry() -> Iterator[None]:
Expand Down Expand Up @@ -334,3 +337,44 @@ def test_run_remote_sync_uses_scoped_roots(tmp_path: Path, monkeypatch: pytest.M
report2 = run_remote_sync()
assert report2 is not None
assert report2 is not report


def test_run_remote_sync_closes_store(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""``run_remote_sync`` must close the store it built (#6989)."""
from infrastructure.filestorage import operations as sync_service

class _ClosingStore(_MemStore):
def __init__(self) -> None:
super().__init__()
self.closed = False

def close(self) -> None:
self.closed = True

stores: list[_ClosingStore] = []

def _factory(_cfg: RemoteSyncConfig) -> _ClosingStore:
store = _ClosingStore()
stores.append(store)
return store

register_object_store("close-wiring", _factory)
try:
sessions = tmp_path / "sessions"
memory = tmp_path / "memory"
sessions.mkdir()
memory.mkdir()
(sessions / "x.jsonl").write_text("{}\n", encoding="utf-8")
roots = (
SyncRoot(name=SyncRootName.SESSIONS, path=sessions),
SyncRoot(name=SyncRootName.MEMORY, path=memory),
)
monkeypatch.setenv(REMOTE_SYNC_ENV, "1")
monkeypatch.setenv(REMOTE_SYNC_BUCKET_ENV, "b")
monkeypatch.setenv(REMOTE_SYNC_PROVIDER_ENV, "close-wiring")
monkeypatch.setattr(sync_service, "syncable_roots", lambda: roots)
report = run_remote_sync()
assert report is not None
assert stores and stores[0].closed is True
Comment on lines +376 to +378

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Failed-sync cleanup lacks coverage

test_run_remote_sync_closes_store only checks a successful sync. It would still pass if cleanup moved out of finally, leaving clients open whenever a sync fails. Add a case where run_sync raises, then assert that the store closes and the original error reaches the caller.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

finally:
unregister_object_store("close-wiring")
23 changes: 23 additions & 0 deletions tests/filestorage/test_vercel_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,3 +327,26 @@ def test_engine_push_restore_via_registry(tmp_path_factory: pytest.TempPathFacto
assert (other_memory / "fact.md").read_text(encoding="utf-8") == "remembered\n"
# Sanity: content_tag matches Vercel etags for single-part puts.
assert content_tag(b'{"turn": 1}\n') == _md5(b'{"turn": 1}\n')


def test_close_releases_owned_httpx_client(monkeypatch: pytest.MonkeyPatch) -> None:
"""Owned clients must close; injected clients must stay open (#6989)."""
closed: list[httpx.Client] = []
real_close = httpx.Client.close

def _tracking_close(self: httpx.Client) -> None:
closed.append(self)
real_close(self)

monkeypatch.setattr(httpx.Client, "close", _tracking_close)
owned = VercelBlobObjectStore(_config(), token=_TOKEN)
owned.close()
assert closed == [owned._client]
owned.close()
assert closed == [owned._client]

injected = httpx.Client(transport=_FakeTransport())
store = VercelBlobObjectStore(_config(), token=_TOKEN, client=injected)
store.close()
assert injected not in closed
injected.close()
Loading