diff --git a/nanobot/channels/manager.py b/nanobot/channels/manager.py index 7110311b..ccac6306 100644 --- a/nanobot/channels/manager.py +++ b/nanobot/channels/manager.py @@ -172,6 +172,7 @@ class ChannelManager: channel=notice.channel, chat_id=notice.chat_id, content=format_restart_completed_message(notice.started_at_raw), + metadata=dict(notice.metadata or {}), ), )) diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 87d4bf64..f46d7dbd 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -28,7 +28,11 @@ async def cmd_stop(ctx: CommandContext) -> OutboundMessage: async def cmd_restart(ctx: CommandContext) -> OutboundMessage: """Restart the process in-place via os.execv.""" msg = ctx.msg - set_restart_notice_to_env(channel=msg.channel, chat_id=msg.chat_id) + set_restart_notice_to_env( + channel=msg.channel, + chat_id=msg.chat_id, + metadata=dict(msg.metadata or {}), + ) async def _do_restart(): await asyncio.sleep(1) diff --git a/nanobot/utils/restart.py b/nanobot/utils/restart.py index 35b8cced..871667f0 100644 --- a/nanobot/utils/restart.py +++ b/nanobot/utils/restart.py @@ -2,12 +2,15 @@ from __future__ import annotations +import json import os import time -from dataclasses import dataclass +from dataclasses import dataclass, field +from typing import Any RESTART_NOTIFY_CHANNEL_ENV = "NANOBOT_RESTART_NOTIFY_CHANNEL" RESTART_NOTIFY_CHAT_ID_ENV = "NANOBOT_RESTART_NOTIFY_CHAT_ID" +RESTART_NOTIFY_METADATA_ENV = "NANOBOT_RESTART_NOTIFY_METADATA" RESTART_STARTED_AT_ENV = "NANOBOT_RESTART_STARTED_AT" @@ -16,6 +19,7 @@ class RestartNotice: channel: str chat_id: str started_at_raw: str + metadata: dict[str, Any] = field(default_factory=dict) def format_restart_completed_message(started_at_raw: str) -> str: @@ -30,11 +34,20 @@ def format_restart_completed_message(started_at_raw: str) -> str: return f"Restart completed{elapsed_suffix}." -def set_restart_notice_to_env(*, channel: str, chat_id: str) -> None: +def set_restart_notice_to_env( + *, channel: str, chat_id: str, metadata: dict[str, Any] | None = None, +) -> None: """Write restart notice env values for the next process.""" os.environ[RESTART_NOTIFY_CHANNEL_ENV] = channel os.environ[RESTART_NOTIFY_CHAT_ID_ENV] = chat_id os.environ[RESTART_STARTED_AT_ENV] = str(time.time()) + if metadata: + try: + os.environ[RESTART_NOTIFY_METADATA_ENV] = json.dumps(metadata, default=str) + except (TypeError, ValueError): + os.environ.pop(RESTART_NOTIFY_METADATA_ENV, None) + else: + os.environ.pop(RESTART_NOTIFY_METADATA_ENV, None) def consume_restart_notice_from_env() -> RestartNotice | None: @@ -42,9 +55,23 @@ def consume_restart_notice_from_env() -> RestartNotice | None: channel = os.environ.pop(RESTART_NOTIFY_CHANNEL_ENV, "").strip() chat_id = os.environ.pop(RESTART_NOTIFY_CHAT_ID_ENV, "").strip() started_at_raw = os.environ.pop(RESTART_STARTED_AT_ENV, "").strip() + metadata_raw = os.environ.pop(RESTART_NOTIFY_METADATA_ENV, "").strip() if not (channel and chat_id): return None - return RestartNotice(channel=channel, chat_id=chat_id, started_at_raw=started_at_raw) + metadata: dict[str, Any] = {} + if metadata_raw: + try: + parsed = json.loads(metadata_raw) + except (TypeError, ValueError): + parsed = None + if isinstance(parsed, dict): + metadata = parsed + return RestartNotice( + channel=channel, + chat_id=chat_id, + started_at_raw=started_at_raw, + metadata=metadata, + ) def should_show_cli_restart_notice(notice: RestartNotice, session_id: str) -> bool: diff --git a/tests/utils/test_restart.py b/tests/utils/test_restart.py index 48124d38..84427215 100644 --- a/tests/utils/test_restart.py +++ b/tests/utils/test_restart.py @@ -16,6 +16,7 @@ from nanobot.utils.restart import ( def test_set_and_consume_restart_notice_env_roundtrip(monkeypatch): monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_CHANNEL", raising=False) monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_CHAT_ID", raising=False) + monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_METADATA", raising=False) monkeypatch.delenv("NANOBOT_RESTART_STARTED_AT", raising=False) set_restart_notice_to_env(channel="feishu", chat_id="oc_123") @@ -25,14 +26,42 @@ def test_set_and_consume_restart_notice_env_roundtrip(monkeypatch): assert notice.channel == "feishu" assert notice.chat_id == "oc_123" assert notice.started_at_raw + assert notice.metadata == {} # Consumed values should be cleared from env. assert consume_restart_notice_from_env() is None assert "NANOBOT_RESTART_NOTIFY_CHANNEL" not in os.environ assert "NANOBOT_RESTART_NOTIFY_CHAT_ID" not in os.environ + assert "NANOBOT_RESTART_NOTIFY_METADATA" not in os.environ assert "NANOBOT_RESTART_STARTED_AT" not in os.environ +def test_restart_notice_preserves_metadata_across_env(monkeypatch): + monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_CHANNEL", raising=False) + monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_CHAT_ID", raising=False) + monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_METADATA", raising=False) + monkeypatch.delenv("NANOBOT_RESTART_STARTED_AT", raising=False) + + set_restart_notice_to_env( + channel="slack", + chat_id="C123", + metadata={"slack": {"thread_ts": "1700.42", "channel_type": "channel"}}, + ) + + notice = consume_restart_notice_from_env() + assert notice is not None + assert notice.metadata == { + "slack": {"thread_ts": "1700.42", "channel_type": "channel"} + } + assert "NANOBOT_RESTART_NOTIFY_METADATA" not in os.environ + + +def test_restart_notice_clears_stale_metadata(monkeypatch): + monkeypatch.setenv("NANOBOT_RESTART_NOTIFY_METADATA", '{"stale": true}') + set_restart_notice_to_env(channel="cli", chat_id="direct") + assert "NANOBOT_RESTART_NOTIFY_METADATA" not in os.environ + + def test_format_restart_completed_message_with_elapsed(monkeypatch): monkeypatch.setattr("nanobot.utils.restart.time.time", lambda: 102.0) assert format_restart_completed_message("100.0") == "Restart completed in 2.0s."