Add process controls and restart command
This commit is contained in:
2
.gitignore
vendored
2
.gitignore
vendored
@@ -5,3 +5,5 @@ sessions/
|
||||
*.acl.json
|
||||
config.json
|
||||
.env
|
||||
workspace/.process-logs/
|
||||
workspace/.background-processes.json
|
||||
|
||||
137
app/ai/processes.py
Normal file
137
app/ai/processes.py
Normal file
@@ -0,0 +1,137 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import signal
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ProcessRecord:
|
||||
name: str
|
||||
pid: int
|
||||
command: list[str]
|
||||
cwd: str
|
||||
log_file: str
|
||||
started_at: str
|
||||
|
||||
|
||||
class BackgroundProcessManager:
|
||||
def __init__(self, workspace_dir: Path) -> None:
|
||||
self.workspace_dir = workspace_dir.resolve()
|
||||
self.log_dir = self.workspace_dir / ".process-logs"
|
||||
self.state_file = self.workspace_dir / ".background-processes.json"
|
||||
self.processes: dict[str, ProcessRecord] = {}
|
||||
self._children: dict[str, asyncio.subprocess.Process] = {}
|
||||
self._load()
|
||||
|
||||
def _load(self) -> None:
|
||||
try:
|
||||
raw = json.loads(self.state_file.read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError):
|
||||
return
|
||||
for name, item in raw.items():
|
||||
try:
|
||||
item = dict(item)
|
||||
item.pop("name", None)
|
||||
self.processes[name] = ProcessRecord(name=name, **item)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
|
||||
def _save(self) -> None:
|
||||
self.workspace_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.state_file.write_text(
|
||||
json.dumps({name: asdict(record) for name, record in self.processes.items()}, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def _alive(self, pid: int) -> bool:
|
||||
if pid in {os.getpid(), 0}:
|
||||
return True
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except (OSError, ProcessLookupError):
|
||||
return False
|
||||
return True
|
||||
|
||||
async def start(self, name: str, command: Any, cwd: str | None = None) -> str:
|
||||
name = name.strip()
|
||||
if not name or not name.replace("-", "").replace("_", "").isalnum():
|
||||
raise ValueError("background_start: name must contain only letters, digits, '-' or '_'")
|
||||
if isinstance(command, str):
|
||||
command = shlex.split(command, posix=(os.name != "nt"))
|
||||
if not isinstance(command, list) or not command or not all(isinstance(x, str) for x in command):
|
||||
raise ValueError("background_start: command must be a non-empty argv list or shell-like string")
|
||||
if name in self.processes and self._alive(self.processes[name].pid):
|
||||
raise ValueError(f"background_start: process {name!r} is already running")
|
||||
|
||||
process_cwd = self.workspace_dir if not cwd else (self.workspace_dir / cwd).resolve()
|
||||
if process_cwd != self.workspace_dir and self.workspace_dir not in process_cwd.parents:
|
||||
raise ValueError("background_start: cwd must be inside workspace_dir")
|
||||
process_cwd.mkdir(parents=True, exist_ok=True)
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
log_file = self.log_dir / f"{name}.log"
|
||||
handle = log_file.open("ab")
|
||||
try:
|
||||
kwargs: dict[str, Any] = {
|
||||
"cwd": str(process_cwd),
|
||||
"stdout": handle,
|
||||
"stderr": asyncio.subprocess.STDOUT,
|
||||
}
|
||||
if os.name != "nt":
|
||||
kwargs["start_new_session"] = True
|
||||
process = await asyncio.create_subprocess_exec(*command, **kwargs)
|
||||
except Exception:
|
||||
handle.close()
|
||||
raise
|
||||
finally:
|
||||
handle.close()
|
||||
|
||||
record = ProcessRecord(
|
||||
name=name,
|
||||
pid=process.pid,
|
||||
command=command,
|
||||
cwd=str(process_cwd),
|
||||
log_file=str(log_file),
|
||||
started_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
self.processes[name] = record
|
||||
self._children[name] = process
|
||||
self._save()
|
||||
return f"Started {name} (pid={process.pid}); log={log_file}"
|
||||
|
||||
async def stop(self, name: str) -> str:
|
||||
record = self.processes.get(name)
|
||||
if record is None:
|
||||
raise ValueError(f"background_stop: unknown process {name!r}")
|
||||
process = self._children.get(name)
|
||||
if process is not None and process.returncode is None:
|
||||
process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=5)
|
||||
except asyncio.TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
elif self._alive(record.pid):
|
||||
try:
|
||||
os.kill(record.pid, signal.SIGTERM)
|
||||
except OSError as exc:
|
||||
raise ValueError(f"background_stop failed: {exc}") from exc
|
||||
self.processes.pop(name, None)
|
||||
self._children.pop(name, None)
|
||||
self._save()
|
||||
return f"Stopped {name} (pid={record.pid})"
|
||||
|
||||
def status(self) -> str:
|
||||
if not self.processes:
|
||||
return "No background processes"
|
||||
lines = []
|
||||
for name, record in sorted(self.processes.items()):
|
||||
state = "running" if self._alive(record.pid) else "exited"
|
||||
lines.append(f"{name}: {state}, pid={record.pid}, log={record.log_file}, command={record.command!r}")
|
||||
return "\n".join(lines)
|
||||
@@ -15,11 +15,15 @@ from openai import AsyncOpenAI
|
||||
from app.ai.parser import AIInvocation
|
||||
from app.ai.prompts import PromptBook
|
||||
from app.ai.page_extract import PageExtractor
|
||||
from app.ai.processes import BackgroundProcessManager
|
||||
from app.ai.tools import (
|
||||
FILE_TOOL_DISPATCH,
|
||||
FILE_TOOL_SCHEMAS,
|
||||
PAGE_EXTRACT_SCHEMA,
|
||||
PYTHON_EXEC_SCHEMA,
|
||||
BACKGROUND_LIST_SCHEMA,
|
||||
BACKGROUND_START_SCHEMA,
|
||||
BACKGROUND_STOP_SCHEMA,
|
||||
WEB_SEARCH_SCHEMA,
|
||||
)
|
||||
from app.ai.web_search import BraveWebSearchClient, format_results
|
||||
@@ -75,6 +79,7 @@ class AIService:
|
||||
self._clients: dict[str, AsyncOpenAI] = {}
|
||||
self.web_search = BraveWebSearchClient(config.web_search)
|
||||
self.page_extractor = PageExtractor(config.page_extract)
|
||||
self.processes = BackgroundProcessManager(config.workspace_dir)
|
||||
|
||||
async def initialize(self) -> None:
|
||||
return None
|
||||
@@ -221,7 +226,12 @@ class AIService:
|
||||
return messages
|
||||
|
||||
async def _build_tool_list(self) -> list[dict]:
|
||||
tools = FILE_TOOL_SCHEMAS + [PYTHON_EXEC_SCHEMA]
|
||||
tools = FILE_TOOL_SCHEMAS + [
|
||||
PYTHON_EXEC_SCHEMA,
|
||||
BACKGROUND_START_SCHEMA,
|
||||
BACKGROUND_STOP_SCHEMA,
|
||||
BACKGROUND_LIST_SCHEMA,
|
||||
]
|
||||
if self.web_search.available:
|
||||
tools.append(WEB_SEARCH_SCHEMA)
|
||||
if self.page_extractor.available:
|
||||
@@ -353,8 +363,18 @@ class AIService:
|
||||
return f"Tool {name!r}: invalid JSON arguments: {exc}"
|
||||
|
||||
try:
|
||||
if name == "python_exec":
|
||||
if name in {"python_exec", "background_start", "background_stop", "background_list"} and not ctx.is_owner:
|
||||
result = "Tool access denied: code execution and process control are owner-only"
|
||||
elif name == "python_exec":
|
||||
result = await self._execute_python_exec(arguments, ctx)
|
||||
elif name == "background_start":
|
||||
result = await self.processes.start(
|
||||
str(arguments.get("name", "")), arguments.get("command"), arguments.get("cwd")
|
||||
)
|
||||
elif name == "background_stop":
|
||||
result = await self.processes.stop(str(arguments.get("name", "")))
|
||||
elif name == "background_list":
|
||||
result = self.processes.status()
|
||||
elif name == "web_search":
|
||||
response = await self.web_search.search(arguments)
|
||||
result = format_results(response, include_heading=True)
|
||||
|
||||
@@ -219,6 +219,41 @@ PYTHON_EXEC_SCHEMA: dict[str, Any] = {
|
||||
},
|
||||
}
|
||||
|
||||
BACKGROUND_START_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "background_start",
|
||||
"description": "Start an owner-approved long-running process without a shell. The process survives bot restarts; output is written to workspace/.process-logs/.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "description": "Stable process name"},
|
||||
"command": {"type": "array", "items": {"type": "string"}, "description": "Executable argv, e.g. [\"python\", \"server.py\"]"},
|
||||
"cwd": {"type": "string", "description": "Optional directory relative to workspace"},
|
||||
},
|
||||
"required": ["name", "command"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
BACKGROUND_STOP_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "background_stop",
|
||||
"description": "Stop a background process previously started by background_start.",
|
||||
"parameters": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]},
|
||||
},
|
||||
}
|
||||
|
||||
BACKGROUND_LIST_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "background_list",
|
||||
"description": "List tracked background processes and their logs.",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
|
||||
WEB_SEARCH_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
|
||||
@@ -14,6 +14,7 @@ from app.commands.core import (
|
||||
shared_access_check,
|
||||
)
|
||||
from app.commands.dispatcher import CommandDispatcher
|
||||
from app.commands.persistent import load_persistent_commands
|
||||
|
||||
__all__ = [
|
||||
"CheckResult",
|
||||
@@ -31,4 +32,5 @@ __all__ = [
|
||||
"require_args_check",
|
||||
"require_reply_check",
|
||||
"shared_access_check",
|
||||
"load_persistent_commands",
|
||||
]
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import signal
|
||||
from textwrap import wrap
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -376,6 +378,17 @@ async def status(ctx: "CommandContext") -> None:
|
||||
await ctx.reply("\n".join(lines))
|
||||
|
||||
|
||||
@command_registry.command(
|
||||
"restart",
|
||||
"Stop the bot process so its supervisor can restart it",
|
||||
usage=".restart",
|
||||
category="system",
|
||||
)
|
||||
async def restart(ctx: "CommandContext") -> None:
|
||||
await ctx.reply("Stopping bot process for restart")
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
@command_registry.command(
|
||||
"id",
|
||||
"Show chat, thread, user, and replied message IDs",
|
||||
|
||||
32
app/commands/persistent.py
Normal file
32
app/commands/persistent.py
Normal file
@@ -0,0 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from app.commands.core import CommandRegistry
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def load_persistent_commands(commands_dir: Path, registry: CommandRegistry) -> list[str]:
|
||||
"""Load command modules from workspace/commands on every bot startup."""
|
||||
commands_dir.mkdir(parents=True, exist_ok=True)
|
||||
loaded: list[str] = []
|
||||
for path in sorted(commands_dir.glob("*.py")):
|
||||
if path.name.startswith("_"):
|
||||
continue
|
||||
module_name = f"bot_persistent_command_{path.stem}"
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(module_name, path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError("cannot create import spec")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
loaded.append(path.name)
|
||||
except Exception:
|
||||
sys.modules.pop(module_name, None)
|
||||
log.exception("Failed to load persistent command module %s", path)
|
||||
return loaded
|
||||
@@ -8,7 +8,7 @@ from typing import Any
|
||||
from pyrogram import Client, idle
|
||||
|
||||
from app.ai import AIService
|
||||
from app.commands import DEFAULT_REGISTRY, CommandRegistry, RuntimeAclStore
|
||||
from app.commands import DEFAULT_REGISTRY, CommandRegistry, RuntimeAclStore, load_persistent_commands
|
||||
from app.config import AccountConfig, AppConfig
|
||||
from app.github_webhook import GitHubWebhookServer
|
||||
from app.handlers import register_common_handlers
|
||||
@@ -37,6 +37,10 @@ class AccountManager:
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.command_registry = DEFAULT_REGISTRY
|
||||
if self.config.ai is not None:
|
||||
loaded = load_persistent_commands(self.config.ai.workspace_dir / "commands", self.command_registry)
|
||||
if loaded:
|
||||
log.info("Loaded persistent commands: %s", ", ".join(loaded))
|
||||
self.ai_service = AIService(self.config.ai) if self.config.ai is not None else None
|
||||
self.states = [AccountState(config=account) for account in self.config.accounts]
|
||||
self.github_webhook_server = GitHubWebhookServer(self, self.config.github_webhook)
|
||||
|
||||
54
app/pyrogram_compat.py
Normal file
54
app/pyrogram_compat.py
Normal file
@@ -0,0 +1,54 @@
|
||||
"""Small compatibility fixes for Kurigram versions with incomplete rich-text parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _plain_text(value: Any) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, (list, tuple)):
|
||||
return "".join(_plain_text(item) for item in value)
|
||||
nested = getattr(value, "text", None)
|
||||
if nested is not None and nested is not value:
|
||||
return _plain_text(nested)
|
||||
return str(value)
|
||||
|
||||
|
||||
def install_kurigram_compat() -> None:
|
||||
"""Patch nested hashtag-like entities in Kurigram 2.2.x.
|
||||
|
||||
Kurigram's parser assumes TextMention/TextHashtag/TextCashtag/TextBotCommand
|
||||
contain TextPlain. Telegram can legally nest formatting there, so the
|
||||
parsed RichText object is passed to ``lstrip`` and message parsing crashes.
|
||||
"""
|
||||
from pyrogram import raw, types
|
||||
from pyrogram.types.messages_and_media.rich_text import RichText
|
||||
|
||||
if getattr(RichText, "_nested_entity_compat", False):
|
||||
return
|
||||
|
||||
original = RichText._parse
|
||||
|
||||
async def patched(client: Any, rich_text: Any, users: dict = {}, chats: dict = {}):
|
||||
try:
|
||||
return await original(client, rich_text, users, chats)
|
||||
except AttributeError as exc:
|
||||
if "lstrip" not in str(exc):
|
||||
raise
|
||||
|
||||
handlers = (
|
||||
(raw.types.TextMention, types.RichTextMention, "@", "username"),
|
||||
(raw.types.TextHashtag, types.RichTextHashtag, "#", "hashtag"),
|
||||
(raw.types.TextCashtag, types.RichTextCashtag, "$", "cashtag"),
|
||||
(raw.types.TextBotCommand, types.RichTextBotCommand, "/", "bot_command"),
|
||||
)
|
||||
for raw_type, rich_type, prefix, field in handlers:
|
||||
if isinstance(rich_text, raw_type):
|
||||
content = await patched(client, rich_text.text, users, chats)
|
||||
return rich_type(text=content, **{field: _plain_text(content).lstrip(prefix)})
|
||||
raise
|
||||
|
||||
RichText._parse = staticmethod(patched)
|
||||
RichText._nested_entity_compat = True
|
||||
2
main.py
2
main.py
@@ -8,6 +8,7 @@ from pathlib import Path
|
||||
|
||||
from app.config import load_app_config
|
||||
from app.manager import AccountManager
|
||||
from app.pyrogram_compat import install_kurigram_compat
|
||||
|
||||
LOG_FORMAT = "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
|
||||
|
||||
@@ -37,6 +38,7 @@ _ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
install_kurigram_compat()
|
||||
config = load_app_config(_ROOT)
|
||||
manager = AccountManager(config)
|
||||
await manager.run()
|
||||
|
||||
6
prompts/ai/parts/40_execution.md
Normal file
6
prompts/ai/parts/40_execution.md
Normal file
@@ -0,0 +1,6 @@
|
||||
Инструменты выполнения доступны только владельцу аккаунта.
|
||||
|
||||
- Для коротких вычислений и проверки кода используй `python_exec`; в `ctx` доступны `chat_id`, `message_id`, `message_text`, `args`, `from_user` и `replied_message`.
|
||||
- Для долгоживущих задач (например, веб-сервера) используй `background_start` с argv-массивом без shell. Логи находятся в `workspace/.process-logs/`; для контроля используй `background_list` и `background_stop`.
|
||||
- Постоянные команды создавай как Python-модули в `workspace/commands/*.py`, импортируя `command_registry` и `CommandContext` из `app.commands`. После перезапуска они загружаются автоматически. По умолчанию команды owner-only; `.shared()` добавляй только по явной просьбе.
|
||||
- Перед изменением существующей команды прочитай файл и сохраняй совместимость с текущим API.
|
||||
@@ -1,4 +1,4 @@
|
||||
kurigram
|
||||
kurigram==2.2.24
|
||||
tgcrypto
|
||||
python-dotenv
|
||||
openai
|
||||
|
||||
24
tests/test_kurigram_compat.py
Normal file
24
tests/test_kurigram_compat.py
Normal file
@@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from app.pyrogram_compat import install_kurigram_compat
|
||||
|
||||
|
||||
class KurigramCompatTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_nested_hashtag_is_parsed_without_lstrip_crash(self) -> None:
|
||||
from pyrogram import raw
|
||||
from pyrogram.types.messages_and_media.rich_text import RichText
|
||||
|
||||
install_kurigram_compat()
|
||||
value = raw.types.TextHashtag(
|
||||
text=raw.types.TextBold(text=raw.types.TextPlain(text="#tag"))
|
||||
)
|
||||
parsed = await RichText._parse(None, value)
|
||||
|
||||
self.assertEqual(parsed.hashtag, "tag")
|
||||
self.assertEqual(parsed.text.text, "#tag")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
23
tests/test_restart_command.py
Normal file
23
tests/test_restart_command.py
Normal file
@@ -0,0 +1,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import signal
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from app.commands.builtins import restart
|
||||
|
||||
|
||||
class RestartCommandTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_restart_confirms_then_stops_current_process(self) -> None:
|
||||
ctx = type("Context", (), {"reply": AsyncMock()})()
|
||||
|
||||
with patch("app.commands.builtins.os.kill") as kill:
|
||||
await restart(ctx)
|
||||
|
||||
ctx.reply.assert_awaited_once_with("Stopping bot process for restart")
|
||||
kill.assert_called_once_with(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user