refactor(agent): unify request context routing

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
+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