mirror of
https://github.com/usestrix/strix.git
synced 2026-11-01 07:02:32 +00:00
feat(agents): soft per-skill turn limits for sub-agents
run_strix_scan(subagent_turn_limits={skill: turns}) gives each sub-agent the
smallest limit among its skills. The agent is warned 10 turns before the limit
and reminded every 5 turns past it, but never stopped.
This commit is contained in:
parent
14cc948c36
commit
b17c0275a9
6 changed files with 118 additions and 3 deletions
|
|
@ -24,6 +24,8 @@ from openai import (
|
|||
|
||||
from strix.config import codex
|
||||
from strix.core.hooks import (
|
||||
SUBAGENT_TURN_LIMITS_KEY,
|
||||
TURN_LIMIT_KEY,
|
||||
BudgetExceededError,
|
||||
BudgetPausedError,
|
||||
SubagentBudgetReservedError,
|
||||
|
|
@ -39,6 +41,7 @@ from strix.core.sessions import (
|
|||
)
|
||||
from strix.llm import request_log
|
||||
from strix.llm.compaction import is_context_overflow, maybe_compact
|
||||
from strix.skills import skill_turn_limit
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -427,6 +430,7 @@ async def spawn_child_agent(
|
|||
name=name,
|
||||
parent_id=parent_id,
|
||||
task=task,
|
||||
skills=skills,
|
||||
initial_input=child_initial_input(
|
||||
name=name,
|
||||
child_id=child_id,
|
||||
|
|
@ -510,6 +514,7 @@ async def respawn_subagents(
|
|||
name=name,
|
||||
parent_id=parent_id,
|
||||
task=str(md.get("task", "")),
|
||||
skills=child_skills,
|
||||
initial_input=[],
|
||||
start_parked=start_parked,
|
||||
event_sink=event_sink,
|
||||
|
|
@ -1161,6 +1166,7 @@ async def _start_child_runner(
|
|||
name: str,
|
||||
parent_id: str | None,
|
||||
task: str,
|
||||
skills: list[str],
|
||||
initial_input: Any,
|
||||
start_parked: bool = False,
|
||||
event_sink: StreamEventSink | None = None,
|
||||
|
|
@ -1174,6 +1180,9 @@ async def _start_child_runner(
|
|||
child_ctx["agent_id"] = child_id
|
||||
child_ctx["parent_id"] = parent_id
|
||||
child_ctx["task"] = task
|
||||
child_ctx[TURN_LIMIT_KEY] = skill_turn_limit(
|
||||
skills, parent_ctx.get(SUBAGENT_TURN_LIMITS_KEY) or {}
|
||||
)
|
||||
|
||||
async def _child_loop() -> None:
|
||||
# A budget stop is a clean scan-wide shutdown, not a child failure: the
|
||||
|
|
|
|||
|
|
@ -25,6 +25,11 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
|
||||
LLM_TURN_KEY = "llm_turn"
|
||||
# ``run_strix_scan(subagent_turn_limits=...)``, and the limit it gives one sub-agent.
|
||||
SUBAGENT_TURN_LIMITS_KEY = "subagent_turn_limits"
|
||||
TURN_LIMIT_KEY = "turn_limit"
|
||||
_TURN_LIMIT_WARN_BEFORE = 10
|
||||
_TURN_LIMIT_REPEAT = 5
|
||||
|
||||
# ``BudgetPolicy`` decides what happens when the accumulated LLM cost reaches
|
||||
# ``max_budget_usd``.
|
||||
|
|
@ -38,6 +43,8 @@ LLM_TURN_KEY = "llm_turn"
|
|||
# same call when the operator raises the limit or resumes.
|
||||
__all__ = [
|
||||
"LLM_TURN_KEY",
|
||||
"SUBAGENT_TURN_LIMITS_KEY",
|
||||
"TURN_LIMIT_KEY",
|
||||
"BudgetExceededError",
|
||||
"BudgetPausedError",
|
||||
"BudgetPolicy",
|
||||
|
|
@ -213,6 +220,7 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
|||
context.context[LLM_TURN_KEY] = int(context.context.get(LLM_TURN_KEY, 0)) + 1
|
||||
try:
|
||||
self._maybe_warn_turns(context, input_items)
|
||||
self._maybe_warn_turn_limit(context, input_items)
|
||||
self._maybe_warn_budget(context, input_items)
|
||||
except Exception:
|
||||
logger.exception("budget/turn warning injection failed")
|
||||
|
|
@ -276,6 +284,40 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
|||
)
|
||||
input_items.append({"role": "user", "content": content})
|
||||
|
||||
def _maybe_warn_turn_limit(
|
||||
self,
|
||||
context: RunContextWrapper[dict[str, Any]],
|
||||
input_items: list[TResponseInputItem],
|
||||
) -> None:
|
||||
"""A soft limit, never a stop: warn 10 turns before it, then every few turns past it."""
|
||||
limit = context.context.get(TURN_LIMIT_KEY)
|
||||
requests = getattr(getattr(context, "usage", None), "requests", None)
|
||||
if not isinstance(limit, int) or not isinstance(requests, int):
|
||||
return
|
||||
turns_used = requests + 1
|
||||
if turns_used == limit - _TURN_LIMIT_WARN_BEFORE:
|
||||
content = (
|
||||
f"[NOTICE] You have used {turns_used} of your {limit}-turn limit. Start "
|
||||
"wrapping up: take only the steps you still need to answer your task, then "
|
||||
"call agent_finish. You may go past the limit only while you are still "
|
||||
"making real progress."
|
||||
)
|
||||
elif turns_used >= limit and (turns_used - limit) % _TURN_LIMIT_REPEAT == 0:
|
||||
content = (
|
||||
f"[NOTICE] You have used {turns_used} turns, past your {limit}-turn limit. "
|
||||
"Call agent_finish now with what you have, listing anything unresolved in "
|
||||
"open_items, unless your next steps are clearly making progress on your task."
|
||||
)
|
||||
else:
|
||||
return
|
||||
logger.info(
|
||||
"Turn-limit notice for agent %s: turn %d, limit %d",
|
||||
context.context.get("agent_id"),
|
||||
turns_used,
|
||||
limit,
|
||||
)
|
||||
input_items.append({"role": "user", "content": content})
|
||||
|
||||
def _maybe_warn_budget(
|
||||
self,
|
||||
context: RunContextWrapper[dict[str, Any]],
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import io
|
|||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
|
|
@ -36,7 +36,12 @@ from strix.core.execution import (
|
|||
from strix.core.execution import (
|
||||
spawn_child_agent as start_child_agent,
|
||||
)
|
||||
from strix.core.hooks import BudgetExceededError, ReportUsageHooks, recomputed_budget_flags
|
||||
from strix.core.hooks import (
|
||||
SUBAGENT_TURN_LIMITS_KEY,
|
||||
BudgetExceededError,
|
||||
ReportUsageHooks,
|
||||
recomputed_budget_flags,
|
||||
)
|
||||
from strix.core.inputs import (
|
||||
build_root_task,
|
||||
build_scan_targets,
|
||||
|
|
@ -222,6 +227,7 @@ async def run_strix_scan(
|
|||
mcp_status_sink: McpStatusSink | None = None,
|
||||
root_finish_tool: Tool = finish_scan,
|
||||
report_tool: Tool = create_vulnerability_report,
|
||||
subagent_turn_limits: Mapping[str, int] | None = None,
|
||||
) -> RunResultBase | None:
|
||||
"""Run or resume one Strix scan against a sandbox.
|
||||
|
||||
|
|
@ -247,6 +253,9 @@ async def run_strix_scan(
|
|||
``root_finish_tool`` is the tool the root agent ends the run with.
|
||||
``report_tool`` is the tool every agent files findings with, in place of
|
||||
``create_vulnerability_report``.
|
||||
``subagent_turn_limits`` maps a skill name to a soft turn limit for sub-agents
|
||||
spawned with that skill (the smallest applies). Agents are warned 10 turns
|
||||
before it and reminded past it, but never stopped.
|
||||
"""
|
||||
|
||||
def report(phase: str) -> None:
|
||||
|
|
@ -568,6 +577,7 @@ async def run_strix_scan(
|
|||
"scan_targets": build_scan_targets(scan_config),
|
||||
"max_context_images": settings.runtime.max_context_images,
|
||||
"supports_images": supports_images,
|
||||
SUBAGENT_TURN_LIMITS_KEY: dict(subagent_turn_limits or {}),
|
||||
}
|
||||
|
||||
root_session = open_agent_session(root_id, agents_db)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import logging
|
||||
import re
|
||||
from collections import Counter
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from typing import TypeGuard
|
||||
|
||||
|
|
@ -270,6 +270,19 @@ def _candidate_skill_files(skill_name: str) -> list[Path]:
|
|||
return _bare_skill_files(skill_name)
|
||||
|
||||
|
||||
def skill_turn_limit(skill_names: list[str], turn_limits: Mapping[str, int]) -> int | None:
|
||||
"""Smallest of ``turn_limits`` (keyed by skill name) among these skills, if any."""
|
||||
files = {path for name in skill_names for path in _candidate_skill_files(name)}
|
||||
return min(
|
||||
(
|
||||
limit
|
||||
for name, limit in turn_limits.items()
|
||||
if files.intersection(_candidate_skill_files(name))
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def load_skills(skill_names: list[str], *, supports_images: bool = True) -> dict[str, str]:
|
||||
"""Load skill markdown bodies (frontmatter stripped) by name.
|
||||
|
||||
|
|
|
|||
|
|
@ -507,3 +507,30 @@ def test_recomputed_budget_flags(
|
|||
expected: tuple[bool, bool],
|
||||
) -> None:
|
||||
assert recomputed_budget_flags(cost, max_budget, interactive=interactive) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_limit_warns_ten_before_and_repeats_past_it() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model")
|
||||
noticed = []
|
||||
for requests in range(40):
|
||||
ctx = _make_warn_context(requests=requests, parent_id="root-1")
|
||||
ctx.context["turn_limit"] = 25
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(ctx, MagicMock(), None, items)
|
||||
if items:
|
||||
noticed.append((requests + 1, items[0]["content"]))
|
||||
|
||||
assert [turn for turn, _ in noticed] == [15, 25, 30, 35, 40]
|
||||
assert "15 of your 25-turn limit" in noticed[0][1]
|
||||
assert "past your 25-turn limit" in noticed[1][1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_turn_limit_no_notice() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model")
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=14, parent_id="root-1"), MagicMock(), None, items
|
||||
)
|
||||
assert items == []
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from strix.skills import (
|
|||
register_skill_dir,
|
||||
registered_skill_dirs,
|
||||
skill_search_dirs,
|
||||
skill_turn_limit,
|
||||
validate_requested_skills,
|
||||
)
|
||||
from strix.utils.resource_paths import get_strix_resource_path
|
||||
|
|
@ -282,3 +283,16 @@ def test_new_skill_files_load() -> None:
|
|||
for name in names:
|
||||
key = name.split("/")[-1]
|
||||
assert loaded.get(key), f"{name} failed to load"
|
||||
|
||||
|
||||
def test_skill_turn_limit_matches_bare_and_qualified_names(tmp_path: Path) -> None:
|
||||
_write_skill(tmp_path, "pr_review", "investigator", "---\nname: investigator\n---\nx")
|
||||
_write_skill(tmp_path, "pr_review", "validator", "---\nname: validator\n---\nx")
|
||||
register_skill_dir(tmp_path)
|
||||
limits = {"pr_review/investigator": 25, "pr_review/validator": 40}
|
||||
|
||||
assert skill_turn_limit(["investigator"], limits) == 25
|
||||
assert skill_turn_limit(["pr_review/validator", "idor"], limits) == 40
|
||||
assert skill_turn_limit(["investigator", "validator"], limits) == 25
|
||||
assert skill_turn_limit(["idor"], limits) is None
|
||||
assert skill_turn_limit(["investigator"], {}) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue