Taming agent/main.py

Ticket PLT-4282. One 543-line entrypoint was carrying three call flows, a prewarm closure, a chat session, and every retry of the same logging boilerplate. Three commits later the same behavior reads as three linear flows. Branch asuraphel/plt-4282-agent-main-refactor, based on origin/preview, not pushed.

543 80
lines in _entrypoint
3
commits, each shippable
1415
agent tests passing
1
duplicate function deleted
0
behavior changes intended
commit 1 · 7df03d6

Delete the duplicate, collapse the boilerplate

Pure mechanics, zero behavior change. The headline find: cancel_pending_generated_audio existed twice, byte-for-byte identical: once on PrewarmResult (never called) and once on AudioGenerator (called, but the wrong layer: an audio class reaching into prewarm task internals). One implementation survives, on the type it actually inspects.

BEFORE: the same 9 lines, twicemain.py + audio_generator.py
# main.py — defined here, called by nobody
class PrewarmResult:
    @staticmethod
    def cancel_pending_generated_audio(prewarm_task):
        if prewarm_task is None or not prewarm_task.done():
            return
        ...

# audio_generator.py — identical copy, this one gets called
class AudioGenerator(ABC):
    @staticmethod
    def cancel_pending_generated_audio(prewarm_task):
        if prewarm_task is None or not prewarm_task.done():
            return
        ...  # needs TYPE_CHECKING import of agent.main (!)

# and at every call site:
AudioGenerator.cancel_pending_generated_audio(prewarm_task)
AFTER: one home, right layermain.py
# main.py — single implementation, on the type it reads
class PrewarmResult:
    """Result of the prewarm task ..."""

    @staticmethod
    def cancel_pending_generated_audio(prewarm_task):
        if prewarm_task is None or not prewarm_task.done():
            return
        ...

# audio_generator.py no longer imports agent.main at all;
# the backwards TYPE_CHECKING dependency is gone with it.

# call sites:
PrewarmResult.cancel_pending_generated_audio(prewarm_task)

The second offender: every startup step was logged three different ways: a [STEP] pair, a [TIMING] line, and a Prometheus observation. Ten repetitions of that pattern, ~80 lines. Now it is one context manager. Metric labels are unchanged, so Grafana dashboards keep working.

BEFORE: 5 lines per step, ×10main.py _entrypoint
logger.info("[AGENT][STEP] pool_open_await_start")
await pool_task
logger.info("[AGENT][STEP] pool_open_await_done")
logger.info("[AGENT][TIMING] pool_opened")
agent_startup_duration.labels(phase="pool_opened").observe(perf_counter() - t0)

logger.info("[AGENT][STEP] settings_await_start")
global_settings = await settings_task
logger.info("[AGENT][STEP] settings_await_done")
logger.info("[AGENT][TIMING] settings_loaded")
agent_startup_duration.labels(phase="settings_loaded").observe(perf_counter() - t0)

# ... same shape for room_connected, settings_resolved,
# telemetry_setup, services_initialized, participant_joined ...
AFTER: the pattern has a namemain.py
@contextlib.contextmanager
def _startup_phase(phase: str, t0: float, **log_ctx) -> Generator[None]:
    logger.info("[AGENT][STEP] {phase}_start", phase=phase, **log_ctx)
    # An exception inside the block skips the done-log and metric
    # on purpose: a failed phase must not record a startup duration.
    yield
    logger.info("[AGENT][STEP] {phase}_done", phase=phase, **log_ctx)
    _mark_startup_phase(phase, t0, **log_ctx)

# usage — control flow is visible again:
with _startup_phase("pool_opened", t0):
    await pool_task
with _startup_phase("settings_loaded", t0):
    global_settings = await settings_task
BEFORE: best-effort room delete, ×3main.py, three sites
try:
    await ctx.api.room.delete_room(api.DeleteRoomRequest(room=ctx.room.name))
    logger.info("[AGENT][WAIT][TIMEOUT] room deleted", ...)
except Exception as e:  # noqa: BLE001 — lifecycle resilience ...
    logger.warning("[AGENT][WAIT][TIMEOUT] failed to delete room", ...)

# ... and again in the dial-failed path ...
# ... and again for unrecognized inbound participants ...
AFTER: one helper, one line per sitemain.py
async def _delete_room_best_effort(ctx: JobContext, **log_ctx) -> None:
    try:
        await ctx.api.room.delete_room(api.DeleteRoomRequest(room=ctx.room.name))
    except Exception as e:  # noqa: BLE001 — room cleanup failure must not abort the call
        logger.warning("[AGENT] failed to delete room", room_name=ctx.room.name, error=str(e), **log_ctx)

# each site:
await _delete_room_best_effort(ctx, call_id=call_id, is_chat_mode=is_chat_mode)
commit 2 · 54be755

Give the chat flow its own module, unhide the prewarm builder

The chat branch was ~140 lines living inside an if inside a with inside _entrypoint, with its own dispatcher build, its own end-of-call handling, its own keep-alive loop. It moved wholesale to agent/lk/chat_entrypoint.py. The prewarm closure got the same treatment: a real function with real parameters instead of five prewarm_* aliases frozen into a closure.

BEFORE: closure with an aliasing dancemain.py, nested in _entrypoint
if not is_chat_mode:
    prewarm_call_settings = call_settings
    prewarm_call_data = call_data
    prewarm_is_telephony = is_telephony
    prewarm_app_settings = app_settings
    prewarm_call_config = call_config

    async def _prewarm_agent() -> PrewarmResult:
        # 75 lines that can see every local of _entrypoint,
        # but must only touch the frozen prewarm_* aliases
        injected = await deps.call_setup_manager.inject_global_variables(
            prewarm_call_settings, prewarm_call_data.custom_scenario
        )
        ...

    prewarm_task = asyncio.create_task(_prewarm_agent())
AFTER: a function with a signaturemain.py, module level
async def _prewarm_agent(
    *,
    ctx: JobContext,
    deps: CallDependencies,
    call_settings: CallSettings,
    call_data: CallData,
    call_config: GlobalSettings,
    app_settings: AppSettings,
    prompt_config: GlobalPromptConfig,
    is_telephony: bool,
) -> PrewarmResult:
    """Build the session and agent while the callee's phone is still ringing."""
    # same 75 lines — but now the inputs are explicit,
    # and nothing else is reachable by accident

The poll loop stayed kept on purpose

The chat disconnect loop (while True: ... await asyncio.sleep(0.5)) looks like a candidate for an asyncio.Event. It isn't. It is level-triggered on purpose: the participant_disconnected handler misses a disconnect that lands before registration, and the loop is also what keeps the job alive (if _entrypoint returns, the agent leaves the room). An Event version would add code and a new hang mode. What did die is the if participant is None re-acquisition block inside it: provably dead (participant is non-None on entry, never reassigned), pyright already flagged it.

BEFORE: dead branch in the loopmain.py chat branch
while True:
    if participant is None:  # pyright: ignore[reportUnnecessaryComparison]
        remote_participants = [
            p for p in ctx.room.remote_participants.values()
            if not p.identity.startswith("agent-")
        ]
        if remote_participants:
            participant = remote_participants[0]

    if participant and participant.disconnect_reason is not None:
        await handle_chat_end(reason=str(participant.disconnect_reason))
        break
    await asyncio.sleep(0.5)
AFTER: loop keeps its two real jobschat_entrypoint.py
# Level-triggered on purpose: the event handler above misses a
# disconnect that lands before registration, and this loop is
# also what keeps the job alive.
while True:
    if participant.disconnect_reason is not None:
        await handle_chat_end(reason=str(participant.disconnect_reason))
        break
    await asyncio.sleep(0.5)
commit 3 · d85b852

Three flows, three functions: the Optional-state weave is gone

The old entrypoint threaded five mutable Optionals (call_settings, call_data, is_chat_mode, is_telephony, prewarm_task) through 300 lines, re-checking combinations of them at every step, which is why it needed a mid-function RuntimeError("call data was never loaded") for a state that "cannot happen". Now the room metadata is classified once into a frozen plan, and each plan variant runs a linear flow. Illegal states are unrepresentable; the RuntimeError is gone because the state it guarded no longer exists.

BEFORE: one function, three interleaved flowsmain.py _entrypoint, 543 lines
call_settings: CallSettings | None = None
call_data: CallData | None = None
is_chat_mode = False
is_telephony = False
prewarm_task: asyncio.Task[PrewarmResult] | None = None

try:
    parsed = CallCreationSettings(**json.loads(ctx.room.metadata))
    ...
except (ValidationError, json.JSONDecodeError, TypeError):
    ...  # fall through with everything still None

if call_settings:                      # flow check #1
    ...
    if not is_chat_mode:               # flow check #2
        ...launch prewarm...

with span_cm:
    if call_settings and call_settings.callee and not is_chat_mode:  # #3
        ...dial...
    participant = await wait_for_participant_with_timeout(...)
    if not call_settings:              # flow check #4
        ...inbound resolution...
    if call_data is None:               # "cannot happen"
        raise RuntimeError("call data was never loaded ...")
    if not is_chat_mode:               # flow check #5
        ...JOIN_STATE log...
    if is_chat_mode:                   # flow check #6
        ...140 lines of chat...
    else:
        ...prewarm await + handoff...
AFTER: classify once, dispatch oncemain.py _entrypoint, ~80 lines
@dataclass(frozen=True)
class _ChatCall:
    call_settings: CallSettings      # never Optional again
    call_data: CallData
    call_config: GlobalSettings
    app_settings: AppSettings
    is_telephony: bool

@dataclass(frozen=True)
class _DispatchedVoiceCall:
    ...  # same fields, plus:
    prewarm_task: asyncio.Task[PrewarmResult]   # always launched

@dataclass(frozen=True)
class _InboundSipCall:
    """No usable room metadata: everything resolves after join."""

_CallPlan = _ChatCall | _DispatchedVoiceCall | _InboundSipCall

# _entrypoint, after shared setup:
plan = await _make_call_plan(ctx=ctx, deps=deps, t0=t0, llm=llm,
                            prompt_config=prompt_config,
                            background_tasks=background_tasks)

with _call_span(ctx, span_settings):
    match plan:
        case _ChatCall():
            await _run_chat_call(...)
        case _DispatchedVoiceCall():
            await _run_dispatched_voice_call(...)
        case _InboundSipCall():
            await _run_inbound_sip_call(...)

Each runner now reads top-to-bottom as the thing it is. The dial + claim-race logic, the prewarm-during-ring orchestration, and every incident-context comment moved verbatim; the mess was the packaging, not the logic.

step_run_chat_call_run_dispatched_voice_call_run_inbound_sip_call
dial + claim race·if callee is set·
prewarm during ring·always (cancelled on dead call)·
wait for participantis_chat_mode=Trueis_chat_mode=Falseis_chat_mode=False
resolve settings post-join··from participant metadata
hand off to sessionrun_chat(...)_handoff_to_session(prewarm)_handoff_to_session(None)

And the repeated cancel pair finally became a function

The cancel_pending_generated_audio + task.cancel() pair appeared three times in the dial/wait failure paths. Now it is _cancel_prewarm(task), with a docstring that says why: stop the prebuild for a call that will never connect.

verification

How it was checked

ruff + pyright clean

Zero errors on every touched file after each commit, including the strict pyright workspace settings. Two new reasoned suppressions only (RUF075 for the intentional skip-on-failure context manager, PLW0717 carried over).

1415 agent tests green

The full agent suite (minus the LTO tests that hit a real LLM) ran after the final commit: 1415 passed, 8 skipped (environment-conditional), 2 xfailed (pre-existing, documented Hamsa STT/TTS bugs unrelated to this change). The targeted main/chat/prewarm test files were also run green after each individual commit.

Behavior invariants audited

Dial gating, wait-timeout failure marking, sentry tag ordering, JOIN_STATE gating, span construction, inbound defaults, and all prewarm cancellation points were compared path-by-path against the old code. Metric phase labels and log queries used by dashboards are unchanged.

Observability preserved

Prometheus agent_startup_duration phase labels identical; a few [STEP] log names were normalized to match their metric name (e.g. pool_open_awaitpool_opened).