mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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.
This commit is contained in:
@@ -199,6 +199,10 @@ ROBOCO_DEFAULT_EMBEDDING_MODEL=qwen3-embedding:0.6b
|
|||||||
# Park-and-retry delays after a rate-limit / auth failure (seconds):
|
# Park-and-retry delays after a rate-limit / auth failure (seconds):
|
||||||
# ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS=60
|
# ROBOCO_KIMI_RATE_LIMIT_RETRY_AFTER_SECONDS=60
|
||||||
# ROBOCO_KIMI_AUTH_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
|
# Cost budgets — optional
|
||||||
|
|||||||
@@ -385,6 +385,10 @@ services:
|
|||||||
# ONE host chain via a symlinked-in credentials/+oauth/ mount rather
|
# ONE host chain via a symlinked-in credentials/+oauth/ mount rather
|
||||||
# than a per-container copy (kimi follows codex's RW mount mode here).
|
# 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_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}
|
ROBOCO_HOST_DATA_DIR: ${ROBOCO_HOST_DATA_DIR:-/opt/roboco/data}
|
||||||
# Reachable base URL for commit-trailer links — set to your host's LAN
|
# Reachable base URL for commit-trailer links — set to your host's LAN
|
||||||
# address or domain so the links in commit bodies resolve.
|
# address or domain so the links in commit bodies resolve.
|
||||||
|
|||||||
@@ -533,6 +533,10 @@ services:
|
|||||||
# ONE host chain via a symlinked-in credentials/+oauth/ mount rather
|
# ONE host chain via a symlinked-in credentials/+oauth/ mount rather
|
||||||
# than a per-container copy (kimi follows codex's RW mount mode here).
|
# 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_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}
|
ROBOCO_HOST_DATA_DIR: ${ROBOCO_HOST_DATA_DIR:-/volume1/roboco/data}
|
||||||
# Public base URL for commit-trailer links. Default 127.0.0.1 produces
|
# 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
|
# unusable links in commit message bodies; set to NAS LAN IP so
|
||||||
|
|||||||
@@ -533,6 +533,10 @@ services:
|
|||||||
# ONE host chain via a symlinked-in credentials/+oauth/ mount rather
|
# ONE host chain via a symlinked-in credentials/+oauth/ mount rather
|
||||||
# than a per-container copy (kimi follows codex's RW mount mode here).
|
# 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_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}
|
ROBOCO_HOST_DATA_DIR: ${ROBOCO_HOST_DATA_DIR:-/volume1/roboco/data}
|
||||||
# Public base URL for commit-trailer links. Default 127.0.0.1 produces
|
# 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
|
# unusable links in commit message bodies; set to NAS LAN IP so
|
||||||
|
|||||||
@@ -2020,6 +2020,23 @@ class Settings(BaseSettings):
|
|||||||
"exit 78); override via ROBOCO_KIMI_AUTH_RETRY_AFTER_SECONDS"
|
"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
|
# An interactive intake/secretary chat the human abandoned (closed the tab
|
||||||
# without confirming/stopping) otherwise leaks its container until the
|
# without confirming/stopping) otherwise leaks its container until the
|
||||||
# orchestrator restarts. The sweeper reaps a live session whose
|
# orchestrator restarts. The sweeper reaps a live session whose
|
||||||
|
|||||||
@@ -3120,28 +3120,16 @@ class AgentOrchestrator:
|
|||||||
# existing-running check above stays first, so a live agent is never
|
# existing-running check above stays first, so a live agent is never
|
||||||
# replaced by this bail. Fail-open: a tracker read error never blocks.
|
# replaced by this bail. Fail-open: a tracker read error never blocks.
|
||||||
route = await self._resolve_agent_route(agent_id, task_id)
|
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)
|
self._mark_task_handled(task_id)
|
||||||
logger.info(
|
logger.info(
|
||||||
"Spawn skipped: provider rate-limited (parked)",
|
skip_reason,
|
||||||
agent_id=agent_id,
|
agent_id=agent_id,
|
||||||
task_id=task_id,
|
task_id=task_id,
|
||||||
provider=route.provider_type.value,
|
provider=route.provider_type.value,
|
||||||
)
|
)
|
||||||
return AgentInstance(
|
return self._offline_route_bail(agent_id, task_id, route, git_context)
|
||||||
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,
|
|
||||||
)
|
|
||||||
|
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
# TOCTOU re-check: another tick may have started this agent during
|
# 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
|
# parked case is already handled above; this guards the window between
|
||||||
# the pre-check and the launch. Fail-open: a tracker read error never
|
# the pre-check and the launch. Fail-open: a tracker read error never
|
||||||
# blocks spawning.
|
# blocks spawning.
|
||||||
if await self._provider_spawn_parked(config.provider_type):
|
skip_reason = await self._spawn_gate_skip_reason(config.provider_type)
|
||||||
self._mark_task_handled(task_id)
|
if skip_reason:
|
||||||
instance.state = AgentState.OFFLINE
|
return self._bail_prepared_instance(
|
||||||
logger.info(
|
instance, agent_id, task_id, config.provider_type, skip_reason
|
||||||
"Spawn skipped: provider rate-limited (parked)",
|
|
||||||
agent_id=agent_id,
|
|
||||||
task_id=task_id,
|
|
||||||
provider=config.provider_type,
|
|
||||||
)
|
)
|
||||||
return instance
|
|
||||||
# Record the task as handled so later dispatchers in the same
|
# Record the task as handled so later dispatchers in the same
|
||||||
# tick don't act on it again. Safe even if _launch_spawn fails
|
# tick don't act on it again. Safe even if _launch_spawn fails
|
||||||
# — the next tick starts fresh.
|
# — the next tick starts fresh.
|
||||||
@@ -10190,6 +10173,96 @@ Start by:
|
|||||||
)
|
)
|
||||||
return False
|
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
|
@staticmethod
|
||||||
def _is_grok_rate_limit_exit(instance: Any, exit_code: int | None) -> bool:
|
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)."""
|
"""True for a one-shot grok container that exited 75 (xAI 429)."""
|
||||||
|
|||||||
@@ -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]
|
||||||
Reference in New Issue
Block a user