Bug: MLE-STAR agent crashes in Data Leakage Checker with TypeError and UnboundLocalError

[88taoqimi]

Name of sample affected

python/agents/machine-learning-engineering

Description of issue

When the use_data_leakage_checker is set to True in the agent's configuration, the MLE-STAR agent crashes during execution. The crash seems to happen due to two different errors under this condition: a TypeError related to a missing function argument, and an UnboundLocalError.

Environment

• OS & Architecture: Ubuntu 24.04.3 LTS (x86_64)
• Python version: 3.12
• Java version: N/A

Reproduction steps or code

1. Configure the agent to use the following settings. This can be done in the relevant .yaml config file (e.g., config/california-housing-prices.yaml). These settings enable the checker and increase concurrency, which makes the bug easy to reproduce.

   # Configuration that reliably triggers the bug
   use_data_leakage_checker: True
   # Concurrency settings used during testing
   num_solutions: 3
   num_model_candidates: 3
   inner_loop_round: 4
   outer_loop_round: 4
   ensemble_loop_round: 5
   num_top_plans: 3

2. Run the agent (e.g., adk run machine_learning_engineering) and start the task.

3. Observe the agent crashing after some time.

Error log

The agent crashes and raises one of the following two errors:

Error A: TypeError
TypeError: replace_leakage_code() missing 1 required positional argument: 'prefix'

Error B: UnboundLocalError
UnboundLocalError: cannot access local variable 'leakage_status' where it is not associated with a value