import asyncio import json import logging import time from collections.abc import AsyncIterator, Callable from typing import Any, Protocol from agent_lab.application.contracts import AgentParams, DebugRunRequest from agent_lab.application.event_agent import EventAgent, EventAgentRequest from agent_lab.application.queues import RuntimeQueues from agent_lab.application.session_store import SessionStore from agent_lab.application.tools import ToolRegistry, build_default_tool_registry from agent_lab.domain.events import EVENT_BLOCK_START, ToolCallEvent from agent_lab.domain.messages import ChatMessage, StreamItem, TokenUsage logger = logging.getLogger(__name__) class ChatClient(Protocol): async def stream_chat( self, messages: list[ChatMessage], tools: list[dict[str, Any]], params: AgentParams, tool_choice: dict[str, Any] | None = None, ) -> AsyncIterator[StreamItem]: ... class DebugRuntime: def __init__( self, chat_client: ChatClient, queues: RuntimeQueues | None = None, registry: ToolRegistry | None = None, session_store: SessionStore | None = None, clock: Callable[[], float] = time.perf_counter, ) -> None: self.chat_client = chat_client self.queues = queues self.registry = registry or build_default_tool_registry() self.session_store = session_store self.clock = clock self._tasks: list[asyncio.Task[None]] = [] def start(self, request: DebugRunRequest) -> RuntimeQueues: queues = self.queues or RuntimeQueues() task = asyncio.create_task(self._produce(request, queues)) self._tasks.append(task) return queues def start_session(self, request: DebugRunRequest) -> RuntimeQueues: queues = self.queues or RuntimeQueues() task = asyncio.create_task(self._produce_session(request, queues)) self._tasks.append(task) return queues async def run(self, request: DebugRunRequest) -> AsyncIterator[dict[str, Any]]: queues = self.start(request) async for message in self.output_messages(queues): yield message await self._wait_for_tasks() async def output_messages( self, queues: RuntimeQueues, ) -> AsyncIterator[dict[str, Any]]: while True: message = await queues.output.get() yield message if message.get("type") in {"done", "error"}: break async def aclose(self) -> None: for task in self._tasks: if not task.done(): task.cancel() await self._wait_for_tasks() async def _wait_for_tasks(self) -> None: if not self._tasks: return await asyncio.gather(*self._tasks, return_exceptions=True) self._tasks = [task for task in self._tasks if not task.done()] async def _produce(self, request: DebugRunRequest, queues: RuntimeQueues) -> None: event_worker = self._start_event_worker(request, queues) try: await self._run_chat_agent(request, queues) except Exception as exc: logger.exception("runtime session failed") await queues.output.put({"type": "error", "message": str(exc)}) finally: if event_worker is not None: await queues.events.put(None) await asyncio.gather(event_worker, return_exceptions=True) async def _produce_session( self, request: DebugRunRequest, queues: RuntimeQueues, ) -> None: event_worker = self._start_event_worker(request, queues) try: await self._run_chat_session(request, queues) except asyncio.CancelledError: raise except Exception as exc: logger.exception("runtime session failed") await queues.output.put({"type": "error", "message": str(exc)}) finally: if event_worker is not None: await queues.events.put(None) await asyncio.gather(event_worker, return_exceptions=True) def _start_event_worker( self, request: DebugRunRequest, queues: RuntimeQueues, ) -> asyncio.Task[None] | None: if request.tool_invocation_mode != "dual_agent": return None event_agent = EventAgent( request.event_agent.enabled_tools, registry=self.registry, chat_client=self.chat_client, params=request.event_agent, ) return asyncio.create_task(self._consume_events(queues, event_agent)) async def _run_chat_agent( self, request: DebugRunRequest, queues: RuntimeQueues, ) -> None: messages = self._build_initial_messages(request) turn_started_at = self.clock() await queues.input.put(ChatMessage(role="user", content=request.user_message)) await queues.output.put({"type": "session_started"}) await self._audit( queues, "session_started", turn_started_at=turn_started_at, tool_invocation_mode=request.tool_invocation_mode, ) event_loops = 0 round_index = 0 while True: await self._drain_input(queues, messages) round_index += 1 started_at = self.clock() ttft_ms: int | None = None prompt_tokens = 0 completion_tokens = 0 total_tokens = 0 cached_tokens = 0 saw_event = False assistant_content: list[str] = [] events: list[ToolCallEvent] = [] tool_replies: list[ChatMessage] = [] message_stream_started = False message_delta_count = 0 configured_events = request.event_agent.enabled_tools event_prompt_events = ( request.event_agent.enabled_tools if event_loops < request.event_agent.max_event_loops else [] ) chat_tools = self._chat_tools_for_round(request, event_prompt_events) raw_chunks: list[dict[str, Any]] = [] chat_messages = self._chat_messages_for_mode( request, messages, event_prompt_events, ) await self._audit( queues, "chat_round_started", turn_started_at=turn_started_at, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, events_enabled=event_prompt_events, configured_events=configured_events, event_generation_enabled=bool(event_prompt_events), event_prompt_events=event_prompt_events, ) await self._audit( queues, "chat_agent_request", turn_started_at=turn_started_at, agent="chat_agent", round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, params=self._params_snapshot(request.chat_agent), messages=self._message_snapshots(chat_messages), tools=chat_tools, ) async for item in self.chat_client.stream_chat( messages=chat_messages, tools=chat_tools, params=request.chat_agent, ): if item.kind == "raw_chunk" and item.raw_chunk is not None: raw_chunks.append(item.raw_chunk) continue if item.kind == "message_delta": if ttft_ms is None: ttft_ms = self._elapsed_ms(started_at) if not message_stream_started: message_stream_started = True await self._audit( queues, "chat_message_stream_started", turn_started_at=turn_started_at, agent="chat_agent", round_index=round_index, ttft_ms=ttft_ms, ) message_delta_count += 1 assistant_content.append(item.content or "") await queues.output.put( {"type": "message_delta", "content": item.content} ) continue if item.kind == "usage" and item.usage is not None: prompt_tokens = item.usage.prompt_tokens completion_tokens = item.usage.completion_tokens total_tokens = item.usage.total_tokens cached_tokens = item.usage.cached_tokens await queues.output.put( {"type": "usage", "usage": item.usage.model_dump()} ) continue event = self._accepted_chat_event( request, item, budget_available=event_loops < request.event_agent.max_event_loops, ) if event is not None: saw_event = True events.append(event) await queues.output.put( {"type": "event", "event": event.model_dump()} ) await self._audit( queues, "chat_event_detected", turn_started_at=turn_started_at, round_index=round_index, event_id=event.id, event_name=event.name, event_source=( "provider_resolved" if request.tool_invocation_mode == "chat_agent_tools" else "text_event" ), tool_invocation_mode=request.tool_invocation_mode, ) if message_stream_started: await self._audit( queues, "chat_message_stream_finished", turn_started_at=turn_started_at, agent="chat_agent", round_index=round_index, delta_count=message_delta_count, content_length=len("".join(assistant_content)), ) await self._audit( queues, "chat_agent_response", turn_started_at=turn_started_at, agent="chat_agent", round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, content="".join(assistant_content), event_names=[event.name for event in events], events=self._event_snapshots(events), usage={ "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens, "cached_tokens": cached_tokens, }, raw_chunks=raw_chunks, ) if events and request.tool_invocation_mode == "chat_agent_tools": self._validate_provider_tool_call_ids(messages, events) if assistant_content or events: assistant_message = ChatMessage( role="assistant", content="".join(assistant_content), tool_calls=( events if request.tool_invocation_mode == "chat_agent_tools" else [] ), ) messages.append(assistant_message) if events: if request.tool_invocation_mode == "chat_agent_tools": tool_replies = await self._execute_provider_tools( events, enabled_names=request.event_agent.enabled_tools, ) messages.extend(tool_replies) else: await queues.events.put( EventAgentRequest( events=events, history=self._event_agent_history(messages), system_prompt=request.event_agent.system_prompt, extra_body=request.event_agent.extra_body, turn_started_at=turn_started_at, ) ) tool_replies = await self._wait_for_tool_replies( queues, [event.id for event in events], ) for reply in tool_replies: await queues.output.put( {"type": "tool_result", "message": reply.model_dump()} ) if request.tool_invocation_mode == "chat_agent_tools": await self._audit( queues, "provider_tools_completed", turn_started_at=turn_started_at, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, event_source="provider_resolved", events=self._event_snapshots(events), result_count=len(tool_replies), ) else: await self._audit( queues, "event_agent_completed", turn_started_at=turn_started_at, round_index=round_index, event_names=[event.name for event in events], result_count=len(tool_replies), result_summary="\n".join( reply.content for reply in tool_replies ), ) elapsed_ms = self._elapsed_ms(started_at) await queues.output.put( { "type": "round_stats", "round_index": round_index, "ttft_ms": ttft_ms, "elapsed_ms": elapsed_ms, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens, "cached_tokens": cached_tokens, "had_event": saw_event, }, ) await self._audit( queues, "chat_round_finished", turn_started_at=turn_started_at, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, had_event=saw_event, elapsed_ms=elapsed_ms, ) if not saw_event: break event_loops += 1 await self._audit( queues, "session_finished", turn_started_at=turn_started_at, round_count=round_index, ) await queues.output.put({"type": "done"}) async def _run_chat_session( self, request: DebugRunRequest, queues: RuntimeQueues, ) -> None: session_id = self._ensure_persisted_session(request) messages = self._build_initial_messages(request) initial_turn_started_at = self.clock() session_message = {"type": "session_started"} if session_id is not None: session_message["session_id"] = session_id await queues.output.put(session_message) await self._audit( queues, "session_started", turn_started_at=initial_turn_started_at, session_id=session_id, tool_invocation_mode=request.tool_invocation_mode, ) await queues.input.put(ChatMessage(role="user", content=request.user_message)) turn_index = self._initial_turn_index(session_id) - 1 next_turn_started_at: float | None = initial_turn_started_at while True: user_message = await queues.input.get() if user_message.role != "user" or user_message.name == "event_agent": messages.append(user_message) continue turn_index += 1 turn_started_at = next_turn_started_at or self.clock() next_turn_started_at = None await self._run_chat_turn( request, queues, messages, user_message, turn_index, session_id, turn_started_at, ) async def _run_chat_turn( self, request: DebugRunRequest, queues: RuntimeQueues, messages: list[ChatMessage], user_message: ChatMessage, turn_index: int, session_id: str | None = None, turn_started_at: float | None = None, ) -> None: turn_started_at = turn_started_at or self.clock() messages.append(user_message) self._start_persisted_turn(session_id, turn_index, user_message) await queues.output.put({"type": "turn_started", "turn_index": turn_index}) await self._audit( queues, "turn_started", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, ) event_loops = 0 round_index = 0 deferred_user_messages: list[ChatMessage] = [] while True: if round_index: await self._drain_input( queues, messages, deferred_user_messages=deferred_user_messages, ) round_index += 1 started_at = self.clock() ttft_ms: int | None = None prompt_tokens = 0 completion_tokens = 0 total_tokens = 0 cached_tokens = 0 saw_event = False assistant_content: list[str] = [] events: list[ToolCallEvent] = [] message_stream_started = False message_delta_count = 0 configured_events = request.event_agent.enabled_tools event_prompt_events = ( request.event_agent.enabled_tools if event_loops < request.event_agent.max_event_loops else [] ) chat_tools = self._chat_tools_for_round(request, event_prompt_events) raw_chunks: list[dict[str, Any]] = [] chat_messages = self._chat_messages_for_mode( request, messages, event_prompt_events, ) await self._audit( queues, "chat_round_started", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, events_enabled=event_prompt_events, configured_events=configured_events, event_generation_enabled=bool(event_prompt_events), event_prompt_events=event_prompt_events, ) await self._audit( queues, "chat_agent_request", turn_started_at=turn_started_at, session_id=session_id, agent="chat_agent", turn_index=turn_index, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, params=self._params_snapshot(request.chat_agent), messages=self._message_snapshots(chat_messages), tools=chat_tools, ) async for item in self.chat_client.stream_chat( messages=chat_messages, tools=chat_tools, params=request.chat_agent, ): if item.kind == "raw_chunk" and item.raw_chunk is not None: raw_chunks.append(item.raw_chunk) continue if item.kind == "message_delta": if ttft_ms is None: ttft_ms = self._elapsed_ms(started_at) if not message_stream_started: message_stream_started = True await self._audit( queues, "chat_message_stream_started", turn_started_at=turn_started_at, session_id=session_id, agent="chat_agent", turn_index=turn_index, round_index=round_index, ttft_ms=ttft_ms, ) message_delta_count += 1 assistant_content.append(item.content or "") await queues.output.put( {"type": "message_delta", "content": item.content} ) continue if item.kind == "usage" and item.usage is not None: prompt_tokens = item.usage.prompt_tokens completion_tokens = item.usage.completion_tokens total_tokens = item.usage.total_tokens cached_tokens = item.usage.cached_tokens await queues.output.put( {"type": "usage", "usage": item.usage.model_dump()} ) continue event = self._accepted_chat_event( request, item, budget_available=event_loops < request.event_agent.max_event_loops, ) if event is not None: saw_event = True events.append(event) await queues.output.put( {"type": "event", "event": event.model_dump()} ) await self._audit( queues, "chat_event_detected", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, round_index=round_index, event_id=event.id, event_name=event.name, event_source=( "provider_resolved" if request.tool_invocation_mode == "chat_agent_tools" else "text_event" ), tool_invocation_mode=request.tool_invocation_mode, ) if message_stream_started: await self._audit( queues, "chat_message_stream_finished", turn_started_at=turn_started_at, session_id=session_id, agent="chat_agent", turn_index=turn_index, round_index=round_index, delta_count=message_delta_count, content_length=len("".join(assistant_content)), ) await self._audit( queues, "chat_agent_response", turn_started_at=turn_started_at, session_id=session_id, agent="chat_agent", turn_index=turn_index, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, content="".join(assistant_content), event_names=[event.name for event in events], events=self._event_snapshots(events), usage={ "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens, "cached_tokens": cached_tokens, }, raw_chunks=raw_chunks, ) if events and request.tool_invocation_mode == "chat_agent_tools": self._validate_provider_tool_call_ids(messages, events) if assistant_content or events: assistant_message = ChatMessage( role="assistant", content="".join(assistant_content), tool_calls=( events if request.tool_invocation_mode == "chat_agent_tools" else [] ), ) messages.append(assistant_message) self._append_persisted_message( session_id, turn_index, assistant_message, ) if events: if request.tool_invocation_mode == "chat_agent_tools": tool_replies = await self._execute_provider_tools( events, enabled_names=request.event_agent.enabled_tools, ) messages.extend(tool_replies) else: await queues.events.put( EventAgentRequest( events=events, history=self._event_agent_history(messages), system_prompt=request.event_agent.system_prompt, extra_body=request.event_agent.extra_body, session_id=session_id, turn_index=turn_index, round_index=round_index, turn_started_at=turn_started_at, ) ) tool_replies = await self._wait_for_tool_replies( queues, [event.id for event in events], ) for reply in tool_replies: await queues.output.put( {"type": "tool_result", "message": reply.model_dump()} ) if request.tool_invocation_mode == "chat_agent_tools": await self._audit( queues, "provider_tools_completed", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, event_source="provider_resolved", events=self._event_snapshots(events), result_count=len(tool_replies), ) else: await self._audit( queues, "event_agent_completed", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, round_index=round_index, event_names=[event.name for event in events], result_count=len(tool_replies), result_summary="\n".join( reply.content for reply in tool_replies ), ) elapsed_ms = self._elapsed_ms(started_at) self._append_persisted_usage( session_id, turn_index, round_index, usage=TokenUsage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens, cached_tokens=cached_tokens, ), ttft_ms=ttft_ms, elapsed_ms=elapsed_ms, ) await queues.output.put( { "type": "round_stats", "round_index": round_index, "ttft_ms": ttft_ms, "elapsed_ms": elapsed_ms, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens, "cached_tokens": cached_tokens, "had_event": saw_event, }, ) await self._audit( queues, "chat_round_finished", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, round_index=round_index, tool_invocation_mode=request.tool_invocation_mode, had_event=saw_event, elapsed_ms=elapsed_ms, ) if not saw_event: break event_loops += 1 for deferred_message in deferred_user_messages: await queues.input.put(deferred_message) await self._audit( queues, "turn_completed", turn_started_at=turn_started_at, session_id=session_id, turn_index=turn_index, round_count=round_index, ) self._complete_persisted_turn(session_id, turn_index) await queues.output.put( { "type": "turn_completed", "turn_index": turn_index, "round_count": round_index, } ) async def _consume_events( self, queues: RuntimeQueues, event_agent: EventAgent, ) -> None: while True: request = await queues.events.get() if request is None: return await self._audit( queues, "event_agent_request", turn_started_at=request.turn_started_at, session_id=request.session_id, agent="event_agent", turn_index=request.turn_index, round_index=request.round_index, params=self._params_snapshot(event_agent.params), events=self._event_snapshots(request.events), history=self._message_snapshots(request.history), system_prompt=request.system_prompt, extra_body=request.extra_body or {}, tools=[ tool for tool in ( event_agent.registry.tool_schema(event.name) for event in request.events ) if tool is not None ], ) replies = await event_agent.handle_many( request.events, history=request.history, system_prompt=request.system_prompt, extra_body=request.extra_body, ) await self._audit( queues, "event_agent_response", turn_started_at=request.turn_started_at, session_id=request.session_id, agent="event_agent", turn_index=request.turn_index, round_index=request.round_index, replies=[reply.model_dump() for reply in replies], raw_model_chunks=event_agent.raw_model_chunks(request.events), ) for reply in replies: await queues.input.put(reply) summary = event_agent.summarize_replies(replies) if summary is not None: await queues.input.put(summary) async def _wait_for_tool_replies( self, queues: RuntimeQueues, tool_call_ids: list[str], ) -> list[ChatMessage]: pending = set(tool_call_ids) replies: dict[str, ChatMessage] = {} deferred: list[ChatMessage] = [] while pending: message = await queues.input.get() if message.role == "tool" and message.tool_call_id in pending: replies[message.tool_call_id] = message pending.remove(message.tool_call_id) continue deferred.append(message) for deferred_message in deferred: await queues.input.put(deferred_message) return [replies[tool_call_id] for tool_call_id in tool_call_ids] def _event_name_only(self, event: ToolCallEvent) -> ToolCallEvent: return event.model_copy(update={"arguments": {}, "raw_arguments": "{}"}) def _build_initial_messages(self, request: DebugRunRequest) -> list[ChatMessage]: messages = [ ChatMessage(role="system", content=prompt) for prompt in request.system_prompts ] if request.chat_agent.system_prompt.strip(): messages.append( ChatMessage(role="system", content=request.chat_agent.system_prompt) ) messages.extend(request.pre_messages) return messages def _chat_messages_for_mode( self, request: DebugRunRequest, messages: list[ChatMessage], enabled_events: list[str], ) -> list[ChatMessage]: if request.tool_invocation_mode == "chat_agent_tools": return list(messages) return self._chat_messages_for_round(messages, enabled_events) def _chat_tools_for_round( self, request: DebugRunRequest, enabled_events: list[str], ) -> list[dict[str, Any]]: if request.tool_invocation_mode != "chat_agent_tools": return [] return self.registry.provider_tool_schemas(enabled_events) def _accepted_chat_event( self, request: DebugRunRequest, item: StreamItem, *, budget_available: bool, ) -> ToolCallEvent | None: if not budget_available: return None if request.tool_invocation_mode == "chat_agent_tools": if ( item.kind == "provider_tool_call" and item.event is not None ): return item.event return None if item.kind == "text_event" and item.event is not None: return self._event_name_only(item.event) return None def _validate_provider_tool_call_ids( self, messages: list[ChatMessage], events: list[ToolCallEvent], ) -> None: declared_ids = { tool_call.id for message in messages for tool_call in message.tool_calls } batch_ids: set[str] = set() for event in events: if event.id in declared_ids or event.id in batch_ids: raise ValueError(f"duplicate provider tool-call ID: {event.id}") batch_ids.add(event.id) async def _execute_provider_tools( self, events: list[ToolCallEvent], *, enabled_names: list[str], ) -> list[ChatMessage]: payloads = await asyncio.gather( *[ self.registry.execute_async( event, enabled_names=enabled_names, ) for event in events ] ) return [ ChatMessage( role="tool", content=json.dumps(payload, ensure_ascii=False), name=event.name, tool_call_id=event.id, ) for event, payload in zip(events, payloads, strict=True) ] def _chat_messages_for_round( self, messages: list[ChatMessage], enabled_events: list[str], ) -> list[ChatMessage]: if self._has_event_instructions(messages): return list(messages) event_prompt = self.registry.chat_event_system_message(enabled_events) if not event_prompt: return list(messages) insert_at = 0 while insert_at < len(messages) and messages[insert_at].role == "system": insert_at += 1 return [ *messages[:insert_at], ChatMessage(role="system", content=event_prompt), *messages[insert_at:], ] def _has_event_instructions(self, messages: list[ChatMessage]) -> bool: return any( message.role == "system" and "Available events:" in message.content and EVENT_BLOCK_START in message.content for message in messages ) def _event_agent_history(self, messages: list[ChatMessage]) -> list[ChatMessage]: return [ message for message in messages if message.role in {"user", "assistant"} and message.name != "event_agent" ] async def _drain_input( self, queues: RuntimeQueues, messages: list[ChatMessage], deferred_user_messages: list[ChatMessage] | None = None, ) -> None: while not queues.input.empty(): message = await queues.input.get() if ( deferred_user_messages is not None and message.role == "user" and message.name != "event_agent" ): deferred_user_messages.append(message) continue messages.append(message) def _elapsed_ms(self, started_at: float) -> int: return round((self.clock() - started_at) * 1000) def _ensure_persisted_session(self, request: DebugRunRequest) -> str | None: if self.session_store is None: return None return self.session_store.ensure_session( request.session_id, title=self._session_title(request.user_message), config=self._session_config_snapshot(request), ) def _initial_turn_index(self, session_id: str | None) -> int: if self.session_store is None or session_id is None: return 1 return self.session_store.next_turn_index(session_id) def _start_persisted_turn( self, session_id: str | None, turn_index: int, user_message: ChatMessage, ) -> None: if self.session_store is None or session_id is None: return self.session_store.start_turn( session_id, turn_index=turn_index, user_message=user_message.content, ) self.session_store.append_message( session_id, turn_index=turn_index, message=user_message, ) def _complete_persisted_turn( self, session_id: str | None, turn_index: int, ) -> None: if self.session_store is None or session_id is None: return self.session_store.complete_turn(session_id, turn_index=turn_index) def _append_persisted_message( self, session_id: str | None, turn_index: int, message: ChatMessage, ) -> None: if self.session_store is None or session_id is None: return self.session_store.append_message( session_id, turn_index=turn_index, message=message, ) def _append_persisted_usage( self, session_id: str | None, turn_index: int, round_index: int, *, usage: TokenUsage, ttft_ms: int | None, elapsed_ms: int, ) -> None: if self.session_store is None or session_id is None: return self.session_store.append_usage( session_id, turn_index=turn_index, round_index=round_index, usage=usage, ttft_ms=ttft_ms, elapsed_ms=elapsed_ms, ) def _session_title(self, user_message: str) -> str: title = " ".join(user_message.split()) if not title: return "Debug session" return title[:80] def _session_config_snapshot(self, request: DebugRunRequest) -> dict[str, Any]: return { "tool_invocation_mode": request.tool_invocation_mode, "system_prompts": list(request.system_prompts), "pre_messages": self._message_snapshots(request.pre_messages), "chat_agent": request.chat_agent.model_dump(), "event_agent": request.event_agent.model_dump(), } def _params_snapshot(self, params: AgentParams) -> dict[str, Any]: return params.model_dump() def _message_snapshots(self, messages: list[ChatMessage] | tuple[ChatMessage, ...]) -> list[dict[str, Any]]: return [message.model_dump() for message in messages] def _event_snapshots(self, events: list[ToolCallEvent]) -> list[dict[str, Any]]: return [event.model_dump() for event in events] async def _audit( self, queues: RuntimeQueues, event: str, *, session_id: str | None = None, turn_started_at: float | None = None, **details: Any, ) -> None: if turn_started_at is not None: details["turn_elapsed_ms"] = max(0, self._elapsed_ms(turn_started_at)) logger.info("audit event=%s details=%s", event, details) if self.session_store is not None and session_id is not None: turn_index = details.get("turn_index") round_index = details.get("round_index") self.session_store.append_audit( session_id, event=event, details=details, turn_index=turn_index if isinstance(turn_index, int) else None, round_index=round_index if isinstance(round_index, int) else None, ) await queues.output.put( { "type": "audit", "event": event, "details": details, } )