refactor(agent): capture subagent runtime before spawn

This commit is contained in:
chengyongru
2026-07-10 17:54:34 +08:00
committed by Xubin Ren
parent 198fd9f869
commit af85c356b8
16 changed files with 261 additions and 132 deletions
+75 -28
View File
@@ -14,27 +14,35 @@ from nanobot.agent.subagent import (
SubagentStatus,
_SubagentHook,
)
from nanobot.agent.tools.context import current_request_context
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMProvider
from nanobot.providers.base import GenerationSettings, LLMProvider
from nanobot.utils.llm_runtime import LLMRuntime
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _manager(tmp_path: Path, **kw) -> SubagentManager:
provider = MagicMock(spec=LLMProvider)
provider.get_default_model.return_value = "test-model"
defaults = dict(
provider=provider,
workspace=tmp_path,
bus=MessageBus(),
model="test-model",
max_tool_result_chars=16_000,
)
defaults.update(kw)
return SubagentManager(**defaults)
def _runtime(*, model: str = "test-model", temperature: float = 0.1) -> LLMRuntime:
provider = MagicMock(spec=LLMProvider)
provider.generation = GenerationSettings(temperature=temperature, max_tokens=4096)
return LLMRuntime.capture(
provider,
model,
context_window_tokens=128_000,
)
def _make_hook_context(**overrides) -> AgentHookContext:
defaults = dict(
iteration=1,
@@ -77,17 +85,16 @@ class TestSubagentStatus:
# ---------------------------------------------------------------------------
# set_provider
# Runtime ownership
# ---------------------------------------------------------------------------
class TestSetProvider:
def test_updates_provider_model_runner(self, tmp_path):
class TestRuntimeOwnership:
def test_manager_has_no_provider_model_mirrors(self, tmp_path):
sm = _manager(tmp_path)
new_provider = MagicMock(spec=LLMProvider)
sm.set_provider(new_provider, "new-model")
assert sm.provider is new_provider
assert sm.model == "new-model"
assert not hasattr(sm, "provider")
assert not hasattr(sm, "model")
assert not hasattr(sm, "context_window_tokens")
assert not hasattr(sm.runner, "provider")
@@ -103,7 +110,7 @@ class TestSpawn:
sm.runner.run = AsyncMock(return_value=AgentRunResult(
final_content="done", messages=[], stop_reason="completed",
))
result = await sm.spawn("do something")
result = await sm.spawn("do something", runtime=_runtime())
assert "started" in result
assert "id:" in result
@@ -116,7 +123,7 @@ class TestSpawn:
return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
sm.runner.run = _slow_run
await sm.spawn("task", session_key="s1")
await sm.spawn("task", runtime=_runtime(), session_key="s1")
assert len(sm._running_tasks) == 1
block.set()
@@ -129,7 +136,7 @@ class TestSpawn:
sm.runner.run = AsyncMock(return_value=AgentRunResult(
final_content="done", messages=[], stop_reason="completed",
))
await sm.spawn("my task")
await sm.spawn("my task", runtime=_runtime())
await _drain_subagent_tasks(sm)
# Status cleaned up after task completes
assert len(sm._task_statuses) == 0
@@ -143,7 +150,7 @@ class TestSpawn:
return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
sm.runner.run = _slow_run
await sm.spawn("task", session_key="s1")
await sm.spawn("task", runtime=_runtime(), session_key="s1")
assert "s1" in sm._session_tasks
assert len(sm._session_tasks["s1"]) == 1
@@ -160,7 +167,7 @@ class TestSpawn:
return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
sm.runner.run = _slow_run
await sm.spawn("task")
await sm.spawn("task", runtime=_runtime())
assert len(sm._session_tasks) == 0
block.set()
@@ -176,7 +183,7 @@ class TestSpawn:
sm.runner.run = _slow_run
long_task = "A" * 50
await sm.spawn(long_task, session_key="s1")
await sm.spawn(long_task, runtime=_runtime(), session_key="s1")
status = next(iter(sm._task_statuses.values()))
assert status.label == long_task[:30] + "..."
@@ -192,7 +199,9 @@ class TestSpawn:
return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
sm.runner.run = _slow_run
await sm.spawn("task", label="Custom Label", session_key="s1")
await sm.spawn(
"task", runtime=_runtime(), label="Custom Label", session_key="s1"
)
status = next(iter(sm._task_statuses.values()))
assert status.label == "Custom Label"
@@ -205,12 +214,47 @@ class TestSpawn:
sm.runner.run = AsyncMock(return_value=AgentRunResult(
final_content="done", messages=[], stop_reason="completed",
))
await sm.spawn("task", session_key="s1")
await sm.spawn("task", runtime=_runtime(), session_key="s1")
await _drain_subagent_tasks(sm)
assert len(sm._running_tasks) == 0
assert len(sm._task_statuses) == 0
assert len(sm._session_tasks) == 0
@pytest.mark.asyncio
async def test_runtime_is_captured_before_background_task_starts(self, tmp_path):
sm = _manager(tmp_path)
runtime = _runtime(temperature=0.2)
entered = asyncio.Event()
release = asyncio.Event()
seen: dict[str, object] = {}
async def observe(spec):
seen["spec_runtime"] = spec.runtime
request_ctx = current_request_context()
seen["context_runtime"] = request_ctx.runtime if request_ctx else None
entered.set()
await release.wait()
return AgentRunResult(
final_content="done",
messages=[],
stop_reason="completed",
)
sm.runner.run = observe
await sm.spawn("task", runtime=runtime, session_key="s1")
runtime.provider.generation = GenerationSettings(
temperature=0.9,
max_tokens=128,
)
await asyncio.wait_for(entered.wait(), timeout=1)
assert seen["spec_runtime"] is runtime
assert seen["context_runtime"] is runtime
assert runtime.generation.temperature == 0.2
release.set()
await _drain_subagent_tasks(sm)
# ---------------------------------------------------------------------------
# _run_subagent
@@ -229,6 +273,7 @@ class TestRunSubagent:
"t1", "do task", "label",
{"channel": "cli", "chat_id": "direct"},
SubagentStatus(task_id="t1", label="label", task_description="do task", started_at=time.monotonic()),
_runtime(),
)
mock_announce.assert_called_once()
assert mock_announce.call_args.args[-2] == "ok"
@@ -244,7 +289,7 @@ class TestRunSubagent:
with patch.object(sm, "_announce_result", new_callable=AsyncMock) as mock_announce:
await sm._run_subagent(
"t1", "do task", "label",
{"channel": "cli", "chat_id": "direct"}, status,
{"channel": "cli", "chat_id": "direct"}, status, _runtime(),
)
assert mock_announce.call_args.args[-2] == "error"
@@ -256,7 +301,7 @@ class TestRunSubagent:
with patch.object(sm, "_announce_result", new_callable=AsyncMock) as mock_announce:
await sm._run_subagent(
"t1", "do task", "label",
{"channel": "cli", "chat_id": "direct"}, status,
{"channel": "cli", "chat_id": "direct"}, status, _runtime(),
)
assert status.phase == "error"
assert "LLM down" in status.error
@@ -272,7 +317,7 @@ class TestRunSubagent:
with patch.object(sm, "_announce_result", new_callable=AsyncMock):
await sm._run_subagent(
"t1", "do task", "label",
{"channel": "cli", "chat_id": "direct"}, status,
{"channel": "cli", "chat_id": "direct"}, status, _runtime(),
)
assert status.phase == "done"
assert status.stop_reason == "completed"
@@ -451,8 +496,9 @@ class TestCancelBySession:
return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
sm.runner.run = _slow_run
await sm.spawn("task1", session_key="s1")
await sm.spawn("task2", session_key="s1")
runtime = _runtime()
await sm.spawn("task1", runtime=runtime, session_key="s1")
await sm.spawn("task2", runtime=runtime, session_key="s1")
assert len(sm._session_tasks.get("s1", set())) == 2
count = await sm.cancel_by_session("s1")
@@ -472,7 +518,7 @@ class TestCancelBySession:
sm.runner.run = AsyncMock(return_value=AgentRunResult(
final_content="done", messages=[], stop_reason="completed",
))
await sm.spawn("task1", session_key="s1")
await sm.spawn("task1", runtime=_runtime(), session_key="s1")
await _drain_subagent_tasks(sm)
count = await sm.cancel_by_session("s1")
@@ -499,8 +545,9 @@ class TestRunningCounts:
return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
sm.runner.run = _slow_run
await sm.spawn("t1", session_key="s1")
await sm.spawn("t2", session_key="s1")
runtime = _runtime()
await sm.spawn("t1", runtime=runtime, session_key="s1")
await sm.spawn("t2", runtime=runtime, session_key="s1")
assert sm.get_running_count() == 2
assert sm.get_running_count_by_session("s1") == 2