**feat: refactor agent loop to state machine with dynamic trajectory checks**

- Replace hard-coded `MAX_REPL_STEPS` death-limit with a dynamic Checkpoint Interval.
- Introduce 4-phase state machine: `PLANNING`, `EXECUTING`, `TRAJECTORY_CHECK`, and `PIVOTING`.
- Add Trajectory Checks allowing the agent to self-assess progress against defined success criteria to extend its execution budget.
- Implement Hard (wipe strategy/kernel) and Soft (retry step) Pivoting logic.
- Establish persistent notebook-style kernel (flushed only on Hard Pivots).
- Upgrade sandbox safety with strict import whitelisting, dunder blocking, and taint tracking.

**bugfix:

- logging_config.py: Add markdown_it to suppressed logger list
- utils.py: Rewrite save_agent_trace for delta-based logging
- edge_rlm.py: Add last_trace_msg_count tracking and pass deltas at all 3 call sites

.env.example

@@ -0,0 +1,11 @@
+# LLM API endpoints
+AGENT_API=http://localhost:8080/v1
+REPL_API=http://localhost:8090/v1
+
+# File paths
+CONTEXT_FILE=context.txt
+TASK_FILE=task.txt
+
+# Agent behaviour
+MAX_REPL_STEPS=20
+MAX_VIRTUAL_CONTEXT_RATIO=0.85

.gitignore

@@ -4,4 +4,5 @@ __pycache__/
 *.pyc
 *.pyo
 *.pyd
+.env
 .fschatignore

edge_rlm.py

@@ -5,7 +5,7 @@ import json
 import argparse
 import io
 import logging
-import types
+import re
 from rich.console import Console
 from rich.panel import Panel
 from rich.markdown import Markdown
@@ -17,17 +17,35 @@ from logging_config import setup_logging
 import utils
 import prompts
 
+
+def _load_dotenv(path=".env"):
+    try:
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#") or "=" not in line:
+                    continue
+                key, _, val = line.partition("=")
+                key, val = key.strip(), val.strip()
+                if (val.startswith('"') and val.endswith('"')) or (val.startswith("'") and val.endswith("'")):
+                    val = val[1:-1]
+                os.environ.setdefault(key, val)
+    except FileNotFoundError:
+        pass
+
+_load_dotenv()
+
 logger = logging.getLogger(__name__)
 console = Console()
 
-# Configuration
+# Configuration (from .env with fallback defaults)
-DEFAULT_AGENT_API = "http://localhost:8080/v1"
+AGENT_API = os.getenv("AGENT_API", "http://localhost:8080/v1")
-DEFAULT_REPL_API = "http://localhost:8090/v1"
+REPL_API = os.getenv("REPL_API", "http://localhost:8090/v1")
 
-DEFAULT_CONTEXT_FILE = "context.txt"
+CONTEXT_FILE = os.getenv("CONTEXT_FILE", "context.txt")
-DEFAULT_TASK_FILE = "task.txt"
+TASK_FILE = os.getenv("TASK_FILE", "task.txt")
-MAX_REPL_STEPS = 20
+MAX_REPL_STEPS = int(os.getenv("MAX_REPL_STEPS", "20"))
-MAX_VIRTUAL_CONTEXT_RATIO = 0.85
+MAX_VIRTUAL_CONTEXT_RATIO = float(os.getenv("MAX_VIRTUAL_CONTEXT_RATIO", "0.85"))
 
 
 class LlamaClient:
@@ -207,24 +225,8 @@ class AgentOutputBuffer:
         self.global_truncated = False
         return value
 
-def run_agent(agent_client, repl_client, context_text, task_text):
+def _make_exec_env(tools, out_buffer):
-    tools = AgentTools(repl_client, context_text)
+    return {
-
-    agent_schema = {
-        "type": "object",
-        "properties": {
-            "thought": {"type": "string", "description": "Reasoning about current state and what to do next."},
-            "action": {"type": "string", "enum": ["execute_python", "final_answer"]},
-            "content": {"type": "string", "description": "Python code or Final Answer text."}
-        },
-        "required": ["thought", "action", "content"]
-    }
-
-    out_buffer = AgentOutputBuffer()
-
-    trace_filepath = utils.init_trace_file()
-
-    exec_env = {
         "RAW_CORPUS": tools.RAW_CORPUS,
         "llm_query": tools.llm_query,
         "re": __import__("re"),
@@ -236,87 +238,179 @@ def run_agent(agent_client, repl_client, context_text, task_text):
         "datetime": __import__("datetime"),
         "difflib": __import__("difflib"),
         "string": __import__("string"),
-
+        "print": out_buffer.custom_print,
-        "print": out_buffer.custom_print
     }
 
-    system_instruction = prompts.get_system_prompt()
-
-    messages = [
-        {"role": "system", "content": system_instruction},
-        {"role": "user", "content": f"USER TASK: {task_text}"}
-    ]
-
-    step = 0
-    while step < MAX_REPL_STEPS:
-        step += 1
-        logger.debug("Step %d of %d", step, MAX_REPL_STEPS)
-
-        modules = []
-        functions = []
-        variables = []
-        ACTIVE_VAR_SNIPPET_LEN = 100
-
-        for name, val in exec_env.items():
-            if name.startswith("__"): continue
-            if name == "print": continue
-
-            if isinstance(val, types.ModuleType):
-                modules.append(name)
-            elif callable(val):
-                functions.append(name)
-            else:
-                type_name = type(val).__name__
-                s_val = str(val)
-                snippet = (s_val[:ACTIVE_VAR_SNIPPET_LEN] + '...') if len(s_val) > ACTIVE_VAR_SNIPPET_LEN else s_val
-                variables.append(f"{name} ({type_name}): {snippet}")
-
-        dynamic_state_msg = (
-            f"[SYSTEM STATE REMINDER]\n"
-            f"Current Step: {step}/{MAX_REPL_STEPS}\n"
-            f"Available Libraries: {', '.join(modules)}\n"
-            f"Available Tools: {', '.join(functions)}\n"
-            f"Active Variables:\n" + ("\n".join([f"  - {v}" for v in variables]) if variables else "  (None)") + "\n---"
-        )
-
-        inference_messages = messages.copy()
-        inference_messages.append({"role": "user", "content": dynamic_state_msg})
 
+def _compress_if_needed(agent_client, messages, inference_messages):
     usage = agent_client.count_tokens(inference_messages)
     logger.debug("Context Usage: %d / %d", usage, agent_client.max_input_tokens)
 
     if usage > agent_client.max_input_tokens:
         logger.warning("Context limit exceeded. Triggering History Compression.")
-
         messages = utils.compress_history(agent_client, messages, keep_last_pairs=2)
-
+        rebuilt = messages + inference_messages[-1:]
-             inference_messages = messages.copy()
+        new_usage = agent_client.count_tokens(rebuilt)
-             inference_messages.append({"role": "user", "content": dynamic_state_msg})
-
-             new_usage = agent_client.count_tokens(inference_messages)
         logger.debug("Context Usage after compression: %d", new_usage)
-
         if new_usage > agent_client.max_input_tokens:
             logger.error("Compression insufficient. Forcing hard truncation.")
             messages.pop(2)
+    return messages
+
+
+def run_agent(agent_client, repl_client, context_text, task_text):
+    tools = AgentTools(repl_client, context_text)
+
+    planning_schema = {
+        "type": "object",
+        "properties": {
+            "goal": {"type": "string", "description": "The high-level objective."},
+            "strategy": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Discrete, programmatic actions.",
+            },
+            "success_criteria": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Success criteria parallel to strategy.",
+            },
+        },
+        "required": ["goal", "strategy", "success_criteria"],
+    }
+
+    executing_schema = {
+        "type": "object",
+        "properties": {
+            "thought": {"type": "string", "description": "Reasoning about current state and what to do next."},
+            "action": {"type": "string", "enum": ["execute_python", "final_answer"]},
+            "content": {"type": "string", "description": "Python code or Final Answer text."},
+            "step_completed": {
+                "type": "boolean",
+                "description": "Set to true when the current strategy step is fully complete and you are ready to advance.",
+            },
+        },
+        "required": ["thought", "action", "content"],
+    }
+
+    trajectory_schema = {
+        "type": "object",
+        "properties": {
+            "progress": {"type": "string", "enum": ["YES", "NO"]},
+            "justification": {"type": "string", "description": "Brief explanation of the assessment."},
+        },
+        "required": ["progress", "justification"],
+    }
+
+    pivot_schema = {
+        "type": "object",
+        "properties": {
+            "diagnosis": {"type": "string", "description": "Why progress stalled."},
+            "new_strategy": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Revised strategy (omit for Soft Pivot).",
+            },
+            "new_success_criteria": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Revised success criteria (omit for Soft Pivot).",
+            },
+        },
+        "required": ["diagnosis"],
+    }
+
+    out_buffer = AgentOutputBuffer()
+    trace_filepath = utils.init_trace_file()
+    exec_env = _make_exec_env(tools, out_buffer)
+
+    system_instruction = prompts.get_system_prompt(max_repl_steps=MAX_REPL_STEPS)
+
+    messages = [
+        {"role": "system", "content": system_instruction},
+        {"role": "user", "content": f"USER TASK: {task_text}"},
+    ]
+
+    state = "PLANNING"
+    plan = None
+    turns_since_checkpoint = 0
+    total_turns = 0
+    last_diagnosis = ""
+    last_trace_msg_count = 0
+
+    while state != "FINISHED":
+        logger.debug("State: %s | Turns since checkpoint: %d / %d", state, turns_since_checkpoint, MAX_REPL_STEPS)
+
+        if state == "PLANNING":
+            planning_prompt = prompts.get_planning_prompt(task_text)
+
+            planning_messages = [
+                {"role": "system", "content": system_instruction},
+                {"role": "user", "content": planning_prompt},
+            ]
+
+            planning_messages = _compress_if_needed(agent_client, planning_messages, planning_messages)
+
+            response_text = agent_client.completion(
+                planning_messages, schema=planning_schema, temperature=0.3
+            )
+
+            try:
+                plan = json.loads(response_text)
+            except json.JSONDecodeError:
+                logger.error("Planning: JSON parse error. Retrying.")
+                messages.append({"role": "user", "content": "System: Invalid PlanObject JSON. Please retry."})
+                continue
+
+            valid, err = utils.validate_plan_object(plan)
+            if not valid:
+                logger.error("Planning: Invalid PlanObject: %s", err)
+                messages.append({"role": "user", "content": f"System: Invalid PlanObject: {err}. Please retry."})
+                continue
+
+            plan["current_step_index"] = 0
+            plan["state"] = "executing"
+            turns_since_checkpoint = 0
+
+            messages.append({"role": "assistant", "content": json.dumps(plan, indent=2, ensure_ascii=False)})
+            messages.append({"role": "user", "content": f"System: Plan accepted. Starting execution of step 1/{len(plan['strategy'])}: {plan['strategy'][0]}"})
+
+            logger.info("Plan accepted: goal=%s, steps=%d", plan["goal"], len(plan["strategy"]))
+            state = "EXECUTING"
+
+        elif state == "EXECUTING":
+            dynamic_state_msg = prompts.get_execution_context(
+                plan, exec_env, turns_since_checkpoint, MAX_REPL_STEPS
+            )
+
             inference_messages = messages.copy()
             inference_messages.append({"role": "user", "content": dynamic_state_msg})
 
-        response_text = agent_client.completion(inference_messages, schema=agent_schema, temperature=0.5)
+            messages = _compress_if_needed(agent_client, messages, inference_messages)
+
+            inference_messages = messages.copy()
+            inference_messages.append({"role": "user", "content": dynamic_state_msg})
+
+            response_text = agent_client.completion(
+                inference_messages, schema=executing_schema, temperature=0.5
+            )
 
             try:
                 response_json = json.loads(response_text)
             except json.JSONDecodeError:
-            logger.error("JSON Parse Error")
+                logger.error("EXECUTING: JSON Parse Error")
                 messages.append({"role": "user", "content": "System: Invalid JSON returned. Please retry."})
                 continue
 
             thought = response_json.get("thought", "")
             action = response_json.get("action", "")
             content = response_json.get("content", "")
+            step_completed = response_json.get("step_completed", False)
 
             if action == "execute_python" and content:
-            content = utils.safeguard_and_repair(agent_client.debug, agent_client, messages, agent_schema, content)
+                content = utils.safeguard_and_repair(
+                    agent_client.debug, agent_client, messages, executing_schema, content
+                )
 
             if agent_client.debug:
                 console.print(Panel(
@@ -326,18 +420,19 @@ def run_agent(agent_client, repl_client, context_text, task_text):
                     border_style="magenta"
                 ))
 
-        messages.append({"role": "assistant", "content": json.dumps(response_json, indent=2, ensure_ascii=False)})
+            messages.append({
+                "role": "assistant",
+                "content": json.dumps(response_json, indent=2, ensure_ascii=False),
+            })
 
             if action == "final_answer":
                 logger.debug("Raw Agent Output: %s", content[:200])
-
                 final_report = utils.generate_final_report(agent_client, task_text, content)
-
                 final_report_md = Markdown(final_report)
                 print("\n\n")
                 console.print(final_report_md)
                 print("\n")
-            break
+                state = "FINISHED"
 
             elif action == "execute_python":
                 if agent_client.debug and content != response_json.get("content"):
@@ -350,7 +445,6 @@ def run_agent(agent_client, repl_client, context_text, task_text):
                     out_buffer.read_and_clear()
                     exec(content, exec_env)
                     observation = out_buffer.read_and_clear()
-
                     if not observation:
                         observation = "Code executed successfully (no output)."
                 except Exception as e:
@@ -366,20 +460,192 @@ def run_agent(agent_client, repl_client, context_text, task_text):
                     ))
                 messages.append({"role": "user", "content": f"Observation:\n{observation}"})
 
+                total_turns += 1
+                turns_since_checkpoint += 1
+
+                if step_completed:
+                    plan["current_step_index"] += 1
+                    if plan["current_step_index"] >= len(plan["strategy"]):
+                        plan["current_step_index"] = len(plan["strategy"]) - 1
+                        messages.append({
+                            "role": "user",
+                            "content": "System: All plan steps completed. Provide your final answer."
+                        })
+                    else:
+                        next_idx = plan["current_step_index"]
+                        messages.append({
+                            "role": "user",
+                            "content": (
+                                f"System: Advancing to step {next_idx + 1}/{len(plan['strategy'])}: "
+                                f"{plan['strategy'][next_idx]}. "
+                                f"Success criteria: {plan['success_criteria'][next_idx]}"
+                            )
+                        })
+                    turns_since_checkpoint = 0
+
+                if turns_since_checkpoint >= MAX_REPL_STEPS:
+                    state = "TRAJECTORY_CHECK"
+                    messages.append({
+                        "role": "user",
+                        "content": "System: Checkpoint reached. Pausing for trajectory review."
+                    })
+
             else:
                 messages.append({"role": "user", "content": f"System: Unknown action '{action}'."})
 
-        utils.save_agent_trace(trace_filepath, messages, step=step)
+            trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
+            utils.save_agent_trace(
+                trace_filepath, trace_delta, step=total_turns,
+                state=state, plan_step=plan["current_step_index"],
+                total_messages=len(messages),
+            )
+            last_trace_msg_count = len(messages)
+
+        elif state == "TRAJECTORY_CHECK":
+            step_idx = plan["current_step_index"]
+            criteria = plan["success_criteria"][step_idx]
+            check_prompt = prompts.get_trajectory_check_prompt(criteria)
+
+            check_messages = [
+                {"role": "system", "content": system_instruction},
+                {"role": "user", "content": check_prompt},
+            ]
+            if len(messages) >= 4:
+                check_messages += messages[-4:]
+            else:
+                check_messages += messages
+
+            response_text = agent_client.completion(
+                check_messages, schema=trajectory_schema, temperature=0.2
+            )
+
+            try:
+                check_json = json.loads(response_text)
+            except json.JSONDecodeError:
+                logger.error("TRAJECTORY_CHECK: JSON parse error. Assuming NO.")
+                check_json = {"progress": "NO", "justification": "JSON parse failure."}
+
+            progress = check_json.get("progress", "NO")
+            justification = check_json.get("justification", "")
+
+            messages.append({
+                "role": "assistant",
+                "content": json.dumps(check_json, indent=2, ensure_ascii=False),
+            })
+
+            if progress == "YES":
+                logger.info("Trajectory check PASSED: %s", justification[:100])
+                turns_since_checkpoint = 0
+                state = "EXECUTING"
+                messages.append({
+                    "role": "user",
+                    "content": (
+                        f"System: Trajectory check passed. Resuming execution. "
+                        f"Turns since checkpoint reset to 0."
+                    )
+                })
+            else:
+                logger.warning("Trajectory check FAILED: %s", justification[:100])
+                last_diagnosis = justification
+                state = "PIVOTING"
+                messages.append({
+                    "role": "user",
+                    "content": "System: Trajectory check failed. Transitioning to pivot."
+                })
+
+            trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
+            utils.save_agent_trace(
+                trace_filepath, trace_delta, step=total_turns,
+                state=state, plan_step=plan["current_step_index"],
+                total_messages=len(messages),
+            )
+            last_trace_msg_count = len(messages)
+
+        elif state == "PIVOTING":
+            pivot_prompt = prompts.get_pivot_prompt(last_diagnosis, task_text)
+
+            pivot_messages = [
+                {"role": "system", "content": system_instruction},
+                {"role": "user", "content": pivot_prompt},
+            ]
+            if len(messages) >= 4:
+                pivot_messages += messages[-4:]
+            else:
+                pivot_messages += messages
+
+            response_text = agent_client.completion(
+                pivot_messages, schema=pivot_schema, temperature=0.4
+            )
+
+            try:
+                pivot_json = json.loads(response_text)
+            except json.JSONDecodeError:
+                logger.error("PIVOTING: JSON parse error. Forcing Soft Pivot.")
+                pivot_json = {"diagnosis": "JSON parse error. Retrying current step."}
+
+            valid, err = utils.validate_pivot_object(pivot_json)
+            if not valid:
+                logger.error("PIVOTING: Invalid PivotObject: %s. Forcing Soft Pivot.", err)
+                pivot_json = {"diagnosis": pivot_json.get("diagnosis", "Validation error: " + err)}
+
+            has_new_strategy = "new_strategy" in pivot_json and "new_success_criteria" in pivot_json
+
+            messages.append({
+                "role": "assistant",
+                "content": json.dumps(pivot_json, indent=2, ensure_ascii=False),
+            })
+
+            if has_new_strategy:
+                logger.info("Hard Pivot: new strategy with %d steps.", len(pivot_json["new_strategy"]))
+                plan = {
+                    "goal": plan["goal"],
+                    "strategy": pivot_json["new_strategy"],
+                    "success_criteria": pivot_json["new_success_criteria"],
+                    "current_step_index": 0,
+                    "state": "executing",
+                }
+                exec_env = _make_exec_env(tools, out_buffer)
+                messages.append({
+                    "role": "user",
+                    "content": (
+                        f"System: Hard Pivot applied. Kernel reset. "
+                        f"New plan has {len(plan['strategy'])} steps. "
+                        f"Starting step 1: {plan['strategy'][0]}"
+                    )
+                })
+            else:
+                logger.info("Soft Pivot: retrying current step with fresh budget.")
+                messages.append({
+                    "role": "user",
+                    "content": (
+                        f"System: Soft Pivot applied. Kernel state preserved. "
+                        f"Resuming step {plan['current_step_index'] + 1}/{len(plan['strategy'])}: "
+                        f"{plan['strategy'][plan['current_step_index']]}"
+                    )
+                })
+
+            turns_since_checkpoint = 0
+            state = "EXECUTING"
+
+            trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
+            utils.save_agent_trace(
+                trace_filepath, trace_delta, step=total_turns,
+                state=state, plan_step=plan["current_step_index"],
+                total_messages=len(messages),
+            )
+            last_trace_msg_count = len(messages)
+
+    logger.info("Agent finished. Total execution turns: %d", total_turns)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="""Edge Recursive Language Model
 
     A sophisticated data extraction and analysis tool that mimics the process of a human data scientist, carefully exploring and structuring a large dataset before performing targeted queries.""")
-    parser.add_argument("--context", default=DEFAULT_CONTEXT_FILE, help="Path to text file to process")
+    parser.add_argument("--context", default=CONTEXT_FILE, help="Path to text file to process")
-    parser.add_argument("--task", default=DEFAULT_TASK_FILE, help="Path to task instruction file")
+    parser.add_argument("--task", default=TASK_FILE, help="Path to task instruction file")
     parser.add_argument("--override_task", help="Direct string override for the task")
-    parser.add_argument("--agent_api", default=DEFAULT_AGENT_API, help="URL for the Main Agent LLM")
+    parser.add_argument("--agent_api", default=AGENT_API, help="URL for the Main Agent LLM")
-    parser.add_argument("--repl_api", default=DEFAULT_REPL_API, help="URL for the Sub-call/REPL LLM")
+    parser.add_argument("--repl_api", default=REPL_API, help="URL for the Sub-call/REPL LLM")
     parser.add_argument("--debug", action="store_true", help="Enable verbose debug logging and JSON log file")
 
     args = parser.parse_args()

logging_config.py

@@ -26,7 +26,7 @@ class JSONFormatter(logging.Formatter):
 
 
 def setup_logging(level=logging.INFO, log_file=None):
-    for lib_name in ("urllib3", "requests", "http.client", "markdown", "Markdown"):
+    for lib_name in ("urllib3", "requests", "http.client", "markdown", "Markdown", "markdown_it"):
         logging.getLogger(lib_name).setLevel(logging.WARNING)
 
     handlers = [

prompts.py

@@ -2,10 +2,11 @@ role = """### ROLE
 You are a Recursive AI Controller operating in a **persistent** Python REPL. Your mission is to answer User Queries by architecting and executing data extraction scripts against a massive text variable named `RAW_CORPUS`.
 """
 
-constraints = """### CRITICAL CONSTRAINTS
+def _constraints(max_repl_steps):
+    return f"""### CRITICAL CONSTRAINTS
 - **BLINDNESS**: You cannot see `RAW_CORPUS` directly. You must "feel" its shape using Python.
 - **MEMORY SAFETY**: Your context window is finite. Summarize findings in Python variables; do not print massive blocks of raw text.
-- **LIMITED ITERATIONS**: You have a limited number of steps to complete your objective, as shown in your SYSTEM STATE REMINDER. Batch as many actions as possible into each step.
+- **TRAJECTORY CHECKS**: You have a checkpoint budget of {max_repl_steps} turns. If you do not demonstrate progress against your plan's success criteria by then, execution will pause for a trajectory review.
 - **JSON FORMATTING**: Always use `print(json.dumps(data, indent=2))` for lists/dicts.
 REPL ENV:
 - `print()`: For sending output to stdout. *Note:* DO NOT print > 1000 char snippets, counts, or summaries to preserve context. **BLINDNESS:** You are blind to function return values unless they are explicitly printed.
@@ -42,7 +43,7 @@ Avoid "Hello World" programming. Do not write one step just to see if it works.
 3.  **Global Scope:** Remember that variables you define are available in future steps. Don't re-calculate them.
 """
 
-outputs = """### YOUR OUTPUTS
+outputs = """### YOUR OUTPUTS (EXECUTING State)
 
 Your outputs must follow this format:
 ```json
@@ -51,14 +52,109 @@ Your outputs must follow this format:
     "properties": {
         "thought": {"type": "string", "description": "Reasoning about previous step, current state and what to do next."},
         "action": {"type": "string", "enum": ["execute_python", "final_answer"]},
-        "content": {"type": "string", "description": "Python code or Final Answer text."}
+        "content": {"type": "string", "description": "Python code or Final Answer text."},
+        "step_completed": {"type": "boolean", "description": "Set to true when the current step's work is done and you are ready to advance to the next strategy step."}
     },
     "required": ["thought", "action", "content"]
     }
 ```
 """
 
-def get_system_prompt():
+def get_system_prompt(max_repl_steps=20):
-    system_prompt = f"{role}\n{workflow_guidelines}\n{constraints}\n{outputs}"
+    system_prompt = f"{role}\n{workflow_guidelines}\n{_constraints(max_repl_steps)}\n{outputs}"
+    return system_prompt
 
-    return(system_prompt)
+
+def get_planning_prompt(task_text):
+    return f"""You are a Data Engineer. Your goal is to: {task_text}
+
+Before writing any code, create a PlanObject. Break down the task into discrete, programmatic actions. For each step, define what success looks like (Success Criteria). You will be audited against these criteria to determine if you are allowed to continue executing.
+
+Your PlanObject must have this exact structure:
+{{
+  "goal": "The high-level objective",
+  "strategy": ["Step 1 description", "Step 2 description", ...],
+  "success_criteria": ["Success looks like for step 1", "Success looks like for step 2", ...]
+}}
+
+Rules:
+- 'strategy' and 'success_criteria' must be parallel arrays of the same length.
+- Each strategy step must be a discrete, programmatic action (e.g., "Parse the text into chapters", "Extract all email addresses").
+- Each success criterion must be a measurable definition of progress.
+- Output ONLY the PlanObject JSON — no other text."""
+
+
+def get_trajectory_check_prompt(success_criteria):
+    return f"""Checkpoint Reached. Based on your defined success criteria for the current step, are you making tangible progress?
+
+Current Step Success Criteria: {success_criteria}
+
+Answer with a JSON object containing:
+- "progress": "YES" if you are making demonstrable progress toward the success criteria
+- "progress": "NO" if you are stuck, hitting errors, or making no forward movement
+- "justification": A brief explanation of your assessment
+
+If YES, you will be granted additional execution turns. If NO, you will be asked to diagnose and pivot."""
+
+
+def get_pivot_prompt(diagnosis, task_text):
+    return f"""You indicated you are stuck.
+
+Diagnosis: {diagnosis}
+
+Original Goal: {task_text}
+
+If the overall strategy is fundamentally flawed, provide a new_strategy and new_success_criteria for a Hard Pivot (this will wipe your active variables and restart the kernel).
+
+If the strategy is sound but your code implementation needs adjusting, omit new_strategy and new_success_criteria to perform a Soft Pivot (your variables and kernel state are preserved).
+
+Output a PivotObject:
+{{
+  "diagnosis": "Why you are stuck",
+  "new_strategy": ["Revised step 1", "Revised step 2", ...],
+  "new_success_criteria": ["Revised criterion 1", "Revised criterion 2", ...]
+}}
+
+For a Soft Pivot, omit new_strategy and new_success_criteria entirely."""
+
+
+def get_execution_context(plan, exec_env, turns_since_checkpoint, max_repl_steps):
+    modules = []
+    functions = []
+    variables = []
+    ACTIVE_VAR_SNIPPET_LEN = 100
+
+    import types
+    for name, val in exec_env.items():
+        if name.startswith("__"):
+            continue
+        if name == "print":
+            continue
+
+        if isinstance(val, types.ModuleType):
+            modules.append(name)
+        elif callable(val):
+            functions.append(name)
+        else:
+            type_name = type(val).__name__
+            s_val = str(val)
+            snippet = (s_val[:ACTIVE_VAR_SNIPPET_LEN] + '...') if len(s_val) > ACTIVE_VAR_SNIPPET_LEN else s_val
+            variables.append(f"{name} ({type_name}): {snippet}")
+
+    step_index = plan["current_step_index"]
+    total_steps = len(plan["strategy"])
+    current_step = plan["strategy"][step_index]
+    criteria = plan["success_criteria"][step_index]
+
+    return (
+        f"[SYSTEM STATE REMINDER]\n"
+        f"Plan Goal: {plan['goal']}\n"
+        f"Current Step ({step_index + 1}/{total_steps}): {current_step}\n"
+        f"Success Criteria: {criteria}\n"
+        f"Turns Since Checkpoint: {turns_since_checkpoint}/{max_repl_steps}\n"
+        f"Available Libraries: {', '.join(modules)}\n"
+        f"Available Tools: {', '.join(functions)}\n"
+        f"Active Variables:\n"
+        + ("\n".join([f"  - {v}" for v in variables]) if variables else "  (None)")
+        + "\n---"
+    )

utils.py

@@ -10,6 +10,13 @@ from rich.panel import Panel
 logger = logging.getLogger(__name__)
 console = Console()
 
+WHITELISTED_IMPORTS = {
+    "json", "math", "datetime", "re", "collections",
+    "statistics", "random", "difflib", "string",
+}
+
+DANGEROUS_BUILTINS = {"eval", "exec", "compile"}
+
 
 def init_trace_file(log_dir="logs"):
     if not os.path.exists(log_dir):
@@ -20,12 +27,15 @@ def init_trace_file(log_dir="logs"):
     logger.debug("Trace logging initialized: %s", filename)
     return filename
 
-def save_agent_trace(filepath, messages, step=0):
+def save_agent_trace(filepath, new_messages, step=0, state=None, plan_step=0, total_messages=0):
     try:
         data_to_save = {
             "timestamp": time.time(),
             "step": step,
-            "context_window": messages
+            "state": state,
+            "plan_step": plan_step,
+            "total_messages": total_messages,
+            "new_messages": new_messages,
         }
 
         with open(filepath, 'a', encoding='utf-8') as f:
@@ -35,6 +45,15 @@ def save_agent_trace(filepath, messages, step=0):
         logger.error("Failed to save trace file: %s", e)
 
 
+def _check_dunder_access(node):
+    if isinstance(node, ast.Attribute) and node.attr.startswith("__") and node.attr.endswith("__"):
+        return True
+    if isinstance(node, ast.Subscript):
+        if isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
+            if node.slice.value.startswith("__") and node.slice.value.endswith("__"):
+                return True
+    return False
+
 def _analyze_code_safety(code_str):
     try:
         tree = ast.parse(code_str)
@@ -45,6 +64,23 @@ def _analyze_code_safety(code_str):
     has_print = False
 
     for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.name not in WHITELISTED_IMPORTS:
+                    return False, f"Import Violation: '{alias.name}' is not whitelisted. Allowed: {', '.join(sorted(WHITELISTED_IMPORTS))}.", node.lineno
+
+        if isinstance(node, ast.ImportFrom):
+            if node.module not in WHITELISTED_IMPORTS:
+                return False, f"Import Violation: '{node.module}' is not whitelisted. Allowed: {', '.join(sorted(WHITELISTED_IMPORTS))}.", node.lineno
+
+        if _check_dunder_access(node):
+            name = None
+            if isinstance(node, ast.Attribute):
+                name = node.attr
+            elif isinstance(node, ast.Subscript):
+                name = getattr(node.slice, 'value', getattr(node.slice, 's', None))
+            return False, f"Dunder Access Violation: Accessing '{name}' is forbidden.", getattr(node, 'lineno', None)
+
         if isinstance(node, ast.Assign):
             if isinstance(node.value, ast.Name) and node.value.id in tainted_vars:
                 for target in node.targets:
@@ -52,6 +88,10 @@ def _analyze_code_safety(code_str):
                         tainted_vars.add(target.id)
 
         if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name):
+                if node.func.id in DANGEROUS_BUILTINS:
+                    return False, f"Builtin Violation: '{node.func.id}()' is forbidden.", node.lineno
+
             if isinstance(node.func, ast.Name) and node.func.id == 'print':
                 has_print = True
                 for arg in node.args:
@@ -230,6 +270,48 @@ Write the final response in natural language (Markdown).
         {"role": "user", "content": user_prompt}
     ])
 
+def validate_plan_object(obj):
+    if not isinstance(obj, dict):
+        return False, "PlanObject must be a JSON object."
+    for key in ("goal", "strategy", "success_criteria"):
+        if key not in obj:
+            return False, f"PlanObject missing required key '{key}'."
+    if not isinstance(obj["goal"], str) or not obj["goal"].strip():
+        return False, "PlanObject 'goal' must be a non-empty string."
+    if not isinstance(obj["strategy"], list) or len(obj["strategy"]) == 0:
+        return False, "PlanObject 'strategy' must be a non-empty list of strings."
+    if not isinstance(obj["success_criteria"], list) or len(obj["success_criteria"]) == 0:
+        return False, "PlanObject 'success_criteria' must be a non-empty list of strings."
+    if len(obj["strategy"]) != len(obj["success_criteria"]):
+        return False, "PlanObject 'strategy' and 'success_criteria' must have the same length."
+    for i, s in enumerate(obj["strategy"]):
+        if not isinstance(s, str) or not s.strip():
+            return False, f"PlanObject 'strategy[{i}]' must be a non-empty string."
+    for i, s in enumerate(obj["success_criteria"]):
+        if not isinstance(s, str) or not s.strip():
+            return False, f"PlanObject 'success_criteria[{i}]' must be a non-empty string."
+    return True, None
+
+
+def validate_pivot_object(obj):
+    if not isinstance(obj, dict):
+        return False, "PivotObject must be a JSON object."
+    if "diagnosis" not in obj:
+        return False, "PivotObject missing required key 'diagnosis'."
+    if not isinstance(obj["diagnosis"], str) or not obj["diagnosis"].strip():
+        return False, "PivotObject 'diagnosis' must be a non-empty string."
+    if "new_strategy" in obj or "new_success_criteria" in obj:
+        if "new_strategy" not in obj or "new_success_criteria" not in obj:
+            return False, "PivotObject: 'new_strategy' and 'new_success_criteria' must both be present or both omitted."
+        if not isinstance(obj["new_strategy"], list) or len(obj["new_strategy"]) == 0:
+            return False, "PivotObject 'new_strategy' must be a non-empty list of strings."
+        if not isinstance(obj["new_success_criteria"], list) or len(obj["new_success_criteria"]) == 0:
+            return False, "PivotObject 'new_success_criteria' must be a non-empty list of strings."
+        if len(obj["new_strategy"]) != len(obj["new_success_criteria"]):
+            return False, "PivotObject 'new_strategy' and 'new_success_criteria' must have the same length."
+    return True, None
+
+
 def load_file(filepath):
     try:
         with open(filepath, 'r', encoding='utf-8') as f: