from __future__ import annotations import sys from pathlib import Path import re from typing import TypedDict from langchain_core.runnables import RunnableConfig from langgraph.graph import END, START, StateGraph from langgraph.types import interrupt from agents import AgentRunner def _default_review_cmd() -> list[str]: """Return the real review_task.py command using the current interpreter.""" review_py = Path(__file__).parent / "review_task.py" return [ sys.executable, str(review_py), "$task_path", "$latest_discuss_path", "$codex_session_id", "$work_dir", ] class WorkflowState(TypedDict): task_id: str user_prompt: str work_dir: str yy: int notes_dir: str task_path: str latest_discuss_path: str codex_session_id: str has_serious_issue: bool review_comment: str should_exit: bool claude_accepts: bool rejection_reason: str def format_yy(value: int) -> str: return f"{value:02d}" def discuss_path(notes_dir: str, task_id: str, yy: int) -> str: return str(Path(notes_dir) / f"discuss{task_id}-{format_yy(yy)}.md") def resolve_work_dir(value: str) -> Path: return Path(value or ".").resolve() def resolve_notes_dir(work_dir: Path, value: str) -> Path: notes_dir = Path(value or "Notes") if not notes_dir.is_absolute(): notes_dir = work_dir / notes_dir return notes_dir.resolve() def next_discuss_index(notes_dir: str, task_id: str) -> int: notes = Path(notes_dir) if not notes.exists(): return 0 pattern = re.compile(rf"^discuss{re.escape(task_id)}-(\d+)\.md$") indexes: list[int] = [] for path in notes.iterdir(): match = pattern.match(path.name) if match: indexes.append(int(match.group(1))) return max(indexes, default=-1) + 1 def latest_discuss_path(notes_dir: str, task_id: str) -> str: notes = Path(notes_dir) if not notes.exists(): return "" pattern = re.compile(rf"^discuss{re.escape(task_id)}-(\d+)\.md$") candidates: list[tuple[int, Path]] = [] for path in notes.iterdir(): match = pattern.match(path.name) if match: candidates.append((int(match.group(1)), path)) if not candidates: return "" return str(max(candidates, key=lambda item: item[0])[1]) def validate_task_file(task_path: str) -> None: path = Path(task_path) if not path.exists(): raise FileNotFoundError( f"Task file does not exist. Current Claude Code session must write it first: {task_path}" ) text = path.read_text(encoding="utf-8").strip() if not text: raise ValueError(f"Task file is empty: {task_path}") # Count level-2 headings (lines starting with ## ) lines = text.split('\n') level2_count = sum(1 for line in lines if line.startswith('## ')) if level2_count < 3: raise ValueError( f"Task file must contain at least 3 level-2 sections (## ...). Found: {level2_count}. " f"File: {task_path}" ) def validate_rejection_file(path: str) -> None: rejection_path = Path(path) if not rejection_path.exists(): raise FileNotFoundError( "Rejection file was not written. Current Claude Code session must " f"write it first: {path}" ) text = rejection_path.read_text(encoding="utf-8").strip() if not text: raise ValueError(f"Rejection file is empty: {path}") required_markers = ["Claude Code Rejection", "拒绝"] if not any(marker in text for marker in required_markers): raise ValueError( f"Rejection file must contain a Claude rejection marker: {path}" ) if len(text) < 30: raise ValueError( f"Rejection file is too short to contain a meaningful reason: {path}" ) def values_from_state(state: WorkflowState) -> dict[str, str]: return { "task_id": state["task_id"], "user_prompt": state["user_prompt"], "work_dir": state["work_dir"], "yy": format_yy(state["yy"]), "notes_dir": state["notes_dir"], "task_path": state["task_path"], "discuss_path": discuss_path(state["notes_dir"], state["task_id"], state["yy"]), "latest_discuss_path": state["latest_discuss_path"], "codex_session_id": state["codex_session_id"], } def build_workflow(*, checkpointer=None): graph = StateGraph(WorkflowState) def initialize(state: WorkflowState) -> WorkflowState: work_dir = resolve_work_dir(state.get("work_dir", "")) notes_dir = resolve_notes_dir(work_dir, state.get("notes_dir") or "Notes") task_id = state["task_id"] task_path = str(notes_dir / f"Task{task_id}.md") yy = next_discuss_index(str(notes_dir), task_id) latest_path = latest_discuss_path(str(notes_dir), task_id) return { **state, "work_dir": str(work_dir), "yy": yy, "notes_dir": str(notes_dir), "task_path": task_path, "latest_discuss_path": latest_path, "codex_session_id": state.get("codex_session_id", ""), "has_serious_issue": False, "review_comment": "", "should_exit": False, "claude_accepts": False, "rejection_reason": "", } def request_plan(state: WorkflowState) -> WorkflowState: interrupt({ "action": "write_plan", "task_path": state["task_path"], "user_prompt": state["user_prompt"], }) validate_task_file(state["task_path"]) return state def review(state: WorkflowState, config: RunnableConfig) -> WorkflowState: cfg = config.get("configurable", {}) codex_cmd: list[str] = cfg.get("codex_review_command") or _default_review_cmd() timeout: int = int(cfg.get("codex_timeout", 600)) runner = AgentRunner(codex_review_command=codex_cmd, timeout=timeout) result = runner.review(values_from_state(state)) updates = { **state, "has_serious_issue": result.has_serious_issue, "review_comment": result.comment, "should_exit": not result.has_serious_issue, } if result.codex_session_id: updates["codex_session_id"] = result.codex_session_id return updates def write_discuss(state: WorkflowState) -> WorkflowState: path = discuss_path(state["notes_dir"], state["task_id"], state["yy"]) comment = state["review_comment"].strip() if not comment: comment = "Codex reported a serious issue but did not provide details." Path(path).write_text( "\n".join([ f"# discuss{state['task_id']}-{format_yy(state['yy'])}", "", "## Codex Review", "", comment, "", ]), encoding="utf-8", ) return { **state, "latest_discuss_path": path, "yy": state["yy"] + 1, } def request_decision(state: WorkflowState) -> WorkflowState: result = interrupt({ "action": "decide", "discuss_path": state["latest_discuss_path"], }) if not isinstance(result, dict): result = {} return { **state, "claude_accepts": bool(result.get("accepts", False)), "rejection_reason": str(result.get("reason", "")), } def request_revision(state: WorkflowState) -> WorkflowState: interrupt({ "action": "revise", "task_path": state["task_path"], "discuss_path": state["latest_discuss_path"], }) validate_task_file(state["task_path"]) return state def request_rejection(state: WorkflowState) -> WorkflowState: path = discuss_path(state["notes_dir"], state["task_id"], state["yy"]) interrupt({ "action": "write_rejection", "output_path": path, "reason_hint": state["rejection_reason"], "required_heading": "## Claude Code Rejection", }) validate_rejection_file(path) return { **state, "latest_discuss_path": path, "yy": state["yy"] + 1, } def review_route(state: WorkflowState) -> str: return "done" if state["should_exit"] else "write_discuss" def decision_route(state: WorkflowState) -> str: return "revise" if state["claude_accepts"] else "reject" graph.add_node("initialize", initialize) graph.add_node("request_plan", request_plan) graph.add_node("review", review) graph.add_node("write_discuss", write_discuss) graph.add_node("request_decision", request_decision) graph.add_node("request_revision", request_revision) graph.add_node("request_rejection", request_rejection) graph.add_edge(START, "initialize") graph.add_edge("initialize", "request_plan") graph.add_edge("request_plan", "review") graph.add_conditional_edges( "review", review_route, {"done": END, "write_discuss": "write_discuss"}, ) graph.add_edge("write_discuss", "request_decision") graph.add_conditional_edges( "request_decision", decision_route, {"revise": "request_revision", "reject": "request_rejection"}, ) graph.add_edge("request_revision", "review") graph.add_edge("request_rejection", "review") return graph.compile(checkpointer=checkpointer)