270 lines
8.8 KiB
Python
270 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import re
|
|
from typing import TypedDict
|
|
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.types import interrupt
|
|
|
|
from agents import AgentRunner
|
|
|
|
|
|
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(runner: AgentRunner, *, 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) -> WorkflowState:
|
|
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)
|