149 lines
4.5 KiB
Python
149 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass
|
|
class CodexResult:
|
|
ok: bool
|
|
final_message: str
|
|
session_id: str | None
|
|
commands: list[str] = field(default_factory=list)
|
|
stderr: str = ""
|
|
returncode: int | None = None
|
|
|
|
|
|
class CodexRunner:
|
|
def __init__(
|
|
self,
|
|
codex_bin: str,
|
|
timeout_seconds: int,
|
|
resume_json_flag_style: str,
|
|
) -> None:
|
|
self.codex_bin = codex_bin
|
|
self.timeout_seconds = timeout_seconds
|
|
self.resume_json_flag_style = resume_json_flag_style
|
|
|
|
def _command(
|
|
self,
|
|
prompt: str,
|
|
sandbox: str,
|
|
session_id: str | None,
|
|
skip_git_repo_check: bool,
|
|
) -> list[str]:
|
|
base = [self.codex_bin, "exec", "--json", "--sandbox", sandbox]
|
|
if skip_git_repo_check:
|
|
base.append("--skip-git-repo-check")
|
|
if not session_id:
|
|
return [*base, prompt]
|
|
|
|
if self.resume_json_flag_style == "after_resume":
|
|
return [
|
|
self.codex_bin,
|
|
"exec",
|
|
"resume",
|
|
session_id,
|
|
"--json",
|
|
"--sandbox",
|
|
sandbox,
|
|
*(["--skip-git-repo-check"] if skip_git_repo_check else []),
|
|
prompt,
|
|
]
|
|
return [*base, "resume", session_id, prompt]
|
|
|
|
async def run(
|
|
self,
|
|
prompt: str,
|
|
cwd: Path,
|
|
sandbox: str,
|
|
session_id: str | None,
|
|
skip_git_repo_check: bool = False,
|
|
) -> CodexResult:
|
|
cmd = self._command(
|
|
prompt=prompt,
|
|
sandbox=sandbox,
|
|
session_id=session_id,
|
|
skip_git_repo_check=skip_git_repo_check,
|
|
)
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
cwd=str(cwd),
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
|
|
try:
|
|
stdout, stderr = await asyncio.wait_for(
|
|
proc.communicate(),
|
|
timeout=self.timeout_seconds,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
proc.kill()
|
|
await proc.wait()
|
|
return CodexResult(
|
|
ok=False,
|
|
final_message=f"Codex timed out after {self.timeout_seconds} seconds.",
|
|
session_id=session_id,
|
|
)
|
|
|
|
parsed = self._parse_jsonl(stdout.decode("utf-8", errors="replace"))
|
|
stderr_text = stderr.decode("utf-8", errors="replace").strip()
|
|
if proc.returncode != 0:
|
|
message = parsed.final_message or stderr_text or "Codex command failed."
|
|
return CodexResult(
|
|
ok=False,
|
|
final_message=message,
|
|
session_id=parsed.session_id or session_id,
|
|
commands=parsed.commands,
|
|
stderr=stderr_text,
|
|
returncode=proc.returncode,
|
|
)
|
|
|
|
return CodexResult(
|
|
ok=True,
|
|
final_message=parsed.final_message or "(Codex finished without text output.)",
|
|
session_id=parsed.session_id or session_id,
|
|
commands=parsed.commands,
|
|
stderr=stderr_text,
|
|
returncode=proc.returncode,
|
|
)
|
|
|
|
@staticmethod
|
|
def _parse_jsonl(text: str) -> CodexResult:
|
|
session_id: str | None = None
|
|
final_messages: list[str] = []
|
|
commands: list[str] = []
|
|
fallback_lines: list[str] = []
|
|
|
|
for line in text.splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
event = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
fallback_lines.append(line)
|
|
continue
|
|
|
|
if event.get("type") == "thread.started":
|
|
session_id = event.get("thread_id") or session_id
|
|
continue
|
|
|
|
item = event.get("item")
|
|
if isinstance(item, dict):
|
|
if item.get("type") == "agent_message" and item.get("text"):
|
|
final_messages.append(str(item["text"]))
|
|
if item.get("type") == "command_execution" and item.get("command"):
|
|
commands.append(str(item["command"]))
|
|
|
|
final_message = final_messages[-1] if final_messages else "\n".join(fallback_lines)
|
|
return CodexResult(
|
|
ok=True,
|
|
final_message=final_message,
|
|
session_id=session_id,
|
|
commands=commands,
|
|
)
|