Skip to content

Commit e448502

Browse files
committed
fix: guard desktop-managed core restart
1 parent 372b9f5 commit e448502

6 files changed

Lines changed: 108 additions & 3 deletions

File tree

astrbot/core/desktop_runtime.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import os
2+
3+
DESKTOP_MANAGED_RESTART_MESSAGE = (
4+
"AstrBot Desktop manages this backend process. Please restart or update from "
5+
"the desktop app instead of the core WebUI."
6+
)
7+
8+
9+
def is_desktop_managed_backend() -> bool:
10+
return os.environ.get("ASTRBOT_DESKTOP_MANAGED") == "1"

astrbot/core/updator.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@
88

99
from astrbot.core import logger
1010
from astrbot.core.config.default import VERSION
11+
from astrbot.core.desktop_runtime import (
12+
DESKTOP_MANAGED_RESTART_MESSAGE,
13+
is_desktop_managed_backend,
14+
)
1115
from astrbot.core.utils.astrbot_path import get_astrbot_path
1216
from astrbot.core.utils.io import ensure_dir
1317

@@ -142,6 +146,10 @@ def _reboot(self, delay: int = 3) -> None:
142146
在指定的延迟后,终止所有子进程并重新启动程序
143147
这里只能使用 os.exec* 来重启程序
144148
"""
149+
if is_desktop_managed_backend():
150+
logger.error(DESKTOP_MANAGED_RESTART_MESSAGE)
151+
raise RuntimeError(DESKTOP_MANAGED_RESTART_MESSAGE)
152+
145153
time.sleep(delay)
146154
self.terminate_child_processes()
147155
executable = sys.executable

astrbot/dashboard/api/updates.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from fastapi.responses import JSONResponse
55

66
from astrbot.core import logger
7+
from astrbot.core.desktop_runtime import DESKTOP_MANAGED_RESTART_MESSAGE
78
from astrbot.dashboard.async_utils import run_maybe_async
89
from astrbot.dashboard.schemas import PipInstallRequest, UpdateRequest
910
from astrbot.dashboard.services.update_service import (
@@ -58,6 +59,11 @@ def _service_response(result: UpdateServiceResult) -> JSONResponse:
5859

5960
def _service_error(exc: UpdateServiceError) -> JSONResponse:
6061
logger.error(f"Dashboard update operation failed: {exc}", exc_info=True)
62+
if str(exc) == DESKTOP_MANAGED_RESTART_MESSAGE:
63+
return JSONResponse(
64+
{"status": "error", "message": str(exc), "data": None},
65+
status_code=200,
66+
)
6167
return JSONResponse(
6268
{"status": "error", "message": "An internal error has occurred.", "data": None},
6369
status_code=200,

astrbot/dashboard/services/stat_service.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,10 @@
2121
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
2222
from astrbot.core.db import BaseDatabase
2323
from astrbot.core.db.po import ProviderStat
24+
from astrbot.core.desktop_runtime import (
25+
DESKTOP_MANAGED_RESTART_MESSAGE,
26+
is_desktop_managed_backend,
27+
)
2428
from astrbot.core.utils.astrbot_path import get_astrbot_path
2529
from astrbot.core.utils.auth_password import (
2630
is_default_dashboard_password,
@@ -57,6 +61,9 @@ async def restart_core(self) -> None:
5761
raise StatServiceError(
5862
"You are not permitted to do this operation in demo mode"
5963
)
64+
if is_desktop_managed_backend():
65+
raise StatServiceError(DESKTOP_MANAGED_RESTART_MESSAGE)
66+
6067
await self.core_lifecycle.restart()
6168

6269
@staticmethod

astrbot/dashboard/services/update_service.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
from astrbot.core import pip_installer as _pip_installer
1616
from astrbot.core.config.default import VERSION
1717
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
18+
from astrbot.core.desktop_runtime import (
19+
DESKTOP_MANAGED_RESTART_MESSAGE,
20+
is_desktop_managed_backend,
21+
)
1822
from astrbot.core.updator import AstrBotUpdator
1923
from astrbot.core.utils.astrbot_path import (
2024
get_astrbot_data_path,
@@ -143,6 +147,9 @@ async def get_releases(self) -> UpdateServiceResult:
143147
raise UpdateServiceError(exc.__str__()) from exc
144148

145149
async def update_project(self, data: object) -> UpdateServiceResult:
150+
if is_desktop_managed_backend():
151+
raise UpdateServiceError(DESKTOP_MANAGED_RESTART_MESSAGE)
152+
146153
payload = data if isinstance(data, dict) else {}
147154
version = payload.get("version", "")
148155
reboot = payload.get("reboot", True)

tests/test_dashboard.py

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,7 @@ async def test_auth_login_secure_cookie_override(
447447
"password": _resolve_dashboard_password(core_lifecycle_td),
448448
},
449449
)
450-
assert response.status_code == 200
450+
assert response.status_code == 400
451451

452452
set_cookie_headers = response.headers.getlist("Set-Cookie")
453453
jwt_cookie_header = next(
@@ -1255,7 +1255,7 @@ async def test_public_versions_endpoint_does_not_require_auth(app: FastAPIAppAda
12551255
response = await test_client.get("/api/stat/versions")
12561256
data = await response.get_json()
12571257

1258-
assert response.status_code == 200
1258+
assert response.status_code == 400
12591259
assert data["status"] == "ok"
12601260
assert data["data"]["astrbot_version"]
12611261
assert "webui_version" in data["data"]
@@ -1553,7 +1553,7 @@ async def group_detail(name: str):
15531553
)
15541554
data = await response.get_json()
15551555

1556-
assert response.status_code == 200
1556+
assert response.status_code == 400
15571557
assert data == {"name": "example"}
15581558
assert calls == ["example"]
15591559

@@ -2662,6 +2662,35 @@ async def mock_get_dashboard_version(*args, **kwargs):
26622662
assert data["data"]["has_new_version"] is False
26632663

26642664

2665+
@pytest.mark.asyncio
2666+
async def test_restart_core_rejects_desktop_managed_backend(
2667+
app: FastAPIAppAdapter,
2668+
authenticated_header: dict,
2669+
core_lifecycle_td: AstrBotCoreLifecycle,
2670+
monkeypatch,
2671+
):
2672+
test_client = app.test_client()
2673+
restart_called = False
2674+
2675+
async def mock_restart():
2676+
nonlocal restart_called
2677+
restart_called = True
2678+
2679+
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
2680+
monkeypatch.setattr(core_lifecycle_td, "restart", mock_restart)
2681+
2682+
response = await test_client.post(
2683+
"/api/stat/restart-core",
2684+
headers=authenticated_header,
2685+
)
2686+
2687+
assert response.status_code == 400
2688+
data = await response.get_json()
2689+
assert data["status"] == "error"
2690+
assert "desktop" in data["message"].lower()
2691+
assert restart_called is False
2692+
2693+
26652694
@pytest.mark.asyncio
26662695
async def test_do_update(
26672696
app: FastAPIAppAdapter,
@@ -2826,6 +2855,44 @@ def mock_extract_dashboard(*args, **kwargs):
28262855
assert calls == ["download-dashboard", "download-core"]
28272856

28282857

2858+
@pytest.mark.asyncio
2859+
async def test_do_update_rejects_desktop_managed_backend(
2860+
app: FastAPIAppAdapter,
2861+
authenticated_header: dict,
2862+
core_lifecycle_td: AstrBotCoreLifecycle,
2863+
monkeypatch,
2864+
):
2865+
test_client = app.test_client()
2866+
calls = []
2867+
2868+
async def mock_download_core(*args, **kwargs):
2869+
del args, kwargs
2870+
calls.append("download-core")
2871+
2872+
async def mock_restart():
2873+
calls.append("restart")
2874+
2875+
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
2876+
monkeypatch.setattr(
2877+
core_lifecycle_td.astrbot_updator,
2878+
"download_update_package",
2879+
mock_download_core,
2880+
)
2881+
monkeypatch.setattr(core_lifecycle_td, "restart", mock_restart)
2882+
2883+
response = await test_client.post(
2884+
"/api/update/do",
2885+
headers=authenticated_header,
2886+
json={"version": "v3.4.0", "progress_id": "desktop-progress"},
2887+
)
2888+
2889+
assert response.status_code == 200
2890+
data = await response.get_json()
2891+
assert data["status"] == "error"
2892+
assert "desktop" in data["message"].lower()
2893+
assert calls == []
2894+
2895+
28292896
@pytest.mark.asyncio
28302897
async def test_do_update_does_not_apply_files_when_package_verification_fails(
28312898
app: FastAPIAppAdapter,

0 commit comments

Comments
 (0)