wehub-resource-sync
1223 lines
45 KiB
Python
1223 lines
45 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import signal
|
|
import sys
|
|
import tempfile
|
|
import textwrap
|
|
import threading
|
|
import time
|
|
import unittest
|
|
|
|
from omp_rpc import RpcClient, RpcCommandError, RpcConcurrencyError, RpcError, host_tool
|
|
|
|
|
|
FAKE_SERVER = textwrap.dedent(
|
|
"""
|
|
import json
|
|
import sys
|
|
import time
|
|
|
|
def usage():
|
|
return {
|
|
"input": 1,
|
|
"output": 1,
|
|
"cacheRead": 0,
|
|
"cacheWrite": 0,
|
|
"totalTokens": 2,
|
|
"cost": {
|
|
"input": 0.0,
|
|
"output": 0.0,
|
|
"cacheRead": 0.0,
|
|
"cacheWrite": 0.0,
|
|
"total": 0.0,
|
|
},
|
|
}
|
|
|
|
def model_info(model_id: str, provider: str = "anthropic"):
|
|
return {
|
|
"id": model_id,
|
|
"name": f"Model {model_id}",
|
|
"api": "anthropic-messages",
|
|
"provider": provider,
|
|
"baseUrl": "https://api.anthropic.com",
|
|
"reasoning": True,
|
|
"input": ["text"],
|
|
"cost": {
|
|
"input": 1.0,
|
|
"output": 2.0,
|
|
"cacheRead": 0.0,
|
|
"cacheWrite": 0.0,
|
|
},
|
|
"contextWindow": 200000,
|
|
"maxTokens": 8192,
|
|
}
|
|
|
|
def assistant_message(text: str):
|
|
return {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": text}],
|
|
"api": "anthropic-messages",
|
|
"provider": model_provider,
|
|
"model": model_id,
|
|
"usage": usage(),
|
|
"stopReason": "stop",
|
|
"timestamp": 1,
|
|
}
|
|
|
|
registered_host_tools = []
|
|
|
|
def current_state():
|
|
return {
|
|
"model": model_info(model_id, model_provider),
|
|
"thinkingLevel": thinking_level,
|
|
"isStreaming": False,
|
|
"isCompacting": False,
|
|
"steeringMode": steering_mode,
|
|
"followUpMode": follow_up_mode,
|
|
"interruptMode": interrupt_mode,
|
|
"sessionId": "fake-session",
|
|
"sessionName": session_name,
|
|
"autoCompactionEnabled": auto_compaction_enabled,
|
|
"messageCount": len(messages),
|
|
"queuedMessageCount": 0,
|
|
"todoPhases": todo_phases,
|
|
"dumpTools": [{"name": "read", "description": "Read files", "parameters": {"type": "object"}}] + registered_host_tools,
|
|
}
|
|
|
|
def emit_prompt_turn(text: str, delay: float = 0.0, include_extra_events: bool = False):
|
|
global last_assistant_text, messages
|
|
print(json.dumps({"type": "agent_start"}), flush=True)
|
|
print(json.dumps({"type": "turn_start"}), flush=True)
|
|
partial = assistant_message("")
|
|
print(json.dumps({"type": "message_start", "message": partial}), flush=True)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "message_update",
|
|
"message": partial,
|
|
"assistantMessageEvent": {
|
|
"type": "text_delta",
|
|
"contentIndex": 0,
|
|
"delta": text,
|
|
"partial": partial,
|
|
},
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
|
|
if delay:
|
|
time.sleep(delay)
|
|
|
|
if include_extra_events:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "tool_execution_start",
|
|
"toolCallId": "tool-1",
|
|
"toolName": "read",
|
|
"args": {"path": "README.md"},
|
|
"intent": "Inspect docs",
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "tool_execution_update",
|
|
"toolCallId": "tool-1",
|
|
"toolName": "read",
|
|
"args": {"path": "README.md"},
|
|
"partialResult": {"bytes": 12},
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "tool_execution_end",
|
|
"toolCallId": "tool-1",
|
|
"toolName": "read",
|
|
"result": {"text": "docs"},
|
|
"isError": False,
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(json.dumps({"type": "auto_compaction_start", "reason": "threshold", "action": "context-full"}), flush=True)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "auto_compaction_end",
|
|
"action": "context-full",
|
|
"result": {
|
|
"summary": "trimmed",
|
|
"shortSummary": "trimmed",
|
|
"firstKeptEntryId": "entry-1",
|
|
"tokensBefore": 123,
|
|
},
|
|
"aborted": False,
|
|
"willRetry": False,
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "auto_retry_start",
|
|
"attempt": 1,
|
|
"maxAttempts": 3,
|
|
"delayMs": 25,
|
|
"errorMessage": "retrying",
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(json.dumps({"type": "auto_retry_end", "success": True, "attempt": 1}), flush=True)
|
|
print(json.dumps({"type": "retry_fallback_applied", "from": "a", "to": "b", "role": "primary"}), flush=True)
|
|
print(json.dumps({"type": "retry_fallback_succeeded", "model": "b", "role": "primary"}), flush=True)
|
|
print(json.dumps({"type": "ttsr_triggered", "rules": [{"id": "rule-1"}]}), flush=True)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "todo_reminder",
|
|
"attempt": 1,
|
|
"maxAttempts": 2,
|
|
"todos": [{"id": "task-1", "content": "Map tools", "status": "pending"}],
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(json.dumps({"type": "todo_auto_clear"}), flush=True)
|
|
|
|
assistant = assistant_message(text)
|
|
print(json.dumps({"type": "message_end", "message": assistant}), flush=True)
|
|
print(json.dumps({"type": "turn_end", "message": assistant, "toolResults": []}), flush=True)
|
|
print(json.dumps({"type": "agent_end", "messages": [assistant]}), flush=True)
|
|
last_assistant_text = text
|
|
messages = [assistant]
|
|
|
|
def respond(request_id, command, data=None, success=True, error=None):
|
|
payload = {"id": request_id, "type": "response", "command": command, "success": success}
|
|
if success and data is not None:
|
|
payload["data"] = data
|
|
if not success:
|
|
payload["error"] = error
|
|
print(json.dumps(payload), flush=True)
|
|
|
|
print(json.dumps({"type": "ready"}), flush=True)
|
|
todo_phases = []
|
|
messages = []
|
|
branch_messages = [{"entryId": "entry-1", "text": "branch message"}]
|
|
model_provider = "anthropic"
|
|
model_id = "claude-sonnet-4-5"
|
|
thinking_level = "medium"
|
|
steering_mode = "one-at-a-time"
|
|
follow_up_mode = "one-at-a-time"
|
|
interrupt_mode = "immediate"
|
|
auto_compaction_enabled = True
|
|
auto_retry_enabled = True
|
|
session_name = "Scratchpad"
|
|
last_assistant_text = None
|
|
|
|
for raw_line in sys.stdin:
|
|
raw_line = raw_line.strip()
|
|
if not raw_line:
|
|
continue
|
|
|
|
command = json.loads(raw_line)
|
|
command_type = command["type"]
|
|
request_id = command.get("id")
|
|
|
|
if command_type == "extension_ui_response":
|
|
emit_prompt_turn("ui acknowledged")
|
|
continue
|
|
|
|
if command_type == "get_state":
|
|
respond(request_id, "get_state", current_state())
|
|
elif command_type == "set_host_tools":
|
|
registered_host_tools = command.get("tools", [])
|
|
respond(
|
|
request_id,
|
|
"set_host_tools",
|
|
{"toolNames": [tool.get("name", "") for tool in registered_host_tools]},
|
|
)
|
|
elif command_type == "set_todos":
|
|
todo_phases = command.get("phases", [])
|
|
respond(request_id, "set_todos", {"todoPhases": todo_phases})
|
|
elif command_type == "get_messages":
|
|
respond(request_id, "get_messages", {"messages": messages})
|
|
elif command_type == "set_host_tools":
|
|
tool_names = [tool.get("name", "") for tool in command.get("tools", [])]
|
|
respond(request_id, "set_host_tools", {"toolNames": tool_names})
|
|
elif command_type == "set_model":
|
|
model_provider = command["provider"]
|
|
model_id = command["modelId"]
|
|
respond(request_id, "set_model", model_info(model_id, model_provider))
|
|
elif command_type == "cycle_model":
|
|
model_id = "claude-sonnet-4-6" if model_id == "claude-sonnet-4-5" else "claude-sonnet-4-5"
|
|
respond(request_id, "cycle_model", {"model": model_info(model_id, model_provider), "thinkingLevel": thinking_level, "isScoped": False})
|
|
elif command_type == "get_available_models":
|
|
respond(
|
|
request_id,
|
|
"get_available_models",
|
|
{
|
|
"models": [
|
|
model_info("claude-sonnet-4-5", "anthropic"),
|
|
model_info("claude-sonnet-4-6", "anthropic"),
|
|
]
|
|
},
|
|
)
|
|
elif command_type == "set_thinking_level":
|
|
thinking_level = command["level"]
|
|
respond(request_id, "set_thinking_level", {})
|
|
elif command_type == "cycle_thinking_level":
|
|
thinking_level = "high" if thinking_level != "high" else "low"
|
|
respond(request_id, "cycle_thinking_level", {"level": thinking_level})
|
|
elif command_type == "set_steering_mode":
|
|
steering_mode = command["mode"]
|
|
respond(request_id, "set_steering_mode", {})
|
|
elif command_type == "set_follow_up_mode":
|
|
follow_up_mode = command["mode"]
|
|
respond(request_id, "set_follow_up_mode", {})
|
|
elif command_type == "set_interrupt_mode":
|
|
interrupt_mode = command["mode"]
|
|
respond(request_id, "set_interrupt_mode", {})
|
|
elif command_type == "compact":
|
|
respond(
|
|
request_id,
|
|
"compact",
|
|
{"summary": "trimmed", "shortSummary": "trimmed", "firstKeptEntryId": "entry-1", "tokensBefore": 123},
|
|
)
|
|
elif command_type == "set_auto_compaction":
|
|
auto_compaction_enabled = command["enabled"]
|
|
respond(request_id, "set_auto_compaction", {})
|
|
elif command_type == "set_auto_retry":
|
|
auto_retry_enabled = command["enabled"]
|
|
respond(request_id, "set_auto_retry", {})
|
|
elif command_type == "abort_retry":
|
|
respond(request_id, "abort_retry", {})
|
|
elif command_type == "bash":
|
|
respond(
|
|
request_id,
|
|
"bash",
|
|
{
|
|
"output": "hello\\n",
|
|
"exitCode": 0,
|
|
"cancelled": False,
|
|
"truncated": False,
|
|
"totalLines": 1,
|
|
"totalBytes": 6,
|
|
"outputLines": 1,
|
|
"outputBytes": 6,
|
|
},
|
|
)
|
|
elif command_type == "abort_bash":
|
|
respond(request_id, "abort_bash", {})
|
|
elif command_type == "get_session_stats":
|
|
respond(
|
|
request_id,
|
|
"get_session_stats",
|
|
{
|
|
"sessionFile": "/tmp/fake-session.jsonl",
|
|
"sessionId": "fake-session",
|
|
"userMessages": 1,
|
|
"assistantMessages": len(messages),
|
|
"toolCalls": 1,
|
|
"toolResults": 1,
|
|
"totalMessages": len(messages) + 1,
|
|
"tokens": {"input": 10, "output": 5, "cacheRead": 0, "cacheWrite": 0, "total": 15},
|
|
"premiumRequests": 0,
|
|
"cost": 0.0,
|
|
},
|
|
)
|
|
elif command_type == "export_html":
|
|
respond(request_id, "export_html", {"path": command.get("outputPath") or "/tmp/session.html"})
|
|
elif command_type == "new_session":
|
|
respond(request_id, "new_session", {"cancelled": False})
|
|
elif command_type == "switch_session":
|
|
respond(request_id, "switch_session", {"cancelled": False})
|
|
elif command_type == "branch":
|
|
branch_messages = [{"entryId": command["entryId"], "text": "branch message"}]
|
|
respond(request_id, "branch", {"text": "branch created", "cancelled": False})
|
|
elif command_type == "get_branch_messages":
|
|
respond(request_id, "get_branch_messages", {"messages": branch_messages})
|
|
elif command_type == "get_last_assistant_text":
|
|
respond(request_id, "get_last_assistant_text", {"text": last_assistant_text})
|
|
elif command_type == "set_session_name":
|
|
session_name = command["name"]
|
|
respond(request_id, "set_session_name", {})
|
|
elif command_type in {"steer", "follow_up", "abort"}:
|
|
respond(request_id, command_type, {})
|
|
elif command_type in {"prompt", "abort_and_prompt"}:
|
|
respond(request_id, command_type, {})
|
|
message = command["message"]
|
|
if message == "needs ui":
|
|
print(json.dumps({"type": "extension_ui_request", "id": "ui-1", "method": "input", "title": "Need input", "placeholder": "value"}), flush=True)
|
|
continue
|
|
if message == "needs confirm":
|
|
print(json.dumps({"type": "extension_ui_request", "id": "ui-2", "method": "confirm", "title": "Confirm", "message": "Continue?"}), flush=True)
|
|
continue
|
|
if message == "needs cancel":
|
|
print(json.dumps({"type": "extension_ui_request", "id": "ui-3", "method": "editor", "title": "Edit", "placeholder": "value"}), flush=True)
|
|
continue
|
|
if message == "needs host tool":
|
|
print(json.dumps({"type": "agent_start"}), flush=True)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "host_tool_call",
|
|
"id": "host-call-1",
|
|
"toolCallId": "toolu_host_1",
|
|
"toolName": "echo_host",
|
|
"arguments": {"message": "hello"},
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
continue
|
|
if message == "notifications":
|
|
print(json.dumps({"type": "extension_error", "extensionPath": "/tmp/ext.py", "event": "run", "error": "boom"}), flush=True)
|
|
print(json.dumps({"type": "unknown_future_event", "value": 1}), flush=True)
|
|
emit_prompt_turn("pong", delay=0.3 if message == "slow" else 0.0, include_extra_events=message == "all events")
|
|
elif command_type == "host_tool_update":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "tool_execution_update",
|
|
"toolCallId": "toolu_host_1",
|
|
"toolName": "echo_host",
|
|
"args": {"message": "hello"},
|
|
"partialResult": command["partialResult"],
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
elif command_type == "host_tool_result":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "tool_execution_end",
|
|
"toolCallId": "toolu_host_1",
|
|
"toolName": "echo_host",
|
|
"result": command["result"],
|
|
"isError": command.get("isError", False),
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(json.dumps({"type": "agent_end", "messages": []}), flush=True)
|
|
else:
|
|
respond(request_id, command_type, success=False, error=f"unsupported: {command_type}")
|
|
"""
|
|
)
|
|
|
|
IDLESS_ERROR_SERVER = textwrap.dedent(
|
|
"""
|
|
import json
|
|
import sys
|
|
|
|
print(json.dumps({"type": "ready"}), flush=True)
|
|
|
|
for raw_line in sys.stdin:
|
|
raw_line = raw_line.strip()
|
|
if not raw_line:
|
|
continue
|
|
|
|
command = json.loads(raw_line)
|
|
if command["type"] == "set_host_tools":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": command.get("id"),
|
|
"type": "response",
|
|
"command": "set_host_tools",
|
|
"success": True,
|
|
"data": {"toolNames": []},
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
continue
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"type": "response",
|
|
"command": command["type"],
|
|
"success": False,
|
|
"error": f"unsupported: {command['type']}",
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
"""
|
|
)
|
|
|
|
LATE_PROMPT_FAILURE_SERVER = textwrap.dedent(
|
|
"""
|
|
import json
|
|
import sys
|
|
|
|
print(json.dumps({"type": "ready"}), flush=True)
|
|
|
|
for raw_line in sys.stdin:
|
|
raw_line = raw_line.strip()
|
|
if not raw_line:
|
|
continue
|
|
|
|
command = json.loads(raw_line)
|
|
request_id = command.get("id")
|
|
if command["type"] == "set_host_tools":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": request_id,
|
|
"type": "response",
|
|
"command": "set_host_tools",
|
|
"success": True,
|
|
"data": {"toolNames": []},
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
continue
|
|
if command["type"] == "prompt":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": request_id,
|
|
"type": "response",
|
|
"command": "prompt",
|
|
"success": True,
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": request_id,
|
|
"type": "response",
|
|
"command": "prompt",
|
|
"success": False,
|
|
"error": "late failure",
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
else:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": request_id,
|
|
"type": "response",
|
|
"command": command["type"],
|
|
"success": True,
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
"""
|
|
)
|
|
|
|
STDERR_SERVER = textwrap.dedent(
|
|
"""
|
|
import json
|
|
import sys
|
|
|
|
sys.stderr.write("first\\n")
|
|
sys.stderr.flush()
|
|
sys.stderr.write("second\\n")
|
|
sys.stderr.flush()
|
|
print(json.dumps({"type": "ready"}), flush=True)
|
|
|
|
for raw_line in sys.stdin:
|
|
raw_line = raw_line.strip()
|
|
if not raw_line:
|
|
continue
|
|
command = json.loads(raw_line)
|
|
if command["type"] == "set_host_tools":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": command.get("id"),
|
|
"type": "response",
|
|
"command": "set_host_tools",
|
|
"success": True,
|
|
"data": {"toolNames": []},
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
"""
|
|
)
|
|
|
|
INVALID_JSON_SERVER = textwrap.dedent(
|
|
"""
|
|
import sys
|
|
|
|
sys.stdout.write('{"type":"ready"}\\n')
|
|
sys.stdout.flush()
|
|
sys.stdout.write('{"type":"broken"\\n')
|
|
sys.stdout.flush()
|
|
"""
|
|
)
|
|
|
|
BROKEN_STARTUP_SERVER = textwrap.dedent(
|
|
"""
|
|
import sys
|
|
|
|
sys.stdout.write('not-json\\n')
|
|
sys.stdout.flush()
|
|
"""
|
|
)
|
|
|
|
|
|
class RpcClientTests(unittest.TestCase):
|
|
def make_client(self, server: str = FAKE_SERVER, **kwargs: object) -> RpcClient:
|
|
return RpcClient(
|
|
command=[sys.executable, "-u", "-c", server],
|
|
startup_timeout=2.0,
|
|
request_timeout=2.0,
|
|
**kwargs,
|
|
)
|
|
|
|
def test_command_builder_supports_common_rpc_options(self) -> None:
|
|
client = RpcClient(
|
|
executable="omp",
|
|
model="openrouter/anthropic/claude-sonnet-4.6",
|
|
cwd="/tmp/workspace",
|
|
thinking="high",
|
|
append_system_prompt="extra instructions",
|
|
provider_session_id="provider-session-1",
|
|
tools=("read", "edit", "write"),
|
|
no_session=True,
|
|
no_skills=True,
|
|
no_rules=True,
|
|
extra_args=("--foo", "bar"),
|
|
)
|
|
|
|
self.assertEqual(
|
|
client.command,
|
|
(
|
|
"omp",
|
|
"--mode",
|
|
"rpc",
|
|
"--model",
|
|
"openrouter/anthropic/claude-sonnet-4.6",
|
|
"--thinking",
|
|
"high",
|
|
"--append-system-prompt",
|
|
"extra instructions",
|
|
"--provider-session-id",
|
|
"provider-session-1",
|
|
"--tools",
|
|
"read,edit,write",
|
|
"--no-session",
|
|
"--no-skills",
|
|
"--no-rules",
|
|
"--no-title",
|
|
"--foo",
|
|
"bar",
|
|
),
|
|
)
|
|
|
|
def test_get_state_and_bash(self) -> None:
|
|
with self.make_client() as client:
|
|
state = client.get_state()
|
|
self.assertEqual(state.session_id, "fake-session")
|
|
self.assertEqual(
|
|
state.model.id if state.model else None, "claude-sonnet-4-5"
|
|
)
|
|
|
|
result = client.bash("echo hello")
|
|
self.assertEqual(result.output, "hello\n")
|
|
self.assertEqual(result.exit_code, 0)
|
|
|
|
def test_prompt_and_wait_returns_assistant_text(self) -> None:
|
|
with self.make_client() as client:
|
|
turn = client.prompt_and_wait("say hello", timeout=2.0)
|
|
self.assertEqual(turn.require_assistant_text(), "pong")
|
|
self.assertGreaterEqual(len(turn.events), 3)
|
|
|
|
def test_custom_tools_are_registered_and_executed_via_rpc(self) -> None:
|
|
def echo_host(args: dict[str, str], context) -> str:
|
|
context.send_update(f"working:{args['message']}")
|
|
return f"host:{args['message']}"
|
|
|
|
with self.make_client(
|
|
custom_tools=(
|
|
host_tool(
|
|
name="echo_host",
|
|
description="Echo from the Python host process",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"message": {"type": "string"}},
|
|
"required": ["message"],
|
|
"additionalProperties": False,
|
|
},
|
|
execute=echo_host,
|
|
),
|
|
)
|
|
) as client:
|
|
state = client.get_state()
|
|
self.assertEqual(state.dump_tools[-1].name, "echo_host")
|
|
|
|
turn = client.prompt_and_wait("needs host tool", timeout=2.0)
|
|
update_events = [
|
|
event
|
|
for event in turn.events
|
|
if getattr(event, "type", None) == "tool_execution_update"
|
|
]
|
|
end_events = [
|
|
event
|
|
for event in turn.events
|
|
if getattr(event, "type", None) == "tool_execution_end"
|
|
]
|
|
|
|
self.assertEqual(len(update_events), 1)
|
|
self.assertEqual(
|
|
update_events[0].partial_result["content"][0]["text"], "working:hello"
|
|
)
|
|
self.assertEqual(len(end_events), 1)
|
|
self.assertEqual(end_events[0].result["content"][0]["text"], "host:hello")
|
|
|
|
def test_extension_ui_round_trip(self) -> None:
|
|
with self.make_client() as client:
|
|
client.prompt("needs ui")
|
|
request = client.next_ui_request(timeout=2.0)
|
|
self.assertEqual(request.method, "input")
|
|
|
|
client.send_ui_value(request.id, "approved")
|
|
client.wait_for_idle(timeout=2.0)
|
|
|
|
def test_install_headless_ui_cancels_interactive_requests(self) -> None:
|
|
seen_methods: list[str] = []
|
|
|
|
with self.make_client() as client:
|
|
client.install_headless_ui(
|
|
on_request=lambda request: seen_methods.append(request.method)
|
|
)
|
|
client.prompt_and_wait("needs ui", timeout=2.0)
|
|
|
|
self.assertEqual(seen_methods, ["input"])
|
|
|
|
def test_ready_and_typed_event_listeners(self) -> None:
|
|
ready_types: list[str] = []
|
|
event_types: list[str] = []
|
|
notification_types: list[str] = []
|
|
client = self.make_client()
|
|
client.on_ready(lambda event: ready_types.append(event.type))
|
|
client.on_notification(
|
|
lambda notification: notification_types.append(notification.type)
|
|
)
|
|
client.on_turn_start(lambda event: event_types.append(event.type))
|
|
client.on_message_update(lambda event: event_types.append(event.type))
|
|
client.on_agent_end(lambda event: event_types.append(event.type))
|
|
|
|
try:
|
|
client.start()
|
|
client.prompt_and_wait("say hello", timeout=2.0)
|
|
finally:
|
|
client.stop()
|
|
|
|
self.assertEqual(ready_types, ["ready"])
|
|
self.assertEqual(event_types, ["turn_start", "message_update", "agent_end"])
|
|
self.assertIn("ready", notification_types)
|
|
self.assertIn("turn_start", notification_types)
|
|
self.assertIn("agent_end", notification_types)
|
|
|
|
def test_set_todos_supports_flat_items(self) -> None:
|
|
with self.make_client() as client:
|
|
phases = client.set_todos(["Map tools", "Exercise edits"])
|
|
|
|
self.assertEqual(len(phases), 1)
|
|
self.assertEqual(phases[0].name, "Todos")
|
|
self.assertEqual(phases[0].tasks[0].content, "Map tools")
|
|
self.assertEqual(phases[0].tasks[1].status, "pending")
|
|
|
|
state = client.get_state()
|
|
self.assertEqual(state.todo_phases[0].tasks[1].content, "Exercise edits")
|
|
|
|
def test_model_mode_and_session_commands(self) -> None:
|
|
with self.make_client() as client:
|
|
model = client.set_model("anthropic", "claude-sonnet-4-6")
|
|
self.assertEqual(model.id, "claude-sonnet-4-6")
|
|
|
|
cycled = client.cycle_model()
|
|
self.assertIsNotNone(cycled)
|
|
self.assertEqual(cycled.model.id, "claude-sonnet-4-5")
|
|
|
|
available = client.get_available_models()
|
|
self.assertEqual(
|
|
[item.id for item in available],
|
|
["claude-sonnet-4-5", "claude-sonnet-4-6"],
|
|
)
|
|
|
|
client.set_thinking_level("high")
|
|
self.assertEqual(client.get_state().thinking_level, "high")
|
|
|
|
cycled_level = client.cycle_thinking_level()
|
|
self.assertIsNotNone(cycled_level)
|
|
self.assertEqual(cycled_level.level, "low")
|
|
|
|
client.set_steering_mode("all")
|
|
client.set_follow_up_mode("all")
|
|
client.set_interrupt_mode("wait")
|
|
client.set_auto_compaction(False)
|
|
client.set_auto_retry(False)
|
|
client.set_session_name("Renamed")
|
|
|
|
state = client.get_state()
|
|
self.assertEqual(state.steering_mode, "all")
|
|
self.assertEqual(state.follow_up_mode, "all")
|
|
self.assertEqual(state.interrupt_mode, "wait")
|
|
self.assertFalse(state.auto_compaction_enabled)
|
|
self.assertEqual(state.session_name, "Renamed")
|
|
|
|
compacted = client.compact()
|
|
self.assertEqual(compacted.summary, "trimmed")
|
|
|
|
stats = client.get_session_stats()
|
|
self.assertEqual(stats.session_id, "fake-session")
|
|
self.assertEqual(stats.tokens.total, 15)
|
|
|
|
exported = client.export_html("/tmp/custom.html")
|
|
self.assertEqual(str(exported), "/tmp/custom.html")
|
|
|
|
new_session = client.new_session()
|
|
switched = client.switch_session("/tmp/session.jsonl")
|
|
self.assertFalse(new_session.cancelled)
|
|
self.assertFalse(switched.cancelled)
|
|
|
|
branch = client.branch("entry-9")
|
|
self.assertEqual(branch.text, "branch created")
|
|
branch_messages = client.get_branch_messages()
|
|
self.assertEqual(branch_messages[0].entry_id, "entry-9")
|
|
|
|
def test_message_and_control_commands(self) -> None:
|
|
with self.make_client() as client:
|
|
turn = client.prompt_and_wait("say hello", timeout=2.0)
|
|
self.assertEqual(turn.require_assistant_text(), "pong")
|
|
self.assertEqual(client.get_last_assistant_text(), "pong")
|
|
|
|
messages = client.get_messages()
|
|
self.assertEqual(len(messages), 1)
|
|
self.assertEqual(messages[0]["role"], "assistant")
|
|
|
|
client.clear_todos()
|
|
self.assertEqual(client.get_todos(), ())
|
|
|
|
client.steer("nudge")
|
|
client.follow_up("later")
|
|
client.abort()
|
|
client.abort_retry()
|
|
client.abort_bash()
|
|
|
|
client.abort_and_prompt("say hello")
|
|
client.wait_for_idle(timeout=2.0)
|
|
self.assertEqual(client.get_last_assistant_text(), "pong")
|
|
|
|
def test_collect_events_returns_turn_events(self) -> None:
|
|
with self.make_client() as client:
|
|
client.prompt("slow")
|
|
events = client.collect_events(timeout=2.0)
|
|
|
|
self.assertGreaterEqual(len(events), 1)
|
|
self.assertEqual(events[-1].type, "agent_end")
|
|
|
|
def test_all_typed_event_listeners_receive_eventful_prompt(self) -> None:
|
|
seen: list[str] = []
|
|
|
|
with self.make_client() as client:
|
|
client.on_event(lambda event: seen.append(f"event:{event.type}"))
|
|
client.on_agent_start(lambda event: seen.append(event.type))
|
|
client.on_turn_end(lambda event: seen.append(event.type))
|
|
client.on_message_start(lambda event: seen.append(event.type))
|
|
client.on_message_end(lambda event: seen.append(event.type))
|
|
client.on_tool_execution_start(lambda event: seen.append(event.type))
|
|
client.on_tool_execution_update(lambda event: seen.append(event.type))
|
|
client.on_tool_execution_end(lambda event: seen.append(event.type))
|
|
client.on_auto_compaction_start(lambda event: seen.append(event.type))
|
|
client.on_auto_compaction_end(lambda event: seen.append(event.type))
|
|
client.on_auto_retry_start(lambda event: seen.append(event.type))
|
|
client.on_auto_retry_end(lambda event: seen.append(event.type))
|
|
client.on_retry_fallback_applied(lambda event: seen.append(event.type))
|
|
client.on_retry_fallback_succeeded(lambda event: seen.append(event.type))
|
|
client.on_ttsr_triggered(lambda event: seen.append(event.type))
|
|
client.on_todo_reminder(lambda event: seen.append(event.type))
|
|
client.on_todo_auto_clear(lambda event: seen.append(event.type))
|
|
|
|
turn = client.prompt_and_wait("all events", timeout=2.0)
|
|
|
|
self.assertEqual(turn.require_assistant_text(), "pong")
|
|
for expected in [
|
|
"agent_start",
|
|
"message_start",
|
|
"message_end",
|
|
"turn_end",
|
|
"tool_execution_start",
|
|
"tool_execution_update",
|
|
"tool_execution_end",
|
|
"auto_compaction_start",
|
|
"auto_compaction_end",
|
|
"auto_retry_start",
|
|
"auto_retry_end",
|
|
"retry_fallback_applied",
|
|
"retry_fallback_succeeded",
|
|
"ttsr_triggered",
|
|
"todo_reminder",
|
|
"todo_auto_clear",
|
|
]:
|
|
self.assertIn(expected, seen)
|
|
|
|
def test_extension_and_unknown_notification_listeners(self) -> None:
|
|
seen_extension_errors: list[str] = []
|
|
seen_unknown: list[str] = []
|
|
|
|
with self.make_client() as client:
|
|
client.on_extension_error(
|
|
lambda event: seen_extension_errors.append(event.error)
|
|
)
|
|
client.on_unknown_notification(
|
|
lambda event: seen_unknown.append(str(event.payload.get("type")))
|
|
)
|
|
client.prompt_and_wait("notifications", timeout=2.0)
|
|
|
|
self.assertEqual(seen_extension_errors, ["boom"])
|
|
self.assertEqual(seen_unknown, ["unknown_future_event"])
|
|
|
|
def test_ui_confirmation_and_cancel_round_trip(self) -> None:
|
|
with self.make_client() as client:
|
|
client.prompt("needs confirm")
|
|
confirm_request = client.next_ui_request(timeout=2.0)
|
|
self.assertEqual(confirm_request.method, "confirm")
|
|
client.send_ui_confirmation(confirm_request.id, True)
|
|
client.wait_for_idle(timeout=2.0)
|
|
|
|
client.prompt("needs cancel")
|
|
editor_request = client.next_ui_request(timeout=2.0)
|
|
self.assertEqual(editor_request.method, "editor")
|
|
client.cancel_ui_request(editor_request.id)
|
|
client.wait_for_idle(timeout=2.0)
|
|
|
|
def test_prompt_lifecycle_collectors_are_single_flight(self) -> None:
|
|
results: list[str] = []
|
|
errors: list[BaseException] = []
|
|
|
|
with self.make_client() as client:
|
|
|
|
def run_prompt() -> None:
|
|
try:
|
|
results.append(
|
|
client.prompt_and_wait(
|
|
"slow", timeout=2.0
|
|
).require_assistant_text()
|
|
)
|
|
except (
|
|
BaseException
|
|
) as exc: # pragma: no cover - defensive thread capture
|
|
errors.append(exc)
|
|
|
|
thread = threading.Thread(target=run_prompt)
|
|
thread.start()
|
|
|
|
deadline = time.time() + 1.0
|
|
while (
|
|
client._prompt_lifecycle.active_operation != "prompt_and_wait"
|
|
and time.time() < deadline
|
|
):
|
|
time.sleep(0.01)
|
|
|
|
self.assertEqual(
|
|
client._prompt_lifecycle.active_operation, "prompt_and_wait"
|
|
)
|
|
with self.assertRaises(RpcConcurrencyError):
|
|
client.collect_events(timeout=1.0)
|
|
|
|
thread.join(timeout=2.0)
|
|
self.assertFalse(thread.is_alive())
|
|
|
|
self.assertEqual(errors, [])
|
|
self.assertEqual(results, ["pong"])
|
|
|
|
def test_listener_mutation_does_not_change_retained_turn(self) -> None:
|
|
with self.make_client() as client:
|
|
client.on_message_end(
|
|
lambda event: event.message["content"].__setitem__(
|
|
0, {"type": "text", "text": "mutated"}
|
|
)
|
|
)
|
|
turn = client.prompt_and_wait("say hello", timeout=2.0)
|
|
messages = client.get_messages()
|
|
|
|
self.assertEqual(turn.require_assistant_text(), "pong")
|
|
self.assertEqual(messages[0]["content"][0]["text"], "pong")
|
|
|
|
def test_id_less_error_responses_are_correlated(self) -> None:
|
|
with self.make_client(server=IDLESS_ERROR_SERVER) as client:
|
|
with self.assertRaises(RpcCommandError) as ctx:
|
|
client.request_raw("unknown")
|
|
|
|
self.assertEqual(ctx.exception.command, "unknown")
|
|
self.assertEqual(ctx.exception.error, "unsupported: unknown")
|
|
|
|
def test_prompt_and_wait_raises_for_late_prompt_failure(self) -> None:
|
|
protocol_errors: list[str] = []
|
|
client = self.make_client(server=LATE_PROMPT_FAILURE_SERVER)
|
|
client.on_protocol_error(lambda error: protocol_errors.append(str(error)))
|
|
|
|
try:
|
|
client.start()
|
|
with self.assertRaises(RpcCommandError) as ctx:
|
|
client.prompt_and_wait("say hello", timeout=2.0)
|
|
finally:
|
|
client.stop()
|
|
|
|
self.assertEqual(ctx.exception.command, "prompt")
|
|
self.assertEqual(ctx.exception.error, "late failure")
|
|
self.assertEqual(len(protocol_errors), 1)
|
|
self.assertIn("late failure", protocol_errors[0])
|
|
self.assertEqual(len(client.protocol_errors), 1)
|
|
|
|
def test_listener_exceptions_are_reported_without_stopping_client(self) -> None:
|
|
listener_errors: list[tuple[str, str | None, str]] = []
|
|
client = self.make_client()
|
|
client.on_notification(
|
|
lambda notification: (
|
|
(_ for _ in ()).throw(RuntimeError("boom"))
|
|
if notification.type == "turn_start"
|
|
else None
|
|
)
|
|
)
|
|
client.on_listener_error(
|
|
lambda event: listener_errors.append(
|
|
(event.listener_kind, event.source_type, str(event.error))
|
|
)
|
|
)
|
|
|
|
try:
|
|
client.start()
|
|
turn = client.prompt_and_wait("say hello", timeout=2.0)
|
|
finally:
|
|
client.stop()
|
|
|
|
self.assertEqual(turn.require_assistant_text(), "pong")
|
|
self.assertEqual(listener_errors, [("notification", "turn_start", "boom")])
|
|
self.assertEqual(len(client.listener_errors), 1)
|
|
self.assertEqual(client.listener_errors[0].listener_kind, "notification")
|
|
|
|
def test_stderr_history_is_bounded(self) -> None:
|
|
client = self.make_client(server=STDERR_SERVER, max_stderr_chunks=1)
|
|
|
|
try:
|
|
client.start()
|
|
finally:
|
|
client.stop()
|
|
|
|
self.assertEqual(client.stderr, "second\n")
|
|
|
|
def test_broken_startup_frame_is_reported(self) -> None:
|
|
client = self.make_client(server=BROKEN_STARTUP_SERVER)
|
|
|
|
with self.assertRaises(RpcError) as ctx:
|
|
client.start()
|
|
|
|
self.assertIn("Frame: 'not-json'", str(ctx.exception))
|
|
|
|
def test_event_history_limit_reports_overflow(self) -> None:
|
|
with self.make_client(max_event_history=2) as client:
|
|
with self.assertRaises(RpcError) as ctx:
|
|
client.prompt_and_wait("say hello", timeout=2.0)
|
|
|
|
self.assertIn("max_event_history", str(ctx.exception))
|
|
|
|
|
|
HANGING_SERVER = textwrap.dedent(
|
|
"""
|
|
import json
|
|
import sys
|
|
|
|
print(json.dumps({"type": "ready"}), flush=True)
|
|
# Read one line (the prompt) and acknowledge it, then never emit agent_end.
|
|
# The client's prompt_and_wait should sit in _wait_for_agent_end forever
|
|
# unless stop() unblocks it.
|
|
line = sys.stdin.readline()
|
|
if line:
|
|
command = json.loads(line)
|
|
if command.get("type") == "prompt":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"id": command["id"],
|
|
"type": "response",
|
|
"command": "prompt",
|
|
"success": True,
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
# Block forever on stdin so the subprocess does not exit on its own.
|
|
sys.stdin.read()
|
|
"""
|
|
)
|
|
|
|
|
|
class StopUnblocksPromptAndWaitTests(unittest.TestCase):
|
|
"""Regression: stop() must wake `_wait_for_agent_end` immediately.
|
|
|
|
Previously, the stdout reader's "if not self._stopping:" guard caused
|
|
`_mark_closed` to be skipped after stop(), so `_closed_error` stayed
|
|
`None` and `_wait_for_agent_end` blocked on its condition variable until
|
|
the prompt timeout. The fix sets `_closed_error` from `stop()` itself.
|
|
"""
|
|
|
|
def test_stop_during_prompt_unblocks_waiter(self) -> None:
|
|
from omp_rpc import RpcProcessExitError
|
|
|
|
client = RpcClient(
|
|
command=[sys.executable, "-u", "-c", HANGING_SERVER],
|
|
startup_timeout=2.0,
|
|
request_timeout=2.0,
|
|
)
|
|
client.start()
|
|
try:
|
|
errors: list[BaseException] = []
|
|
|
|
def run_prompt() -> None:
|
|
try:
|
|
# 30s is more than enough to let stop() race in; if the
|
|
# bug regresses, the worker hangs the full 30s.
|
|
client.prompt_and_wait("hang", timeout=30.0)
|
|
except BaseException as exc:
|
|
errors.append(exc)
|
|
|
|
thread = threading.Thread(target=run_prompt)
|
|
thread.start()
|
|
|
|
# Wait until the prompt is in flight.
|
|
deadline = time.time() + 2.0
|
|
while (
|
|
client._prompt_lifecycle.active_operation != "prompt_and_wait"
|
|
and time.time() < deadline
|
|
):
|
|
time.sleep(0.01)
|
|
self.assertEqual(
|
|
client._prompt_lifecycle.active_operation, "prompt_and_wait"
|
|
)
|
|
|
|
t0 = time.time()
|
|
client.stop()
|
|
thread.join(timeout=2.0)
|
|
elapsed = time.time() - t0
|
|
|
|
self.assertFalse(
|
|
thread.is_alive(), "prompt_and_wait did not return after stop()"
|
|
)
|
|
self.assertLess(
|
|
elapsed, 2.0, f"stop() took {elapsed:.2f}s to unblock prompt_and_wait"
|
|
)
|
|
self.assertEqual(len(errors), 1)
|
|
self.assertIsInstance(errors[0], RpcProcessExitError)
|
|
finally:
|
|
# stop() is idempotent; safe to call again on cleanup paths.
|
|
client.stop()
|
|
|
|
|
|
class TerminatesProcessGroupTests(unittest.TestCase):
|
|
"""Regression: stop() must reap descendants the agent spawned, not only
|
|
the omp leader.
|
|
|
|
A `bun test` launched by the agent's `bash` tool runs as a grandchild of
|
|
the omp process. Before the fix, stop() signalled only the leader pid, so
|
|
such grandchildren reparented to the container init and kept running —
|
|
once ballooning to tens of GB of RAM. omp is now spawned in its own
|
|
session and stop() tears down the whole process group.
|
|
"""
|
|
|
|
@unittest.skipUnless(hasattr(os, "killpg"), "POSIX process groups only")
|
|
def test_stop_kills_grandchild_spawned_by_server(self) -> None:
|
|
work = tempfile.mkdtemp()
|
|
self.addCleanup(shutil.rmtree, work, ignore_errors=True)
|
|
pid_file = os.path.join(work, "gc.pid")
|
|
beat_file = os.path.join(work, "gc.beat")
|
|
gc_script = os.path.join(work, "gc.py")
|
|
with open(gc_script, "w", encoding="utf-8") as handle:
|
|
handle.write(
|
|
textwrap.dedent(
|
|
f"""
|
|
import os, time
|
|
with open({pid_file!r}, "w") as f:
|
|
f.write(str(os.getpid()))
|
|
while True:
|
|
with open({beat_file!r}, "w") as f:
|
|
f.write(str(time.time()))
|
|
time.sleep(0.02)
|
|
"""
|
|
)
|
|
)
|
|
|
|
def _reap_leaked_grandchild() -> None:
|
|
try:
|
|
with open(pid_file, encoding="utf-8") as f:
|
|
os.kill(int(f.read()), signal.SIGKILL)
|
|
except (OSError, ValueError):
|
|
pass
|
|
|
|
self.addCleanup(_reap_leaked_grandchild)
|
|
|
|
# Fake omp server: spawn the long-lived grandchild, signal ready, then
|
|
# idle until torn down (sleep past stdin EOF so the group is still
|
|
# alive when stop() fires).
|
|
server = textwrap.dedent(
|
|
f"""
|
|
import json, subprocess, sys, time
|
|
subprocess.Popen([sys.executable, {gc_script!r}])
|
|
print(json.dumps({{"type": "ready"}}), flush=True)
|
|
for _line in sys.stdin:
|
|
pass
|
|
time.sleep(30)
|
|
"""
|
|
)
|
|
|
|
client = RpcClient(
|
|
command=[sys.executable, "-u", "-c", server],
|
|
startup_timeout=2.0,
|
|
request_timeout=2.0,
|
|
)
|
|
client.start()
|
|
try:
|
|
deadline = time.time() + 2.0
|
|
while time.time() < deadline and not os.path.exists(pid_file):
|
|
time.sleep(0.02)
|
|
self.assertTrue(os.path.exists(pid_file), "grandchild never started")
|
|
with open(pid_file, encoding="utf-8") as f:
|
|
os.kill(int(f.read()), 0) # alive before teardown
|
|
finally:
|
|
client.stop()
|
|
|
|
# The grandchild writes `time.time()` every 20ms. Once the group is
|
|
# killed it stops writing, so the file contents stay frozen. Compare
|
|
# contents (not mtime) to stay independent of filesystem timestamp
|
|
# resolution.
|
|
time.sleep(0.2)
|
|
with open(beat_file, encoding="utf-8") as f:
|
|
first = f.read()
|
|
time.sleep(0.3)
|
|
with open(beat_file, encoding="utf-8") as f:
|
|
second = f.read()
|
|
self.assertEqual(
|
|
second,
|
|
first,
|
|
"grandchild kept running after stop() — process group leaked",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|