From ef9392d5fd5b985c94ac12c29ab041d6de72617d Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sat, 13 Jun 2026 02:13:06 +0000
Subject: [PATCH] Add `og-veil restart` command
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds a `restart` subcommand that stops the background server (if running),
waits for the old process to exit, then starts a fresh one. Useful after
`og-veil update` to load the new version without typing `stop && og-veil`.

- New `wait_until_stopped(pid)` daemon helper polls the live process (the
  pidfile is already removed by stop_background) so the new server doesn't
  collide with the old one on the same port.
- `restart` mirrors `serve`'s flags (host/port/tee-id/expected-pcr/pii-scrub)
  and reuses the saved login, so it never prompts.
- Update hint, docstrings, and README now point at `og-veil restart`.
- Tests cover the stop→wait→start ordering, the no-server-running path, the
  lingering-process error, and the wait helper itself.

https://claude.ai/code/session_01TYNjSYZ6QoreCr2dcURYSH
---
 README.md            |  6 +++--
 tests/test_cli.py    | 39 ++++++++++++++++++++++++++++
 tests/test_daemon.py | 17 ++++++++++++
 uv.lock              |  2 +-
 veil/cli.py          | 62 +++++++++++++++++++++++++++++++++++++++++---
 veil/daemon.py       | 26 +++++++++++++++++++
 6 files changed, 146 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 2b25625..2011fbf 100644
--- a/README.md
+++ b/README.md
@@ -61,8 +61,9 @@ That's it. Every response is verified before you see it — check the
 on the body). Streaming works too; it's verified before the first token replays.
 
 Useful commands: `og-veil test` (send a one-off prompt to check the path),
-`og-veil stop`, `og-veil status`, `og-veil env` (re-prints the env vars),
-`og-veil models` (list available models), `og-veil update`, `og-veil logout`.
+`og-veil stop`, `og-veil restart` (after an update), `og-veil status`,
+`og-veil env` (re-prints the env vars), `og-veil models` (list available
+models), `og-veil update`, `og-veil logout`.
 
 ### Use it with Hermes Agent
 
@@ -144,6 +145,7 @@ the local OpenAI-compatible server.
 |---------|--------------|
 | `og-veil` | Set up on first run, then serve (detached). The one command you need. |
 | `og-veil stop` | Stop the background server. |
+| `og-veil restart` | Stop and start the background server — e.g. after `og-veil update`. |
 | `og-veil status` | Login + network config + whether the server is running. |
 | `og-veil test ["prompt"]` | Send a one-off prompt to the running server and print the verified reply. |
 | `og-veil update` | Update og-veil to the latest version. |
diff --git a/tests/test_cli.py b/tests/test_cli.py
index b4ffd62..d842b41 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -129,3 +129,42 @@ def test_update_surfaces_failure():
         result = CliRunner().invoke(cli.main, ["update"])
     assert result.exit_code != 0
     assert "update failed" in result.output
+
+
+def test_restart_stops_waits_then_starts():
+    with (
+        mock.patch("veil.daemon.stop_background", return_value=4321) as stop,
+        mock.patch("veil.daemon.wait_until_stopped", return_value=True) as wait,
+        mock.patch.object(cli, "_start_server") as start,
+    ):
+        result = CliRunner().invoke(cli.main, ["restart"])
+    assert stop.called
+    assert wait.call_args.args[0] == 4321, "should wait on the stopped pid"
+    assert start.called and start.call_args.kwargs["foreground"] is False
+    assert "Stopped background server (pid 4321)" in result.output
+    assert result.exit_code == 0
+
+
+def test_restart_starts_fresh_when_nothing_running():
+    with (
+        mock.patch("veil.daemon.stop_background", return_value=None),
+        mock.patch("veil.daemon.wait_until_stopped") as wait,
+        mock.patch.object(cli, "_start_server") as start,
+    ):
+        result = CliRunner().invoke(cli.main, ["restart"])
+    assert not wait.called, "no running server → nothing to wait for"
+    assert start.called
+    assert "No background server was running" in result.output
+    assert result.exit_code == 0
+
+
+def test_restart_errors_if_old_process_lingers():
+    with (
+        mock.patch("veil.daemon.stop_background", return_value=99),
+        mock.patch("veil.daemon.wait_until_stopped", return_value=False),
+        mock.patch.object(cli, "_start_server") as start,
+    ):
+        result = CliRunner().invoke(cli.main, ["restart"])
+    assert not start.called, "must not start a new server while the old one lingers"
+    assert result.exit_code != 0
+    assert "did not exit" in result.output
diff --git a/tests/test_daemon.py b/tests/test_daemon.py
index bec6cc3..a4f0e72 100644
--- a/tests/test_daemon.py
+++ b/tests/test_daemon.py
@@ -64,3 +64,20 @@ def test_stop_command(home):
         result = CliRunner().invoke(cli.main, ["stop"])
     assert "4321" in result.output
     assert result.exit_code == 0
+
+
+def test_wait_until_stopped_returns_true_once_process_gone():
+    # Alive on the first poll, gone on the second.
+    with (
+        mock.patch("os.kill", side_effect=[None, OSError]),
+        mock.patch("time.sleep"),
+    ):
+        assert daemon.wait_until_stopped(4321, timeout=1.0, interval=0.0) is True
+
+
+def test_wait_until_stopped_times_out_if_process_lingers():
+    with (
+        mock.patch("os.kill"),  # always alive
+        mock.patch("time.sleep"),
+    ):
+        assert daemon.wait_until_stopped(4321, timeout=0.0) is False
diff --git a/uv.lock b/uv.lock
index b4967bf..3fb649d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2329,7 +2329,7 @@ wheels = [
 
 [[package]]
 name = "opengradient-veil"
-version = "0.2.5"
+version = "0.2.6"
 source = { editable = "." }
 dependencies = [
     { name = "click" },
diff --git a/veil/cli.py b/veil/cli.py
index 16f6118..5aec652 100644
--- a/veil/cli.py
+++ b/veil/cli.py
@@ -2,8 +2,8 @@
 
 The common path is a single command: run ``og-veil`` and it logs you in on first
 use, then starts the local server in the background. Individual steps (``serve``,
-``login``, ``stop``, ``status``, ``env``, ``models``, ``test``, ``update``,
-``logout``) are available on their own too.
+``login``, ``stop``, ``restart``, ``status``, ``env``, ``models``, ``test``,
+``update``, ``logout``) are available on their own too.
 """
 
 from __future__ import annotations
@@ -191,6 +191,62 @@ def stop() -> None:
         click.secho(f"✓ Stopped background server (pid {pid}).", fg="green")
 
 
+@main.command()
+@click.option("--host", default=None, help="Bind host (default 127.0.0.1 / OG_VEIL_HOST).")
+@click.option("--port", type=int, default=None, help="Bind port (default 11434 / OG_VEIL_PORT).")
+@click.option("--tee-id", default=None, help="Pin a specific tee_id from the registry.")
+@click.option(
+    "--expected-pcr", default=None, help="Refuse any TEE whose registry pcrHash differs from this."
+)
+@click.option(
+    "--pii-scrub",
+    is_flag=True,
+    default=False,
+    help="Redact high-impact PII (email, SSN, bank numbers; addresses with the [pii] extra) "
+    "from prompts locally before they leave this machine.",
+)
+def restart(
+    host: str | None,
+    port: int | None,
+    tee_id: str | None,
+    expected_pcr: str | None,
+    pii_scrub: bool,
+) -> None:
+    """Stop the background server (if running) and start it again.
+
+    Handy after ``og-veil update`` to load the new version. Login is reused, so
+    this never prompts. Flags work the same as ``og-veil serve``; without them
+    the server restarts with its environment-based config.
+    """
+    from veil.daemon import stop_background, wait_until_stopped
+
+    pid = stop_background()
+    if pid is None:
+        click.echo("No background server was running — starting a fresh one.")
+    else:
+        click.secho(f"✓ Stopped background server (pid {pid}).", fg="green")
+        if not wait_until_stopped(pid):
+            raise click.ClickException(
+                f"the previous server (pid {pid}) did not exit — try `og-veil stop` again"
+            )
+
+    config = ServerConfig.from_env()
+    if host:
+        config.host = host
+    if port:
+        config.port = port
+    if tee_id:
+        config.pinned_tee_id = tee_id if tee_id.startswith("0x") else "0x" + tee_id
+    if expected_pcr:
+        config.expected_pcr_hash = (
+            expected_pcr if expected_pcr.startswith("0x") else "0x" + expected_pcr
+        ).lower()
+    if pii_scrub:
+        config.pii_scrub = True
+
+    _start_server(config, foreground=False)
+
+
 @main.command(name="env")
 def env_cmd() -> None:
     """Print the env vars to point your agent at OpenGradient Veil."""
@@ -338,7 +394,7 @@ def update() -> None:
         raise click.ClickException(
             f"update failed: {exc}\nTry manually, e.g.:  uv tool upgrade opengradient-veil"
         )
-    click.secho("✓ Updated. Restart the server to pick it up:  og-veil stop && og-veil", fg="green")
+    click.secho("✓ Updated. Restart the server to pick it up:  og-veil restart", fg="green")
 
 
 @main.command()
diff --git a/veil/daemon.py b/veil/daemon.py
index 2d1d841..ca681d3 100644
--- a/veil/daemon.py
+++ b/veil/daemon.py
@@ -76,3 +76,29 @@ def stop_background() -> int | None:
         pass
     pid_path().unlink(missing_ok=True)
     return pid
+
+
+def _pid_alive(pid: int) -> bool:
+    try:
+        os.kill(pid, 0)  # signal 0 just checks the process exists
+    except OSError:
+        return False
+    return True
+
+
+def wait_until_stopped(pid: int, timeout: float = 5.0, interval: float = 0.1) -> bool:
+    """Block until ``pid`` has exited, or ``timeout`` seconds elapse.
+
+    Returns True once the process is gone, False if it was still alive at the
+    deadline. Used by ``og-veil restart`` so the new server doesn't collide with
+    the old one still occupying the same port. The pidfile is already gone by
+    this point, so we poll the process directly.
+    """
+    import time
+
+    deadline = time.monotonic() + timeout
+    while _pid_alive(pid):
+        if time.monotonic() >= deadline:
+            return False
+        time.sleep(interval)
+    return True