refactor(agent): capture subagent runtime before spawn
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user