Skip to content
31 changes: 31 additions & 0 deletions src/google/adk/agents/base_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,23 @@ def clone(
cloned_agent.parent_agent = None
return cloned_agent

AGENT_LIFECYCLE_KEY: ClassVar[str] = 'agent_lifecycle'
"""Key under ``Event.custom_metadata`` marking an agent lifecycle event as
``'start'`` or ``'finish'`` (emitted only when
``RunConfig.emit_agent_lifecycle_events`` is enabled).
See https://github.com/google/adk-python/issues/6267."""

def _create_agent_lifecycle_event(
self, ctx: InvocationContext, phase: str
) -> Event:
"""Build a lightweight lifecycle marker event authored by this agent."""
return Event(
author=self.name,
invocation_id=ctx.invocation_id,
branch=ctx.branch,
custom_metadata={self.AGENT_LIFECYCLE_KEY: phase},
)

async def run_async(
self,
parent_context: InvocationContext,
Expand All @@ -295,22 +312,36 @@ async def run_async(
"""

ctx = self._create_invocation_context(parent_context)
emit_lifecycle = bool(
ctx.run_config and ctx.run_config.emit_agent_lifecycle_events
)
async with _instrumentation.record_agent_invocation(ctx, self):
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return

# Emit the lifecycle "start" only once the agent is actually going to run
# (i.e. not short-circuited by before_agent_callback), and pair it with a
# "finish" on every exit path below.
if emit_lifecycle:
yield self._create_agent_lifecycle_event(ctx, 'start')

async with Aclosing(self._run_async_impl(ctx)) as agen:
async for event in agen:
yield event

if ctx.end_invocation:
if emit_lifecycle:
yield self._create_agent_lifecycle_event(ctx, 'finish')
return

if event := await self._handle_after_agent_callback(ctx):
yield event

if emit_lifecycle:
yield self._create_agent_lifecycle_event(ctx, 'finish')

@override
async def _run_impl(
self,
Expand Down
14 changes: 14 additions & 0 deletions src/google/adk/agents/run_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,20 @@ class RunConfig(BaseModel):
),
)

emit_agent_lifecycle_events: bool = False
"""Whether each agent invocation emits lightweight lifecycle marker events.

When True, every agent that runs yields a ``start`` event before its logic and
a ``finish`` event after, authored by the agent and carrying
``Event.custom_metadata['agent_lifecycle'] == 'start' | 'finish'`` (plus the
agent's ``branch``). This lets consumers bracket agent/node execution exactly —
including per ``LoopAgent`` iteration and parallel-branch boundaries — without
inferring boundaries from ``author``/``branch`` transitions.

Defaults to False; existing event streams are unchanged. See
https://github.com/google/adk-python/issues/6267.
"""

max_llm_calls: int = 500
"""
A limit on the total number of llm calls for a given run.
Expand Down
42 changes: 42 additions & 0 deletions tests/unittests/agents/test_base_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from google.adk.agents.base_agent import BaseAgentState
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.run_config import RunConfig
from google.adk.apps.app import ResumabilityConfig
from google.adk.events.event import Event
from google.adk.plugins.base_plugin import BasePlugin
Expand Down Expand Up @@ -197,6 +198,47 @@ async def test_run_async(request: pytest.FixtureRequest):
assert events[0].content.parts[0].text == 'Hello, world!'


@pytest.mark.asyncio
async def test_no_agent_lifecycle_events_by_default(
request: pytest.FixtureRequest,
):
# Backward-compat: without opting in, the stream is unchanged.
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = await _create_parent_invocation_context(
request.function.__name__, agent
)

events = [e async for e in agent.run_async(parent_ctx)]

assert len(events) == 1
assert all(
(e.custom_metadata or {}).get(BaseAgent.AGENT_LIFECYCLE_KEY) is None
for e in events
)


@pytest.mark.asyncio
async def test_run_async_emits_agent_lifecycle_events(
request: pytest.FixtureRequest,
):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = await _create_parent_invocation_context(
request.function.__name__, agent
)
parent_ctx.run_config = RunConfig(emit_agent_lifecycle_events=True)

events = [e async for e in agent.run_async(parent_ctx)]

# start -> agent's own event -> finish, all authored by the agent.
phases = [
(e.custom_metadata or {}).get(BaseAgent.AGENT_LIFECYCLE_KEY)
for e in events
]
assert phases == ['start', None, 'finish']
assert all(e.author == agent.name for e in events)
assert events[1].content.parts[0].text == 'Hello, world!'


@pytest.mark.asyncio
async def test_run_async_with_branch(request: pytest.FixtureRequest):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
Expand Down
Loading