From 64678dc7053d69a6d84dcb3ca69f71bc4669a6cb Mon Sep 17 00:00:00 2001 From: Shuxin Lin Date: Mon, 30 Mar 2026 14:25:29 -0400 Subject: [PATCH 1/2] Add validation for agent selection and dependency references in planner Signed-off-by: Shuxin Lin --- src/workflow/planner.py | 56 +++++++++++++++++++++++++++++------------ 1 file changed, 40 insertions(+), 16 deletions(-) diff --git a/src/workflow/planner.py b/src/workflow/planner.py index f3e1cb91..97092768 100644 --- a/src/workflow/planner.py +++ b/src/workflow/planner.py @@ -77,22 +77,37 @@ def parse_plan(raw: str) -> Plan: except json.JSONDecodeError: args[n] = {} - steps = [ - PlanStep( - step_number=n, - task=tasks[n], - server=servers.get(n, ""), - tool=tools.get(n, ""), - tool_args=args.get(n, {}), - dependencies=( - [] - if deps_raw.get(n, "None").strip().lower() == "none" - else [int(x) for x in _DEP_NUM_RE.findall(deps_raw.get(n, ""))] - ), - expected_output=outputs.get(n, ""), + steps = [] + for n in sorted(tasks): + raw_dep = deps_raw.get(n, "None").strip() + + if raw_dep.lower() == "none": + dependencies = [] + else: + dependencies = [int(x) for x in _DEP_NUM_RE.findall(raw_dep)] + + # Make sure dependency references only point to earlier valid steps. + if not dependencies: + raise ValueError(f"Invalid dependency format for step {n}: {raw_dep}") + + for dep in dependencies: + if dep < 1 or dep >= n: + raise ValueError( + f"Invalid dependency reference for step {n}: #S{dep}" + ) + + steps.append( + PlanStep( + step_number=n, + task=tasks[n], + server=servers.get(n, ""), + tool=tools.get(n, ""), + tool_args=args.get(n, {}), + dependencies=dependencies, + expected_output=outputs.get(n, ""), + ) ) - for n in sorted(tasks) - ] + return Plan(steps=steps, raw=raw) @@ -121,4 +136,13 @@ def generate_plan( ) prompt = _PLAN_PROMPT.format(servers=servers_text, question=question) raw = self._llm.generate(prompt) - return parse_plan(raw) + plan = parse_plan(raw) + + # Make sure the planner only uses servers that are actually available. + for step in plan.steps: + if step.server not in server_descriptions: + raise ValueError( + f"Planner selected unknown server '{step.server}' at step {step.step_number}" + ) + + return plan \ No newline at end of file From f765c27a87c561478613a61b93fa7da85e086f8d Mon Sep 17 00:00:00 2001 From: Shuxin Lin Date: Mon, 30 Mar 2026 14:31:29 -0400 Subject: [PATCH 2/2] fix: remove redundant server validation from generate_plan (executor handles it) Signed-off-by: Shuxin Lin --- src/workflow/planner.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/workflow/planner.py b/src/workflow/planner.py index 97092768..c6cf2dcf 100644 --- a/src/workflow/planner.py +++ b/src/workflow/planner.py @@ -136,13 +136,4 @@ def generate_plan( ) prompt = _PLAN_PROMPT.format(servers=servers_text, question=question) raw = self._llm.generate(prompt) - plan = parse_plan(raw) - - # Make sure the planner only uses servers that are actually available. - for step in plan.steps: - if step.server not in server_descriptions: - raise ValueError( - f"Planner selected unknown server '{step.server}' at step {step.step_number}" - ) - - return plan \ No newline at end of file + return parse_plan(raw) \ No newline at end of file