test/acp/test_agent_lifecycle.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
import asyncio
from collections.abc import Sequence
from typing import Any, Self
import acp
import pytest
from acp import schema
from acp.core import ClientSideConnection
from ag2 import Agent, Context
from ag2.acp import ACPAgent, SessionConfig
from ag2.acp.guard import serve
from ag2.acp.testing import RecordingClient, connect, duplex
from ag2.config import LLMClient, ModelConfig
from ag2.events import BaseEvent, ModelResponse, ToolCallEvent
from ag2.history import MemoryStorage
from ag2.testing import TestConfig
class _HeldClient(LLMClient):
"""Parks inside the LLM call and records whether the turn was cancelled there."""
def __init__(self, entered: asyncio.Event, cancelled: asyncio.Event) -> None:
self.entered = entered
self.cancelled = cancelled
async def __call__(self, messages: Sequence[BaseEvent], context: Context, **kwargs: Any) -> ModelResponse:
self.entered.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
self.cancelled.set()
raise
raise AssertionError("the held turn was never meant to finish") # pragma: no cover
class _HeldConfig(ModelConfig):
"""A config whose turns hang until something cancels them."""
def __init__(self) -> None:
self.entered = asyncio.Event()
self.cancelled = asyncio.Event()
def copy(self) -> Self:
return self
def create(self) -> _HeldClient:
return _HeldClient(self.entered, self.cancelled)
def create_files_client(self) -> None:
raise NotImplementedError
class _StaticAgent:
"""An ACP agent object rather than a per-connection factory — ``serve`` takes both.
Records whether teardown reached it: an object the caller built is the
caller's to close, so it should not.
"""
def __init__(self) -> None:
self.initialized = False
self.closed = False
async def initialize(self, protocol_version: int, **kwargs: Any) -> schema.InitializeResponse:
self.initialized = True
return schema.InitializeResponse(protocol_version=min(protocol_version, acp.PROTOCOL_VERSION))
async def aclose(self) -> None:
self.closed = True
@pytest.mark.asyncio
class TestSessionsDoNotOutliveTheirConnection:
async def test_the_registry_is_empty_after_the_client_disconnects(self) -> None:
server = ACPAgent(Agent("workie", config=TestConfig("ok")))
async with connect(server) as (conn, _):
await conn.new_session(cwd="/tmp")
assert len(server.sessions) == 1
assert len(server.sessions) == 0
async def test_every_session_of_a_connection_goes_at_once(self) -> None:
server = ACPAgent(Agent("workie", config=TestConfig("ok")))
async with connect(server) as (conn, _):
for _ in range(3):
await conn.new_session(cwd="/tmp")
assert len(server.sessions) == 3
assert len(server.sessions) == 0
async def test_a_reconnecting_client_cannot_name_the_old_connection_s_session(self) -> None:
server = ACPAgent(Agent("workie", config=TestConfig("ok")))
async with connect(server) as (first, _):
stale = (await first.new_session(cwd="/tmp")).session_id
async with connect(server) as (second, _):
with pytest.raises(acp.RequestError):
await second.prompt(session_id=stale, prompt=[acp.text_block("still there?")])
async def test_one_connection_s_teardown_leaves_another_s_sessions_alone(self) -> None:
"""Scopes are per-connection, so teardown has to be too."""
server = ACPAgent(Agent("workie", config=TestConfig("ok")))
async with connect(server) as (first, _):
kept = (await first.new_session(cwd="/tmp")).session_id
survivor = server.sessions
async with connect(server) as (second, _):
await second.new_session(cwd="/tmp")
doomed = server.sessions
assert len(doomed) == 0
assert len(survivor) == 1
assert (await survivor.get(kept)).session_id == kept
@pytest.mark.asyncio
class TestTeardownStopsLiveWork:
async def test_a_turn_still_running_at_disconnect_is_cancelled(self) -> None:
config = _HeldConfig()
server = ACPAgent(Agent("workie", config=config))
async with connect(server) as (conn, _):
session = await conn.new_session(cwd="/tmp")
turn = asyncio.create_task(conn.prompt(session_id=session.session_id, prompt=[acp.text_block("slow")]))
await asyncio.wait_for(config.entered.wait(), timeout=5)
assert config.cancelled.is_set()
# The request dies with the connection instead of being answered.
with pytest.raises(ConnectionError):
await asyncio.wait_for(turn, timeout=5)
async def test_a_busy_session_is_still_dropped_from_the_registry(self) -> None:
"""A turn in flight must not be able to hold its connection's scope open."""
config = _HeldConfig()
server = ACPAgent(Agent("workie", config=config))
async with connect(server) as (conn, _):
session = await conn.new_session(cwd="/tmp")
turn = asyncio.create_task(conn.prompt(session_id=session.session_id, prompt=[acp.text_block("slow")]))
await asyncio.wait_for(config.entered.wait(), timeout=5)
assert len(server.sessions) == 0
with pytest.raises(ConnectionError):
await asyncio.wait_for(turn, timeout=5)
@pytest.mark.asyncio
class TestDurableHistoryIsDropped:
async def test_a_session_s_history_does_not_outlive_its_connection(self) -> None:
"""The default storage is in-memory; an injected one keeps orphans visible."""
storage = MemoryStorage()
server = ACPAgent(Agent("workie", config=TestConfig("answered")), sessions=SessionConfig(storage=storage))
async with connect(server) as (conn, _):
session_id = (await conn.new_session(cwd="/tmp")).session_id
await conn.prompt(session_id=session_id, prompt=[acp.text_block("hello")])
store = server.sessions
session = await store.get(session_id)
assert await store.stream(session).history.get_events() != []
assert list(await store.stream(session).history.get_events()) == []
async def test_only_the_closing_connection_s_history_is_dropped(self) -> None:
storage = MemoryStorage()
agent = Agent("workie", config=TestConfig("answered", "answered"))
server = ACPAgent(agent, sessions=SessionConfig(storage=storage))
async with connect(server) as (first, _):
kept_id = (await first.new_session(cwd="/tmp")).session_id
await first.prompt(session_id=kept_id, prompt=[acp.text_block("hello")])
survivor = server.sessions
kept = await survivor.get(kept_id)
async with connect(server) as (second, _):
doomed_id = (await second.new_session(cwd="/tmp")).session_id
await second.prompt(session_id=doomed_id, prompt=[acp.text_block("hello")])
doomed_store = server.sessions
doomed = await doomed_store.get(doomed_id)
assert list(await doomed_store.stream(doomed).history.get_events()) == []
assert await survivor.stream(kept).history.get_events() != []
async def test_retained_history_outlives_its_connection(self) -> None:
"""The opt-in mirror of the first test: with ``retain_history`` the disconnect only releases."""
storage = MemoryStorage()
server = ACPAgent(
Agent("workie", config=TestConfig("answered")),
sessions=SessionConfig(storage=storage, retain_history=True),
)
async with connect(server) as (conn, _):
session_id = (await conn.new_session(cwd="/tmp")).session_id
await conn.prompt(session_id=session_id, prompt=[acp.text_block("hello")])
store = server.sessions
session = await store.get(session_id)
assert len(store) == 0
assert await store.stream(session).history.get_events() != []
@pytest.mark.asyncio
class TestServeOverStreams:
"""``run_stdio``'s own path: the Client goes away, nothing cancels the server."""
async def test_eof_from_the_client_releases_the_scope(self) -> None:
server = ACPAgent(Agent("workie", config=TestConfig("ok")))
(agent_r, agent_w), (client_r, client_w) = await duplex()
served = asyncio.create_task(serve(server.bind, agent_r, agent_w))
conn = ClientSideConnection(lambda _agent: RecordingClient(), client_w, client_r)
try:
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION)
await conn.new_session(cwd="/tmp")
assert len(server.sessions) == 1
client_w.close() # EOF on the agent's input: the Client hung up
await asyncio.wait_for(served, timeout=5)
assert len(server.sessions) == 0
finally:
await conn.close()
agent_w.close()
async def test_an_agent_object_the_caller_built_is_not_closed(self) -> None:
"""``serve`` only releases the scope it created; a passed-in agent is not one."""
agent = _StaticAgent()
(agent_r, agent_w), (client_r, client_w) = await duplex()
served = asyncio.create_task(serve(agent, agent_r, agent_w))
conn = ClientSideConnection(lambda _agent: RecordingClient(), client_w, client_r)
try:
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION)
client_w.close()
await asyncio.wait_for(served, timeout=5)
assert agent.initialized is True
assert agent.closed is False
finally:
await conn.close()
agent_w.close()
@pytest.mark.asyncio
class TestBackgroundWorkBelongsToTheSession:
"""A stream carries an inbox and background tasks, not only history.
A fresh stream per turn would strand work that finished after its own turn
and hide it from teardown.
"""
@staticmethod
def _spawning_agent(ran: list[str]) -> Agent:
agent = Agent(
"workie",
config=TestConfig(ToolCallEvent(name="spawn", arguments="{}"), "ok", "ok"),
)
@agent.tool
async def spawn(ctx: Context) -> str:
"""Start work that deliberately outlives this turn."""
async def later() -> None:
await asyncio.sleep(0.05)
ran.append("done")
ctx.spawn_background(later())
return "spawned"
return agent
async def test_the_session_keeps_one_stream_across_turns(self) -> None:
server = ACPAgent(Agent("workie", config=TestConfig("ok", "ok")))
async with connect(server) as (conn, _):
created = await conn.new_session(cwd="/tmp")
await conn.prompt(session_id=created.session_id, prompt=[acp.text_block("one")])
session = await server.sessions.get(created.session_id)
first = server.sessions.stream(session)
await conn.prompt(session_id=created.session_id, prompt=[acp.text_block("two")])
assert server.sessions.stream(session) is first
async def test_background_work_is_tracked_where_teardown_can_see_it(self) -> None:
ran: list[str] = []
server = ACPAgent(self._spawning_agent(ran))
async with connect(server) as (conn, _):
created = await conn.new_session(cwd="/tmp")
await conn.prompt(session_id=created.session_id, prompt=[acp.text_block("go")])
session = await server.sessions.get(created.session_id)
assert server.sessions.stream(session)._background_tasks
async def test_closing_a_session_stops_its_background_work(self) -> None:
"""Otherwise a closed conversation keeps reaching the outside world."""
ran: list[str] = []
server = ACPAgent(self._spawning_agent(ran))
async with connect(server) as (conn, _):
created = await conn.new_session(cwd="/tmp")
await conn.prompt(session_id=created.session_id, prompt=[acp.text_block("go")])
await server.sessions.close(created.session_id)
await asyncio.sleep(0.2)
assert ran == []
async def test_a_turn_does_not_leave_its_forwarder_behind(self) -> None:
"""The stream outlives the turn now, so the subscriber must be released.
Left attached, the second turn would deliver every update twice.
"""
server = ACPAgent(Agent("workie", config=TestConfig("first", "second")))
async with connect(server) as (conn, recorder):
created = await conn.new_session(cwd="/tmp")
await conn.prompt(session_id=created.session_id, prompt=[acp.text_block("one")])
await conn.prompt(session_id=created.session_id, prompt=[acp.text_block("two")])
# Two turns, so two message updates. A leaked forwarder would make the
# second turn deliver its update once per prompt the session had run.
texts = [
u.content.text for u in recorder.updates_for(created.session_id) if isinstance(u, schema.AgentMessageChunk)
]
assert len(texts) == 2