diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 88f77b49f9..bdaceae0e1 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -7,7 +7,10 @@ import base64 import json import logging import os +import tempfile +import threading from collections.abc import AsyncIterable, AsyncIterator, Generator, Mapping, Sequence +from contextlib import suppress from typing import Protocol, cast from agent_framework import ( @@ -139,38 +142,61 @@ class InMemoryFunctionApprovalStorage: class FileBasedFunctionApprovalStorage: - """A simple file-based storage for function approval requests.""" + """A simple file-based storage for function approval requests. + + Concurrent writes from multiple threads in the same process are + serialized by a ``threading.Lock``, and the on-disk JSON file is + updated atomically (write to a temp file, then ``os.replace``) so a + crash mid-write cannot leave a partially written file behind. + """ def __init__(self, storage_path: str) -> None: self._storage_path = storage_path + self._lock = threading.Lock() def _create_storage_file_if_not_exists_sync(self) -> None: - """Lazy-create the storage file (and its parent directory) if it does not already exist.""" - if not os.path.exists(self._storage_path): - os.makedirs(os.path.dirname(self._storage_path), exist_ok=True) - with open(self._storage_path, "w") as f: - json.dump({}, f) + """Lazy-create the storage file (and its parent directory) if it does not already exist. + + Uses exclusive-create mode (``"x"``) so a concurrent creator cannot + be truncated by an ``open(..., "w")`` after a stale existence check. + """ + os.makedirs(os.path.dirname(self._storage_path) or ".", exist_ok=True) + with suppress(FileExistsError), open(self._storage_path, "x") as f: + json.dump({}, f) + + def _atomic_write(self, data: dict[str, Any]) -> None: + """Atomically replace the storage file with the serialized ``data``.""" + directory = os.path.dirname(self._storage_path) or "." + # Serialize first so any error doesn't leave a partial file behind. + serialized = json.dumps(data) + fd, tmp_path = tempfile.mkstemp(prefix=".approvals-", suffix=".tmp", dir=directory) + try: + with os.fdopen(fd, "w") as tmp: + tmp.write(serialized) + os.replace(tmp_path, self._storage_path) + except BaseException: + with suppress(OSError): + os.unlink(tmp_path) + raise def _save_sync(self, approval_request_id: str, request: Content) -> None: - self._create_storage_file_if_not_exists_sync() - with open(self._storage_path, "r+") as f: - data = json.load(f) + with self._lock: + self._create_storage_file_if_not_exists_sync() + with open(self._storage_path) as f: + data = json.load(f) if approval_request_id in data: raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.") data[approval_request_id] = request.to_dict() - # Serialize to a string first so any error doesn't leave the file in a partially written state. - serialized = json.dumps(data) - f.seek(0) - f.write(serialized) - f.truncate() + self._atomic_write(data) def _load_sync(self, approval_request_id: str) -> Content: - self._create_storage_file_if_not_exists_sync() - with open(self._storage_path) as f: - data = json.load(f) - if approval_request_id not in data: - raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.") - return Content.from_dict(data[approval_request_id]) + with self._lock: + self._create_storage_file_if_not_exists_sync() + with open(self._storage_path) as f: + data = json.load(f) + if approval_request_id not in data: + raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.") + return Content.from_dict(data[approval_request_id]) async def save_approval_request(self, approval_request_id: str, request: Content) -> None: await asyncio.to_thread(self._save_sync, approval_request_id, request) @@ -243,7 +269,7 @@ class ResponsesHostServer(ResponsesAgentServerHost): self._agent = agent self._approval_storage = ( - FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH.lstrip("/")) + FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH) if self.config.is_hosted else InMemoryFunctionApprovalStorage() )