Reuse the existing run input for command images and ordinary queue/steer delivery. Align Unicode command separators, bound goal observation, preserve plan input encoding, and handle explicitly expected Windows cancellation disconnects in runtime tests.
188 lines
7.7 KiB
Python
188 lines
7.7 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 contextlib import closing
|
|
from pathlib import Path
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from deepseek_harness import CommandExecution, DeepSeekHarness, Notification, UserQuestionRequest
|
|
from deepseek_harness.api import final_response
|
|
|
|
|
|
class WaitResult(BaseModel):
|
|
"""The empty session/wait result."""
|
|
|
|
|
|
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 os, sys
|
|
sys.stdin.reconfigure(encoding=sys.argv[1])
|
|
if sys.stdin.isatty():
|
|
try:
|
|
text = input() + '\\n'
|
|
except EOFError:
|
|
text = ''
|
|
else:
|
|
line = bytearray()
|
|
while True:
|
|
value = os.read(0, 1)
|
|
if not value:
|
|
break
|
|
line.extend(value)
|
|
if value == b'\\n':
|
|
break
|
|
text = line.decode(sys.stdin.encoding or 'utf-8')
|
|
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 after the plan command completes."""
|
|
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()
|
|
with closing(harness.client.subscribe_session_notifications(session.id)) as notifications:
|
|
execution = session.run(f"/plan {args.task}")
|
|
if not isinstance(execution, CommandExecution):
|
|
raise RuntimeError("The SDK profile did not execute the requested command")
|
|
print(execution.result.text, flush=True)
|
|
if execution.result.kind == "error":
|
|
raise RuntimeError(execution.result.text)
|
|
harness.client.request(
|
|
"session/wait", {"sessionId": session.id}, response_model=WaitResult,
|
|
on_notification=lambda notice: display(notice, session.id), notification_subscription=notifications,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|