11from __future__ import annotations
22
3+ import asyncio
34from collections .abc import AsyncIterator
45from contextlib import asynccontextmanager
56from typing import TYPE_CHECKING , Any
@@ -26,6 +27,20 @@ def __init__(self, runtime: AgentRuntime, *, auto_step: bool = True) -> None:
2627 self .runtime = runtime
2728 self ._step_open = False
2829 self ._auto_step = bool (auto_step )
30+ self ._auto_opened_step = False
31+ self ._auto_opened_step_id : str | None = None
32+
33+ def _schedule_close_auto_step (self ) -> None :
34+ """
35+ Best-effort: close an auto-opened step without forcing callers to await.
36+ """
37+ if not (self ._step_open and self ._auto_opened_step ):
38+ return
39+ try :
40+ loop = asyncio .get_running_loop ()
41+ except RuntimeError :
42+ return
43+ loop .create_task (self .end_step ())
2944
3045 @classmethod
3146 def attach (
@@ -46,15 +61,25 @@ def attach(
4661 return cls (runtime = runtime )
4762
4863 def begin_step (self , goal : str , step_index : int | None = None ) -> str :
64+ # If we previously auto-opened a verification step, close it before starting a real step.
65+ if self ._step_open and self ._auto_opened_step :
66+ self ._schedule_close_auto_step ()
67+ self ._auto_opened_step = False
68+ self ._auto_opened_step_id = None
4969 self ._step_open = True
5070 return self .runtime .begin_step (goal , step_index = step_index )
5171
5272 async def end_step (self , ** kwargs : Any ) -> dict [str , Any ]:
5373 self ._step_open = False
74+ self ._auto_opened_step = False
75+ self ._auto_opened_step_id = None
5476 return await self .runtime .emit_step_end (** kwargs )
5577
5678 @asynccontextmanager
5779 async def step (self , goal : str , step_index : int | None = None ) -> AsyncIterator [None ]:
80+ # Async form can safely close any auto-opened step before starting.
81+ if self ._step_open and self ._auto_opened_step :
82+ await self .end_step ()
5883 self .begin_step (goal , step_index = step_index )
5984 try :
6085 yield
@@ -64,11 +89,61 @@ async def step(self, goal: str, step_index: int | None = None) -> AsyncIterator[
6489 async def snapshot (self , ** kwargs : Any ):
6590 return await self .runtime .snapshot (** kwargs )
6691
92+ async def record_action (self , action : str , * , url : str | None = None ) -> None :
93+ """
94+ Sidecar helper: let the host framework report the action it performed.
95+
96+ This improves trace readability and (when artifacts are enabled) enriches the action timeline.
97+ """
98+ await self .runtime .record_action (action , url = url )
99+
67100 def check (self , predicate , label : str , required : bool = False ):
68101 if not self ._step_open :
69102 if not self ._auto_step :
70103 raise RuntimeError (
71104 f"No active step. Call dbg.begin_step(...) or use 'async with dbg.step(...)' before check(label={ label !r} )."
72105 )
73106 self .begin_step (f"verify:{ label } " )
74- return self .runtime .check (predicate , label , required = required )
107+ self ._auto_opened_step = True
108+ self ._auto_opened_step_id = getattr (self .runtime , "step_id" , None )
109+
110+ base = self .runtime .check (predicate , label , required = required )
111+
112+ # Auto-close auto-opened verification steps after the check completes.
113+ if not self ._auto_opened_step :
114+ return base
115+
116+ dbg = self
117+ opened_step_id = self ._auto_opened_step_id
118+
119+ class _AutoCloseAssertionHandle :
120+ def __init__ (self , inner ):
121+ self ._inner = inner
122+
123+ def once (self ) -> bool :
124+ ok = self ._inner .once ()
125+ if (
126+ dbg ._step_open
127+ and dbg ._auto_opened_step
128+ and (
129+ opened_step_id is None
130+ or getattr (dbg .runtime , "step_id" , None ) == opened_step_id
131+ )
132+ ):
133+ dbg ._schedule_close_auto_step ()
134+ return ok
135+
136+ async def eventually (self , ** kwargs : Any ) -> bool :
137+ ok = await self ._inner .eventually (** kwargs )
138+ if (
139+ dbg ._step_open
140+ and dbg ._auto_opened_step
141+ and (
142+ opened_step_id is None
143+ or getattr (dbg .runtime , "step_id" , None ) == opened_step_id
144+ )
145+ ):
146+ await dbg .end_step ()
147+ return ok
148+
149+ return _AutoCloseAssertionHandle (base )
0 commit comments