From 95e7d5df7ca8ae70f7742741bc96abdcb334c91b Mon Sep 17 00:00:00 2001 From: Renn F Date: Wed, 29 Jul 2026 06:14:22 +0200 Subject: [PATCH] fix(kimi): cap concurrent Kimi agents to protect the shared auth chain Every Kimi container redeems the same rotating refresh-token chain; Moonshot rotates with a short reuse grace, so two containers refreshing near-simultaneously fork the chain and a later stale redemption revokes the whole family - fleet-wide re-login (observed twice in production, each after paired spawns). With one consumer at a time refreshes are strictly sequential and the chain stays coherent, so the spawn gate now skips-and-retries Kimi spawns past ROBOCO_KIMI_MAX_CONCURRENT (default 1), sharing the provider-parked bail path. The compose files also gain the four Kimi tunables their environment blocks silently dropped - documented .env overrides never reached the orchestrator container. --- .env.example | 4 + docker-compose.registry.yml | 4 + docker-compose.yaml | 4 + docker-compose.yml | 4 + roboco/config.py | 17 ++ roboco/runtime/orchestrator.py | 123 ++++++-- .../unit/runtime/test_kimi_concurrency_cap.py | 262 ++++++++++++++++++ 7 files changed, 393 insertions(+), 25 deletions(-) create mode 100644 tests/unit/runtime/test_kimi_concurrency_cap.py diff --git a/.env.example b/.env.example index 487eff56..e25940ee 100644 --- a/.env.example +++ b/.env.example @@ -199,6 +199,10 @@ ROBOCO_DEFAULT_EMBEDDING_MODEL=qwen3-embedding:0.6b # Park-and-retry delays after a rate-limit / auth failure (seconds): # ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS=60 # ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS=60 +# Max concurrent live Kimi containers (default 1). Raising this risks two +# containers refreshing the shared credential chain at once, forking it and +# revoking fleet-wide Kimi auth — only raise if you understand that risk. +# ROBOCO_KIMI_MAX_CONCURRENT=1 # ============================================================================= # Cost budgets — optional diff --git a/docker-compose.registry.yml b/docker-compose.registry.yml index 83943b80..5740011f 100644 --- a/docker-compose.registry.yml +++ b/docker-compose.registry.yml @@ -385,6 +385,10 @@ services: # ONE host chain via a symlinked-in credentials/+oauth/ mount rather # than a per-container copy (kimi follows codex's RW mount mode here). ROBOCO_HOST_KIMI_DIR: ${ROBOCO_HOST_KIMI_DIR:-${HOME}/.kimi-code} + ROBOCO_KIMI_CLI_MODEL: ${ROBOCO_KIMI_CLI_MODEL:-kimi-code/k3} + ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS: ${ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS:-60} + ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS: ${ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS:-60} + ROBOCO_KIMI_MAX_CONCURRENT: ${ROBOCO_KIMI_MAX_CONCURRENT:-1} ROBOCO_HOST_DATA_DIR: ${ROBOCO_HOST_DATA_DIR:-/opt/roboco/data} # Reachable base URL for commit-trailer links — set to your host's LAN # address or domain so the links in commit bodies resolve. diff --git a/docker-compose.yaml b/docker-compose.yaml index 5ff01e6c..f6cc11df 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -533,6 +533,10 @@ services: # ONE host chain via a symlinked-in credentials/+oauth/ mount rather # than a per-container copy (kimi follows codex's RW mount mode here). ROBOCO_HOST_KIMI_DIR: ${ROBOCO_HOST_KIMI_DIR:-/home/renzof/.kimi-code} + ROBOCO_KIMI_CLI_MODEL: ${ROBOCO_KIMI_CLI_MODEL:-kimi-code/k3} + ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS: ${ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS:-60} + ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS: ${ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS:-60} + ROBOCO_KIMI_MAX_CONCURRENT: ${ROBOCO_KIMI_MAX_CONCURRENT:-1} ROBOCO_HOST_DATA_DIR: ${ROBOCO_HOST_DATA_DIR:-/volume1/roboco/data} # Public base URL for commit-trailer links. Default 127.0.0.1 produces # unusable links in commit message bodies; set to NAS LAN IP so diff --git a/docker-compose.yml b/docker-compose.yml index 5ff01e6c..f6cc11df 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -533,6 +533,10 @@ services: # ONE host chain via a symlinked-in credentials/+oauth/ mount rather # than a per-container copy (kimi follows codex's RW mount mode here). ROBOCO_HOST_KIMI_DIR: ${ROBOCO_HOST_KIMI_DIR:-/home/renzof/.kimi-code} + ROBOCO_KIMI_CLI_MODEL: ${ROBOCO_KIMI_CLI_MODEL:-kimi-code/k3} + ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS: ${ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS:-60} + ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS: ${ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS:-60} + ROBOCO_KIMI_MAX_CONCURRENT: ${ROBOCO_KIMI_MAX_CONCURRENT:-1} ROBOCO_HOST_DATA_DIR: ${ROBOCO_HOST_DATA_DIR:-/volume1/roboco/data} # Public base URL for commit-trailer links. Default 127.0.0.1 produces # unusable links in commit message bodies; set to NAS LAN IP so diff --git a/roboco/config.py b/roboco/config.py index f317ee50..f007acef 100644 --- a/roboco/config.py +++ b/roboco/config.py @@ -2020,6 +2020,23 @@ class Settings(BaseSettings): "exit 78); override via ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS" ), ) + # Every Kimi container shares ONE OAuth refresh-token chain (the + # read-write ~/.kimi-code mount above). Moonshot rotates refresh tokens + # with a short reuse-grace; two containers refreshing near-simultaneously + # fork the chain, and a later redemption of a stale ancestor triggers + # family revocation — fleet-wide Kimi auth dies and the CEO must + # re-login. One consumer at a time keeps refreshes strictly sequential + # (the proven-stable regime, live-incident-verified 2026-07-29). + kimi_max_concurrent: int = Field( + default=1, + ge=1, + description=( + "Max concurrent live KIMI agent containers. Raising this risks " + "forking the shared OAuth refresh-token chain and triggering " + "fleet-wide Kimi re-login; override via " + "ROBOCO_KIMI_MAX_CONCURRENT only if you understand that risk" + ), + ) # An interactive intake/secretary chat the human abandoned (closed the tab # without confirming/stopping) otherwise leaks its container until the # orchestrator restarts. The sweeper reaps a live session whose diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index e9235b50..6363086b 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -3120,28 +3120,16 @@ class AgentOrchestrator: # existing-running check above stays first, so a live agent is never # replaced by this bail. Fail-open: a tracker read error never blocks. route = await self._resolve_agent_route(agent_id, task_id) - if await self._provider_spawn_parked(route.provider_type.value): + skip_reason = await self._spawn_gate_skip_reason(route.provider_type.value) + if skip_reason: self._mark_task_handled(task_id) logger.info( - "Spawn skipped: provider rate-limited (parked)", + skip_reason, agent_id=agent_id, task_id=task_id, provider=route.provider_type.value, ) - return AgentInstance( - agent_id=agent_id, - state=AgentState.OFFLINE, - config=AgentConfig( - agent_id=agent_id, - blueprint_path=Path(), # not launching — no blueprint written - model=route.model_name, - provider_type=route.provider_type.value, - provider_base_url=route.base_url, - provider_auth_token=route.auth_token, - git_context=git_context, - ), - current_task_id=task_id, - ) + return self._offline_route_bail(agent_id, task_id, route, git_context) async with self._lock: # TOCTOU re-check: another tick may have started this agent during @@ -3157,16 +3145,11 @@ class AgentOrchestrator: # parked case is already handled above; this guards the window between # the pre-check and the launch. Fail-open: a tracker read error never # blocks spawning. - if await self._provider_spawn_parked(config.provider_type): - self._mark_task_handled(task_id) - instance.state = AgentState.OFFLINE - logger.info( - "Spawn skipped: provider rate-limited (parked)", - agent_id=agent_id, - task_id=task_id, - provider=config.provider_type, + skip_reason = await self._spawn_gate_skip_reason(config.provider_type) + if skip_reason: + return self._bail_prepared_instance( + instance, agent_id, task_id, config.provider_type, skip_reason ) - return instance # Record the task as handled so later dispatchers in the same # tick don't act on it again. Safe even if _launch_spawn fails # — the next tick starts fresh. @@ -10190,6 +10173,96 @@ Start by: ) return False + @staticmethod + def _provider_concurrency_cap(provider_type: str | None) -> int | None: + """Max concurrent live containers for *provider_type*, or None if uncapped. + + Only KIMI is capped: every Kimi container shares one OAuth + refresh-token chain (see ``settings.kimi_max_concurrent``) — a second + concurrent container risks forking it and revoking fleet-wide Kimi + auth. No other provider has this constraint. + """ + if provider_type == ModelProvider.KIMI.value: + return settings.kimi_max_concurrent + return None + + def _live_provider_instance_count(self, provider_type: str) -> int: + """Count ``_instances`` entries for *provider_type* still holding a slot. + + Mirrors ``_existing_running_instance``'s liveness definition: any + state other than OFFLINE/WAITING_LONG occupies (or is about to + occupy) a container. + """ + return sum( + 1 + for inst in self._instances.values() + if inst.config is not None + and inst.config.provider_type == provider_type + and inst.state not in (AgentState.OFFLINE, AgentState.WAITING_LONG) + ) + + def _provider_spawn_at_capacity(self, provider_type: str | None) -> bool: + """True when *provider_type* is at its concurrency cap, if any.""" + cap = self._provider_concurrency_cap(provider_type) + if cap is None or provider_type is None: + return False + return self._live_provider_instance_count(provider_type) >= cap + + async def _spawn_gate_skip_reason(self, provider_type: str | None) -> str | None: + """Why ``spawn_agent`` should bail before launch, or None to proceed. + + Checked in order: provider-parked (rate-limited/overloaded), then the + provider's concurrency cap (currently kimi-only, see + ``_provider_concurrency_cap``). The returned string is both the skip + reason and the log event name — same shape both callers already used. + """ + if await self._provider_spawn_parked(provider_type): + return "Spawn skipped: provider rate-limited (parked)" + if self._provider_spawn_at_capacity(provider_type): + return "Spawn skipped: kimi concurrency cap reached" + return None + + def _bail_prepared_instance( + self, + instance: AgentInstance, + agent_id: str, + task_id: str | None, + provider_type: str | None, + reason: str, + ) -> AgentInstance: + """OFFLINE a just-``_prepare_agent_spawn``'d instance and log why the + launch was skipped (used by both the parked and concurrency-cap + rare-race checks after prepare).""" + self._mark_task_handled(task_id) + instance.state = AgentState.OFFLINE + logger.info(reason, agent_id=agent_id, task_id=task_id, provider=provider_type) + return instance + + def _offline_route_bail( + self, + agent_id: str, + task_id: str | None, + route: Any, + git_context: SpawnGitContext | None, + ) -> AgentInstance: + """Build the unregistered OFFLINE instance returned by a pre-prepare + spawn bail (parked or at-capacity) — no blueprint/settings/image work + has run, so nothing needs undoing.""" + return AgentInstance( + agent_id=agent_id, + state=AgentState.OFFLINE, + config=AgentConfig( + agent_id=agent_id, + blueprint_path=Path(), # not launching — no blueprint written + model=route.model_name, + provider_type=route.provider_type.value, + provider_base_url=route.base_url, + provider_auth_token=route.auth_token, + git_context=git_context, + ), + current_task_id=task_id, + ) + @staticmethod def _is_grok_rate_limit_exit(instance: Any, exit_code: int | None) -> bool: """True for a one-shot grok container that exited 75 (xAI 429).""" diff --git a/tests/unit/runtime/test_kimi_concurrency_cap.py b/tests/unit/runtime/test_kimi_concurrency_cap.py new file mode 100644 index 00000000..a417e66c --- /dev/null +++ b/tests/unit/runtime/test_kimi_concurrency_cap.py @@ -0,0 +1,262 @@ +"""Kimi per-provider spawn concurrency cap. + +Every Kimi container shares ONE OAuth refresh-token chain (a host-mounted, +read-write ``~/.kimi-code`` — see ``roboco.llm.providers.kimi``'s module +docstring). Moonshot rotates refresh tokens with a short reuse-grace: two +containers refreshing near-simultaneously fork the chain, and a later +redemption of a stale ancestor triggers family revocation — fleet-wide Kimi +auth dies and the CEO must re-login (a live incident, 2026-07-29). With one +Kimi consumer at a time refreshes stay strictly sequential. + +``settings.kimi_max_concurrent`` (default 1) caps live Kimi containers, +enforced in ``AgentOrchestrator.spawn_agent`` at the exact same chokepoints as +the provider-parked check (see ``test_parked_spawn_shortcut.py`` for the +pre-prepare / post-prepare shape this mirrors). No other provider is capped. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from roboco.config import settings +from roboco.models.runtime import AgentInstance +from roboco.runtime import orchestrator as orch_module +from roboco.runtime.orchestrator import AgentConfig, AgentOrchestrator, AgentState + + +def _make_orchestrator() -> AgentOrchestrator: + # __new__ + skip __init__: avoid all constructor I/O (mirrors + # test_parked_spawn_shortcut.py's _make_orchestrator). + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._instances = {} + orch._lock = asyncio.Lock() + orch._tick_handled_tasks = set() + orch._bg_tasks = set() + orch._running = True + return orch + + +def _live_instance( + agent_id: str, provider_type: str, state: AgentState +) -> AgentInstance: + cfg = AgentConfig( + agent_id=agent_id, blueprint_path=Path(), provider_type=provider_type + ) + return AgentInstance(agent_id=agent_id, state=state, config=cfg) + + +def _wire(monitor: dict[str, Any], provider_value: str) -> Any: + """Build the mock wiring closure capturing call counts in ``monitor``.""" + + async def _readiness_gate(_aid: str, _tid: str | None) -> None: + return None + + async def _git_context(_gc: Any, _tid: str | None) -> None: + return None + + async def _route(_aid: str, _tid: str | None = None) -> Any: + monitor["route_calls"] += 1 + return SimpleNamespace( + provider_type=SimpleNamespace(value=provider_value), + model_name="kimi-code/k3" if provider_value == "kimi" else "opus", + base_url=None, + auth_token=None, + ) + + async def _prepare(*_a: Any, **_k: Any) -> Any: + monitor["prepare_calls"] += 1 + cfg = SimpleNamespace(provider_type=provider_value, model="model") + inst = AgentInstance( + agent_id="spawning-agent", state=AgentState.STARTING, config=None + ) + return cfg, inst, None + + return _readiness_gate, _git_context, _route, _prepare + + +def _wire_orch( + orch: AgentOrchestrator, + monkeypatch: pytest.MonkeyPatch, + monitor: dict[str, Any], + provider_value: str, +) -> None: + _rg, _gc, _route, _prepare = _wire(monitor, provider_value) + monkeypatch.setattr(orch, "_readiness_gate", _rg) + monkeypatch.setattr(orch, "_resolve_spawn_git_context", _gc) + monkeypatch.setattr(orch, "_resolve_agent_route", _route) + monkeypatch.setattr(orch, "_prepare_agent_spawn", _prepare) + monkeypatch.setattr(orch, "_provider_spawn_parked", AsyncMock(return_value=False)) + + +def _capture_info_logs( + monkeypatch: pytest.MonkeyPatch, +) -> list[tuple[str, dict[str, Any]]]: + calls: list[tuple[str, dict[str, Any]]] = [] + + def _info(event: str, **kwargs: Any) -> None: + calls.append((event, kwargs)) + + monkeypatch.setattr(orch_module.logger, "info", _info) + return calls + + +# --------------------------------------------------------------------------- +# Helper-level unit tests +# --------------------------------------------------------------------------- + + +def test_provider_concurrency_cap_kimi_only() -> None: + assert ( + AgentOrchestrator._provider_concurrency_cap("kimi") + == settings.kimi_max_concurrent + ) + assert AgentOrchestrator._provider_concurrency_cap("anthropic") is None + assert AgentOrchestrator._provider_concurrency_cap("openai") is None + assert AgentOrchestrator._provider_concurrency_cap(None) is None + + +def test_live_provider_instance_count_and_capacity() -> None: + orch = _make_orchestrator() + assert orch._live_provider_instance_count("kimi") == 0 + assert orch._provider_spawn_at_capacity("kimi") is False + + orch._instances["be-dev-1"] = _live_instance("be-dev-1", "kimi", AgentState.ACTIVE) + assert orch._live_provider_instance_count("kimi") == 1 + assert orch._provider_spawn_at_capacity("kimi") is True # default cap 1 + + # OFFLINE/WAITING_LONG never occupy a slot (mirrors _existing_running_instance). + orch._instances["be-dev-1"].state = AgentState.OFFLINE + assert orch._live_provider_instance_count("kimi") == 0 + assert orch._provider_spawn_at_capacity("kimi") is False + + # An uncapped provider is never at capacity, however many live instances. + orch._instances["fe-dev-1"] = _live_instance( + "fe-dev-1", "anthropic", AgentState.ACTIVE + ) + assert orch._provider_spawn_at_capacity("anthropic") is False + + +# --------------------------------------------------------------------------- +# spawn_agent integration: exercises the exact chokepoint every kimi spawn +# path crosses (mirrors test_parked_spawn_shortcut.py). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_second_concurrent_kimi_spawn_skipped_while_first_live( + monkeypatch: pytest.MonkeyPatch, +) -> None: + orch = _make_orchestrator() + orch._instances["be-dev-1"] = _live_instance("be-dev-1", "kimi", AgentState.ACTIVE) + monitor = {"route_calls": 0, "prepare_calls": 0} + _wire_orch(orch, monkeypatch, monitor, "kimi") + info_calls = _capture_info_logs(monkeypatch) + + result = await orch.spawn_agent(agent_id="fe-dev-1", task_id="task-9") + + assert monitor["prepare_calls"] == 0 + assert isinstance(result, AgentInstance) + assert result.state is AgentState.OFFLINE + assert "fe-dev-1" not in orch._instances + assert "task-9" in orch._tick_handled_tasks + assert any( + event == "Spawn skipped: kimi concurrency cap reached" + and kwargs.get("agent_id") == "fe-dev-1" + and kwargs.get("task_id") == "task-9" + and kwargs.get("provider") == "kimi" + for event, kwargs in info_calls + ), info_calls + + +@pytest.mark.asyncio +async def test_non_kimi_spawn_unaffected_by_kimi_cap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A kimi agent already at cap must never block an unrelated provider.""" + orch = _make_orchestrator() + orch._instances["be-dev-1"] = _live_instance("be-dev-1", "kimi", AgentState.ACTIVE) + monitor = {"route_calls": 0, "prepare_calls": 0} + _wire_orch(orch, monkeypatch, monitor, "anthropic") + + launched: list[bool] = [] + + async def _launch(*_a: Any, **_k: Any) -> AgentInstance: + launched.append(True) + return AgentInstance(agent_id="fe-dev-1", state=AgentState.ACTIVE, config=None) + + monkeypatch.setattr(orch, "_launch_spawn", _launch) + + await orch.spawn_agent(agent_id="fe-dev-1", task_id="task-9") + + assert monitor["prepare_calls"] == 1 + assert launched == [True] + + +@pytest.mark.asyncio +async def test_cap_honors_env_override(monkeypatch: pytest.MonkeyPatch) -> None: + """ROBOCO_KIMI_MAX_CONCURRENT populates settings.kimi_max_concurrent at + construction; monkeypatching the singleton attribute is this suite's + established stand-in for an env override (see test_agent_image_registry.py).""" + monkeypatch.setattr(orch_module.settings, "kimi_max_concurrent", 2) + orch = _make_orchestrator() + orch._instances["be-dev-1"] = _live_instance("be-dev-1", "kimi", AgentState.ACTIVE) + orch._instances["be-dev-2"] = _live_instance("be-dev-2", "kimi", AgentState.ACTIVE) + monitor = {"route_calls": 0, "prepare_calls": 0} + _wire_orch(orch, monkeypatch, monitor, "kimi") + + launched: list[bool] = [] + + async def _launch(*_a: Any, **_k: Any) -> AgentInstance: + launched.append(True) + return AgentInstance(agent_id="fe-dev-1", state=AgentState.ACTIVE, config=None) + + monkeypatch.setattr(orch, "_launch_spawn", _launch) + + # Two live kimi agents, cap raised to 2: a third spawn is still skipped. + result = await orch.spawn_agent(agent_id="fe-dev-1", task_id="task-9") + assert monitor["prepare_calls"] == 0 + assert result.state is AgentState.OFFLINE + assert launched == [] + + # Drop to one live kimi agent: now under the raised cap, the spawn proceeds. + orch._instances["be-dev-2"].state = AgentState.OFFLINE + orch._tick_handled_tasks.clear() + await orch.spawn_agent(agent_id="fe-dev-1", task_id="task-10") + assert monitor["prepare_calls"] == 1 + assert launched == [True] + + +@pytest.mark.asyncio +async def test_spawn_proceeds_once_prior_kimi_instance_is_gone( + monkeypatch: pytest.MonkeyPatch, +) -> None: + orch = _make_orchestrator() + orch._instances["be-dev-1"] = _live_instance("be-dev-1", "kimi", AgentState.ACTIVE) + monitor = {"route_calls": 0, "prepare_calls": 0} + _wire_orch(orch, monkeypatch, monitor, "kimi") + + launched: list[bool] = [] + + async def _launch(*_a: Any, **_k: Any) -> AgentInstance: + launched.append(True) + return AgentInstance(agent_id="fe-dev-1", state=AgentState.ACTIVE, config=None) + + monkeypatch.setattr(orch, "_launch_spawn", _launch) + + # At cap: skipped. + result = await orch.spawn_agent(agent_id="fe-dev-1", task_id="task-9") + assert result.state is AgentState.OFFLINE + assert monitor["prepare_calls"] == 0 + + # The first kimi instance goes OFFLINE (finished/parked/reaped) — the + # slot is free, the next spawn proceeds normally. + orch._instances["be-dev-1"].state = AgentState.OFFLINE + orch._tick_handled_tasks.clear() + await orch.spawn_agent(agent_id="fe-dev-1", task_id="task-10") + assert monitor["prepare_calls"] == 1 + assert launched == [True]