agent_search_gateway.daemon
Foreground Unix-socket daemon and graceful shutdown lifecycle.
1"""Foreground Unix-socket daemon and graceful shutdown lifecycle.""" 2 3import asyncio 4import logging 5import os 6import stat 7import time 8from collections.abc import Awaitable, Callable, Mapping 9from contextlib import suppress 10from pathlib import Path 11from typing import Protocol 12 13from .config import load_toml, resolve_config 14from .errors import ConfigFailure, ErrorCode, GatewayError 15from .models import ( 16 ErrorResponse, 17 KeywordSearchRequest, 18 LLMSearchRequest, 19 LLMSearchScope, 20 PaperSearchRequest, 21 Request, 22 Response, 23 ShutdownRequest, 24 SuccessResponse, 25 URLFetchRequest, 26) 27from .observability import DebugLoggingSession, elapsed_ms, log_event 28from .paths import RuntimePaths 29from .protocol import NDJSONDecoder, encode_response 30from .providers.academic.defaults import ( 31 build_default_academic_registry, 32 build_default_oa_resolver_registry, 33) 34from .providers.defaults import build_default_registry 35from .request_ids import RequestIdFactory, RequestIdRegistry, bind_request_id, generate_request_id 36from .runtime import Runtime 37from .socket_probe import SocketState, probe_unix_socket 38 39_SHUTDOWN_GRACE_SECONDS = 10.0 40_SOCKET_PROBE_TIMEOUT_SECONDS = 2.0 41 42 43class _SearchOrchestratorLike(Protocol): 44 async def keyword_search(self, query: str, *, request_id: str) -> str: ... 45 46 async def llm_search( 47 self, 48 prompt: str, 49 *, 50 request_id: str, 51 scope: LLMSearchScope = "web", 52 ) -> str: ... 53 54 55class _PaperSearchOrchestratorLike(Protocol): 56 async def paper_search(self, query: str, *, request_id: str) -> str: ... 57 58 59class _FetchOrchestratorLike(Protocol): 60 async def url_fetch(self, url: str, focus: str | None = None) -> str: ... 61 62 63class RuntimeLike(Protocol): 64 @property 65 def search_orchestrator(self) -> _SearchOrchestratorLike: ... 66 67 @property 68 def paper_search_orchestrator(self) -> _PaperSearchOrchestratorLike: ... 69 70 @property 71 def fetch_orchestrator(self) -> _FetchOrchestratorLike: ... 72 73 async def aclose(self) -> None: ... 74 75 76RuntimeFactory = Callable[[], RuntimeLike] 77ShutdownWaiter = Callable[[tuple[asyncio.Task[object], ...], float], Awaitable[None]] 78MonotonicClock = Callable[[], float] 79 80 81async def _default_shutdown_waiter( 82 tasks: tuple[asyncio.Task[object], ...], 83 timeout: float, 84) -> None: 85 if not tasks: 86 return 87 _, pending = await asyncio.wait(tasks, timeout=timeout) 88 if pending: 89 raise TimeoutError 90 91 92def _command_name( 93 request: KeywordSearchRequest | PaperSearchRequest | LLMSearchRequest | URLFetchRequest, 94) -> str: 95 if isinstance(request, KeywordSearchRequest): 96 return "keyword-search" 97 if isinstance(request, PaperSearchRequest): 98 return "paper-search" 99 if isinstance(request, LLMSearchRequest): 100 return "llm-search" 101 return "url-fetch" 102 103 104class ForegroundDaemon: 105 def __init__( 106 self, 107 paths: RuntimePaths, 108 *, 109 runtime_factory: RuntimeFactory | None = None, 110 environ: Mapping[str, str] | None = None, 111 logger: logging.Logger | None = None, 112 shutdown_waiter: ShutdownWaiter = _default_shutdown_waiter, 113 request_id_factory: RequestIdFactory = generate_request_id, 114 debug: bool = False, 115 logging_session: DebugLoggingSession | None = None, 116 monotonic: MonotonicClock = time.monotonic, 117 ) -> None: 118 self.paths = paths 119 self.ready = asyncio.Event() 120 self.stopped = asyncio.Event() 121 self._logger = logger or logging.getLogger(__name__) 122 self._environ = dict(os.environ if environ is None else environ) 123 self._runtime_factory = runtime_factory 124 self._shutdown_waiter = shutdown_waiter 125 self._request_ids = RequestIdRegistry(paths.results_dir, factory=request_id_factory) 126 self._debug = debug 127 self._logging_session = logging_session 128 self._monotonic = monotonic 129 self._session_started = False 130 self._session_stopped = False 131 self._runtime: RuntimeLike | None = None 132 self._server: asyncio.AbstractServer | None = None 133 self._socket_identity: tuple[int, int] | None = None 134 self._state_lock = asyncio.Lock() 135 self._shutting_down = False 136 self._shutdown_task: asyncio.Task[None] | None = None 137 self._shutdown_response_ready = asyncio.Event() 138 self._active_workflows: set[asyncio.Task[object]] = set() 139 140 @property 141 def shutting_down(self) -> bool: 142 return self._shutting_down 143 144 async def start(self) -> None: 145 self.paths.socket_file.parent.mkdir(parents=True, exist_ok=True) 146 self.paths.results_dir.mkdir(parents=True, exist_ok=True) 147 try: 148 await self._prepare_socket_path() 149 self._runtime = self._create_runtime() 150 try: 151 self._server = await asyncio.start_unix_server( 152 self._handle_connection, 153 path=self.paths.socket_file, 154 ) 155 except OSError as exc: 156 await self._close_runtime_after_startup_failure() 157 raise ConfigFailure( 158 ErrorCode.CONFIG_ERROR, 159 f"Failed to bind daemon socket: {self.paths.socket_file}", 160 ) from exc 161 socket_stat = self.paths.socket_file.stat() 162 self._socket_identity = (socket_stat.st_dev, socket_stat.st_ino) 163 if self._debug: 164 log_event( 165 self._logger, 166 logging.INFO, 167 "session_started", 168 pid=os.getpid(), 169 debug=True, 170 ) 171 self._session_started = True 172 self.ready.set() 173 await self.stopped.wait() 174 except (asyncio.CancelledError, KeyboardInterrupt): 175 await self.stop_for_test() 176 raise 177 178 async def _prepare_socket_path(self) -> None: 179 path = self.paths.socket_file 180 probe = await probe_unix_socket( 181 path, 182 timeout_seconds=_SOCKET_PROBE_TIMEOUT_SECONDS, 183 connector=asyncio.open_unix_connection, 184 ) 185 if probe.state is SocketState.MISSING: 186 return 187 if probe.state is SocketState.LIVE: 188 raise ConfigFailure( 189 ErrorCode.CONFIG_ERROR, 190 f"Daemon is already running at: {path}", 191 ) 192 if probe.state is SocketState.TIMEOUT: 193 raise ConfigFailure( 194 ErrorCode.CONFIG_ERROR, 195 f"Daemon socket did not respond in time: {path}", 196 ) 197 if probe.state is SocketState.NOT_SOCKET: 198 raise ConfigFailure( 199 ErrorCode.CONFIG_ERROR, 200 f"Daemon socket path is not a Unix socket: {path}", 201 ) 202 if probe.state is SocketState.OS_ERROR: 203 raise ConfigFailure( 204 ErrorCode.CONFIG_ERROR, 205 f"Unable to inspect daemon socket: {path}", 206 ) 207 208 identity = probe.identity 209 if probe.state is not SocketState.REFUSED or identity is None: 210 raise RuntimeError("unexpected daemon socket probe state") 211 try: 212 current = path.lstat() 213 except FileNotFoundError: 214 return 215 if not stat.S_ISSOCK(current.st_mode) or (current.st_dev, current.st_ino) != identity: 216 raise ConfigFailure( 217 ErrorCode.CONFIG_ERROR, 218 f"Daemon socket changed during startup: {path}", 219 ) 220 try: 221 path.unlink() 222 except FileNotFoundError: 223 return 224 except OSError as exc: 225 raise ConfigFailure( 226 ErrorCode.CONFIG_ERROR, 227 f"Failed to remove stale daemon socket: {path}", 228 ) from exc 229 230 def _create_runtime(self) -> RuntimeLike: 231 try: 232 if self._runtime_factory is not None: 233 return self._runtime_factory() 234 registry = build_default_registry() 235 academic_registry = build_default_academic_registry() 236 resolver_registry = build_default_oa_resolver_registry() 237 data = load_toml(self.paths.config_file) 238 config = resolve_config( 239 data, 240 registry, 241 self._environ, 242 academic_registry=academic_registry, 243 oa_resolver_registry=resolver_registry, 244 ) 245 if self._logging_session is not None: 246 web_secrets = ( 247 provider.secret 248 for provider in config.web.providers 249 if provider.secret is not None 250 ) 251 llm_secrets = (provider.secret for provider in config.llm.providers) 252 academic_secrets = ( 253 secret 254 for provider in config.academic.providers 255 for secret in (provider.api_key, provider.contact_email) 256 if secret is not None 257 ) 258 resolver_secrets = ( 259 secret 260 for resolver in (config.oa_resolver,) 261 if resolver is not None 262 for secret in (resolver.api_key, resolver.contact_email) 263 if secret is not None 264 ) 265 self._logging_session.add_secrets( 266 web_secrets, 267 llm_secrets, 268 academic_secrets, 269 resolver_secrets, 270 ) 271 return Runtime.build( 272 config, 273 self.paths, 274 registry=registry, 275 academic_registry=academic_registry, 276 oa_resolver_registry=resolver_registry, 277 ) 278 except ConfigFailure: 279 raise 280 except Exception as exc: 281 raise ConfigFailure( 282 ErrorCode.CONFIG_ERROR, 283 "Failed to initialize daemon runtime", 284 ) from exc 285 286 async def _close_runtime_after_startup_failure(self) -> None: 287 runtime = self._runtime 288 self._runtime = None 289 if runtime is None: 290 return 291 try: 292 await runtime.aclose() 293 except Exception: 294 self._logger.error("daemon runtime close failed during startup cleanup") 295 296 async def _handle_connection( 297 self, 298 reader: asyncio.StreamReader, 299 writer: asyncio.StreamWriter, 300 ) -> None: 301 decoder = NDJSONDecoder() 302 try: 303 while data := await reader.read(4096): 304 for decoded in decoder.feed(data): 305 if isinstance(decoded, ErrorResponse): 306 response: Response = decoded 307 else: 308 response = await self._dispatch(decoded) 309 writer.write(encode_response(response)) 310 await writer.drain() 311 except (ConnectionError, BrokenPipeError): 312 return 313 finally: 314 writer.close() 315 with suppress(ConnectionError, BrokenPipeError): 316 await writer.wait_closed() 317 318 async def _dispatch(self, request: Request) -> Response: 319 if isinstance(request, ShutdownRequest): 320 await self._begin_shutdown() 321 await self._shutdown_response_ready.wait() 322 return SuccessResponse("Daemon stopped.") 323 if not isinstance( 324 request, 325 (KeywordSearchRequest, PaperSearchRequest, LLMSearchRequest, URLFetchRequest), 326 ): 327 return ErrorResponse(ErrorCode.BAD_REQUEST, "Unknown request type") 328 329 is_search = isinstance( 330 request, 331 (KeywordSearchRequest, PaperSearchRequest, LLMSearchRequest), 332 ) 333 with ( 334 self._request_ids.reserve(may_write_search_result=is_search) as request_id, 335 bind_request_id(request_id), 336 ): 337 return await self._dispatch_business(request, request_id=request_id) 338 339 async def _dispatch_business( 340 self, 341 request: KeywordSearchRequest | PaperSearchRequest | LLMSearchRequest | URLFetchRequest, 342 *, 343 request_id: str, 344 ) -> Response: 345 command = _command_name(request) 346 started = self._monotonic() 347 log_event(self._logger, logging.DEBUG, "workflow_started", command=command) 348 current = asyncio.current_task() 349 if current is None: 350 return self._internal_workflow_failure(command, started, RuntimeError("missing task")) 351 352 async with self._state_lock: 353 if self._shutting_down: 354 log_event( 355 self._logger, 356 logging.DEBUG, 357 "workflow_rejected", 358 command=command, 359 elapsed_ms=elapsed_ms(self._monotonic, started), 360 error_code=ErrorCode.DAEMON_SHUTTING_DOWN.value, 361 ) 362 return ErrorResponse( 363 ErrorCode.DAEMON_SHUTTING_DOWN, 364 "Daemon is shutting down", 365 ) 366 self._active_workflows.add(current) 367 368 try: 369 text = await self._invoke_workflow(request, request_id=request_id) 370 except asyncio.CancelledError: 371 log_event( 372 self._logger, 373 logging.DEBUG, 374 "workflow_cancelled", 375 command=command, 376 elapsed_ms=elapsed_ms(self._monotonic, started), 377 ) 378 raise 379 except GatewayError as exc: 380 log_event( 381 self._logger, 382 logging.DEBUG, 383 "workflow_failed", 384 command=command, 385 elapsed_ms=elapsed_ms(self._monotonic, started), 386 error_code=exc.code.value, 387 error_type=type(exc).__name__, 388 ) 389 return ErrorResponse(exc.code, exc.message) 390 except Exception as exc: 391 return self._internal_workflow_failure(command, started, exc) 392 else: 393 log_event( 394 self._logger, 395 logging.DEBUG, 396 "workflow_completed", 397 command=command, 398 elapsed_ms=elapsed_ms(self._monotonic, started), 399 ) 400 return SuccessResponse(text) 401 finally: 402 async with self._state_lock: 403 self._active_workflows.discard(current) 404 405 async def _invoke_workflow( 406 self, 407 request: KeywordSearchRequest | PaperSearchRequest | LLMSearchRequest | URLFetchRequest, 408 *, 409 request_id: str, 410 ) -> str: 411 runtime = self._require_runtime() 412 if isinstance(request, KeywordSearchRequest): 413 return await runtime.search_orchestrator.keyword_search( 414 request.query, 415 request_id=request_id, 416 ) 417 if isinstance(request, PaperSearchRequest): 418 return await runtime.paper_search_orchestrator.paper_search( 419 request.query, 420 request_id=request_id, 421 ) 422 if isinstance(request, LLMSearchRequest): 423 return await runtime.search_orchestrator.llm_search( 424 request.prompt, 425 request_id=request_id, 426 scope=request.scope, 427 ) 428 return await runtime.fetch_orchestrator.url_fetch(request.url, request.focus) 429 430 def _internal_workflow_failure( 431 self, 432 command: str, 433 started: float, 434 exc: Exception, 435 ) -> ErrorResponse: 436 if self._debug: 437 log_event( 438 self._logger, 439 logging.ERROR, 440 "workflow_failed", 441 command=command, 442 elapsed_ms=elapsed_ms(self._monotonic, started), 443 error_type=type(exc).__name__, 444 exc_info=exc, 445 ) 446 else: 447 self._logger.error( 448 "unexpected daemon workflow failure type=%s", 449 type(exc).__name__, 450 ) 451 return ErrorResponse(ErrorCode.PROTOCOL_ERROR, "Internal daemon error") 452 453 def _require_runtime(self) -> RuntimeLike: 454 if self._runtime is None: 455 raise RuntimeError("daemon runtime is not initialized") 456 return self._runtime 457 458 async def _begin_shutdown(self) -> asyncio.Task[None]: 459 async with self._state_lock: 460 if self._shutdown_task is None: 461 self._shutting_down = True 462 self._shutdown_task = asyncio.create_task(self._shutdown_coordinator()) 463 return self._shutdown_task 464 465 async def _shutdown_coordinator(self) -> None: 466 server: asyncio.AbstractServer | None = None 467 try: 468 async with self._state_lock: 469 active = tuple(self._active_workflows) 470 try: 471 await self._shutdown_waiter(active, _SHUTDOWN_GRACE_SECONDS) 472 except TimeoutError: 473 for task in active: 474 if not task.done(): 475 task.cancel() 476 if active: 477 await asyncio.gather(*active, return_exceptions=True) 478 server = await self._cleanup_before_response() 479 self._shutdown_response_ready.set() 480 if server is not None: 481 try: 482 await server.wait_closed() 483 except Exception: 484 self._logger.error("daemon socket server close failed") 485 finally: 486 self._shutdown_response_ready.set() 487 self.stopped.set() 488 489 async def _cleanup_before_response(self) -> asyncio.AbstractServer | None: 490 self._emit_session_stopped() 491 runtime = self._runtime 492 self._runtime = None 493 if runtime is not None: 494 try: 495 await runtime.aclose() 496 except Exception: 497 self._logger.error("daemon runtime close failed") 498 499 server = self._server 500 self._server = None 501 if server is not None: 502 try: 503 server.close() 504 except Exception: 505 self._logger.error("daemon socket server close failed") 506 507 if self._owns_socket(self.paths.socket_file): 508 try: 509 self.paths.socket_file.unlink() 510 except FileNotFoundError: 511 pass 512 except OSError: 513 self._logger.error("daemon socket unlink failed") 514 return server 515 516 def _emit_session_stopped(self) -> None: 517 if not self._session_started or self._session_stopped: 518 return 519 self._session_stopped = True 520 log_event( 521 self._logger, 522 logging.INFO, 523 "session_stopped", 524 pid=os.getpid(), 525 debug=self._debug, 526 ) 527 528 def _owns_socket(self, path: Path) -> bool: 529 identity = self._socket_identity 530 if identity is None: 531 return False 532 try: 533 stat = path.stat() 534 except FileNotFoundError: 535 return False 536 return (stat.st_dev, stat.st_ino) == identity 537 538 async def stop_for_test(self) -> None: 539 task = await self._begin_shutdown() 540 await task
class
RuntimeLike(typing.Protocol):
64class RuntimeLike(Protocol): 65 @property 66 def search_orchestrator(self) -> _SearchOrchestratorLike: ... 67 68 @property 69 def paper_search_orchestrator(self) -> _PaperSearchOrchestratorLike: ... 70 71 @property 72 def fetch_orchestrator(self) -> _FetchOrchestratorLike: ... 73 74 async def aclose(self) -> None: ...
Base class for protocol classes.
Protocol classes are defined as::
class Proto(Protocol):
def meth(self) -> int:
...
Such classes are primarily used with static type checkers that recognize structural subtyping (static duck-typing).
For example::
class C:
def meth(self) -> int:
return 0
def func(x: Proto) -> int:
return x.meth()
func(C()) # Passes static type check
See PEP 544 for details. Protocol classes decorated with @typing.runtime_checkable act as simple-minded runtime protocols that check only the presence of given attributes, ignoring their type signatures. Protocol classes can be generic, they are defined as::
class GenProto[T](Protocol):
def meth(self) -> T:
...
RuntimeLike(*args, **kwargs)
1739def _no_init_or_replace_init(self, *args, **kwargs): 1740 cls = type(self) 1741 1742 if cls._is_protocol: 1743 raise TypeError('Protocols cannot be instantiated') 1744 1745 # Already using a custom `__init__`. No need to calculate correct 1746 # `__init__` to call. This can lead to RecursionError. See bpo-45121. 1747 if cls.__init__ is not _no_init_or_replace_init: 1748 return 1749 1750 # Initially, `__init__` of a protocol subclass is set to `_no_init_or_replace_init`. 1751 # The first instantiation of the subclass will call `_no_init_or_replace_init` which 1752 # searches for a proper new `__init__` in the MRO. The new `__init__` 1753 # replaces the subclass' old `__init__` (ie `_no_init_or_replace_init`). Subsequent 1754 # instantiation of the protocol subclass will thus use the new 1755 # `__init__` and no longer call `_no_init_or_replace_init`. 1756 for base in cls.__mro__: 1757 init = base.__dict__.get('__init__', _no_init_or_replace_init) 1758 if init is not _no_init_or_replace_init: 1759 cls.__init__ = init 1760 break 1761 else: 1762 # should not happen 1763 cls.__init__ = object.__init__ 1764 1765 cls.__init__(self, *args, **kwargs)
RuntimeFactory =
collections.abc.Callable[[], RuntimeLike]
ShutdownWaiter =
collections.abc.Callable[[tuple[_asyncio.Task[object], ...], float], collections.abc.Awaitable[None]]
MonotonicClock =
collections.abc.Callable[[], float]
class
ForegroundDaemon:
105class ForegroundDaemon: 106 def __init__( 107 self, 108 paths: RuntimePaths, 109 *, 110 runtime_factory: RuntimeFactory | None = None, 111 environ: Mapping[str, str] | None = None, 112 logger: logging.Logger | None = None, 113 shutdown_waiter: ShutdownWaiter = _default_shutdown_waiter, 114 request_id_factory: RequestIdFactory = generate_request_id, 115 debug: bool = False, 116 logging_session: DebugLoggingSession | None = None, 117 monotonic: MonotonicClock = time.monotonic, 118 ) -> None: 119 self.paths = paths 120 self.ready = asyncio.Event() 121 self.stopped = asyncio.Event() 122 self._logger = logger or logging.getLogger(__name__) 123 self._environ = dict(os.environ if environ is None else environ) 124 self._runtime_factory = runtime_factory 125 self._shutdown_waiter = shutdown_waiter 126 self._request_ids = RequestIdRegistry(paths.results_dir, factory=request_id_factory) 127 self._debug = debug 128 self._logging_session = logging_session 129 self._monotonic = monotonic 130 self._session_started = False 131 self._session_stopped = False 132 self._runtime: RuntimeLike | None = None 133 self._server: asyncio.AbstractServer | None = None 134 self._socket_identity: tuple[int, int] | None = None 135 self._state_lock = asyncio.Lock() 136 self._shutting_down = False 137 self._shutdown_task: asyncio.Task[None] | None = None 138 self._shutdown_response_ready = asyncio.Event() 139 self._active_workflows: set[asyncio.Task[object]] = set() 140 141 @property 142 def shutting_down(self) -> bool: 143 return self._shutting_down 144 145 async def start(self) -> None: 146 self.paths.socket_file.parent.mkdir(parents=True, exist_ok=True) 147 self.paths.results_dir.mkdir(parents=True, exist_ok=True) 148 try: 149 await self._prepare_socket_path() 150 self._runtime = self._create_runtime() 151 try: 152 self._server = await asyncio.start_unix_server( 153 self._handle_connection, 154 path=self.paths.socket_file, 155 ) 156 except OSError as exc: 157 await self._close_runtime_after_startup_failure() 158 raise ConfigFailure( 159 ErrorCode.CONFIG_ERROR, 160 f"Failed to bind daemon socket: {self.paths.socket_file}", 161 ) from exc 162 socket_stat = self.paths.socket_file.stat() 163 self._socket_identity = (socket_stat.st_dev, socket_stat.st_ino) 164 if self._debug: 165 log_event( 166 self._logger, 167 logging.INFO, 168 "session_started", 169 pid=os.getpid(), 170 debug=True, 171 ) 172 self._session_started = True 173 self.ready.set() 174 await self.stopped.wait() 175 except (asyncio.CancelledError, KeyboardInterrupt): 176 await self.stop_for_test() 177 raise 178 179 async def _prepare_socket_path(self) -> None: 180 path = self.paths.socket_file 181 probe = await probe_unix_socket( 182 path, 183 timeout_seconds=_SOCKET_PROBE_TIMEOUT_SECONDS, 184 connector=asyncio.open_unix_connection, 185 ) 186 if probe.state is SocketState.MISSING: 187 return 188 if probe.state is SocketState.LIVE: 189 raise ConfigFailure( 190 ErrorCode.CONFIG_ERROR, 191 f"Daemon is already running at: {path}", 192 ) 193 if probe.state is SocketState.TIMEOUT: 194 raise ConfigFailure( 195 ErrorCode.CONFIG_ERROR, 196 f"Daemon socket did not respond in time: {path}", 197 ) 198 if probe.state is SocketState.NOT_SOCKET: 199 raise ConfigFailure( 200 ErrorCode.CONFIG_ERROR, 201 f"Daemon socket path is not a Unix socket: {path}", 202 ) 203 if probe.state is SocketState.OS_ERROR: 204 raise ConfigFailure( 205 ErrorCode.CONFIG_ERROR, 206 f"Unable to inspect daemon socket: {path}", 207 ) 208 209 identity = probe.identity 210 if probe.state is not SocketState.REFUSED or identity is None: 211 raise RuntimeError("unexpected daemon socket probe state") 212 try: 213 current = path.lstat() 214 except FileNotFoundError: 215 return 216 if not stat.S_ISSOCK(current.st_mode) or (current.st_dev, current.st_ino) != identity: 217 raise ConfigFailure( 218 ErrorCode.CONFIG_ERROR, 219 f"Daemon socket changed during startup: {path}", 220 ) 221 try: 222 path.unlink() 223 except FileNotFoundError: 224 return 225 except OSError as exc: 226 raise ConfigFailure( 227 ErrorCode.CONFIG_ERROR, 228 f"Failed to remove stale daemon socket: {path}", 229 ) from exc 230 231 def _create_runtime(self) -> RuntimeLike: 232 try: 233 if self._runtime_factory is not None: 234 return self._runtime_factory() 235 registry = build_default_registry() 236 academic_registry = build_default_academic_registry() 237 resolver_registry = build_default_oa_resolver_registry() 238 data = load_toml(self.paths.config_file) 239 config = resolve_config( 240 data, 241 registry, 242 self._environ, 243 academic_registry=academic_registry, 244 oa_resolver_registry=resolver_registry, 245 ) 246 if self._logging_session is not None: 247 web_secrets = ( 248 provider.secret 249 for provider in config.web.providers 250 if provider.secret is not None 251 ) 252 llm_secrets = (provider.secret for provider in config.llm.providers) 253 academic_secrets = ( 254 secret 255 for provider in config.academic.providers 256 for secret in (provider.api_key, provider.contact_email) 257 if secret is not None 258 ) 259 resolver_secrets = ( 260 secret 261 for resolver in (config.oa_resolver,) 262 if resolver is not None 263 for secret in (resolver.api_key, resolver.contact_email) 264 if secret is not None 265 ) 266 self._logging_session.add_secrets( 267 web_secrets, 268 llm_secrets, 269 academic_secrets, 270 resolver_secrets, 271 ) 272 return Runtime.build( 273 config, 274 self.paths, 275 registry=registry, 276 academic_registry=academic_registry, 277 oa_resolver_registry=resolver_registry, 278 ) 279 except ConfigFailure: 280 raise 281 except Exception as exc: 282 raise ConfigFailure( 283 ErrorCode.CONFIG_ERROR, 284 "Failed to initialize daemon runtime", 285 ) from exc 286 287 async def _close_runtime_after_startup_failure(self) -> None: 288 runtime = self._runtime 289 self._runtime = None 290 if runtime is None: 291 return 292 try: 293 await runtime.aclose() 294 except Exception: 295 self._logger.error("daemon runtime close failed during startup cleanup") 296 297 async def _handle_connection( 298 self, 299 reader: asyncio.StreamReader, 300 writer: asyncio.StreamWriter, 301 ) -> None: 302 decoder = NDJSONDecoder() 303 try: 304 while data := await reader.read(4096): 305 for decoded in decoder.feed(data): 306 if isinstance(decoded, ErrorResponse): 307 response: Response = decoded 308 else: 309 response = await self._dispatch(decoded) 310 writer.write(encode_response(response)) 311 await writer.drain() 312 except (ConnectionError, BrokenPipeError): 313 return 314 finally: 315 writer.close() 316 with suppress(ConnectionError, BrokenPipeError): 317 await writer.wait_closed() 318 319 async def _dispatch(self, request: Request) -> Response: 320 if isinstance(request, ShutdownRequest): 321 await self._begin_shutdown() 322 await self._shutdown_response_ready.wait() 323 return SuccessResponse("Daemon stopped.") 324 if not isinstance( 325 request, 326 (KeywordSearchRequest, PaperSearchRequest, LLMSearchRequest, URLFetchRequest), 327 ): 328 return ErrorResponse(ErrorCode.BAD_REQUEST, "Unknown request type") 329 330 is_search = isinstance( 331 request, 332 (KeywordSearchRequest, PaperSearchRequest, LLMSearchRequest), 333 ) 334 with ( 335 self._request_ids.reserve(may_write_search_result=is_search) as request_id, 336 bind_request_id(request_id), 337 ): 338 return await self._dispatch_business(request, request_id=request_id) 339 340 async def _dispatch_business( 341 self, 342 request: KeywordSearchRequest | PaperSearchRequest | LLMSearchRequest | URLFetchRequest, 343 *, 344 request_id: str, 345 ) -> Response: 346 command = _command_name(request) 347 started = self._monotonic() 348 log_event(self._logger, logging.DEBUG, "workflow_started", command=command) 349 current = asyncio.current_task() 350 if current is None: 351 return self._internal_workflow_failure(command, started, RuntimeError("missing task")) 352 353 async with self._state_lock: 354 if self._shutting_down: 355 log_event( 356 self._logger, 357 logging.DEBUG, 358 "workflow_rejected", 359 command=command, 360 elapsed_ms=elapsed_ms(self._monotonic, started), 361 error_code=ErrorCode.DAEMON_SHUTTING_DOWN.value, 362 ) 363 return ErrorResponse( 364 ErrorCode.DAEMON_SHUTTING_DOWN, 365 "Daemon is shutting down", 366 ) 367 self._active_workflows.add(current) 368 369 try: 370 text = await self._invoke_workflow(request, request_id=request_id) 371 except asyncio.CancelledError: 372 log_event( 373 self._logger, 374 logging.DEBUG, 375 "workflow_cancelled", 376 command=command, 377 elapsed_ms=elapsed_ms(self._monotonic, started), 378 ) 379 raise 380 except GatewayError as exc: 381 log_event( 382 self._logger, 383 logging.DEBUG, 384 "workflow_failed", 385 command=command, 386 elapsed_ms=elapsed_ms(self._monotonic, started), 387 error_code=exc.code.value, 388 error_type=type(exc).__name__, 389 ) 390 return ErrorResponse(exc.code, exc.message) 391 except Exception as exc: 392 return self._internal_workflow_failure(command, started, exc) 393 else: 394 log_event( 395 self._logger, 396 logging.DEBUG, 397 "workflow_completed", 398 command=command, 399 elapsed_ms=elapsed_ms(self._monotonic, started), 400 ) 401 return SuccessResponse(text) 402 finally: 403 async with self._state_lock: 404 self._active_workflows.discard(current) 405 406 async def _invoke_workflow( 407 self, 408 request: KeywordSearchRequest | PaperSearchRequest | LLMSearchRequest | URLFetchRequest, 409 *, 410 request_id: str, 411 ) -> str: 412 runtime = self._require_runtime() 413 if isinstance(request, KeywordSearchRequest): 414 return await runtime.search_orchestrator.keyword_search( 415 request.query, 416 request_id=request_id, 417 ) 418 if isinstance(request, PaperSearchRequest): 419 return await runtime.paper_search_orchestrator.paper_search( 420 request.query, 421 request_id=request_id, 422 ) 423 if isinstance(request, LLMSearchRequest): 424 return await runtime.search_orchestrator.llm_search( 425 request.prompt, 426 request_id=request_id, 427 scope=request.scope, 428 ) 429 return await runtime.fetch_orchestrator.url_fetch(request.url, request.focus) 430 431 def _internal_workflow_failure( 432 self, 433 command: str, 434 started: float, 435 exc: Exception, 436 ) -> ErrorResponse: 437 if self._debug: 438 log_event( 439 self._logger, 440 logging.ERROR, 441 "workflow_failed", 442 command=command, 443 elapsed_ms=elapsed_ms(self._monotonic, started), 444 error_type=type(exc).__name__, 445 exc_info=exc, 446 ) 447 else: 448 self._logger.error( 449 "unexpected daemon workflow failure type=%s", 450 type(exc).__name__, 451 ) 452 return ErrorResponse(ErrorCode.PROTOCOL_ERROR, "Internal daemon error") 453 454 def _require_runtime(self) -> RuntimeLike: 455 if self._runtime is None: 456 raise RuntimeError("daemon runtime is not initialized") 457 return self._runtime 458 459 async def _begin_shutdown(self) -> asyncio.Task[None]: 460 async with self._state_lock: 461 if self._shutdown_task is None: 462 self._shutting_down = True 463 self._shutdown_task = asyncio.create_task(self._shutdown_coordinator()) 464 return self._shutdown_task 465 466 async def _shutdown_coordinator(self) -> None: 467 server: asyncio.AbstractServer | None = None 468 try: 469 async with self._state_lock: 470 active = tuple(self._active_workflows) 471 try: 472 await self._shutdown_waiter(active, _SHUTDOWN_GRACE_SECONDS) 473 except TimeoutError: 474 for task in active: 475 if not task.done(): 476 task.cancel() 477 if active: 478 await asyncio.gather(*active, return_exceptions=True) 479 server = await self._cleanup_before_response() 480 self._shutdown_response_ready.set() 481 if server is not None: 482 try: 483 await server.wait_closed() 484 except Exception: 485 self._logger.error("daemon socket server close failed") 486 finally: 487 self._shutdown_response_ready.set() 488 self.stopped.set() 489 490 async def _cleanup_before_response(self) -> asyncio.AbstractServer | None: 491 self._emit_session_stopped() 492 runtime = self._runtime 493 self._runtime = None 494 if runtime is not None: 495 try: 496 await runtime.aclose() 497 except Exception: 498 self._logger.error("daemon runtime close failed") 499 500 server = self._server 501 self._server = None 502 if server is not None: 503 try: 504 server.close() 505 except Exception: 506 self._logger.error("daemon socket server close failed") 507 508 if self._owns_socket(self.paths.socket_file): 509 try: 510 self.paths.socket_file.unlink() 511 except FileNotFoundError: 512 pass 513 except OSError: 514 self._logger.error("daemon socket unlink failed") 515 return server 516 517 def _emit_session_stopped(self) -> None: 518 if not self._session_started or self._session_stopped: 519 return 520 self._session_stopped = True 521 log_event( 522 self._logger, 523 logging.INFO, 524 "session_stopped", 525 pid=os.getpid(), 526 debug=self._debug, 527 ) 528 529 def _owns_socket(self, path: Path) -> bool: 530 identity = self._socket_identity 531 if identity is None: 532 return False 533 try: 534 stat = path.stat() 535 except FileNotFoundError: 536 return False 537 return (stat.st_dev, stat.st_ino) == identity 538 539 async def stop_for_test(self) -> None: 540 task = await self._begin_shutdown() 541 await task
ForegroundDaemon( paths: agent_search_gateway.paths.RuntimePaths, *, runtime_factory: Callable[[], RuntimeLike] | None = None, environ: Mapping[str, str] | None = None, logger: logging.Logger | None = None, shutdown_waiter: Callable[[tuple[_asyncio.Task[object], ...], float], Awaitable[None]] = <function _default_shutdown_waiter>, request_id_factory: Callable[[], str] = <function generate_request_id>, debug: bool = False, logging_session: agent_search_gateway.observability.DebugLoggingSession | None = None, monotonic: Callable[[], float] = <built-in function monotonic>)
106 def __init__( 107 self, 108 paths: RuntimePaths, 109 *, 110 runtime_factory: RuntimeFactory | None = None, 111 environ: Mapping[str, str] | None = None, 112 logger: logging.Logger | None = None, 113 shutdown_waiter: ShutdownWaiter = _default_shutdown_waiter, 114 request_id_factory: RequestIdFactory = generate_request_id, 115 debug: bool = False, 116 logging_session: DebugLoggingSession | None = None, 117 monotonic: MonotonicClock = time.monotonic, 118 ) -> None: 119 self.paths = paths 120 self.ready = asyncio.Event() 121 self.stopped = asyncio.Event() 122 self._logger = logger or logging.getLogger(__name__) 123 self._environ = dict(os.environ if environ is None else environ) 124 self._runtime_factory = runtime_factory 125 self._shutdown_waiter = shutdown_waiter 126 self._request_ids = RequestIdRegistry(paths.results_dir, factory=request_id_factory) 127 self._debug = debug 128 self._logging_session = logging_session 129 self._monotonic = monotonic 130 self._session_started = False 131 self._session_stopped = False 132 self._runtime: RuntimeLike | None = None 133 self._server: asyncio.AbstractServer | None = None 134 self._socket_identity: tuple[int, int] | None = None 135 self._state_lock = asyncio.Lock() 136 self._shutting_down = False 137 self._shutdown_task: asyncio.Task[None] | None = None 138 self._shutdown_response_ready = asyncio.Event() 139 self._active_workflows: set[asyncio.Task[object]] = set()
async def
start(self) -> None:
145 async def start(self) -> None: 146 self.paths.socket_file.parent.mkdir(parents=True, exist_ok=True) 147 self.paths.results_dir.mkdir(parents=True, exist_ok=True) 148 try: 149 await self._prepare_socket_path() 150 self._runtime = self._create_runtime() 151 try: 152 self._server = await asyncio.start_unix_server( 153 self._handle_connection, 154 path=self.paths.socket_file, 155 ) 156 except OSError as exc: 157 await self._close_runtime_after_startup_failure() 158 raise ConfigFailure( 159 ErrorCode.CONFIG_ERROR, 160 f"Failed to bind daemon socket: {self.paths.socket_file}", 161 ) from exc 162 socket_stat = self.paths.socket_file.stat() 163 self._socket_identity = (socket_stat.st_dev, socket_stat.st_ino) 164 if self._debug: 165 log_event( 166 self._logger, 167 logging.INFO, 168 "session_started", 169 pid=os.getpid(), 170 debug=True, 171 ) 172 self._session_started = True 173 self.ready.set() 174 await self.stopped.wait() 175 except (asyncio.CancelledError, KeyboardInterrupt): 176 await self.stop_for_test() 177 raise