test_debug_runtime.py 87 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727
  1. import asyncio
  2. import importlib
  3. import json
  4. import logging
  5. from collections.abc import AsyncIterator
  6. from typing import Any
  7. import pytest
  8. from agent_lab.application.contracts import AgentParams, DebugRunRequest, EventAgentParams
  9. from agent_lab.application.event_agent import EventAgentRequest
  10. from agent_lab.application.events import EventBatchExecutor, ResultPolicy
  11. from agent_lab.application.runtime import DebugRuntime
  12. from agent_lab.application.tools import ToolDefinition, ToolExecutionContext, ToolRegistry
  13. from agent_lab.domain.events import ToolCallEvent
  14. from agent_lab.domain.messages import ChatMessage, StreamItem, TokenUsage
  15. from agent_lab.infrastructure.sqlite_store import SQLiteSessionStore
  16. def _runtime_queues_class():
  17. module = importlib.import_module("agent_lab.application.queues")
  18. return module.RuntimeQueues
  19. async def _collect_outputs(stream: AsyncIterator[dict[str, Any]]) -> list[dict[str, Any]]:
  20. return [message async for message in stream]
  21. def _without_audit(outputs: list[dict[str, Any]]) -> list[dict[str, Any]]:
  22. return [message for message in outputs if message["type"] != "audit"]
  23. def _message_types(outputs: list[dict[str, Any]]) -> list[str]:
  24. return [message["type"] for message in _without_audit(outputs)]
  25. def _event_tool_call_from_tools(
  26. tools: list[dict],
  27. messages: list[ChatMessage],
  28. ) -> StreamItem:
  29. tool_name = tools[0]["function"]["name"]
  30. content = ""
  31. for role in ("assistant", "user"):
  32. content = next(
  33. (
  34. message.content
  35. for message in reversed(messages)
  36. if message.role == role and message.content.strip()
  37. ),
  38. "",
  39. )
  40. if content:
  41. break
  42. arguments = {
  43. "message": content,
  44. "query": content,
  45. "title": content,
  46. }
  47. return StreamItem.provider_tool_call(
  48. ToolCallEvent(
  49. id="event_agent_call_1",
  50. name=tool_name,
  51. arguments=arguments,
  52. raw_arguments=json.dumps(arguments),
  53. )
  54. )
  55. async def _next_non_audit(
  56. stream: AsyncIterator[dict[str, Any]],
  57. ) -> dict[str, Any]:
  58. while True:
  59. message = await anext(stream)
  60. if message["type"] != "audit":
  61. return message
  62. async def _next_non_audit_from_queue(queues: Any) -> dict[str, Any]:
  63. while True:
  64. message = await queues.output.get()
  65. if message["type"] != "audit":
  66. return message
  67. class RecordingQueue(asyncio.Queue):
  68. def __init__(self, name: str, log: list[tuple[str, str, str]]) -> None:
  69. super().__init__()
  70. self.name = name
  71. self.log = log
  72. async def put(self, item: Any) -> None:
  73. self.log.append((self.name, "put", self._describe(item)))
  74. await super().put(item)
  75. async def get(self) -> Any:
  76. item = await super().get()
  77. self.log.append((self.name, "get", self._describe(item)))
  78. return item
  79. def _describe(self, item: Any) -> str:
  80. if isinstance(item, ChatMessage):
  81. if item.role == "tool":
  82. return f"tool:{item.tool_call_id}"
  83. return item.role
  84. if isinstance(item, ToolCallEvent):
  85. return f"event:{item.name}:{item.id}"
  86. if isinstance(item, EventAgentRequest):
  87. events = ",".join(f"{event.name}:{event.id}" for event in item.events)
  88. return f"event_request:{events}"
  89. if isinstance(item, dict):
  90. return f"output:{item.get('type')}"
  91. return type(item).__name__
  92. class FakeChatClient:
  93. def __init__(self) -> None:
  94. self.calls = 0
  95. async def stream_chat(
  96. self,
  97. messages: list[ChatMessage],
  98. tools: list[dict],
  99. params: AgentParams,
  100. tool_choice: dict[str, Any] | None = None,
  101. ) -> AsyncIterator[StreamItem]:
  102. if tools:
  103. yield StreamItem.raw_response_chunk(
  104. {
  105. "choices": [
  106. {
  107. "delta": {
  108. "tool_calls": [
  109. {
  110. "index": 0,
  111. "function": {"name": tools[0]["function"]["name"]},
  112. }
  113. ]
  114. },
  115. "finish_reason": None,
  116. }
  117. ]
  118. }
  119. )
  120. yield _event_tool_call_from_tools(tools, messages)
  121. return
  122. self.calls += 1
  123. if self.calls == 1:
  124. yield StreamItem.raw_response_chunk(
  125. {
  126. "choices": [
  127. {
  128. "delta": {"content": "<agent_events>handoff_note</agent_events>"},
  129. "finish_reason": None,
  130. }
  131. ]
  132. }
  133. )
  134. yield StreamItem.text_event(
  135. ToolCallEvent(
  136. id="call_1",
  137. name="handoff_note",
  138. arguments={},
  139. raw_arguments="{}",
  140. )
  141. )
  142. return
  143. assert not any(message.role == "tool" for message in messages)
  144. assert any(message.name == "event_agent" for message in messages)
  145. yield StreamItem.message_delta("final answer")
  146. class WrongSourceEventChatClient:
  147. def __init__(self, item: StreamItem) -> None:
  148. self.item = item
  149. self.calls = 0
  150. async def stream_chat(
  151. self,
  152. messages: list[ChatMessage],
  153. tools: list[dict],
  154. params: AgentParams,
  155. tool_choice: dict[str, Any] | None = None,
  156. ) -> AsyncIterator[StreamItem]:
  157. if tools:
  158. yield _event_tool_call_from_tools(tools, messages)
  159. return
  160. self.calls += 1
  161. if self.calls == 1:
  162. yield self.item
  163. return
  164. yield StreamItem.message_delta("unexpected continuation")
  165. class IncrementingClock:
  166. def __init__(self, current: float = 100.0, step: float = 0.01) -> None:
  167. self.current = current
  168. self.step = step
  169. def __call__(self) -> float:
  170. self.current += self.step
  171. return self.current
  172. class StrictHistoryChatClient:
  173. def __init__(self) -> None:
  174. self.calls = 0
  175. self.second_call_messages: list[ChatMessage] = []
  176. async def stream_chat(
  177. self,
  178. messages: list[ChatMessage],
  179. tools: list[dict],
  180. params: AgentParams,
  181. tool_choice: dict[str, Any] | None = None,
  182. ) -> AsyncIterator[StreamItem]:
  183. if tools:
  184. yield StreamItem.raw_response_chunk(
  185. {
  186. "choices": [
  187. {
  188. "delta": {
  189. "tool_calls": [
  190. {
  191. "index": 0,
  192. "function": {"name": tools[0]["function"]["name"]},
  193. }
  194. ]
  195. },
  196. "finish_reason": None,
  197. }
  198. ]
  199. }
  200. )
  201. yield _event_tool_call_from_tools(tools, messages)
  202. return
  203. self.calls += 1
  204. if self.calls == 1:
  205. yield StreamItem.raw_response_chunk(
  206. {
  207. "choices": [
  208. {
  209. "delta": {"content": "<agent_events>handoff_note</agent_events>"},
  210. "finish_reason": None,
  211. }
  212. ]
  213. }
  214. )
  215. yield StreamItem.text_event(
  216. ToolCallEvent(
  217. id="call_1",
  218. name="handoff_note",
  219. arguments={},
  220. raw_arguments="{}",
  221. )
  222. )
  223. return
  224. self.second_call_messages = list(messages)
  225. yield StreamItem.message_delta("final answer")
  226. class MultiEventChatClient:
  227. def __init__(self) -> None:
  228. self.calls = 0
  229. self.second_call_messages: list[ChatMessage] = []
  230. async def stream_chat(
  231. self,
  232. messages: list[ChatMessage],
  233. tools: list[dict],
  234. params: AgentParams,
  235. tool_choice: dict[str, Any] | None = None,
  236. ) -> AsyncIterator[StreamItem]:
  237. if tools:
  238. yield _event_tool_call_from_tools(tools, messages)
  239. return
  240. self.calls += 1
  241. if self.calls == 1:
  242. yield StreamItem.message_delta("Checking events.")
  243. yield StreamItem.text_event(
  244. ToolCallEvent(
  245. id="call_1",
  246. name="handoff_note",
  247. arguments={"message": "ignored chat argument"},
  248. raw_arguments='{"message":"ignored chat argument"}',
  249. )
  250. )
  251. yield StreamItem.text_event(
  252. ToolCallEvent(
  253. id="call_2",
  254. name="audit_note",
  255. arguments={"message": "ignored chat argument"},
  256. raw_arguments='{"message":"ignored chat argument"}',
  257. )
  258. )
  259. return
  260. self.second_call_messages = list(messages)
  261. yield StreamItem.message_delta("Final answer.")
  262. class ToolCapturingChatClient:
  263. def __init__(self) -> None:
  264. self.tools: list[dict[str, Any]] = []
  265. self.messages: list[ChatMessage] = []
  266. async def stream_chat(
  267. self,
  268. messages: list[ChatMessage],
  269. tools: list[dict],
  270. params: AgentParams,
  271. tool_choice: dict[str, Any] | None = None,
  272. ) -> AsyncIterator[StreamItem]:
  273. self.messages = list(messages)
  274. self.tools = list(tools)
  275. yield StreamItem.message_delta("final answer")
  276. class EventLoopLimitChatClient:
  277. def __init__(self) -> None:
  278. self.calls = 0
  279. self.tools_by_call: list[list[dict[str, Any]]] = []
  280. self.messages_by_call: list[list[ChatMessage]] = []
  281. async def stream_chat(
  282. self,
  283. messages: list[ChatMessage],
  284. tools: list[dict],
  285. params: AgentParams,
  286. tool_choice: dict[str, Any] | None = None,
  287. ) -> AsyncIterator[StreamItem]:
  288. if tools:
  289. yield _event_tool_call_from_tools(tools, messages)
  290. return
  291. self.calls += 1
  292. self.tools_by_call.append(list(tools))
  293. self.messages_by_call.append(list(messages))
  294. if self.calls == 1:
  295. yield StreamItem.text_event(
  296. ToolCallEvent(
  297. id="call_1",
  298. name="handoff_note",
  299. arguments={},
  300. raw_arguments="{}",
  301. )
  302. )
  303. return
  304. yield StreamItem.message_delta("final after event limit")
  305. class RoundStatsChatClient:
  306. async def stream_chat(
  307. self,
  308. messages: list[ChatMessage],
  309. tools: list[dict],
  310. params: AgentParams,
  311. tool_choice: dict[str, Any] | None = None,
  312. ) -> AsyncIterator[StreamItem]:
  313. yield StreamItem.message_delta("hello")
  314. yield StreamItem.usage_item(
  315. TokenUsage(
  316. prompt_tokens=10,
  317. completion_tokens=20,
  318. total_tokens=30,
  319. cached_tokens=5,
  320. )
  321. )
  322. class EventRoundStatsChatClient:
  323. def __init__(self) -> None:
  324. self.calls = 0
  325. async def stream_chat(
  326. self,
  327. messages: list[ChatMessage],
  328. tools: list[dict],
  329. params: AgentParams,
  330. tool_choice: dict[str, Any] | None = None,
  331. ) -> AsyncIterator[StreamItem]:
  332. if tools:
  333. yield StreamItem.raw_response_chunk(
  334. {
  335. "choices": [
  336. {
  337. "delta": {
  338. "tool_calls": [
  339. {
  340. "index": 0,
  341. "function": {
  342. "name": tools[0]["function"]["name"],
  343. },
  344. }
  345. ]
  346. },
  347. "finish_reason": None,
  348. }
  349. ]
  350. }
  351. )
  352. yield _event_tool_call_from_tools(tools, messages)
  353. return
  354. self.calls += 1
  355. if self.calls == 1:
  356. yield StreamItem.raw_response_chunk(
  357. {
  358. "choices": [
  359. {
  360. "delta": {
  361. "content": "<agent_events>handoff_note</agent_events>",
  362. },
  363. "finish_reason": None,
  364. }
  365. ]
  366. }
  367. )
  368. yield StreamItem.text_event(
  369. ToolCallEvent(
  370. id="call_1",
  371. name="handoff_note",
  372. arguments={},
  373. raw_arguments="{}",
  374. )
  375. )
  376. yield StreamItem.usage_item(
  377. TokenUsage(prompt_tokens=3, completion_tokens=0, total_tokens=3)
  378. )
  379. return
  380. yield StreamItem.raw_response_chunk(
  381. {
  382. "choices": [
  383. {
  384. "delta": {"content": "final answer"},
  385. "finish_reason": None,
  386. }
  387. ]
  388. }
  389. )
  390. yield StreamItem.message_delta("final answer")
  391. yield StreamItem.usage_item(
  392. TokenUsage(prompt_tokens=4, completion_tokens=6, total_tokens=10)
  393. )
  394. class ContextBoundaryChatClient:
  395. def __init__(self) -> None:
  396. self.calls = 0
  397. self.event_agent_messages: list[list[ChatMessage]] = []
  398. async def stream_chat(
  399. self,
  400. messages: list[ChatMessage],
  401. tools: list[dict],
  402. params: AgentParams,
  403. tool_choice: dict[str, Any] | None = None,
  404. ) -> AsyncIterator[StreamItem]:
  405. if tools:
  406. self.event_agent_messages.append(list(messages))
  407. yield _event_tool_call_from_tools(tools, messages)
  408. return
  409. self.calls += 1
  410. if self.calls == 1:
  411. yield StreamItem.message_delta("I will check.")
  412. yield StreamItem.text_event(
  413. ToolCallEvent(
  414. id="call_1",
  415. name="handoff_note",
  416. arguments={},
  417. raw_arguments="{}",
  418. )
  419. )
  420. return
  421. yield StreamItem.message_delta("final answer")
  422. class TwoTurnSessionChatClient:
  423. def __init__(self) -> None:
  424. self.calls = 0
  425. self.messages_by_call: list[list[ChatMessage]] = []
  426. async def stream_chat(
  427. self,
  428. messages: list[ChatMessage],
  429. tools: list[dict],
  430. params: AgentParams,
  431. tool_choice: dict[str, Any] | None = None,
  432. ) -> AsyncIterator[StreamItem]:
  433. if tools:
  434. yield _event_tool_call_from_tools(tools, messages)
  435. return
  436. self.calls += 1
  437. self.messages_by_call.append(list(messages))
  438. if self.calls in {1, 3}:
  439. yield StreamItem.text_event(
  440. ToolCallEvent(
  441. id=f"call_{self.calls}",
  442. name="handoff_note",
  443. arguments={},
  444. raw_arguments="{}",
  445. )
  446. )
  447. return
  448. yield StreamItem.message_delta(f"final answer {self.calls}")
  449. class SlowAfterEventChatClient:
  450. def __init__(self) -> None:
  451. self.calls = 0
  452. self.event_seen = asyncio.Event()
  453. self.release_stream = asyncio.Event()
  454. async def stream_chat(
  455. self,
  456. messages: list[ChatMessage],
  457. tools: list[dict],
  458. params: AgentParams,
  459. tool_choice: dict[str, Any] | None = None,
  460. ) -> AsyncIterator[StreamItem]:
  461. if tools:
  462. yield _event_tool_call_from_tools(tools, messages)
  463. return
  464. self.calls += 1
  465. if self.calls == 1:
  466. yield StreamItem.message_delta("Need event.")
  467. yield StreamItem.text_event(
  468. ToolCallEvent(
  469. id="call_1",
  470. name="handoff_note",
  471. arguments={},
  472. raw_arguments="{}",
  473. )
  474. )
  475. self.event_seen.set()
  476. await self.release_stream.wait()
  477. return
  478. yield StreamItem.message_delta("final answer")
  479. def test_runtime_queues_exposes_input_output_and_events_queues():
  480. RuntimeQueues = _runtime_queues_class()
  481. queues = RuntimeQueues()
  482. assert isinstance(queues.input, asyncio.Queue)
  483. assert isinstance(queues.output, asyncio.Queue)
  484. assert isinstance(queues.events, asyncio.Queue)
  485. assert queues.input is not queues.output
  486. assert queues.input is not queues.events
  487. assert queues.output is not queues.events
  488. @pytest.mark.asyncio
  489. async def test_runtime_routes_chat_events_through_event_agent_then_continues_chat():
  490. request = DebugRunRequest(
  491. user_message="debug this",
  492. system_prompts=["You are a debugger."],
  493. pre_messages=[],
  494. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  495. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=3),
  496. )
  497. client = FakeChatClient()
  498. runtime = DebugRuntime(client)
  499. outputs = [message async for message in runtime.run(request)]
  500. business_outputs = _without_audit(outputs)
  501. assert client.calls == 2
  502. assert _message_types(outputs) == [
  503. "session_started",
  504. "event",
  505. "tool_result",
  506. "round_stats",
  507. "message_delta",
  508. "round_stats",
  509. "done",
  510. ]
  511. assert business_outputs[1]["event"]["name"] == "handoff_note"
  512. assert business_outputs[4]["content"] == "final answer"
  513. @pytest.mark.asyncio
  514. @pytest.mark.parametrize(
  515. "item",
  516. [
  517. StreamItem.provider_tool_call(
  518. ToolCallEvent(
  519. id="provider_call_1",
  520. name="handoff_note",
  521. arguments={"message": "wrong source"},
  522. raw_arguments='{"message":"wrong source"}',
  523. )
  524. ),
  525. StreamItem.event(
  526. ToolCallEvent(
  527. id="legacy_event_1",
  528. name="handoff_note",
  529. arguments={},
  530. raw_arguments="{}",
  531. )
  532. ),
  533. ],
  534. ids=["provider_tool_call", "legacy_event"],
  535. )
  536. async def test_runtime_ignores_non_text_event_sources(item: StreamItem):
  537. request = DebugRunRequest(
  538. user_message="debug this",
  539. system_prompts=[],
  540. pre_messages=[],
  541. chat_agent=AgentParams(model="fake-model"),
  542. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
  543. )
  544. client = WrongSourceEventChatClient(item)
  545. outputs = [message async for message in DebugRuntime(client).run(request)]
  546. assert client.calls == 1
  547. assert "event" not in _message_types(outputs)
  548. assert "tool_result" not in _message_types(outputs)
  549. @pytest.mark.asyncio
  550. async def test_runtime_emits_audit_events_and_backend_logs(caplog):
  551. request = DebugRunRequest(
  552. user_message="debug this",
  553. system_prompts=[],
  554. pre_messages=[],
  555. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  556. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
  557. )
  558. runtime = DebugRuntime(FakeChatClient())
  559. with caplog.at_level(logging.INFO, logger="agent_lab.application.runtime"):
  560. outputs = [message async for message in runtime.run(request)]
  561. audit_events = [
  562. message["event"]
  563. for message in outputs
  564. if message["type"] == "audit"
  565. ]
  566. assert audit_events == [
  567. "session_started",
  568. "chat_round_started",
  569. "chat_agent_request",
  570. "chat_event_detected",
  571. "chat_agent_response",
  572. "event_batch_started",
  573. "event_agent_request",
  574. "event_agent_response",
  575. "event_agent_completed",
  576. "event_batch_results",
  577. "event_policy_decision",
  578. "chat_round_finished",
  579. "chat_round_started",
  580. "chat_agent_request",
  581. "chat_message_stream_started",
  582. "chat_message_stream_finished",
  583. "chat_agent_response",
  584. "chat_round_finished",
  585. "session_finished",
  586. ]
  587. assert "chat_event_detected" in caplog.text
  588. assert "event_agent_completed" in caplog.text
  589. @pytest.mark.asyncio
  590. async def test_runtime_audits_chat_message_stream_boundaries_in_output_order():
  591. request = DebugRunRequest(
  592. user_message="debug this",
  593. system_prompts=[],
  594. pre_messages=[],
  595. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  596. event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
  597. )
  598. runtime = DebugRuntime(RoundStatsChatClient())
  599. outputs = [message async for message in runtime.run(request)]
  600. ordered_labels = [
  601. message["event"] if message["type"] == "audit" else message["type"]
  602. for message in outputs
  603. ]
  604. assert ordered_labels.index("chat_message_stream_started") < ordered_labels.index(
  605. "message_delta"
  606. )
  607. assert ordered_labels.index("message_delta") < ordered_labels.index(
  608. "chat_message_stream_finished"
  609. )
  610. assert ordered_labels.index("chat_message_stream_finished") < ordered_labels.index(
  611. "chat_agent_response"
  612. )
  613. stream_started = next(
  614. message
  615. for message in outputs
  616. if message.get("event") == "chat_message_stream_started"
  617. )
  618. stream_finished = next(
  619. message
  620. for message in outputs
  621. if message.get("event") == "chat_message_stream_finished"
  622. )
  623. assert stream_started["details"]["agent"] == "chat_agent"
  624. assert stream_started["details"]["round_index"] == 1
  625. assert stream_finished["details"]["delta_count"] == 1
  626. assert stream_finished["details"]["content_length"] == len("hello")
  627. @pytest.mark.asyncio
  628. async def test_runtime_audit_events_include_turn_relative_elapsed_time():
  629. request = DebugRunRequest(
  630. user_message="debug this",
  631. system_prompts=[],
  632. pre_messages=[],
  633. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  634. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
  635. )
  636. runtime = DebugRuntime(FakeChatClient(), clock=IncrementingClock())
  637. queues = runtime.start_session(request)
  638. outputs: list[dict[str, Any]] = []
  639. while not any(message["type"] == "turn_completed" for message in outputs):
  640. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  641. await runtime.aclose()
  642. turn_audits = [
  643. message
  644. for message in outputs
  645. if message["type"] == "audit"
  646. ]
  647. elapsed_values = [
  648. message["details"].get("turn_elapsed_ms") for message in turn_audits
  649. ]
  650. assert turn_audits
  651. assert all(isinstance(value, int) for value in elapsed_values)
  652. assert all(value >= 0 for value in elapsed_values)
  653. assert elapsed_values == sorted(elapsed_values)
  654. assert next(
  655. message
  656. for message in turn_audits
  657. if message["event"] == "event_agent_request"
  658. )["details"]["turn_elapsed_ms"] >= 0
  659. @pytest.mark.asyncio
  660. async def test_runtime_audit_includes_model_params_prompts_results_and_usage():
  661. request = DebugRunRequest(
  662. user_message="debug this",
  663. system_prompts=["You are a debugger."],
  664. pre_messages=[],
  665. chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
  666. event_agent=EventAgentParams(
  667. model="event-model",
  668. temperature=0.4,
  669. max_tokens=80,
  670. enabled_tools=["handoff_note"],
  671. max_event_loops=1,
  672. ),
  673. )
  674. runtime = DebugRuntime(EventRoundStatsChatClient())
  675. outputs = [message async for message in runtime.run(request)]
  676. audits = [message for message in outputs if message["type"] == "audit"]
  677. chat_request = next(
  678. message for message in audits if message["event"] == "chat_agent_request"
  679. )
  680. assert chat_request["details"]["agent"] == "chat_agent"
  681. assert chat_request["details"]["params"]["model"] == "chat-model"
  682. assert chat_request["details"]["params"]["temperature"] == 0.1
  683. assert chat_request["details"]["params"]["max_tokens"] == 200
  684. assert chat_request["details"]["tools"] == []
  685. assert chat_request["details"]["messages"][0] == {
  686. "role": "system",
  687. "content": "You are a debugger.",
  688. "name": None,
  689. "tool_call_id": None,
  690. "tool_calls": [],
  691. }
  692. chat_response = next(
  693. message for message in audits if message["event"] == "chat_agent_response"
  694. )
  695. assert chat_response["details"]["event_names"] == ["handoff_note"]
  696. assert chat_response["details"]["raw_chunks"] == [
  697. {
  698. "choices": [
  699. {
  700. "delta": {"content": "<agent_events>handoff_note</agent_events>"},
  701. "finish_reason": None,
  702. }
  703. ]
  704. }
  705. ]
  706. assert chat_response["details"]["usage"] == {
  707. "prompt_tokens": 3,
  708. "completion_tokens": 0,
  709. "total_tokens": 3,
  710. "cached_tokens": 0,
  711. }
  712. event_request = next(
  713. message for message in audits if message["event"] == "event_agent_request"
  714. )
  715. assert event_request["details"]["agent"] == "event_agent"
  716. assert event_request["details"]["params"]["model"] == "event-model"
  717. assert event_request["details"]["params"]["temperature"] == 0.4
  718. assert event_request["details"]["events"][0]["name"] == "handoff_note"
  719. assert event_request["details"]["tools"][0]["function"]["name"] == "handoff_note"
  720. assert "Authorization" not in str(event_request["details"])
  721. event_response = next(
  722. message for message in audits if message["event"] == "event_agent_response"
  723. )
  724. assert event_response["details"]["replies"][0]["role"] == "tool"
  725. assert '"tool": "handoff_note"' in event_response["details"]["replies"][0]["content"]
  726. assert event_response["details"]["raw_model_chunks"][0]["event_name"] == "handoff_note"
  727. assert event_response["details"]["raw_model_chunks"][0]["chunks"] == []
  728. @pytest.mark.asyncio
  729. async def test_runtime_round_started_separates_available_events_from_round_budget():
  730. request = DebugRunRequest(
  731. user_message="debug this",
  732. system_prompts=[],
  733. pre_messages=[],
  734. chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
  735. event_agent=EventAgentParams(
  736. enabled_tools=["handoff_note"],
  737. max_event_loops=1,
  738. ),
  739. )
  740. runtime = DebugRuntime(EventRoundStatsChatClient())
  741. outputs = [message async for message in runtime.run(request)]
  742. round_starts = [
  743. message
  744. for message in outputs
  745. if message.get("event") == "chat_round_started"
  746. ]
  747. assert round_starts[0]["details"]["events_enabled"] == ["handoff_note"]
  748. assert round_starts[0]["details"]["configured_events"] == ["handoff_note"]
  749. assert round_starts[0]["details"]["event_generation_enabled"] is True
  750. assert round_starts[0]["details"]["event_prompt_events"] == ["handoff_note"]
  751. assert round_starts[1]["details"]["events_enabled"] == []
  752. assert round_starts[1]["details"]["configured_events"] == ["handoff_note"]
  753. assert round_starts[1]["details"]["event_generation_enabled"] is False
  754. assert round_starts[1]["details"]["event_prompt_events"] == []
  755. @pytest.mark.asyncio
  756. async def test_runtime_event_agent_history_excludes_chat_agent_system_context():
  757. request = DebugRunRequest(
  758. user_message="debug this",
  759. system_prompts=["ChatAgent root prompt."],
  760. pre_messages=[
  761. ChatMessage(role="system", content="ChatAgent pre system."),
  762. ChatMessage(role="user", content="earlier user"),
  763. ChatMessage(role="assistant", content="earlier assistant"),
  764. ],
  765. chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
  766. event_agent=EventAgentParams(
  767. model="event-model",
  768. enabled_tools=["handoff_note"],
  769. max_event_loops=1,
  770. ),
  771. )
  772. client = ContextBoundaryChatClient()
  773. runtime = DebugRuntime(client)
  774. outputs = [message async for message in runtime.run(request)]
  775. event_request = next(
  776. message for message in outputs if message.get("event") == "event_agent_request"
  777. )
  778. assert [
  779. (message["role"], message["content"])
  780. for message in event_request["details"]["history"]
  781. ] == [
  782. ("user", "earlier user"),
  783. ("assistant", "earlier assistant"),
  784. ("user", "debug this"),
  785. ("assistant", "I will check."),
  786. ]
  787. assert not any(
  788. message["role"] == "system"
  789. for message in event_request["details"]["history"]
  790. )
  791. assert client.event_agent_messages == []
  792. @pytest.mark.asyncio
  793. async def test_runtime_session_event_agent_history_excludes_internal_replies():
  794. request = DebugRunRequest(
  795. user_message="debug this",
  796. system_prompts=["ChatAgent root prompt."],
  797. pre_messages=[ChatMessage(role="assistant", content="prior answer")],
  798. chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
  799. event_agent=EventAgentParams(
  800. model="event-model",
  801. enabled_tools=["handoff_note"],
  802. max_event_loops=1,
  803. ),
  804. )
  805. runtime = DebugRuntime(ContextBoundaryChatClient())
  806. queues = runtime.start_session(request)
  807. outputs: list[dict[str, Any]] = []
  808. while not any(message["type"] == "turn_completed" for message in outputs):
  809. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  810. await runtime.aclose()
  811. event_request = next(
  812. message for message in outputs if message.get("event") == "event_agent_request"
  813. )
  814. assert [
  815. (message["role"], message["content"], message["name"])
  816. for message in event_request["details"]["history"]
  817. ] == [
  818. ("assistant", "prior answer", None),
  819. ("user", "debug this", None),
  820. ("assistant", "I will check.", None),
  821. ]
  822. assert not any(
  823. message["name"] == "event_agent"
  824. for message in event_request["details"]["history"]
  825. )
  826. @pytest.mark.asyncio
  827. async def test_runtime_outputs_event_as_soon_as_chat_stream_detects_it():
  828. RuntimeQueues = _runtime_queues_class()
  829. queues = RuntimeQueues()
  830. client = SlowAfterEventChatClient()
  831. runtime = DebugRuntime(client, queues=queues)
  832. request = DebugRunRequest(
  833. user_message="debug this",
  834. system_prompts=[],
  835. pre_messages=[],
  836. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  837. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
  838. )
  839. runtime.start(request)
  840. try:
  841. assert await _next_non_audit_from_queue(queues) == {"type": "session_started"}
  842. assert await _next_non_audit_from_queue(queues) == {
  843. "type": "message_delta",
  844. "content": "Need event.",
  845. }
  846. await asyncio.wait_for(client.event_seen.wait(), timeout=1)
  847. event_message = await asyncio.wait_for(
  848. _next_non_audit_from_queue(queues),
  849. timeout=0.2,
  850. )
  851. assert event_message["type"] == "event"
  852. assert event_message["event"]["name"] == "handoff_note"
  853. finally:
  854. client.release_stream.set()
  855. await runtime.aclose()
  856. @pytest.mark.asyncio
  857. async def test_runtime_batches_round_events_before_continuing_chat_agent():
  858. def resolve_from_history(
  859. event: ToolCallEvent,
  860. context: ToolExecutionContext,
  861. ) -> dict[str, Any]:
  862. return {"message": context.history[-1].content, "event": event.name}
  863. registry = ToolRegistry(
  864. [
  865. ToolDefinition(
  866. name="handoff_note",
  867. description="Send a handoff note.",
  868. parameters={
  869. "type": "object",
  870. "properties": {"message": {"type": "string"}},
  871. "required": ["message"],
  872. },
  873. handler=lambda event: {
  874. "tool": event.name,
  875. "message": event.arguments["message"],
  876. },
  877. argument_resolver=resolve_from_history,
  878. ),
  879. ToolDefinition(
  880. name="audit_note",
  881. description="Send an audit note.",
  882. parameters={
  883. "type": "object",
  884. "properties": {"message": {"type": "string"}},
  885. "required": ["message"],
  886. },
  887. handler=lambda event: {
  888. "tool": event.name,
  889. "message": event.arguments["message"],
  890. },
  891. argument_resolver=resolve_from_history,
  892. ),
  893. ]
  894. )
  895. request = DebugRunRequest(
  896. user_message="debug this",
  897. system_prompts=[],
  898. pre_messages=[],
  899. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  900. event_agent=EventAgentParams(
  901. enabled_tools=["handoff_note", "audit_note"],
  902. max_event_loops=2,
  903. ),
  904. )
  905. client = MultiEventChatClient()
  906. runtime = DebugRuntime(client, registry=registry)
  907. outputs = [message async for message in runtime.run(request)]
  908. assert _message_types(outputs) == [
  909. "session_started",
  910. "message_delta",
  911. "event",
  912. "event",
  913. "tool_result",
  914. "tool_result",
  915. "round_stats",
  916. "message_delta",
  917. "round_stats",
  918. "done",
  919. ]
  920. assert client.calls == 2
  921. assert [message.role for message in client.second_call_messages] == [
  922. "system",
  923. "user",
  924. "assistant",
  925. "user",
  926. ]
  927. assistant_message = client.second_call_messages[2]
  928. assert assistant_message.content == "Checking events."
  929. assert not any(message.role == "tool" for message in client.second_call_messages)
  930. assert client.second_call_messages[-1].content == (
  931. "EventAgent results:\n"
  932. '{"tool": "handoff_note", "message": "Checking events."}\n'
  933. '{"tool": "audit_note", "message": "Checking events."}'
  934. )
  935. assert client.second_call_messages[-1].name == "event_agent"
  936. @pytest.mark.asyncio
  937. async def test_runtime_start_returns_queues_for_downstream_output_consumer():
  938. request = DebugRunRequest(
  939. user_message="debug this",
  940. system_prompts=[],
  941. pre_messages=[],
  942. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  943. event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
  944. )
  945. runtime = DebugRuntime(RoundStatsChatClient())
  946. queues = runtime.start(request)
  947. outputs: list[dict[str, Any]] = []
  948. while True:
  949. message = await asyncio.wait_for(queues.output.get(), timeout=1)
  950. outputs.append(message)
  951. if message["type"] == "done":
  952. break
  953. assert _message_types(outputs) == [
  954. "session_started",
  955. "message_delta",
  956. "usage",
  957. "round_stats",
  958. "done",
  959. ]
  960. @pytest.mark.asyncio
  961. async def test_runtime_finalizes_chat_after_reaching_event_loop_limit():
  962. request = DebugRunRequest(
  963. user_message="debug this",
  964. system_prompts=[],
  965. pre_messages=[],
  966. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  967. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
  968. )
  969. client = EventLoopLimitChatClient()
  970. runtime = DebugRuntime(client)
  971. outputs = [message async for message in runtime.run(request)]
  972. business_outputs = _without_audit(outputs)
  973. assert client.calls == 2
  974. assert client.tools_by_call[0] == []
  975. assert client.tools_by_call[1] == []
  976. assert "Available events:" in client.messages_by_call[0][0].content
  977. assert not any(
  978. "Available events:" in message.content
  979. for message in client.messages_by_call[1]
  980. if message.role == "system"
  981. )
  982. assert _message_types(outputs) == [
  983. "session_started",
  984. "event",
  985. "tool_result",
  986. "round_stats",
  987. "message_delta",
  988. "round_stats",
  989. "done",
  990. ]
  991. assert business_outputs[4]["content"] == "final after event limit"
  992. @pytest.mark.asyncio
  993. async def test_runtime_buffers_upstream_user_input_until_after_matching_tool_reply():
  994. RuntimeQueues = _runtime_queues_class()
  995. queues = RuntimeQueues()
  996. request = DebugRunRequest(
  997. user_message="debug this",
  998. system_prompts=[],
  999. pre_messages=[],
  1000. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1001. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
  1002. )
  1003. client = StrictHistoryChatClient()
  1004. runtime = DebugRuntime(client, queues=queues)
  1005. stream = runtime.run(request)
  1006. assert await _next_non_audit(stream) == {"type": "session_started"}
  1007. event_message = await _next_non_audit(stream)
  1008. assert event_message["type"] == "event"
  1009. await queues.input.put(ChatMessage(role="user", content="follow-up while tool runs"))
  1010. remaining = [message async for message in stream]
  1011. assert remaining[-1] == {"type": "done"}
  1012. assert [message.role for message in client.second_call_messages] == [
  1013. "system",
  1014. "user",
  1015. "assistant",
  1016. "user",
  1017. "user",
  1018. ]
  1019. assert not any(message.role == "tool" for message in client.second_call_messages)
  1020. assert client.second_call_messages[3].content.startswith("EventAgent results:\n")
  1021. assert client.second_call_messages[3].name == "event_agent"
  1022. assert client.second_call_messages[4].content == "follow-up while tool runs"
  1023. @pytest.mark.asyncio
  1024. async def test_runtime_continues_when_event_agent_tool_handler_raises():
  1025. def fail_tool(event: ToolCallEvent) -> dict[str, Any]:
  1026. raise RuntimeError("boom")
  1027. registry = ToolRegistry(
  1028. [
  1029. ToolDefinition(
  1030. name="handoff_note",
  1031. description="Broken handoff tool.",
  1032. parameters={"type": "object"},
  1033. handler=fail_tool,
  1034. )
  1035. ]
  1036. )
  1037. request = DebugRunRequest(
  1038. user_message="debug this",
  1039. system_prompts=[],
  1040. pre_messages=[],
  1041. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1042. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
  1043. )
  1044. runtime = DebugRuntime(FakeChatClient(), registry=registry)
  1045. outputs = await asyncio.wait_for(
  1046. _collect_outputs(runtime.run(request)),
  1047. timeout=1,
  1048. )
  1049. business_outputs = _without_audit(outputs)
  1050. assert _message_types(outputs) == [
  1051. "session_started",
  1052. "event",
  1053. "tool_result",
  1054. "round_stats",
  1055. "message_delta",
  1056. "round_stats",
  1057. "done",
  1058. ]
  1059. assert json.loads(business_outputs[2]["message"]["content"]) == {
  1060. "tool": "handoff_note",
  1061. "error": "tool handler failed: boom",
  1062. }
  1063. @pytest.mark.asyncio
  1064. async def test_runtime_uses_event_and_input_queues_for_event_agent_handoff():
  1065. RuntimeQueues = _runtime_queues_class()
  1066. queue_log: list[tuple[str, str, str]] = []
  1067. queues = RuntimeQueues(
  1068. input=RecordingQueue("input", queue_log),
  1069. output=RecordingQueue("output", queue_log),
  1070. events=RecordingQueue("events", queue_log),
  1071. )
  1072. request = DebugRunRequest(
  1073. user_message="debug this",
  1074. system_prompts=["You are a debugger."],
  1075. pre_messages=[],
  1076. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1077. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=3),
  1078. )
  1079. runtime = DebugRuntime(FakeChatClient(), queues=queues)
  1080. outputs = [message async for message in runtime.run(request)]
  1081. assert _message_types(outputs) == [
  1082. "session_started",
  1083. "event",
  1084. "tool_result",
  1085. "round_stats",
  1086. "message_delta",
  1087. "round_stats",
  1088. "done",
  1089. ]
  1090. assert queue_log.index(("input", "put", "user")) < queue_log.index(
  1091. ("input", "get", "user")
  1092. )
  1093. assert queue_log.index(
  1094. ("events", "put", "event_request:handoff_note:call_1")
  1095. ) < queue_log.index(("events", "get", "event_request:handoff_note:call_1"))
  1096. assert queue_log.index(
  1097. ("events", "get", "event_request:handoff_note:call_1")
  1098. ) < queue_log.index(("input", "put", "tool:call_1"))
  1099. assert queue_log.index(("input", "put", "tool:call_1")) < queue_log.index(
  1100. ("input", "get", "tool:call_1")
  1101. )
  1102. @pytest.mark.asyncio
  1103. async def test_runtime_run_consumes_output_queue_in_stream_order():
  1104. RuntimeQueues = _runtime_queues_class()
  1105. queue_log: list[tuple[str, str, str]] = []
  1106. queues = RuntimeQueues(
  1107. input=RecordingQueue("input", queue_log),
  1108. output=RecordingQueue("output", queue_log),
  1109. events=RecordingQueue("events", queue_log),
  1110. )
  1111. request = DebugRunRequest(
  1112. user_message="debug this",
  1113. system_prompts=["You are a debugger."],
  1114. pre_messages=[],
  1115. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1116. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=3),
  1117. )
  1118. runtime = DebugRuntime(FakeChatClient(), queues=queues)
  1119. outputs = [message async for message in runtime.run(request)]
  1120. assert _message_types(outputs) == [
  1121. "session_started",
  1122. "event",
  1123. "tool_result",
  1124. "round_stats",
  1125. "message_delta",
  1126. "round_stats",
  1127. "done",
  1128. ]
  1129. output_puts = [
  1130. entry[2] for entry in queue_log if entry[0] == "output" and entry[1] == "put"
  1131. ]
  1132. output_gets = [
  1133. entry[2] for entry in queue_log if entry[0] == "output" and entry[1] == "get"
  1134. ]
  1135. assert [message for message in output_puts if message != "output:audit"] == [
  1136. "output:session_started",
  1137. "output:event",
  1138. "output:tool_result",
  1139. "output:round_stats",
  1140. "output:message_delta",
  1141. "output:round_stats",
  1142. "output:done",
  1143. ]
  1144. assert output_gets == output_puts
  1145. @pytest.mark.asyncio
  1146. async def test_runtime_continues_with_event_summary_without_tool_call_history():
  1147. request = DebugRunRequest(
  1148. user_message="debug this",
  1149. system_prompts=["You are a debugger."],
  1150. pre_messages=[],
  1151. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1152. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
  1153. )
  1154. client = StrictHistoryChatClient()
  1155. runtime = DebugRuntime(client)
  1156. outputs = [message async for message in runtime.run(request)]
  1157. assert client.calls == 2
  1158. assert [message.role for message in client.second_call_messages] == [
  1159. "system",
  1160. "system",
  1161. "user",
  1162. "assistant",
  1163. "user",
  1164. ]
  1165. assistant_message = client.second_call_messages[3]
  1166. assert assistant_message.content == ""
  1167. assert "handoff_note" in client.second_call_messages[1].content
  1168. assert not any(message.role == "tool" for message in client.second_call_messages)
  1169. assert client.second_call_messages[4].content.startswith("EventAgent results:\n")
  1170. assert client.second_call_messages[4].name == "event_agent"
  1171. assert outputs[-1] == {"type": "done"}
  1172. @pytest.mark.asyncio
  1173. async def test_runtime_passes_event_catalog_system_message_without_chat_tools():
  1174. registry = ToolRegistry(
  1175. [
  1176. ToolDefinition(
  1177. name="handoff_note",
  1178. description="Registry-owned handoff tool.",
  1179. parameters={
  1180. "type": "object",
  1181. "properties": {
  1182. "message": {"type": "string"},
  1183. "priority": {"type": "number"},
  1184. },
  1185. "required": ["message"],
  1186. },
  1187. handler=lambda event: {"tool": event.name, "message": "handled"},
  1188. )
  1189. ]
  1190. )
  1191. request = DebugRunRequest(
  1192. user_message="debug this",
  1193. system_prompts=[],
  1194. pre_messages=[],
  1195. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1196. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
  1197. )
  1198. client = ToolCapturingChatClient()
  1199. runtime = DebugRuntime(client, registry=registry)
  1200. outputs = [message async for message in runtime.run(request)]
  1201. assert outputs[-1] == {"type": "done"}
  1202. assert client.tools == []
  1203. assert client.messages[0].role == "system"
  1204. assert "Available events:" in client.messages[0].content
  1205. assert "- handoff_note: Registry-owned handoff tool." in client.messages[0].content
  1206. assert "message" not in client.messages[0].content
  1207. assert "priority" not in client.messages[0].content
  1208. @pytest.mark.asyncio
  1209. async def test_runtime_emits_round_stats_with_clock_and_usage_after_model_turn():
  1210. request = DebugRunRequest(
  1211. user_message="debug this",
  1212. system_prompts=[],
  1213. pre_messages=[],
  1214. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1215. event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
  1216. )
  1217. ticks = iter(
  1218. [
  1219. 0.9,
  1220. 0.91,
  1221. 1.0,
  1222. 1.01,
  1223. 1.02,
  1224. 1.123,
  1225. 1.2,
  1226. 1.25,
  1227. 1.3,
  1228. 1.456,
  1229. 1.7,
  1230. 1.8,
  1231. ]
  1232. )
  1233. runtime = DebugRuntime(RoundStatsChatClient(), clock=lambda: next(ticks))
  1234. outputs = [message async for message in runtime.run(request)]
  1235. business_outputs = _without_audit(outputs)
  1236. assert _message_types(outputs) == [
  1237. "session_started",
  1238. "message_delta",
  1239. "usage",
  1240. "round_stats",
  1241. "done",
  1242. ]
  1243. assert business_outputs[3] == {
  1244. "type": "round_stats",
  1245. "round_index": 1,
  1246. "ttft_ms": 123,
  1247. "elapsed_ms": 456,
  1248. "prompt_tokens": 10,
  1249. "completion_tokens": 20,
  1250. "total_tokens": 30,
  1251. "cached_tokens": 5,
  1252. "had_event": False,
  1253. }
  1254. @pytest.mark.asyncio
  1255. async def test_runtime_emits_round_stats_for_each_chat_call_in_event_handoff():
  1256. request = DebugRunRequest(
  1257. user_message="debug this",
  1258. system_prompts=[],
  1259. pre_messages=[],
  1260. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1261. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
  1262. )
  1263. ticks = iter(
  1264. [
  1265. 1.9,
  1266. 1.91,
  1267. 2.0,
  1268. 2.01,
  1269. 2.02,
  1270. 2.03,
  1271. 2.04,
  1272. 2.05,
  1273. 2.06,
  1274. 2.07,
  1275. 2.08,
  1276. 2.09,
  1277. 2.1,
  1278. 2.25,
  1279. 2.26,
  1280. 3.0,
  1281. 3.01,
  1282. 3.02,
  1283. 3.05,
  1284. 3.1,
  1285. 3.15,
  1286. 3.18,
  1287. 3.2,
  1288. 3.23,
  1289. 3.24,
  1290. 3.25,
  1291. 3.26,
  1292. ]
  1293. )
  1294. client = EventRoundStatsChatClient()
  1295. runtime = DebugRuntime(client, clock=lambda: next(ticks))
  1296. outputs = [message async for message in runtime.run(request)]
  1297. stats = [message for message in outputs if message["type"] == "round_stats"]
  1298. assert client.calls == 2
  1299. assert stats == [
  1300. {
  1301. "type": "round_stats",
  1302. "round_index": 1,
  1303. "ttft_ms": None,
  1304. "elapsed_ms": 250,
  1305. "prompt_tokens": 3,
  1306. "completion_tokens": 0,
  1307. "total_tokens": 3,
  1308. "cached_tokens": 0,
  1309. "had_event": True,
  1310. },
  1311. {
  1312. "type": "round_stats",
  1313. "round_index": 2,
  1314. "ttft_ms": 50,
  1315. "elapsed_ms": 200,
  1316. "prompt_tokens": 4,
  1317. "completion_tokens": 6,
  1318. "total_tokens": 10,
  1319. "cached_tokens": 0,
  1320. "had_event": False,
  1321. },
  1322. ]
  1323. @pytest.mark.asyncio
  1324. async def test_runtime_session_resets_event_budget_for_each_user_turn():
  1325. request = DebugRunRequest(
  1326. user_message="first turn",
  1327. system_prompts=[],
  1328. pre_messages=[],
  1329. chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
  1330. event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
  1331. )
  1332. client = TwoTurnSessionChatClient()
  1333. runtime = DebugRuntime(client)
  1334. queues = runtime.start_session(request)
  1335. outputs: list[dict[str, Any]] = []
  1336. while len([message for message in outputs if message["type"] == "turn_completed"]) < 1:
  1337. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1338. await queues.input.put(ChatMessage(role="user", content="second turn"))
  1339. while len([message for message in outputs if message["type"] == "turn_completed"]) < 2:
  1340. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1341. await runtime.aclose()
  1342. business_types = [message["type"] for message in _without_audit(outputs)]
  1343. assert business_types.count("turn_started") == 2
  1344. assert business_types.count("turn_completed") == 2
  1345. assert client.calls == 4
  1346. assert "Available events:" in client.messages_by_call[0][0].content
  1347. assert "Available events:" in client.messages_by_call[2][0].content
  1348. @pytest.mark.asyncio
  1349. async def test_runtime_persists_session_turn_messages_audit_and_usage(tmp_path):
  1350. store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
  1351. request = DebugRunRequest(
  1352. session_id="session-1",
  1353. user_message="persist this",
  1354. system_prompts=["You are a debugger."],
  1355. pre_messages=[],
  1356. chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
  1357. event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
  1358. )
  1359. runtime = DebugRuntime(RoundStatsChatClient(), session_store=store)
  1360. queues = runtime.start_session(request)
  1361. outputs: list[dict[str, Any]] = []
  1362. while not any(message["type"] == "turn_completed" for message in outputs):
  1363. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1364. await runtime.aclose()
  1365. assert _without_audit(outputs)[0] == {
  1366. "type": "session_started",
  1367. "session_id": "session-1",
  1368. }
  1369. assert store.get_session("session-1")["config"]["chat_agent"]["model"] == "chat-model"
  1370. assert [
  1371. (message["turn_index"], message["role"], message["content"])
  1372. for message in store.list_messages("session-1")
  1373. ] == [
  1374. (1, "user", "persist this"),
  1375. (1, "assistant", "hello"),
  1376. ]
  1377. audit_events = [
  1378. audit["event"]
  1379. for audit in store.list_audit_logs("session-1")
  1380. ]
  1381. assert "session_started" in audit_events
  1382. assert "chat_agent_request" in audit_events
  1383. assert "turn_completed" in audit_events
  1384. usage = store.usage_summary("session-1")
  1385. assert usage["calls"][0]["total_tokens"] == 30
  1386. assert usage["turns"][0]["turn_index"] == 1
  1387. assert usage["session"]["total_tokens"] == 30
  1388. @pytest.mark.asyncio
  1389. async def test_runtime_continues_persisted_turn_indexes_for_existing_session(tmp_path):
  1390. store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
  1391. session_id = store.create_session(title="existing", config={})
  1392. store.start_turn(session_id, turn_index=1, user_message="old")
  1393. request = DebugRunRequest(
  1394. session_id=session_id,
  1395. user_message="new",
  1396. system_prompts=[],
  1397. pre_messages=[],
  1398. chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
  1399. event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
  1400. )
  1401. runtime = DebugRuntime(RoundStatsChatClient(), session_store=store)
  1402. queues = runtime.start_session(request)
  1403. outputs: list[dict[str, Any]] = []
  1404. while not any(message["type"] == "turn_completed" for message in outputs):
  1405. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1406. await runtime.aclose()
  1407. assert [
  1408. (message["turn_index"], message["content"])
  1409. for message in store.list_messages(session_id)
  1410. ] == [
  1411. (2, "new"),
  1412. (2, "hello"),
  1413. ]
  1414. class ScriptedChatClient:
  1415. def __init__(self, rounds: list[list[StreamItem]]) -> None:
  1416. self.rounds = rounds
  1417. self.calls = 0
  1418. self.messages_by_call: list[list[ChatMessage]] = []
  1419. self.tools_by_call: list[list[dict[str, Any]]] = []
  1420. async def stream_chat(
  1421. self,
  1422. messages: list[ChatMessage],
  1423. tools: list[dict],
  1424. params: AgentParams,
  1425. tool_choice: dict[str, Any] | None = None,
  1426. ) -> AsyncIterator[StreamItem]:
  1427. self.messages_by_call.append(list(messages))
  1428. self.tools_by_call.append(list(tools))
  1429. round_items = self.rounds[self.calls]
  1430. self.calls += 1
  1431. for item in round_items:
  1432. yield item
  1433. def _direct_request(
  1434. *,
  1435. enabled_tools: list[str],
  1436. max_event_loops: int = 1,
  1437. session_id: str | None = None,
  1438. ) -> DebugRunRequest:
  1439. return DebugRunRequest(
  1440. session_id=session_id,
  1441. user_message="debug direct tools",
  1442. system_prompts=["You are a debugger."],
  1443. pre_messages=[],
  1444. chat_agent=AgentParams(model="chat-model"),
  1445. event_agent=EventAgentParams(
  1446. enabled_tools=enabled_tools,
  1447. max_event_loops=max_event_loops,
  1448. ),
  1449. tool_invocation_mode="chat_agent_tools",
  1450. )
  1451. @pytest.mark.asyncio
  1452. async def test_direct_mode_uses_enabled_schemas_and_valid_ordered_provider_transcript():
  1453. executed: list[tuple[str, dict[str, Any], str]] = []
  1454. def handle(event: ToolCallEvent) -> dict[str, Any]:
  1455. executed.append((event.name, event.arguments, event.raw_arguments))
  1456. return {"tool": event.name, "value": event.arguments["value"]}
  1457. registry = ToolRegistry(
  1458. [
  1459. ToolDefinition(
  1460. name="first_tool",
  1461. description="First direct tool.",
  1462. parameters={
  1463. "type": "object",
  1464. "properties": {"value": {"type": "string"}},
  1465. "required": ["value"],
  1466. },
  1467. handler=handle,
  1468. ),
  1469. ToolDefinition(
  1470. name="second_tool",
  1471. description="Second direct tool.",
  1472. parameters={
  1473. "type": "object",
  1474. "properties": {"value": {"type": "string"}},
  1475. "required": ["value"],
  1476. },
  1477. handler=handle,
  1478. ),
  1479. ]
  1480. )
  1481. calls = [
  1482. ToolCallEvent(
  1483. id="provider-1",
  1484. name="first_tool",
  1485. arguments={"value": "one"},
  1486. raw_arguments='{"value":"one"}',
  1487. ),
  1488. ToolCallEvent(
  1489. id="provider-2",
  1490. name="second_tool",
  1491. arguments={"value": "two"},
  1492. raw_arguments='{"value":"two"}',
  1493. ),
  1494. ]
  1495. ignored_text_event = ToolCallEvent(
  1496. id="text-ignored",
  1497. name="first_tool",
  1498. arguments={},
  1499. raw_arguments="{}",
  1500. )
  1501. client = ScriptedChatClient(
  1502. [
  1503. [
  1504. StreamItem.message_delta("Visible before tools."),
  1505. StreamItem.provider_tool_call(calls[0]),
  1506. StreamItem.text_event(ignored_text_event),
  1507. StreamItem.provider_tool_call(calls[1]),
  1508. ],
  1509. [StreamItem.message_delta("Final answer.")],
  1510. ]
  1511. )
  1512. outputs = await _collect_outputs(
  1513. DebugRuntime(client, registry=registry).run(
  1514. _direct_request(enabled_tools=["first_tool", "second_tool"])
  1515. )
  1516. )
  1517. assert [tool["function"]["name"] for tool in client.tools_by_call[0]] == [
  1518. "first_tool",
  1519. "second_tool",
  1520. ]
  1521. assert client.tools_by_call[1] == []
  1522. assert not any(
  1523. "Available events:" in message.content
  1524. for message in client.messages_by_call[0]
  1525. if message.role == "system"
  1526. )
  1527. transcript = client.messages_by_call[1]
  1528. assert [message.role for message in transcript[-4:]] == [
  1529. "user",
  1530. "assistant",
  1531. "tool",
  1532. "tool",
  1533. ]
  1534. assert transcript[-3].content == "Visible before tools."
  1535. assert transcript[-3].tool_calls == calls
  1536. assert [message.tool_call_id for message in transcript[-2:]] == [
  1537. "provider-1",
  1538. "provider-2",
  1539. ]
  1540. assert [message.name for message in transcript[-2:]] == [
  1541. "first_tool",
  1542. "second_tool",
  1543. ]
  1544. assert executed == [
  1545. ("first_tool", {"value": "one"}, '{"value":"one"}'),
  1546. ("second_tool", {"value": "two"}, '{"value":"two"}'),
  1547. ]
  1548. emitted_events = [
  1549. message["event"] for message in outputs if message["type"] == "event"
  1550. ]
  1551. assert emitted_events == [call.model_dump() for call in calls]
  1552. assert "event_agent_request" not in [
  1553. message.get("event") for message in outputs if message["type"] == "audit"
  1554. ]
  1555. assert not any(
  1556. message.get("event", {}).get("id") == "text-ignored"
  1557. for message in outputs
  1558. if message["type"] == "event"
  1559. )
  1560. assert all(
  1561. message["details"]["tool_invocation_mode"] == "chat_agent_tools"
  1562. for message in outputs
  1563. if message["type"] == "audit"
  1564. and message["event"] in {"chat_round_started", "chat_agent_request"}
  1565. )
  1566. @pytest.mark.asyncio
  1567. async def test_direct_mode_does_not_execute_unknown_disabled_or_failed_handlers():
  1568. executed: list[str] = []
  1569. def disabled_handler(event: ToolCallEvent) -> dict[str, Any]:
  1570. executed.append(event.name)
  1571. return {"tool": event.name}
  1572. def failing_handler(event: ToolCallEvent) -> dict[str, Any]:
  1573. executed.append(event.name)
  1574. raise RuntimeError("direct boom")
  1575. registry = ToolRegistry(
  1576. [
  1577. ToolDefinition(
  1578. name="enabled_tool",
  1579. description="Enabled.",
  1580. parameters={"type": "object"},
  1581. handler=lambda event: {"tool": event.name},
  1582. ),
  1583. ToolDefinition(
  1584. name="disabled_tool",
  1585. description="Disabled.",
  1586. parameters={"type": "object"},
  1587. handler=disabled_handler,
  1588. ),
  1589. ToolDefinition(
  1590. name="failing_tool",
  1591. description="Fails.",
  1592. parameters={"type": "object"},
  1593. handler=failing_handler,
  1594. ),
  1595. ]
  1596. )
  1597. client = ScriptedChatClient(
  1598. [
  1599. [
  1600. StreamItem.provider_tool_call(
  1601. ToolCallEvent(
  1602. id="unknown-1",
  1603. name="unknown_tool",
  1604. arguments={"kept": True},
  1605. raw_arguments='{"kept":true}',
  1606. )
  1607. ),
  1608. StreamItem.provider_tool_call(
  1609. ToolCallEvent(
  1610. id="disabled-1",
  1611. name="disabled_tool",
  1612. arguments={},
  1613. raw_arguments="{}",
  1614. )
  1615. ),
  1616. StreamItem.provider_tool_call(
  1617. ToolCallEvent(
  1618. id="failed-1",
  1619. name="failing_tool",
  1620. arguments={},
  1621. raw_arguments="{}",
  1622. )
  1623. ),
  1624. ],
  1625. [StreamItem.message_delta("continued")],
  1626. ]
  1627. )
  1628. outputs = await _collect_outputs(
  1629. DebugRuntime(client, registry=registry).run(
  1630. _direct_request(enabled_tools=["enabled_tool", "failing_tool"])
  1631. )
  1632. )
  1633. assert [tool["function"]["name"] for tool in client.tools_by_call[0]] == [
  1634. "enabled_tool",
  1635. "failing_tool",
  1636. ]
  1637. assert executed == ["failing_tool"]
  1638. results = [
  1639. json.loads(message["message"]["content"])
  1640. for message in outputs
  1641. if message["type"] == "tool_result"
  1642. ]
  1643. assert results == [
  1644. {"tool": "unknown_tool", "error": "unknown tool"},
  1645. {"tool": "disabled_tool", "error": "tool disabled"},
  1646. {"tool": "failing_tool", "error": "tool handler failed: direct boom"},
  1647. ]
  1648. @pytest.mark.asyncio
  1649. async def test_direct_mode_ignores_provider_calls_after_event_budget_exhaustion():
  1650. executed: list[str] = []
  1651. registry = ToolRegistry(
  1652. [
  1653. ToolDefinition(
  1654. name="once_tool",
  1655. description="Run once.",
  1656. parameters={"type": "object"},
  1657. handler=lambda event: executed.append(event.id) or {"tool": event.name},
  1658. )
  1659. ]
  1660. )
  1661. client = ScriptedChatClient(
  1662. [
  1663. [
  1664. StreamItem.provider_tool_call(
  1665. ToolCallEvent(
  1666. id="accepted",
  1667. name="once_tool",
  1668. arguments={},
  1669. raw_arguments="{}",
  1670. )
  1671. )
  1672. ],
  1673. [
  1674. StreamItem.provider_tool_call(
  1675. ToolCallEvent(
  1676. id="ignored",
  1677. name="once_tool",
  1678. arguments={},
  1679. raw_arguments="{}",
  1680. )
  1681. ),
  1682. StreamItem.message_delta("budget exhausted"),
  1683. ],
  1684. ]
  1685. )
  1686. outputs = await _collect_outputs(
  1687. DebugRuntime(client, registry=registry).run(
  1688. _direct_request(enabled_tools=["once_tool"], max_event_loops=1)
  1689. )
  1690. )
  1691. assert client.tools_by_call == [
  1692. [registry.tool_schema("once_tool")],
  1693. [],
  1694. ]
  1695. assert executed == ["accepted"]
  1696. assert [
  1697. message["event"]["id"]
  1698. for message in outputs
  1699. if message["type"] == "event"
  1700. ] == ["accepted"]
  1701. assert outputs[-1] == {"type": "done"}
  1702. @pytest.mark.asyncio
  1703. async def test_direct_mode_session_turns_reset_budget_and_snapshot_mode(tmp_path):
  1704. store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
  1705. registry = ToolRegistry(
  1706. [
  1707. ToolDefinition(
  1708. name="turn_tool",
  1709. description="Per-turn tool.",
  1710. parameters={"type": "object"},
  1711. handler=lambda event: {"tool": event.name, "call_id": event.id},
  1712. )
  1713. ]
  1714. )
  1715. client = ScriptedChatClient(
  1716. [
  1717. [
  1718. StreamItem.provider_tool_call(
  1719. ToolCallEvent(
  1720. id="turn-1-call",
  1721. name="turn_tool",
  1722. arguments={},
  1723. raw_arguments="{}",
  1724. )
  1725. )
  1726. ],
  1727. [StreamItem.message_delta("turn one done")],
  1728. [
  1729. StreamItem.provider_tool_call(
  1730. ToolCallEvent(
  1731. id="turn-2-call",
  1732. name="turn_tool",
  1733. arguments={},
  1734. raw_arguments="{}",
  1735. )
  1736. )
  1737. ],
  1738. [StreamItem.message_delta("turn two done")],
  1739. ]
  1740. )
  1741. runtime = DebugRuntime(client, registry=registry, session_store=store)
  1742. request = _direct_request(
  1743. enabled_tools=["turn_tool"],
  1744. max_event_loops=1,
  1745. session_id="direct-session",
  1746. )
  1747. queues = runtime.start_session(request)
  1748. outputs: list[dict[str, Any]] = []
  1749. while sum(message["type"] == "turn_completed" for message in outputs) < 1:
  1750. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1751. await queues.input.put(ChatMessage(role="user", content="second direct turn"))
  1752. while sum(message["type"] == "turn_completed" for message in outputs) < 2:
  1753. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1754. await runtime.aclose()
  1755. assert [bool(tools) for tools in client.tools_by_call] == [True, False, True, False]
  1756. assert [message.role for message in client.messages_by_call[1][-3:]] == [
  1757. "user",
  1758. "assistant",
  1759. "tool",
  1760. ]
  1761. second_turn_transcript = client.messages_by_call[3]
  1762. assert not any(message.name == "event_agent" for message in second_turn_transcript)
  1763. assert [
  1764. message.tool_call_id
  1765. for message in second_turn_transcript
  1766. if message.role == "tool"
  1767. ] == ["turn-1-call", "turn-2-call"]
  1768. session = store.get_session("direct-session")
  1769. assert session is not None
  1770. assert session["config"]["tool_invocation_mode"] == "chat_agent_tools"
  1771. request_audits = [
  1772. audit
  1773. for audit in store.list_audit_logs("direct-session")
  1774. if audit["event"] == "chat_agent_request"
  1775. ]
  1776. assert request_audits
  1777. assert all(
  1778. audit["details"]["tool_invocation_mode"] == "chat_agent_tools"
  1779. for audit in request_audits
  1780. )
  1781. def _single_tool_registry(
  1782. handler: Any,
  1783. *,
  1784. name: str = "safe_tool",
  1785. ) -> ToolRegistry:
  1786. return ToolRegistry(
  1787. [
  1788. ToolDefinition(
  1789. name=name,
  1790. description="Tool used by execution-safety tests.",
  1791. parameters={"type": "object"},
  1792. handler=handler,
  1793. )
  1794. ]
  1795. )
  1796. def _tool_call(call_id: str, *, name: str = "safe_tool") -> ToolCallEvent:
  1797. return ToolCallEvent(
  1798. id=call_id,
  1799. name=name,
  1800. arguments={},
  1801. raw_arguments="{}",
  1802. )
  1803. @pytest.mark.asyncio
  1804. async def test_direct_mode_rejects_provider_call_id_reused_across_one_shot_rounds():
  1805. executed: list[str] = []
  1806. registry = _single_tool_registry(
  1807. lambda event: executed.append(event.id) or {"tool": event.name}
  1808. )
  1809. client = ScriptedChatClient(
  1810. [
  1811. [StreamItem.provider_tool_call(_tool_call("duplicate-call"))],
  1812. [StreamItem.provider_tool_call(_tool_call("duplicate-call"))],
  1813. ]
  1814. )
  1815. outputs = await _collect_outputs(
  1816. DebugRuntime(client, registry=registry).run(
  1817. _direct_request(enabled_tools=["safe_tool"], max_event_loops=2)
  1818. )
  1819. )
  1820. assert executed == ["duplicate-call"]
  1821. assert [
  1822. message["message"]["tool_call_id"]
  1823. for message in outputs
  1824. if message["type"] == "tool_result"
  1825. ] == ["duplicate-call"]
  1826. assert outputs[-1]["type"] == "error"
  1827. assert "duplicate provider tool-call ID: duplicate-call" in outputs[-1]["message"]
  1828. @pytest.mark.asyncio
  1829. async def test_direct_mode_rejects_provider_call_id_reused_in_session_rounds():
  1830. executed: list[str] = []
  1831. registry = _single_tool_registry(
  1832. lambda event: executed.append(event.id) or {"tool": event.name}
  1833. )
  1834. client = ScriptedChatClient(
  1835. [
  1836. [StreamItem.provider_tool_call(_tool_call("session-duplicate"))],
  1837. [StreamItem.provider_tool_call(_tool_call("session-duplicate"))],
  1838. ]
  1839. )
  1840. runtime = DebugRuntime(client, registry=registry)
  1841. queues = runtime.start_session(
  1842. _direct_request(enabled_tools=["safe_tool"], max_event_loops=2)
  1843. )
  1844. outputs: list[dict[str, Any]] = []
  1845. while not any(message["type"] == "error" for message in outputs):
  1846. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  1847. await runtime.aclose()
  1848. assert executed == ["session-duplicate"]
  1849. assert [
  1850. message["message"]["tool_call_id"]
  1851. for message in outputs
  1852. if message["type"] == "tool_result"
  1853. ] == ["session-duplicate"]
  1854. assert "duplicate provider tool-call ID: session-duplicate" in outputs[-1][
  1855. "message"
  1856. ]
  1857. @pytest.mark.asyncio
  1858. async def test_direct_mode_rejects_duplicate_provider_call_ids_within_batch():
  1859. executed: list[str] = []
  1860. registry = _single_tool_registry(
  1861. lambda event: executed.append(event.id) or {"tool": event.name}
  1862. )
  1863. client = ScriptedChatClient(
  1864. [
  1865. [
  1866. StreamItem.provider_tool_call(_tool_call("same-batch")),
  1867. StreamItem.provider_tool_call(_tool_call("same-batch")),
  1868. ]
  1869. ]
  1870. )
  1871. outputs = await _collect_outputs(
  1872. DebugRuntime(client, registry=registry).run(
  1873. _direct_request(enabled_tools=["safe_tool"])
  1874. )
  1875. )
  1876. assert executed == []
  1877. assert outputs[-1]["type"] == "error"
  1878. assert "duplicate provider tool-call ID: same-batch" in outputs[-1]["message"]
  1879. @pytest.mark.asyncio
  1880. async def test_direct_mode_executes_batch_concurrently_but_replies_in_provider_order():
  1881. both_started = asyncio.Event()
  1882. first_release = asyncio.Event()
  1883. second_release = asyncio.Event()
  1884. second_finished = asyncio.Event()
  1885. started: list[str] = []
  1886. async def handler(event: ToolCallEvent) -> dict[str, Any]:
  1887. started.append(event.id)
  1888. if len(started) == 2:
  1889. both_started.set()
  1890. if event.id == "first-call":
  1891. await first_release.wait()
  1892. else:
  1893. await second_release.wait()
  1894. second_finished.set()
  1895. return {"tool": event.name, "call_id": event.id}
  1896. registry = _single_tool_registry(handler)
  1897. client = ScriptedChatClient(
  1898. [
  1899. [
  1900. StreamItem.provider_tool_call(
  1901. _tool_call("first-call").model_copy(
  1902. update={"arguments": {"order": 1}, "raw_arguments": '{"order":1}'}
  1903. )
  1904. ),
  1905. StreamItem.provider_tool_call(
  1906. _tool_call("second-call").model_copy(
  1907. update={"arguments": {"order": 2}, "raw_arguments": '{"order":2}'}
  1908. )
  1909. ),
  1910. ],
  1911. [StreamItem.message_delta("finished")],
  1912. ]
  1913. )
  1914. runtime = DebugRuntime(client, registry=registry)
  1915. run_task = asyncio.create_task(
  1916. _collect_outputs(
  1917. runtime.run(_direct_request(enabled_tools=["safe_tool"]))
  1918. )
  1919. )
  1920. overlapped = False
  1921. try:
  1922. try:
  1923. await asyncio.wait_for(both_started.wait(), timeout=1)
  1924. except TimeoutError:
  1925. pass
  1926. else:
  1927. overlapped = True
  1928. second_release.set()
  1929. await asyncio.wait_for(second_finished.wait(), timeout=1)
  1930. assert not run_task.done()
  1931. finally:
  1932. first_release.set()
  1933. second_release.set()
  1934. outputs = await asyncio.wait_for(run_task, timeout=1)
  1935. assert overlapped is True
  1936. assert started == ["first-call", "second-call"]
  1937. assert [
  1938. message["message"]["tool_call_id"]
  1939. for message in outputs
  1940. if message["type"] == "tool_result"
  1941. ] == ["first-call", "second-call"]
  1942. def _dual_budget_request() -> DebugRunRequest:
  1943. return DebugRunRequest(
  1944. user_message="dual budget",
  1945. system_prompts=[],
  1946. pre_messages=[],
  1947. chat_agent=AgentParams(model="chat-model"),
  1948. event_agent=EventAgentParams(
  1949. enabled_tools=["safe_tool"],
  1950. max_event_loops=1,
  1951. ),
  1952. )
  1953. @pytest.mark.asyncio
  1954. async def test_dual_mode_ignores_text_event_after_one_shot_budget_exhaustion():
  1955. executed: list[str] = []
  1956. registry = _single_tool_registry(
  1957. lambda event: executed.append(event.id) or {"tool": event.name}
  1958. )
  1959. client = ScriptedChatClient(
  1960. [
  1961. [StreamItem.text_event(_tool_call("accepted-text-event"))],
  1962. [
  1963. StreamItem.text_event(_tool_call("ignored-text-event")),
  1964. StreamItem.message_delta("final dual answer"),
  1965. ],
  1966. ]
  1967. )
  1968. outputs = await _collect_outputs(
  1969. DebugRuntime(client, registry=registry).run(_dual_budget_request())
  1970. )
  1971. assert executed == ["accepted-text-event"]
  1972. assert [
  1973. message["event"]["id"]
  1974. for message in outputs
  1975. if message["type"] == "event"
  1976. ] == ["accepted-text-event"]
  1977. assert outputs[-1] == {"type": "done"}
  1978. @pytest.mark.asyncio
  1979. async def test_dual_mode_session_main_path_ignores_text_event_after_budget_exhaustion():
  1980. executed: list[str] = []
  1981. registry = _single_tool_registry(
  1982. lambda event: executed.append(event.id) or {"tool": event.name}
  1983. )
  1984. client = ScriptedChatClient(
  1985. [
  1986. [StreamItem.text_event(_tool_call("session-accepted"))],
  1987. [
  1988. StreamItem.text_event(_tool_call("session-ignored")),
  1989. StreamItem.message_delta("session final"),
  1990. ],
  1991. ]
  1992. )
  1993. runtime = DebugRuntime(client, registry=registry)
  1994. queues = runtime.start_session(_dual_budget_request())
  1995. outputs: list[dict[str, Any]] = []
  1996. while not any(
  1997. message["type"] in {"turn_completed", "error"} for message in outputs
  1998. ):
  1999. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  2000. await runtime.aclose()
  2001. assert executed == ["session-accepted"]
  2002. assert [
  2003. message["event"]["id"]
  2004. for message in outputs
  2005. if message["type"] == "event"
  2006. ] == ["session-accepted"]
  2007. assert outputs[-1]["type"] == "turn_completed"
  2008. def _policy_request(
  2009. *,
  2010. enabled_tools: list[str],
  2011. mode: str = "dual_agent",
  2012. session_id: str | None = None,
  2013. ) -> DebugRunRequest:
  2014. return DebugRunRequest(
  2015. session_id=session_id,
  2016. user_message="exercise event result policy",
  2017. system_prompts=[],
  2018. pre_messages=[],
  2019. chat_agent=AgentParams(model="chat-model"),
  2020. event_agent=EventAgentParams(
  2021. enabled_tools=enabled_tools,
  2022. max_event_loops=1,
  2023. ),
  2024. tool_invocation_mode=mode,
  2025. )
  2026. @pytest.mark.asyncio
  2027. async def test_dual_silent_success_does_not_enqueue_legacy_summary_or_call_chat_again():
  2028. registry = ToolRegistry(
  2029. [
  2030. ToolDefinition(
  2031. name="example.silent",
  2032. description="Complete silently.",
  2033. parameters={"type": "object"},
  2034. handler=lambda event: {"tool": event.name, "status": "applied"},
  2035. result_policy=ResultPolicy.SILENT_SUCCESS,
  2036. )
  2037. ]
  2038. )
  2039. client = ScriptedChatClient(
  2040. [
  2041. [
  2042. StreamItem.message_delta("Applying it now."),
  2043. StreamItem.text_event(_tool_call("silent-1", name="example.silent")),
  2044. ]
  2045. ]
  2046. )
  2047. outputs = await _collect_outputs(
  2048. DebugRuntime(client, registry=registry).run(
  2049. _policy_request(enabled_tools=["example.silent"])
  2050. )
  2051. )
  2052. assert client.calls == 1
  2053. assert outputs[-1] == {"type": "done"}
  2054. assert not any(
  2055. message.get("details", {}).get("result_summary")
  2056. for message in outputs
  2057. if message.get("event") == "event_agent_completed"
  2058. )
  2059. @pytest.mark.asyncio
  2060. async def test_direct_template_success_emits_plugin_confirmation_without_model_call():
  2061. schedule = ToolCallEvent(
  2062. id="schedule-1",
  2063. name="calendar.schedule.create",
  2064. arguments={
  2065. "title": "Design review",
  2066. "start_at": "2026-07-14T09:30:00+08:00",
  2067. "timezone": "Asia/Shanghai",
  2068. },
  2069. raw_arguments=(
  2070. '{"title":"Design review","start_at":"2026-07-14T09:30:00+08:00",'
  2071. '"timezone":"Asia/Shanghai"}'
  2072. ),
  2073. )
  2074. client = ScriptedChatClient(
  2075. [
  2076. [
  2077. StreamItem.message_delta("I will add it."),
  2078. StreamItem.provider_tool_call(schedule),
  2079. ]
  2080. ]
  2081. )
  2082. outputs = await _collect_outputs(
  2083. DebugRuntime(client).run(
  2084. _policy_request(
  2085. enabled_tools=["calendar.schedule.create"],
  2086. mode="chat_agent_tools",
  2087. )
  2088. )
  2089. )
  2090. assert client.calls == 1
  2091. assert [
  2092. message["content"]
  2093. for message in outputs
  2094. if message["type"] == "message_delta"
  2095. ] == [
  2096. "I will add it.",
  2097. "Scheduled Design review for 2026-07-14T09:30:00+08:00 (Asia/Shanghai).",
  2098. ]
  2099. assert outputs[-1] == {"type": "done"}
  2100. @pytest.mark.asyncio
  2101. async def test_direct_silent_failure_requests_exactly_one_corrective_follow_up():
  2102. def fail(event: ToolCallEvent) -> dict[str, Any]:
  2103. raise RuntimeError("device offline")
  2104. registry = ToolRegistry(
  2105. [
  2106. ToolDefinition(
  2107. name="example.silent",
  2108. description="Complete silently on success.",
  2109. parameters={"type": "object"},
  2110. handler=fail,
  2111. result_policy=ResultPolicy.SILENT_SUCCESS,
  2112. )
  2113. ]
  2114. )
  2115. client = ScriptedChatClient(
  2116. [
  2117. [
  2118. StreamItem.message_delta("Trying it."),
  2119. StreamItem.provider_tool_call(
  2120. _tool_call("silent-failure", name="example.silent")
  2121. ),
  2122. ],
  2123. [StreamItem.message_delta("I could not apply that change.")],
  2124. ]
  2125. )
  2126. outputs = await _collect_outputs(
  2127. DebugRuntime(client, registry=registry).run(
  2128. _policy_request(
  2129. enabled_tools=["example.silent"],
  2130. mode="chat_agent_tools",
  2131. )
  2132. )
  2133. )
  2134. assert client.calls == 2
  2135. assert [
  2136. message["content"]
  2137. for message in outputs
  2138. if message["type"] == "message_delta"
  2139. ] == ["Trying it.", "I could not apply that change."]
  2140. @pytest.mark.asyncio
  2141. async def test_mixed_builtin_batch_emits_template_then_terminates_without_llm():
  2142. calls = [
  2143. ToolCallEvent(
  2144. id="terminate-1",
  2145. name="session.terminate",
  2146. arguments={},
  2147. raw_arguments="{}",
  2148. ),
  2149. ToolCallEvent(
  2150. id="schedule-1",
  2151. name="calendar.schedule.create",
  2152. arguments={
  2153. "title": "Design review",
  2154. "start_at": "2026-07-14T09:30:00+08:00",
  2155. "timezone": "Asia/Shanghai",
  2156. },
  2157. raw_arguments=(
  2158. '{"title":"Design review","start_at":"2026-07-14T09:30:00+08:00",'
  2159. '"timezone":"Asia/Shanghai"}'
  2160. ),
  2161. ),
  2162. ToolCallEvent(
  2163. id="search-1",
  2164. name="knowledge.web.search",
  2165. arguments={"query": "event batch executors"},
  2166. raw_arguments='{"query":"event batch executors"}',
  2167. ),
  2168. ToolCallEvent(
  2169. id="volume-1",
  2170. name="device.volume.adjust",
  2171. arguments={"mode": "absolute", "value": 30},
  2172. raw_arguments='{"mode":"absolute","value":30}',
  2173. ),
  2174. ]
  2175. client = ScriptedChatClient(
  2176. [
  2177. [
  2178. StreamItem.message_delta("Goodbye, I will finish those first."),
  2179. *[StreamItem.provider_tool_call(call) for call in calls],
  2180. ]
  2181. ]
  2182. )
  2183. outputs = await _collect_outputs(
  2184. DebugRuntime(client).run(
  2185. _policy_request(
  2186. enabled_tools=[call.name for call in calls],
  2187. mode="chat_agent_tools",
  2188. )
  2189. )
  2190. )
  2191. assert client.calls == 1
  2192. assert [
  2193. message["message"]["tool_call_id"]
  2194. for message in outputs
  2195. if message["type"] == "tool_result"
  2196. ] == [call.id for call in calls]
  2197. assert [
  2198. message["content"]
  2199. for message in outputs
  2200. if message["type"] == "message_delta"
  2201. ][-1] == (
  2202. "Scheduled Design review for 2026-07-14T09:30:00+08:00 "
  2203. "(Asia/Shanghai)."
  2204. )
  2205. decision = next(
  2206. message
  2207. for message in outputs
  2208. if message.get("event") == "event_policy_decision"
  2209. )
  2210. assert decision["details"]["terminate"] is True
  2211. assert decision["details"]["llm_follow_up"] is False
  2212. assert any(
  2213. message.get("event") == "terminal_completed" for message in outputs
  2214. )
  2215. assert outputs[-1] == {"type": "done"}
  2216. @pytest.mark.asyncio
  2217. async def test_reusable_session_terminate_completes_turn_and_session_task():
  2218. client = ScriptedChatClient(
  2219. [
  2220. [
  2221. StreamItem.message_delta("Goodbye."),
  2222. StreamItem.provider_tool_call(
  2223. ToolCallEvent(
  2224. id="terminate-session",
  2225. name="session.terminate",
  2226. arguments={},
  2227. raw_arguments="{}",
  2228. )
  2229. ),
  2230. ]
  2231. ]
  2232. )
  2233. runtime = DebugRuntime(client)
  2234. queues = runtime.start_session(
  2235. _policy_request(
  2236. enabled_tools=["session.terminate"],
  2237. mode="chat_agent_tools",
  2238. session_id="reusable-session",
  2239. )
  2240. )
  2241. outputs: list[dict[str, Any]] = []
  2242. while not any(message["type"] == "done" for message in outputs):
  2243. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  2244. await asyncio.wait_for(runtime._wait_for_tasks(), timeout=1)
  2245. business_types = [message["type"] for message in _without_audit(outputs)]
  2246. assert business_types[-2:] == ["turn_completed", "done"]
  2247. @pytest.mark.asyncio
  2248. async def test_dual_text_event_scope_uses_round_and_does_not_replay_stale_history():
  2249. handled_messages: list[str] = []
  2250. def resolve_from_history(
  2251. event: ToolCallEvent,
  2252. context: ToolExecutionContext,
  2253. ) -> dict[str, Any]:
  2254. del event
  2255. return {
  2256. "message": next(
  2257. message.content
  2258. for message in reversed(context.history)
  2259. if message.role == "assistant" and message.content
  2260. )
  2261. }
  2262. registry = ToolRegistry(
  2263. [
  2264. ToolDefinition(
  2265. name="example.history",
  2266. description="Capture the current assistant history.",
  2267. parameters={
  2268. "type": "object",
  2269. "properties": {"message": {"type": "string"}},
  2270. "required": ["message"],
  2271. },
  2272. handler=lambda event: handled_messages.append(
  2273. event.arguments["message"]
  2274. )
  2275. or {"message": event.arguments["message"]},
  2276. argument_resolver=resolve_from_history,
  2277. )
  2278. ]
  2279. )
  2280. repeated = _tool_call("event_1", name="example.history")
  2281. client = ScriptedChatClient(
  2282. [
  2283. [StreamItem.message_delta("first history"), StreamItem.text_event(repeated)],
  2284. [StreamItem.message_delta("second history"), StreamItem.text_event(repeated)],
  2285. [StreamItem.message_delta("final answer")],
  2286. ]
  2287. )
  2288. request = DebugRunRequest(
  2289. user_message="capture each round",
  2290. system_prompts=[],
  2291. pre_messages=[],
  2292. chat_agent=AgentParams(model="chat-model"),
  2293. event_agent=EventAgentParams(
  2294. enabled_tools=["example.history"],
  2295. max_event_loops=2,
  2296. ),
  2297. )
  2298. outputs = await _collect_outputs(
  2299. DebugRuntime(client, registry=registry).run(request)
  2300. )
  2301. assert handled_messages == ["first history", "second history"]
  2302. scopes = [
  2303. message["details"]["scope"]
  2304. for message in outputs
  2305. if message.get("event") == "event_batch_started"
  2306. ]
  2307. assert scopes[0].endswith(":round-1")
  2308. assert scopes[1].endswith(":round-2")
  2309. @pytest.mark.asyncio
  2310. @pytest.mark.parametrize("mode", ["dual_agent", "chat_agent_tools"])
  2311. async def test_runtime_releases_batch_scope_after_results(monkeypatch, mode: str):
  2312. released: list[Any] = []
  2313. original_release = EventBatchExecutor.release_scope
  2314. def record_release(self: EventBatchExecutor, scope: Any) -> None:
  2315. released.append(scope)
  2316. original_release(self, scope)
  2317. monkeypatch.setattr(EventBatchExecutor, "release_scope", record_release)
  2318. registry = ToolRegistry(
  2319. [
  2320. ToolDefinition(
  2321. name="example.silent",
  2322. description="Complete without another model call.",
  2323. parameters={"type": "object"},
  2324. handler=lambda event: {"event_id": event.id},
  2325. result_policy=ResultPolicy.SILENT_SUCCESS,
  2326. )
  2327. ]
  2328. )
  2329. event = _tool_call("release-1", name="example.silent")
  2330. item = (
  2331. StreamItem.text_event(event)
  2332. if mode == "dual_agent"
  2333. else StreamItem.provider_tool_call(event)
  2334. )
  2335. client = ScriptedChatClient([[item]])
  2336. await _collect_outputs(
  2337. DebugRuntime(client, registry=registry).run(
  2338. _policy_request(enabled_tools=["example.silent"], mode=mode)
  2339. )
  2340. )
  2341. assert len(released) == 1
  2342. assert str(released[0]).endswith(":round-1")
  2343. @pytest.mark.asyncio
  2344. @pytest.mark.parametrize("mode", ["dual_agent", "chat_agent_tools"])
  2345. async def test_tool_only_terminate_emits_plugin_farewell_once(mode: str):
  2346. event = ToolCallEvent(
  2347. id="terminate-only",
  2348. name="session.terminate",
  2349. arguments={},
  2350. raw_arguments="{}",
  2351. )
  2352. item = (
  2353. StreamItem.text_event(event)
  2354. if mode == "dual_agent"
  2355. else StreamItem.provider_tool_call(event)
  2356. )
  2357. client = ScriptedChatClient([[item]])
  2358. outputs = await _collect_outputs(
  2359. DebugRuntime(client).run(
  2360. _policy_request(enabled_tools=["session.terminate"], mode=mode)
  2361. )
  2362. )
  2363. assert client.calls == 1
  2364. assert [
  2365. message["content"]
  2366. for message in outputs
  2367. if message["type"] == "message_delta"
  2368. ] == ["Goodbye."]
  2369. assert outputs[-1] == {"type": "done"}
  2370. @pytest.mark.asyncio
  2371. async def test_tool_only_terminate_ends_reusable_session_after_plugin_farewell():
  2372. client = ScriptedChatClient(
  2373. [
  2374. [
  2375. StreamItem.provider_tool_call(
  2376. ToolCallEvent(
  2377. id="terminate-session-only",
  2378. name="session.terminate",
  2379. arguments={},
  2380. raw_arguments="{}",
  2381. )
  2382. )
  2383. ]
  2384. ]
  2385. )
  2386. runtime = DebugRuntime(client)
  2387. queues = runtime.start_session(
  2388. _policy_request(
  2389. enabled_tools=["session.terminate"],
  2390. mode="chat_agent_tools",
  2391. session_id="terminal-only-session",
  2392. )
  2393. )
  2394. outputs: list[dict[str, Any]] = []
  2395. while not any(message["type"] == "done" for message in outputs):
  2396. outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
  2397. await asyncio.wait_for(runtime._wait_for_tasks(), timeout=1)
  2398. assert [
  2399. message["content"]
  2400. for message in outputs
  2401. if message["type"] == "message_delta"
  2402. ] == ["Goodbye."]
  2403. assert [message["type"] for message in _without_audit(outputs)][-2:] == [
  2404. "turn_completed",
  2405. "done",
  2406. ]
  2407. @pytest.mark.asyncio
  2408. @pytest.mark.parametrize("mode", ["dual_agent", "chat_agent_tools"])
  2409. async def test_web_search_has_first_answer_tool_result_and_one_second_answer(mode: str):
  2410. search = ToolCallEvent(
  2411. id="search-two-answer",
  2412. name="knowledge.web.search",
  2413. arguments={"query": "event batch safety"},
  2414. raw_arguments='{"query":"event batch safety"}',
  2415. )
  2416. item = (
  2417. StreamItem.text_event(search)
  2418. if mode == "dual_agent"
  2419. else StreamItem.provider_tool_call(search)
  2420. )
  2421. client = ScriptedChatClient(
  2422. [
  2423. [StreamItem.message_delta("I will check."), item],
  2424. [StreamItem.message_delta("Grounded search update.")],
  2425. ]
  2426. )
  2427. outputs = await _collect_outputs(
  2428. DebugRuntime(client).run(
  2429. _policy_request(enabled_tools=["knowledge.web.search"], mode=mode)
  2430. )
  2431. )
  2432. business = _without_audit(outputs)
  2433. visible = [
  2434. (index, message["content"])
  2435. for index, message in enumerate(business)
  2436. if message["type"] == "message_delta"
  2437. ]
  2438. tool_index = next(
  2439. index for index, message in enumerate(business) if message["type"] == "tool_result"
  2440. )
  2441. assert client.calls == 2
  2442. assert [content for _, content in visible] == [
  2443. "I will check.",
  2444. "Grounded search update.",
  2445. ]
  2446. assert visible[0][0] < tool_index < visible[1][0]
  2447. assert "sources" in business[tool_index]["message"]["content"]