Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(agent): unify request context routing

Browse Source
This commit is contained in:
chengyongru
2026-07-10 17:54:34 +08:00
committed by Xubin Ren
parent bb3b449e09
commit 42d7ad34a4
20 changed files with 363 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/tools/test_long_task.py
+45 -37
View File
@@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.context import RequestContext
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.long_task import (
CompleteGoalTool,
LongTaskTool,
@@ -24,15 +24,16 @@ from nanobot.session.webui_turns import WebuiTurnCoordinator
def _tools(sm: SessionManager) -> tuple[LongTaskTool, CompleteGoalTool]:
lt = LongTaskTool(sessions=sm)
cg = CompleteGoalTool(sessions=sm)
rc = RequestContext(
return lt, cg
def _request_context(chat_id: str = "c1") -> RequestContext:
return RequestContext(
channel="websocket",
chat_id="c1",
session_key="websocket:c1",
chat_id=chat_id,
session_key=f"websocket:{chat_id}",
metadata={},
)
lt.set_context(rc)
cg.set_context(rc)
return lt, cg
@pytest.mark.asyncio
@@ -40,7 +41,8 @@ async def test_long_task_records_goal_metadata(tmp_path):
sm = SessionManager(tmp_path)
lt, _cg = _tools(sm)
out = await lt.execute(goal="Do the thing", ui_summary="thing")
with request_context(_request_context()):
out = await lt.execute(goal="Do the thing", ui_summary="thing")
assert "Goal recorded" in out
sess = sm.get_or_create("websocket:c1")
@@ -56,8 +58,9 @@ async def test_long_task_rejects_second_active_goal(tmp_path):
sm = SessionManager(tmp_path)
lt, _cg = _tools(sm)
await lt.execute(goal="First")
out = await lt.execute(goal="Second")
with request_context(_request_context()):
await lt.execute(goal="First")
out = await lt.execute(goal="Second")
assert "already active" in out
@@ -66,8 +69,9 @@ async def test_complete_goal_closes_active_goal(tmp_path):
sm = SessionManager(tmp_path)
lt, cg = _tools(sm)
await lt.execute(goal="X")
out = await cg.execute(recap="Done.")
with request_context(_request_context()):
await lt.execute(goal="X")
out = await cg.execute(recap="Done.")
assert "marked complete" in out
sess = sm.get_or_create("websocket:c1")
@@ -84,19 +88,23 @@ async def test_goal_tools_keep_request_context_per_task(tmp_path):
ctx_a = RequestContext(channel="websocket", chat_id="a", session_key="websocket:a")
ctx_b = RequestContext(channel="websocket", chat_id="b", session_key="websocket:b")
lt.set_context(ctx_a)
task_a = asyncio.create_task(lt.execute(goal="Goal A"))
lt.set_context(ctx_b)
task_b = asyncio.create_task(lt.execute(goal="Goal B"))
async def start_goal(ctx: RequestContext, goal: str) -> str:
with request_context(ctx):
return await lt.execute(goal=goal)
task_a = asyncio.create_task(start_goal(ctx_a, "Goal A"))
task_b = asyncio.create_task(start_goal(ctx_b, "Goal B"))
await asyncio.gather(task_a, task_b)
assert sm.get_or_create("websocket:a").metadata[GOAL_STATE_KEY]["objective"] == "Goal A"
assert sm.get_or_create("websocket:b").metadata[GOAL_STATE_KEY]["objective"] == "Goal B"
cg.set_context(ctx_a)
done_a = asyncio.create_task(cg.execute(recap="Done A"))
cg.set_context(ctx_b)
done_b = asyncio.create_task(cg.execute(recap="Done B"))
async def complete_goal(ctx: RequestContext, recap: str) -> str:
with request_context(ctx):
return await cg.execute(recap=recap)
done_a = asyncio.create_task(complete_goal(ctx_a, "Done A"))
done_b = asyncio.create_task(complete_goal(ctx_b, "Done B"))
await asyncio.gather(done_a, done_b)
assert sm.get_or_create("websocket:a").metadata[GOAL_STATE_KEY]["recap"] == "Done A"
@@ -104,19 +112,19 @@ async def test_goal_tools_keep_request_context_per_task(tmp_path):
@pytest.mark.asyncio
async def test_goal_tools_context_isolated_across_tool_types(tmp_path):
"""LongTaskTool and CompleteGoalTool must not share routing context."""
async def test_goal_tools_share_authoritative_request_context(tmp_path):
"""Both goal tools resolve routing from the same request snapshot."""
sm = SessionManager(tmp_path)
lt = LongTaskTool(sessions=sm)
cg = CompleteGoalTool(sessions=sm)
ctx = RequestContext(channel="websocket", chat_id="a", session_key="websocket:a")
lt.set_context(ctx)
assert cg._request_ctx.get() is None
with request_context(ctx):
assert lt._session() is sm.get_or_create("websocket:a")
assert cg._session() is sm.get_or_create("websocket:a")
cg.set_context(ctx)
assert lt._request_ctx.get() is ctx
assert cg._request_ctx.get() is ctx
assert lt._session() is None
assert cg._session() is None
@pytest.mark.asyncio
@@ -137,9 +145,8 @@ async def test_long_task_publishes_goal_state_ws_after_save(tmp_path):
session_key="websocket:chat-99",
metadata={},
)
lt.set_context(rc)
await lt.execute(goal="Objective alpha", ui_summary="alpha")
with request_context(rc):
await lt.execute(goal="Objective alpha", ui_summary="alpha")
bus.publish_outbound.assert_awaited_once()
call = bus.publish_outbound.await_args.args[0]
@@ -172,12 +179,11 @@ async def test_complete_goal_publishes_inactive_goal_state_ws(tmp_path):
session_key="websocket:chat-z",
metadata={},
)
lt.set_context(rc)
await lt.execute(goal="X")
with request_context(rc):
await lt.execute(goal="X")
bus.publish_outbound.reset_mock()
cg.set_context(rc)
await cg.execute(recap="Done.")
bus.publish_outbound.reset_mock()
await cg.execute(recap="Done.")
bus.publish_outbound.assert_awaited_once()
call = bus.publish_outbound.await_args.args[0]
@@ -190,7 +196,8 @@ async def test_complete_goal_without_active_is_noop_message(tmp_path):
sm = SessionManager(tmp_path)
_lt, cg = _tools(sm)
out = await cg.execute(recap="n/a")
with request_context(_request_context()):
out = await cg.execute(recap="n/a")
assert "No active" in out
@@ -198,7 +205,8 @@ async def test_complete_goal_without_active_is_noop_message(tmp_path):
async def test_long_task_skips_ws_publish_without_bus(tmp_path):
sm = SessionManager(tmp_path)
lt, _cg = _tools(sm)
out = await lt.execute(goal="Solo", ui_summary="s")
with request_context(_request_context()):
out = await lt.execute(goal="Solo", ui_summary="s")
assert "Goal recorded" in out
tests/agent/tools/test_self_tool.py
+21 -8
View File
@@ -4,11 +4,12 @@ from __future__ import annotations
import time
from pathlib import Path
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import pytest
from pydantic import BaseModel
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.self import MyTool
# ---------------------------------------------------------------------------
@@ -1121,14 +1122,26 @@ class TestLastUsageInSummary:
# ---------------------------------------------------------------------------
# set_context (audit session tracking)
# request context (audit session tracking)
# ---------------------------------------------------------------------------
class TestSetContext:
class TestRequestContext:
def test_set_context_stores_channel_and_chat_id(self):
from nanobot.agent.tools.context import RequestContext
def test_audit_reads_bound_session(self):
tool = _make_tool()
tool.set_context(RequestContext(channel="feishu", chat_id="oc_abc123"))
assert tool._channel == "feishu"
assert tool._chat_id == "oc_abc123"
ctx = RequestContext(
channel="feishu",
chat_id="oc_abc123",
session_key="feishu:oc_abc123",
)
with patch("nanobot.agent.tools.self.logger.info") as info:
with request_context(ctx):
tool._audit("modify", "temperature = 0.2")
info.assert_called_once_with(
"self.{} | {} | session:{}",
"modify",
"temperature = 0.2",
"feishu:oc_abc123",
)
tests/agent/tools/test_subagent_tools.py
+9 -10
View File
@@ -159,19 +159,18 @@ async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
mgr.runner.run = AsyncMock(side_effect=fake_run)
from nanobot.agent.tools.context import RequestContext
from nanobot.agent.tools.context import RequestContext, request_context
tool = SpawnTool(mgr)
tool.set_context(RequestContext(channel="test", chat_id="c1", session_key="test:c1"))
with request_context(RequestContext(channel="test", chat_id="c1", session_key="test:c1")):
# First spawn succeeds
result = await tool.execute(task="first task")
assert "started" in result
# First spawn succeeds
result = await tool.execute(task="first task")
assert "started" in result
# Second spawn should be rejected (default limit is 1)
result = await tool.execute(task="second task")
assert "Cannot spawn subagent" in result
assert "concurrency limit reached" in result
# Second spawn should be rejected (default limit is 1)
result = await tool.execute(task="second task")
assert "Cannot spawn subagent" in result
assert "concurrency limit reached" in result
# Release the first subagent
release.set()