174 lines
7.2 KiB
Python
174 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Submit a plan command, review its plan, and observe the Session's continuing work."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from deepseek_harness import DeepSeekHarness, Notification, UserQuestionRequest
|
|
from deepseek_harness.api import final_response
|
|
|
|
|
|
class _PlanReviews:
|
|
"""Own cancellable readers with the parent's stdin encoding and join review callbacks before exit."""
|
|
|
|
# A buffered reader can consume later answers from a pipe before it exits.
|
|
_READ_LINE = """import codecs, os, sys
|
|
sys.stdin.reconfigure(encoding=sys.argv[1])
|
|
if sys.stdin.isatty():
|
|
try:
|
|
text = input() + '\\n'
|
|
except EOFError:
|
|
text = ''
|
|
else:
|
|
decoder = codecs.getincrementaldecoder(sys.stdin.encoding or 'utf-8')()
|
|
line = []
|
|
while True:
|
|
value = os.read(0, 1)
|
|
decoded = decoder.decode(value, final=not value)
|
|
if decoded:
|
|
line.append(decoded)
|
|
if not value or '\\n' in decoded:
|
|
break
|
|
text = ''.join(line)
|
|
sys.stdout.buffer.write(text.encode('utf-8'))
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._closed = threading.Event()
|
|
self._lock = threading.Lock()
|
|
self._question_lock = threading.Lock()
|
|
self._readers: list[subprocess.Popen[bytes]] = []
|
|
self._callbacks: set[threading.Thread] = set()
|
|
|
|
def __enter__(self) -> _PlanReviews:
|
|
return self
|
|
|
|
def __exit__(self, _exc_type, _exc, _tb) -> None:
|
|
with self._lock:
|
|
self._closed.set()
|
|
readers = tuple(self._readers)
|
|
callbacks = tuple(self._callbacks)
|
|
for reader in readers:
|
|
self._terminate(reader)
|
|
for reader in readers:
|
|
reader.wait()
|
|
for callback in callbacks:
|
|
if callback is not threading.current_thread():
|
|
callback.join()
|
|
|
|
@staticmethod
|
|
def _terminate(reader: subprocess.Popen[bytes]) -> None:
|
|
if reader.poll() is None:
|
|
try:
|
|
reader.kill()
|
|
except ProcessLookupError as error:
|
|
# The reader exited between poll and kill.
|
|
pass
|
|
|
|
def _read_line(self, prompt: str, cancelled: threading.Event) -> str | None:
|
|
with self._lock:
|
|
if self._closed.is_set() or cancelled.is_set():
|
|
return None
|
|
reader = subprocess.Popen(
|
|
[sys.executable, "-c", self._READ_LINE, sys.stdin.encoding or "utf-8"], stdin=sys.stdin,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
)
|
|
self._readers.append(reader)
|
|
try:
|
|
print(prompt, end="", flush=True)
|
|
while True:
|
|
if self._closed.is_set() or cancelled.is_set():
|
|
self._terminate(reader)
|
|
try:
|
|
output, errors = reader.communicate(timeout=0.05)
|
|
break
|
|
except subprocess.TimeoutExpired:
|
|
continue
|
|
if self._closed.is_set() or cancelled.is_set():
|
|
return None
|
|
if reader.returncode != 0:
|
|
raise RuntimeError(f"Plan review input reader failed: {errors.decode(errors='replace').strip()}")
|
|
return output.decode("utf-8").removesuffix("\n").removesuffix("\r") if output else None
|
|
finally:
|
|
self._terminate(reader)
|
|
reader.communicate()
|
|
with self._lock:
|
|
self._readers.remove(reader)
|
|
|
|
def answer(self, request: UserQuestionRequest, cancelled: threading.Event) -> dict:
|
|
"""Serialize plan reviews; cancellation or empty stdin dismisses the current question."""
|
|
with self._lock:
|
|
if self._closed.is_set() or cancelled.is_set():
|
|
return {"cancelled": True}
|
|
self._callbacks = {callback for callback in self._callbacks if callback.is_alive()}
|
|
self._callbacks.add(threading.current_thread())
|
|
while not self._question_lock.acquire(timeout=0.05):
|
|
if self._closed.is_set() or cancelled.is_set():
|
|
return {"cancelled": True}
|
|
try:
|
|
if self._closed.is_set() or cancelled.is_set():
|
|
return {"cancelled": True}
|
|
if len(request.questions) != 1 or request.questions[0].get("intent", {}).get("kind") != "plan-review":
|
|
raise RuntimeError("This example handles plan reviews only")
|
|
question = request.questions[0]
|
|
print(question["detail"], flush=True)
|
|
while True:
|
|
choice = self._read_line("[a] Approve, [k] Keep planning, [d] Dismiss: ", cancelled)
|
|
if choice is None or choice.strip().lower() == "d":
|
|
return {"cancelled": True}
|
|
if choice.strip().lower() == "a":
|
|
return {"answers": [{"id": question["id"], "selected": ["Approve"]}]}
|
|
if choice.strip().lower() == "k":
|
|
feedback = self._read_line("Feedback: ", cancelled)
|
|
if feedback is None:
|
|
return {"cancelled": True}
|
|
return {"answers": [{"id": question["id"], "selected": ["Keep planning"], "custom": feedback}]}
|
|
finally:
|
|
self._question_lock.release()
|
|
|
|
|
|
def display(notification: Notification, session_id: str) -> None:
|
|
"""Print assistant text and report an unsuccessful root Session turn."""
|
|
if notification.method == "session.event":
|
|
event = notification.payload.get("event")
|
|
if isinstance(event, dict) and event.get("type") == "assistant/message":
|
|
text = final_response([event])
|
|
if text:
|
|
print(text)
|
|
elif (isinstance(event, dict) and event.get("type") == "turn/end"
|
|
and notification.payload.get("sessionId") == session_id):
|
|
reason = event["data"]["reason"]
|
|
if reason["kind"] != "completed":
|
|
raise RuntimeError(f"Agent turn ended: {reason}")
|
|
|
|
|
|
def main() -> None:
|
|
"""Keep the runtime and question handler alive until the plan activity finishes."""
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--dsh-home", required=True, type=Path)
|
|
parser.add_argument("--workspace", type=Path, default=Path.cwd())
|
|
parser.add_argument("--model", default=os.environ.get("DSH_MODEL", "deepseek-v4-flash"))
|
|
parser.add_argument("task", help="Task to plan; implementation waits for your approval")
|
|
args = parser.parse_args()
|
|
if not args.task.strip():
|
|
parser.error("task must be non-blank")
|
|
with _PlanReviews() as reviews, DeepSeekHarness(
|
|
dsh_home=str(args.dsh_home.resolve()), cwd=str(args.workspace.resolve()),
|
|
profile="sdk", model=args.model, on_user_question=reviews.answer,
|
|
env={"DSH_PERMISSION_MODE": "workspace-write"},
|
|
) as harness:
|
|
session = harness.start_session()
|
|
result = session.run(f"/plan {args.task}", on_notification=lambda notice: display(notice, session.id))
|
|
if result.finish_reason != "completed":
|
|
raise RuntimeError(f"Plan activity ended with {result.finish_reason}: {result.final_response}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|