52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""Create the opt-in Worker controller without restricting normal startup."""
|
|
|
|
import json
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
|
|
def _control_directory() -> Path:
|
|
return Path(os.getenv("DEER_FLOW_WORKER_CONTROL_HOME", ".worker-pool")).resolve()
|
|
|
|
|
|
def _persisted_feature_enabled(directory: Path) -> bool:
|
|
def read(path):
|
|
try:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
return {}
|
|
|
|
state = read(directory / "settings.json")
|
|
diagnostics = read(directory / "diagnostics-settings.json")
|
|
return bool(state.get("settings", {}).get("enabled") or diagnostics.get("enabled"))
|
|
|
|
|
|
def ensure_controller(app):
|
|
pool = getattr(app.state, "worker_pool", None)
|
|
if pool is not None:
|
|
return pool
|
|
from .controller import Controller
|
|
from .diagnostics import Diagnostics
|
|
from .settings import SettingsStore
|
|
|
|
directory = getattr(app.state, "worker_pool_directory", _control_directory())
|
|
pool = Controller(SettingsStore(directory))
|
|
pool.diagnostics = Diagnostics(pool, app)
|
|
app.state.worker_pool = pool
|
|
return pool
|
|
|
|
|
|
@asynccontextmanager
|
|
async def managed_controller(app):
|
|
directory = _control_directory()
|
|
app.state.worker_pool_directory = directory
|
|
app.state.worker_pool = None
|
|
pool = ensure_controller(app) if _persisted_feature_enabled(directory) else None
|
|
try:
|
|
yield pool
|
|
finally:
|
|
pool = getattr(app.state, "worker_pool", None)
|
|
if pool is not None:
|
|
await pool.close()
|