refactor：重构流程来简化状态更新和工具调用逻辑

Rasic2 · Rasic2 · commit 3ec6cab89f77 · 2026-02-06T15:13:21.000+08:00
diff --git a/agents/matmaster_agent/core_agents/public_agents/job_agents/agent.py b/agents/matmaster_agent/core_agents/public_agents/job_agents/agent.py
@@ -25,7 +25,6 @@
     SubmitRenderAgent,
 )
 from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
-from agents.matmaster_agent.locales import i18n
 from agents.matmaster_agent.logger import PrefixFilter
 from agents.matmaster_agent.state import CURRENT_STEP
 from agents.matmaster_agent.utils.event_utils import (
@@ -143,7 +142,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
             # Only Query Job Result
             step_title = ctx.session.state.get('step_title', {}).get(
                 'title',
-                f"{i18n.t(ctx.session.state['separate_card_info'])} {ctx.session.state['plan_index'] + 1}: {current_step_tool_name}",
+                current_step_tool_name,
             )
             for matmaster_flow_event in context_function_event(
                 ctx,
diff --git a/agents/matmaster_agent/core_agents/public_agents/job_agents/result_core_agent/agent.py b/agents/matmaster_agent/core_agents/public_agents/job_agents/result_core_agent/agent.py
@@ -25,12 +25,14 @@
 from agents.matmaster_agent.core_agents.public_agents.job_agents.result_core_agent.prompt import (
     ResultCoreAgentDescription,
 )
+from agents.matmaster_agent.flow_agents.step_utils import get_current_step
 from agents.matmaster_agent.logger import PrefixFilter
 from agents.matmaster_agent.services.job import (
     get_job_detail,
     parse_and_prepare_err,
     parse_and_prepare_results,
 )
+from agents.matmaster_agent.state import CURRENT_STEP, CURRENT_STEP_STATUS
 from agents.matmaster_agent.utils.event_utils import (
     all_text_event,
     context_function_event,
@@ -115,14 +117,8 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
             if status != 'Running':
                 # 更新状态
                 plan_status = 'success' if status == 'Finished' else 'failed'
-                update_plan = copy.deepcopy(ctx.session.state['plan'])
-                logger.info(
-                    f'{ctx.session.id} plan_index = {ctx.session.state['plan_index']}'
-                )
-                update_plan['steps'][ctx.session.state['plan_index']][
-                    'status'
-                ] = plan_status
-
+                post_execution_step = copy.deepcopy(get_current_step(ctx))
+                post_execution_step[CURRENT_STEP_STATUS] = plan_status
                 update_long_running_jobs = copy.deepcopy(
                     ctx.session.state['long_running_jobs']
                 )
@@ -131,7 +127,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
                     ctx,
                     state_delta={
                         'long_running_jobs': update_long_running_jobs,
-                        'plan': update_plan,
+                        CURRENT_STEP: post_execution_step,
                     },
                 )
 
diff --git a/agents/matmaster_agent/flow_agents/agent.py b/agents/matmaster_agent/flow_agents/agent.py
@@ -86,14 +86,17 @@
 from agents.matmaster_agent.flow_agents.scene_agent.constant import SCENE_AGENT
 from agents.matmaster_agent.flow_agents.scene_agent.prompt import SCENE_INSTRUCTION
 from agents.matmaster_agent.flow_agents.scene_agent.schema import SceneSchema
-from agents.matmaster_agent.flow_agents.schema import FlowStatusEnum
 from agents.matmaster_agent.flow_agents.step_title_agent.callback import (
     filter_llm_contents,
 )
 from agents.matmaster_agent.flow_agents.step_title_agent.prompt import (
     STEP_TITLE_INSTRUCTION,
 )
 from agents.matmaster_agent.flow_agents.step_title_agent.schema import StepTitleSchema
+from agents.matmaster_agent.flow_agents.step_utils import (
+    get_current_step,
+    is_job_submitted_step,
+)
 from agents.matmaster_agent.flow_agents.step_validation_agent.prompt import (
     STEP_VALIDATION_INSTRUCTION,
 )
@@ -103,9 +106,7 @@
 from agents.matmaster_agent.flow_agents.thinking_agent.agent import ThinkingAgent
 from agents.matmaster_agent.flow_agents.thinking_agent.constant import THINKING_AGENT
 from agents.matmaster_agent.flow_agents.utils import (
-    check_plan,
     get_tools_list,
-    is_plan_confirmed,
     scenes_contain_query_job_status,
     should_bypass_confirmation,
 )
@@ -345,28 +346,20 @@ def _build_execution_agent_for_plan(
             before_model_callback=filter_llm_contents,
             after_model_callback=MatMasterLlmConfig.opik_tracer.after_model_callback,
         )
-        plan_steps = ctx.session.state.get('plan', {}).get('steps', [])
-        agent_names = []
-        for step in plan_steps:
-            tool_name = step.get('tool_name')
-            if not tool_name:
-                continue
-            belonging_agent = ALL_TOOLS.get(tool_name, {}).get('belonging_agent')
-            if belonging_agent and belonging_agent not in agent_names:
-                agent_names.append(belonging_agent)
-
-        sub_agents = [
-            AGENT_CLASS_MAPPING[agent_name](MatMasterLlmConfig)
-            for agent_name in agent_names
-            if agent_name in AGENT_CLASS_MAPPING
-        ]
+        current_step = get_current_step(ctx)
+        tool_name = current_step.get('tool_name')
+        belonging_agent = ALL_TOOLS.get(tool_name, {}).get('belonging_agent')
 
         execution_agent = MatMasterSupervisorAgent(
             name='execution_agent',
             model=MatMasterLlmConfig.default_litellm_model,
             description='根据 materials_plan 返回的计划进行总结',
             instruction='',
-            sub_agents=sub_agents + [step_title_agent] + [step_validation_agent],
+            sub_agents=[
+                AGENT_CLASS_MAPPING[belonging_agent](MatMasterLlmConfig),
+                step_title_agent,
+                step_validation_agent,
+            ],
         )
         track_adk_agent_recursive(execution_agent, MatMasterLlmConfig.opik_tracer)
         return execution_agent
@@ -711,8 +704,6 @@ async def _run_step_make_agent(
     async def _run_plan_execute_agent(
         self, ctx: InvocationContext
     ) -> AsyncGenerator[Event, None]:
-        # 重置 scenes
-        yield update_state_event(ctx, state_delta={'scenes': []})
         # 执行计划
         self._execution_agent = self._build_execution_agent_for_plan(ctx)
         if self._execution_agent:
@@ -841,26 +832,6 @@ async def _run_summary_agent(
             yield matmaster_flow_event
         yield update_state_event(ctx, state_delta={'matmaster_flow_active': None})
 
-        # 渲染追问组件
-        follow_up_list = await get_random_questions(i18n=i18n)
-        for generate_follow_up_event in context_function_event(
-            ctx,
-            self.name,
-            'matmaster_generate_follow_up',
-            {},
-            ModelRole,
-            {
-                'follow_up_result': json.dumps(
-                    {
-                        'invocation_id': ctx.invocation_id,
-                        'title': i18n.t('MoreQuestions'),
-                        'list': follow_up_list,
-                    }
-                )
-            },
-        ):
-            yield generate_follow_up_event
-
     async def _run_research_flow(
         self, ctx: InvocationContext
     ) -> AsyncGenerator[Event, None]:
@@ -891,14 +862,8 @@ async def _run_research_flow(
             yield _scene_event
 
         while True:
-            # 制定计划（1. 无计划；2. 计划已完成；3. 计划失败；4. 用户未确认计划）
-            # 仅查询任务状态时跳过 thinking（查任务状态不 thinking）
-            skip_thinking = scenes_contain_query_job_status(ctx)
-            if check_plan(ctx) in [
-                FlowStatusEnum.NO_PLAN,
-                FlowStatusEnum.COMPLETE,
-                FlowStatusEnum.FAILED,
-            ] or not is_plan_confirmed(ctx):
+            if not is_job_submitted_step(ctx):
+                skip_thinking = scenes_contain_query_job_status(ctx)
                 async for _step_make_event in self._run_step_make_agent(
                     ctx,
                     UPDATE_USER_CONTENT,
@@ -907,27 +872,50 @@ async def _run_research_flow(
                 ):
                     yield _step_make_event
 
-            # 计划未确认，暂停往下执行
-            # if is_plan_confirmed(ctx):
             async for _plan_execute_event in self._run_plan_execute_agent(ctx):
                 yield _plan_execute_event
 
-            # 回顾历史执行
-            user_request = ctx.user_content.parts[0].text
-            history_steps = ctx.session.state[HISTORY_STEPS]
-            session_files = await get_session_files(ctx.session.id)
-            self.all_finished_agent.instruction = create_all_finished_instruction(
-                user_request, history_steps, session_files
-            )
-            async for _all_finished_event in self.all_finished_agent.run_async(ctx):
-                yield _all_finished_event
+            # 检查是否为等待异步任务执行完成的阶段
+            if not is_job_submitted_step(ctx):
+                # 回顾历史执行
+                user_request = ctx.user_content.parts[0].text
+                history_steps = ctx.session.state[HISTORY_STEPS]
+                session_files = await get_session_files(ctx.session.id)
+                self.all_finished_agent.instruction = create_all_finished_instruction(
+                    user_request, history_steps, session_files
+                )
+                async for _all_finished_event in self.all_finished_agent.run_async(ctx):
+                    yield _all_finished_event
 
-            if ctx.session.state[FINISHED_STATE]['finished']:
+                if ctx.session.state[FINISHED_STATE]['finished']:
+                    break
+            else:
                 break
 
-        # 总结计划
-        async for _plan_summary_event in self._run_summary_agent(ctx):
-            yield _plan_summary_event
+        if not is_job_submitted_step(ctx):
+            # 总结计划
+            async for _plan_summary_event in self._run_summary_agent(ctx):
+                yield _plan_summary_event
+
+        # 渲染追问组件
+        follow_up_list = await get_random_questions(i18n=i18n)
+        for generate_follow_up_event in context_function_event(
+            ctx,
+            self.name,
+            'matmaster_generate_follow_up',
+            {},
+            ModelRole,
+            {
+                'follow_up_result': json.dumps(
+                    {
+                        'invocation_id': ctx.invocation_id,
+                        'title': i18n.t('MoreQuestions'),
+                        'list': follow_up_list,
+                    }
+                )
+            },
+        ):
+            yield generate_follow_up_event
 
     async def _run_async_impl(
         self, ctx: InvocationContext
diff --git a/agents/matmaster_agent/flow_agents/execution_agent/agent.py b/agents/matmaster_agent/flow_agents/execution_agent/agent.py
@@ -17,6 +17,10 @@
     MATMASTER_SUPERVISOR_AGENT,
 )
 from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
+from agents.matmaster_agent.flow_agents.step_utils import (
+    get_current_step,
+    is_job_submitted_step,
+)
 from agents.matmaster_agent.flow_agents.step_validation_agent.prompt import (
     STEP_VALIDATION_INSTRUCTION,
 )
@@ -303,14 +307,8 @@ async def _prepare_retry_other_tool(
 
     @override
     async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]:
-        plan = ctx.session.state['plan']
-        logger.info(f'{ctx.session.id} plan = {plan}')
-
-        current_step = ctx.session.state[CURRENT_STEP]
-        current_step_status = current_step['status']
-
         # 制造工具调用上下文，已提交的任务跳过该步骤
-        if current_step_status != PlanStepStatusEnum.SUBMITTED:
+        if not is_job_submitted_step(ctx):
             async for (
                 _construct_function_call_event
             ) in self._construct_function_call_ctx(ctx):
@@ -320,7 +318,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
         async for _core_execution_event in self._core_execution_agent(ctx):
             yield _core_execution_event
 
-        post_execution_step = ctx.session.state[CURRENT_STEP]
+        post_execution_step = get_current_step(ctx)
         # 工具调用结果返回【成功】
         if post_execution_step['status'] == PlanStepStatusEnum.SUCCESS:
             # 对成功的工具调用结果进行校验
@@ -330,7 +328,6 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
                     ctx
                 ):
                     yield _tool_result_validation_event
-
         # 异步任务，直接退出当前函数
         elif post_execution_step['status'] == PlanStepStatusEnum.SUBMITTED:
             return
diff --git a/agents/matmaster_agent/flow_agents/step_utils.py b/agents/matmaster_agent/flow_agents/step_utils.py
@@ -0,0 +1,14 @@
+from google.adk.agents import InvocationContext
+
+from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
+from agents.matmaster_agent.state import CURRENT_STEP, CURRENT_STEP_STATUS
+
+
+def get_current_step(ctx: InvocationContext):
+    return ctx.session.state.get(CURRENT_STEP, {})
+
+
+def is_job_submitted_step(ctx: InvocationContext) -> bool:
+    return (
+        get_current_step(ctx).get(CURRENT_STEP_STATUS) == PlanStepStatusEnum.SUBMITTED
+    )
diff --git a/agents/matmaster_agent/state.py b/agents/matmaster_agent/state.py
@@ -11,6 +11,7 @@
 MULTI_PLANS = 'multi_plans'
 PLAN_CONFIRM = 'plan_confirm'
 CURRENT_STEP = 'current_step'
+CURRENT_STEP_STATUS = 'status'
 HISTORY_STEPS = 'history_steps'
 FINISHED_STATE = 'finished_state'
 
