"""Task attempt orchestration for Beaver Task mode.""" from __future__ import annotations from time import perf_counter from typing import Any, Callable from beaver.engine import AgentRunResult from beaver.engine.context import SkillContext from beaver.prompts.main_agent import normalize_main_agent_prompt_locale from .evidence import EvidenceBuilder, TaskEvidencePacket, render_task_evidence from .models import TaskRecord from .planner import TaskExecutionPlan class TaskAttemptOrchestrator: """Own the execution order inside one Task attempt.""" def __init__( self, *, loaded: Any, create_loop: Callable[[], Any], make_provider_bundle_for_task: Callable[[Any, dict[str, Any]], Any], ) -> None: self.loaded = loaded self.create_loop = create_loop self.make_provider_bundle_for_task = make_provider_bundle_for_task async def run( self, *, message: str, runner: Any, kwargs: dict[str, Any], task: TaskRecord, ) -> AgentRunResult: task_service = self._require_loaded(self.loaded, "task_service") task_execution_planner = self._require_loaded(self.loaded, "task_execution_planner") session_manager = self._require_loaded(self.loaded, "session_manager") base_execution_context = kwargs.get("execution_context") prompt_locale = kwargs.get("prompt_locale") or task.metadata.get("prompt_locale") output_language_instruction = self._output_language_instruction(prompt_locale) provider_bundle = kwargs.get("provider_bundle") or self.make_provider_bundle_for_task(self.loaded, kwargs) kwargs = dict(kwargs) kwargs.pop("team_provider_bundle_factory", None) kwargs["provider_bundle"] = provider_bundle attempt_index = int(task.metadata.get("latest_attempt_index") or 0) + 1 task_service.start_run(task.task_id, user_message=message, attempt_index=attempt_index) pre_skill_context = self._build_skill_selection_context( task=task, user_message=message, attempt_index=attempt_index, ) preselected_skills, pre_skill_latency_ms = await self._assemble_task_attempt_skills( task_description=pre_skill_context, provider_bundle=provider_bundle, thinking_enabled=kwargs.get("thinking_enabled"), include_skill_assembly=bool(kwargs.get("include_skill_assembly", True)), pinned_skill_contexts=kwargs.get("pinned_skill_contexts"), ) if pre_skill_latency_ms: kwargs["pre_run_latency_ms"] = self._merge_latency_ms( kwargs.get("pre_run_latency_ms"), {"pre_skill_assembly_ms": pre_skill_latency_ms}, ) plan = await task_execution_planner.plan( task=task, user_message=message, attempt_index=attempt_index, provider_bundle=provider_bundle, skill_summaries=self._skill_summaries_for_planner(preselected_skills), tool_hints=self._tool_hints_for_skills(preselected_skills), activated_skills=preselected_skills, ) self._append_task_observation( session_manager, task.session_id, event_type="task_execution_planned", payload={ "task_id": task.task_id, "attempt_index": attempt_index, **plan.to_event_payload(), }, ) if plan.is_team: plan = TaskExecutionPlan.single( "legacy_planner_team_ignored", planner_adaptation=plan.planner_adaptation, ) outcome_metadata = { "task_outcome": "single", "incomplete_node_ids": [], "node_statuses": {}, "evidence_gaps": {}, } attempt_kwargs = dict(kwargs) attempt_kwargs.update( { "task_id": task.task_id, "task_mode": True, "attempt_index": attempt_index, "allow_candidate_generation": False, "pinned_skill_contexts": preselected_skills, "include_skill_assembly": False, } ) attempt_kwargs["execution_context"] = self._join_context( base_execution_context, output_language_instruction, ) attempt_kwargs["skill_selection_context"] = self._build_skill_selection_context( task=task, user_message=message, attempt_index=attempt_index, plan=plan, ) result = await runner(message, **attempt_kwargs) self._append_task_observation( session_manager, task.session_id, event_type="task_synthesis_completed", payload={ "task_id": task.task_id, "attempt_index": attempt_index, "main_run_id": result.run_id, "plan_mode": plan.mode, "strategy": plan.graph.strategy if plan.graph else None, **outcome_metadata, }, ) task = task_service.append_run( task.task_id, result.run_id, skill_names=self._skill_names_for_run(result.run_id), ) evidence_packet = self._build_task_evidence_packet( session_manager=session_manager, task=task, attempt_index=attempt_index, result=result, ) evidence_text = render_task_evidence(evidence_packet) evidence_debug = { "evidence_run_ids": [ item.run_id for item in [evidence_packet.main_run, *evidence_packet.team_runs] if item is not None ], "evidence_session_ids": [ item.session_id for item in [evidence_packet.main_run, *evidence_packet.team_runs] if item is not None ], "tool_result_count": sum( len(item.tool_results) for item in [evidence_packet.main_run, *evidence_packet.team_runs] if item is not None ), "evidence_length": len(evidence_text), } session_manager.update_latest_assistant_event_payload( result.session_id, result.run_id, { "task_id": task.task_id, "task_status": task.status, "evidence_status": "recorded", }, ) session_manager.append_message( result.session_id, run_id=result.run_id, role="system", event_type="task_evidence_recorded", event_payload={ "task_id": task.task_id, "attempt_index": attempt_index, "evidence_debug": evidence_debug, }, content=None, context_visible=False, ) result.task_id = task.task_id result.task_status = task.status result.validation_result = None return result async def _assemble_task_attempt_skills( self, *, task_description: str, provider_bundle: Any, thinking_enabled: bool | None, include_skill_assembly: bool, pinned_skill_contexts: Any, ) -> tuple[list[SkillContext], float]: started = perf_counter() selected = self._coerce_skill_contexts(pinned_skill_contexts) if include_skill_assembly: skill_assembler = self._require_loaded(self.loaded, "skill_assembler") runtime = provider_bundle.auxiliary_runtime or provider_bundle.main_runtime assembled = await skill_assembler.assemble( task_description=task_description, provider=provider_bundle.auxiliary_provider or provider_bundle.main_provider, model=getattr(runtime, "model", None), embedding_runtime=getattr(provider_bundle, "embedding_runtime", None), thinking_enabled=thinking_enabled, ) selected = self._merge_skill_contexts( selected, list(getattr(assembled, "activated_skills", []) or []), ) return selected, (perf_counter() - started) * 1000 @staticmethod def _coerce_skill_contexts(value: Any) -> list[SkillContext]: if not isinstance(value, list): return [] return [item for item in value if isinstance(item, SkillContext)] @staticmethod def _merge_skill_contexts(left: list[SkillContext], right: list[SkillContext]) -> list[SkillContext]: merged: list[SkillContext] = [] seen: set[str] = set() for skill in [*left, *right]: if skill.name in seen: continue seen.add(skill.name) merged.append(skill) return merged @staticmethod def _skill_summaries_for_planner(skills: list[SkillContext]) -> list[str]: summaries: list[str] = [] for skill in skills: content = " ".join((skill.content or "").split()) if len(content) > 240: content = content[:237].rstrip() + "..." summaries.append(f"{skill.name}: {content}" if content else skill.name) return summaries @staticmethod def _tool_hints_for_skills(skills: list[SkillContext]) -> list[str]: result: list[str] = [] for skill in skills: for hint in skill.tool_hints: if hint and hint not in result: result.append(hint) return result @staticmethod def _require_loaded(loaded: Any, field_name: str) -> Any: value = getattr(loaded, field_name) if value is None: raise RuntimeError(f"Engine loader did not provide required dependency {field_name!r}") return value @staticmethod def _merge_latency_ms(current: Any, updates: dict[str, float]) -> dict[str, float]: merged: dict[str, float] = {} if isinstance(current, dict): for key, value in current.items(): if isinstance(value, (int, float)): merged[str(key)] = float(value) for key, value in updates.items(): merged[key] = merged.get(key, 0.0) + float(value) return merged @staticmethod def _output_language_instruction(prompt_locale: str | None) -> str: locale = normalize_main_agent_prompt_locale(prompt_locale) if locale == "en": return ( "Output language: English. Use English for user-facing task titles, summaries, plans, " "and final answers unless the user explicitly requests another language." ) if locale == "zh-Hant": return ( "輸出語言：繁體中文。除非使用者明確要求其他語言，所有面向使用者的任務標題、摘要、" "計劃與最終回答都使用繁體中文。" ) return ( "输出语言：简体中文。除非用户明确要求其他语言，所有面向用户的任务标题、摘要、" "计划与最终回答都使用简体中文。" ) def _skill_names_for_run(self, run_id: str) -> list[str]: store = getattr(self.loaded, "run_memory_store", None) if store is None: return [] for record in store.list_runs(): if record.run_id == run_id: return [receipt.skill_name for receipt in record.activated_skills] return [] @staticmethod def _build_skill_selection_context( *, task: TaskRecord, user_message: str, attempt_index: int, plan: TaskExecutionPlan | None = None, ) -> str: phase = f"attempt_{attempt_index}" if task.feedback and task.feedback[-1].get("acceptance_type") == "revise": phase = f"revision_attempt_{attempt_index}" elif plan is not None and plan.is_team: phase = f"team_synthesis_attempt_{attempt_index}" sections = [ f"Task goal:\n{task.goal or task.description}", f"Task description:\n{task.description}", f"Current user request:\n{user_message}", f"Execution phase:\n{phase}", f"Task status:\n{task.status}", ] if task.constraints: sections.append("Known constraints:\n" + "\n".join(f"- {item}" for item in task.constraints)) if task.skill_names: sections.append( "Previously activated skills (reuse bias, not pinned):\n" + "\n".join(f"- {item}" for item in task.skill_names) ) else: sections.append("Previously activated skills:\nNone") if task.feedback: history_lines = [] for item in task.feedback[-5:]: kind = item.get("acceptance_type") or item.get("feedback_type") comment = item.get("comment") or "" run_id = item.get("run_id") or "" history_lines.append(f"- {kind} run={run_id}: {comment}".strip()) sections.append("Task acceptance history:\n" + "\n".join(history_lines)) if plan is not None: plan_lines = [ f"mode: {plan.mode}", f"reason: {plan.reason}", ] if plan.final_synthesis_instruction: plan_lines.append(f"final synthesis instruction: {plan.final_synthesis_instruction}") if plan.graph is not None: plan_lines.append(f"strategy: {plan.graph.strategy}") plan_lines.append( "nodes:\n" + "\n".join( f"- {node.node_id}: {node.task}" for node in plan.graph.nodes ) ) sections.append("Execution plan:\n" + "\n".join(plan_lines)) sections.append( "Skill selection instruction:\n" "Prefer reusing previously activated skills when they still match the Task. " "Select new skills only if the current request, revision, or execution plan needs a different capability. " "If no published skill matches, return [] and let the run continue without skills." ) return "\n\n".join(section for section in sections if section.strip()) @staticmethod def _append_task_observation( session_manager: Any, session_id: str, *, event_type: str, payload: dict[str, Any], ) -> None: session_manager.append_message( session_id, role="system", event_type=event_type, event_payload=payload, content=payload.get("reason") or payload.get("error"), context_visible=False, ) @staticmethod def _join_context(*parts: str | None) -> str: return "\n\n".join(part.strip() for part in parts if part and part.strip()) def _build_task_evidence_packet( self, *, session_manager: Any, task: TaskRecord, attempt_index: int, result: AgentRunResult, ) -> TaskEvidencePacket: main_run = EvidenceBuilder(session_manager).build_run_evidence( result.session_id, result.run_id, result.output_text, result.finish_reason, ) return TaskEvidencePacket( task_id=task.task_id, attempt_index=attempt_index, main_run=main_run, team_runs=[], team_node_results=[], final_output=result.output_text, )