fix(agent): address session and streaming concurrency bugs

tests/agent/test_runner_governance.py

@@ -105,6 +105,60 @@ def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch
     assert trimmed[0]["role"] == "system"
     non_system = [m for m in trimmed if m["role"] != "system"]
     assert non_system[0]["role"] == "user", f"Expected user after system, got {non_system[0]['role']}"
+
+
+def test_snip_history_reserves_budget_for_tool_definitions(monkeypatch):
+    from nanobot.agent.runner import AgentRunSpec, AgentRunner
+
+    provider = MagicMock()
+    tools = MagicMock()
+    tools.get_definitions.return_value = [{"type": "function", "function": {"name": "large_tool"}}]
+    runner = AgentRunner(provider)
+    messages = [
+        {"role": "system", "content": "system"},
+        {"role": "user", "content": "old user"},
+        {"role": "assistant", "content": "old assistant"},
+        {"role": "user", "content": "recent one"},
+        {"role": "assistant", "content": "recent answer"},
+        {"role": "user", "content": "recent two"},
+    ]
+    spec = AgentRunSpec(
+        initial_messages=messages,
+        tools=tools,
+        model="test-model",
+        max_iterations=1,
+        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
+        context_window_tokens=2000,
+        context_block_limit=500,
+    )
+
+    def _estimate(_provider, _model, estimate_messages, estimate_tools):
+        if estimate_messages == messages:
+            return 1000, None
+        assert estimate_messages == [{"role": "system", "content": "system"}]
+        assert estimate_tools == tools.get_definitions.return_value
+        return 350, None
+
+    monkeypatch.setattr("nanobot.agent.runner.estimate_prompt_tokens_chain", _estimate)
+    token_sizes = {
+        "system": 50,
+        "old user": 200,
+        "old assistant": 200,
+        "recent one": 200,
+        "recent answer": 200,
+        "recent two": 200,
+    }
+    monkeypatch.setattr(
+        "nanobot.agent.runner.estimate_message_tokens",
+        lambda msg: token_sizes.get(str(msg.get("content")), 40),
+    )
+
+    trimmed = runner._snip_history(spec, messages)
+
+    contents = [message.get("content") for message in trimmed]
+    assert contents == ["system", "recent two"]
+
+
 async def test_backfill_missing_tool_results_inserts_error():
     """Orphaned tool_use (no matching tool_result) should get a synthetic error."""
     from nanobot.agent.runner import AgentRunner, _BACKFILL_CONTENT

tests/agent/test_runner_injections.py

@@ -554,6 +554,31 @@ async def test_pending_queue_cleanup_on_dispatch(tmp_path):
     assert msg.session_key not in loop._pending_queues
 
 
+@pytest.mark.asyncio
+async def test_waiting_dispatch_does_not_replace_active_pending_queue(tmp_path):
+    """A queued dispatch must not steal the active task's injection queue."""
+    from nanobot.bus.events import InboundMessage
+
+    loop = _make_loop(tmp_path)
+    session_key = "cli:c"
+    lock = loop._session_locks.setdefault(session_key, asyncio.Lock())
+    await lock.acquire()
+    active_pending = asyncio.Queue(maxsize=1)
+    loop._pending_queues[session_key] = active_pending
+
+    waiting = asyncio.create_task(
+        loop._dispatch(InboundMessage(channel="cli", sender_id="u", chat_id="c", content="queued"))
+    )
+    await asyncio.sleep(0.05)
+
+    assert loop._pending_queues[session_key] is active_pending
+
+    waiting.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await waiting
+    lock.release()
+
+
 @pytest.mark.asyncio
 async def test_followup_routed_to_pending_queue(tmp_path):
     """Unified-session follow-ups should route into the active pending queue."""

tests/agent/test_unified_session.py

@@ -474,6 +474,32 @@ class TestStopCommandWithUnifiedSession:
         assert task.cancelled() or task.done()
         assert "Stopped 1 task" in result.content
 
+    @pytest.mark.asyncio
+    async def test_stop_command_uses_effective_key_without_session_override(self, tmp_path: Path):
+        """Priority /stop must cancel the unified session even before dispatch rewrites the message."""
+        from nanobot.agent.loop import UNIFIED_SESSION_KEY
+        from nanobot.command.builtin import cmd_stop
+
+        loop = _make_loop(tmp_path, unified_session=True)
+
+        async def long_running():
+            await asyncio.sleep(10)
+
+        task = asyncio.create_task(long_running())
+        loop._active_tasks[UNIFIED_SESSION_KEY] = [task]
+        msg = InboundMessage(
+            channel="telegram",
+            chat_id="123456",
+            sender_id="user1",
+            content="/stop",
+        )
+        ctx = CommandContext(msg=msg, session=None, key=UNIFIED_SESSION_KEY, raw="/stop", loop=loop)
+
+        result = await cmd_stop(ctx)
+
+        assert task.cancelled() or task.done()
+        assert "Stopped 1 task" in result.content
+
     @pytest.mark.asyncio
     async def test_stop_command_cross_channel_in_unified_mode(self, tmp_path: Path):
         """In unified mode, /stop from one channel cancels tasks from another channel."""
@@ -504,4 +530,4 @@ class TestStopCommandWithUnifiedSession:
         result = await cmd_stop(ctx)
 
         # Both tasks should be cancelled
-        assert "Stopped 2 task" in result.content
+        assert "Stopped 2 task" in result.content

tests/agent/tools/test_long_task.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -72,6 +73,33 @@ async def test_complete_goal_closes_active_goal(tmp_path):
     assert blob["recap"] == "Done."
 
 
+@pytest.mark.asyncio
+async def test_goal_tools_keep_request_context_per_task(tmp_path):
+    sm = SessionManager(tmp_path)
+    lt = LongTaskTool(sessions=sm)
+    cg = CompleteGoalTool(sessions=sm)
+    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"))
+    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"))
+    await asyncio.gather(done_a, done_b)
+
+    assert sm.get_or_create("websocket:a").metadata[GOAL_STATE_KEY]["recap"] == "Done A"
+    assert sm.get_or_create("websocket:b").metadata[GOAL_STATE_KEY]["recap"] == "Done B"
+
+
 @pytest.mark.asyncio
 async def test_long_task_publishes_goal_state_ws_after_save(tmp_path):
     bus = MagicMock()