fix(orchestrator): keep RunLoop generator alive for multi-turn sessions

[Leoyzen]

Summary

Fix 4 interdependent bugs causing message loss in multi-turn OpenCode server scenarios (issue #229, sub-issue of #221).

E1: _consume_run() kills generator after turn 1 (root cause)

_consume_run() broke on StreamCompleteEvent | RunErrorEvent and called gen.aclose() in the finally block, killing the RunHandle.start() generator after turn 1. Turn 2's followup() messages were lost because ProtocolChannel.close() cleared the feedback queue.

Fix: Replaced break+aclose with a drain-only loop (async for _event in gen: pass). The generator stays alive for the entire run lifetime, closed via RunHandle.close() during session cleanup.

E2: wait_for_completion() waits on wrong event (cascade from E1)

wait_for_completion() waited on complete_event (set when generator exits), not _turn_complete_event (per-turn completion). With E1 fix, the generator stays alive between turns, so complete_event never fires within the timeout — wait_for_completion() blocks until timeout (300s), then cancel_run_for_session() kills the run.

Fix: Use asyncio.wait() with both _turn_complete_event and complete_event as FIRST_COMPLETED. Do NOT clear _turn_complete_event — rely on _execute_turn()'s existing clear at turn start to avoid racing with turn completion.

E3: cancel_run_for_session() kills idle RunHandle (cascade from E1)

cancel_run_for_session() called run_handle.cancel() unconditionally. With E1 fix, RunHandle is IDLE between turns — cancel() force-cancels the _consume_run task, killing the generator. Subsequent messages need a new RunHandle, but if current_run_id still points to the dead run, routing breaks.

Fix: Guard with if run_handle._run_state != RunState.RUNNING: return False. Log at WARNING level. Return type changed from None to bool.

D1: OpenCode event bridge reuses turn 1's assistant message ID (independent)

_before_consumer_loop() initialized EventProcessorContext (with assistant_msg_id, assistant_msg) and set _message_registered[session_id] = False only ONCE. On turn 2, _message_registered stayed True → skipped append_message_to_session → turn 2 reused turn 1's assistant_msg_id → frontend merges/overwrites.

Fix: On RunStartedEvent when _message_registered is True (subsequent turn), reset all per-turn state: pop _pending_message_ids for new message ID, create fresh ctx.assistant_msg, pop _pending_message_metadata for model info, reset EventProcessorContext mutable fields, set _message_registered = False.

Changes

File	Changes
src/agentpool/orchestrator/session_controller_runs.py	E1: drain-only loop; E2: asyncio.wait dual events; E3: RunState guard + bool return
src/agentpool_server/opencode_server/opencode_event_bridge.py	D1: per-turn state reset on RunStartedEvent
tests/orchestrator/test_cancellation.py	Removed xfail from E1/E3 tests, rewrote E1 test for fixed behavior
tests/e2e/test_opencode_multiturn_redflag.py	E1 test: check parts instead of text; E2 test: xfail (TestModel text part issue)

Test Results

Suite	Result
test_cancellation.py	10 passed, 7 skipped (pre-existing), 0 failures
Orchestrator + OpenCode + Sessions	275 passed, 0 failures
E2E red-flag tests	5 passed, 1 xfailed (E2 TestModel), 1 failed (C1 pre-existing)
Ruff	All checks passed
Mypy	Success: no issues found

Closes #229

[gemini-code-assist]

Code Review

This pull request updates the orchestrator to support multi-turn runs by keeping the run generator alive across turns instead of closing it prematurely on stream completion. It also ensures that session cancellation is skipped when the run is idle, and updates wait_for_completion to wait for either per-turn or full run completion. Additionally, the event bridge now resets per-turn state on subsequent turns. Feedback is provided to wrap the asyncio.wait block in a try...finally block to prevent task leakage on cancellation, and to extract duplicated assistant message creation logic into a helper method.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If wait_for_completion is cancelled by the caller (e.g., due to a timeout or session shutdown), the await asyncio.wait call will raise asyncio.CancelledError. Since there is no try...finally block wrapping the wait, the pending tasks turn_done and run_done will be leaked as active background tasks. Wrapping the wait in a try...finally block ensures that both tasks are properly cancelled upon any exit path.

Suggested change

	turn_done = asyncio.ensure_future(run_handle._turn_complete_event.wait())
	run_done = asyncio.ensure_future(run_handle.complete_event.wait())
	done, pending = await asyncio.wait(
	{turn_done, run_done},
	timeout=timeout,
	return_when=asyncio.FIRST_COMPLETED,
	)
	for p in pending:
	p.cancel()
	if not done:
	raise TimeoutError
	turn_done = asyncio.create_task(run_handle._turn_complete_event.wait())
	run_done = asyncio.create_task(run_handle.complete_event.wait())
	try:
	done, pending = await asyncio.wait(
	{turn_done, run_done},
	timeout=timeout,
	return_when=asyncio.FIRST_COMPLETED,
	)
	if not done:
	raise TimeoutError
	finally:
	turn_done.cancel()
	run_done.cancel()

[gemini-code-assist]

The logic to construct a new assistant message (retrieving/generating the message ID, resolving default model info, fetching session state, resolving pending metadata, and calling MessageWithParts.assistant) is duplicated verbatim between _before_consumer_loop (lines 181-214) and _handle_event (lines 437-469). This duplication reduces maintainability and increases the risk of inconsistent updates. Consider extracting this logic into a helper method.

Suggested change

	assistant_msg_id = self._pending_message_ids.pop(session_id, None)
	if assistant_msg_id is None:
	assistant_msg_id = identifier.ascending("message")
	agent_name = "agentpool"
	model_id, provider_id = self.server_state.resolve_default_model_info()
	session_state = self.session_pool.sessions.get_session(session_id)
	if session_state is not None:
	agent_name = session_state.agent_name
	pending_meta = self._pending_message_metadata.pop(session_id, None)
	if pending_meta is not None:
	pending_model_id = pending_meta.get("model_id")
	if pending_model_id is not None:
	model_id = pending_model_id
	pending_provider_id = pending_meta.get("provider_id")
	if pending_provider_id is not None:
	provider_id = pending_provider_id
	ctx.assistant_msg_id = assistant_msg_id
	ctx.assistant_msg = MessageWithParts.assistant(
	message_id=assistant_msg_id,
	session_id=session_id,
	time=MessageTime(created=now_ms()),
	agent_name=agent_name,
	model_id=model_id,
	parent_id=session_id,
	provider_id=provider_id,
	path=MessagePath(
	cwd=self.server_state.working_dir,
	root=self.server_state.working_dir,
	),
	mode=agent_name,
	)
	assistant_msg_id, assistant_msg = self._create_assistant_message(session_id)
	ctx.assistant_msg_id = assistant_msg_id
	ctx.assistant_msg = assistant_msg