Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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
28 changes: 20 additions & 8 deletions src/infra/task/worker_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,12 @@

启动面(任务执行在独立进程里依赖的进程内服务,对齐 API lifespan 的最小集):

- **initialize_settings(首要)**:从数据库加载生效设置(DB > env > 默认),
再做分布式配置校验——与 API 进程看到同一份设置。缺了这一步 worker 只跑
pydantic 默认值,与 API 的 DB 生效值分叉:生产开
SESSION_EVENT_CHUNK_STORAGE_ENABLED 时 API 写 chunk、worker 走 legacy 内联,
事件被 chunk 视图的 merge 覆写蒸发(2026-09-09 生产 P0,刷新即丢历史);
SANDBOX_PLATFORM 同理回落默认,本地 daemon 身份段(#499)不再注入;
- 事件循环滞后监控:观测「重任务饿死循环」的眼睛;
- task_manager pubsub 监听:分布式取消信号必须到达**执行方**,不启动则用户
取消不了跑在本进程上的任务;
Expand All @@ -30,23 +36,23 @@
import signal
from typing import Any, Callable

# 协作者在模块级具名:测试按 monkeypatch worker_main.<name> 注入替身。
from src.infra.distributed_validation import validate_distributed_runtime_settings # noqa: E402
from src.infra.llm.pubsub import get_model_config_pubsub # noqa: E402
from src.infra.logging import get_logger
from src.infra.monitoring.event_loop import start_event_loop_lag_monitor
from src.kernel.config import settings
from src.infra.pricing.pubsub import get_pricing_pubsub # noqa: E402
from src.infra.settings.pubsub import get_settings_pubsub # noqa: E402
from src.infra.tool.cache_pubsub import get_tool_cache_pubsub # noqa: E402
from src.infra.tool.mcp_global import get_mcp_cache_pubsub # noqa: E402
from src.kernel.config import initialize_settings, settings

from .arq_runtime import get_arq_runtime
from .arq_worker import worker_shutdown, worker_startup
from .manager import get_task_manager

logger = get_logger(__name__)

# 协作者在模块级具名:测试按 monkeypatch worker_main.<name> 注入替身。
from src.infra.llm.pubsub import get_model_config_pubsub # noqa: E402
from src.infra.pricing.pubsub import get_pricing_pubsub # noqa: E402
from src.infra.settings.pubsub import get_settings_pubsub # noqa: E402
from src.infra.tool.cache_pubsub import get_tool_cache_pubsub # noqa: E402
from src.infra.tool.mcp_global import get_mcp_cache_pubsub # noqa: E402


def get_memory_pubsub() -> Any:
"""记忆 pubsub 惰性解析(ENABLE_MEMORY 才需要,模块级导入会拖重依赖)。"""
Expand Down Expand Up @@ -111,6 +117,12 @@ async def _amain(stop: asyncio.Event) -> None:

logger.info("standalone arq worker starting (queue=%s)", settings.ARQ_QUEUE_NAME)
try:
# 设置加载必须先于一切服务启动(对齐 API lifespan):worker 与 API
# 看到同一份 DB > env > 默认的生效设置。失败快速退出(k8s
# CrashLoopBackOff 显性暴露),绝不带默认值吞事件。
await initialize_settings()
logger.info("Settings initialized from database")
validate_distributed_runtime_settings(settings)
await start_event_loop_lag_monitor()
await task_manager.start_pubsub_listener()
await _start_cache_listeners()
Expand Down
64 changes: 64 additions & 0 deletions tests/infra/task/test_worker_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,14 @@ async def fake_worker_startup(ctx: dict) -> None:
async def fake_worker_shutdown(ctx: dict) -> None:
shutdown_calls.append("worker_shutdown")

async def fake_initialize_settings() -> None:
return None

def fake_validate(s) -> None:
return None

monkeypatch.setattr(worker_main, "initialize_settings", fake_initialize_settings)
monkeypatch.setattr(worker_main, "validate_distributed_runtime_settings", fake_validate)
monkeypatch.setattr(worker_main, "worker_startup", fake_worker_startup)
monkeypatch.setattr(worker_main, "worker_shutdown", fake_worker_shutdown)

Expand Down Expand Up @@ -124,6 +132,62 @@ async def release_stop(stop: asyncio.Event) -> None:
assert wired.startup_calls == ["worker_startup"] # 分布式校验 + loop_bridge 登记


async def test_amain_loads_db_settings_before_any_service(wired, monkeypatch):
"""设置契约(2026-09-09 生产 P0):worker 必须先 initialize_settings 再起服务。

worker 只跑 pydantic 默认值时与 API 的 DB 生效值分叉:生产开了
SESSION_EVENT_CHUNK_STORAGE_ENABLED,API 写 chunk、worker 走 legacy 内联,
事件被 chunk 视图的 merge 覆写蒸发(刷新即丢历史);SANDBOX_PLATFORM
同理回落默认,本地 daemon 身份段(#499)不再注入。分布式校验随设置
加载之后执行,对齐 API lifespan(main.py initialize_settings →
validate_distributed_runtime_settings)。
"""
order: list[str] = []

async def fake_initialize_settings() -> None:
order.append("initialize_settings")

def fake_validate(s) -> None:
order.append("validate")

async def fake_lag_monitor() -> None:
order.append("lag_monitor")

monkeypatch.setattr(worker_main, "initialize_settings", fake_initialize_settings)
monkeypatch.setattr(worker_main, "validate_distributed_runtime_settings", fake_validate)
monkeypatch.setattr(worker_main, "start_event_loop_lag_monitor", fake_lag_monitor)

async def release_stop(stop: asyncio.Event) -> None:
await asyncio.sleep(0.01)
stop.set()

stop = asyncio.Event()
done = asyncio.create_task(release_stop(stop))
await worker_main._amain(stop)
await done

assert order[:2] == ["initialize_settings", "validate"]
assert order.index("initialize_settings") < order.index("lag_monitor")
assert wired.runtime.started_with == {"force": True}


async def test_amain_settings_init_failure_stops_startup(wired, monkeypatch):
"""设置加载失败必须快速失败(k8s CrashLoopBackOff 显性暴露),绝不带
默认值继续起 worker 吞事件。"""

async def failing_initialize() -> None:
raise RuntimeError("db unreachable")

monkeypatch.setattr(worker_main, "initialize_settings", failing_initialize)

stop = asyncio.Event()
with pytest.raises(RuntimeError):
await worker_main._amain(stop)

assert wired.runtime.started_with is None # 未起 worker
assert wired.startup_calls == []


async def test_amain_shutdown_stops_listeners_and_runtime_in_order(wired, monkeypatch):
"""退出面契约:先停 worker(等任务收尾),再停监听,最后 worker_shutdown。"""
monkeypatch.setattr(worker_main.settings, "ENABLE_MEMORY", False)
Expand Down
Loading