onw/tray.py

"""onw AMD tray: the resident icon (Windows task tray / Ubuntu top-bar indicator).

  onw tray      icon + menu: open the onw window (models / server / settings), chat UI, check page, copy the
                OpenAI API URL, start / stop the server. With no model downloaded yet it opens the model manager.

The server itself runs as a separate process managed by onw.app (so the window and the tray share it). The tray
also hosts the manager (onw.manager: downloads, model list, settings) that the window talks to, so downloads keep
going while the window is closed.
"""
import json, os, subprocess, sys, threading, time, webbrowser

from . import app, __version__


def open_window(tab="server", updated=None):
    exe = sys.executable
    if os.name == "nt" and exe.lower().endswith("python.exe"):
        exe = exe[:-10] + "pythonw.exe"
    subprocess.Popen([exe, "-m", "onw", "window", "--tab", tab] + (["--updated", updated] if updated else []))


def icon_image(state):
    from . import logo
    return logo.draw(64, simple=True, state=state)          # the small-size logo + a status dot (grey when stopped)


def copy_text(text):
    """Copy to the clipboard: clip (Windows), wl-copy (Wayland) or xclip / xsel (X11). False when none is there."""
    if os.name == "nt":
        cmds = [(["clip"], text.encode("utf-16-le"))]
    else:
        cmds = [(["wl-copy"], text.encode()), (["xclip", "-selection", "clipboard"], text.encode()),
                (["xsel", "--clipboard", "--input"], text.encode())]
    for cmd, data in cmds:
        try:
            if subprocess.run(cmd, input=data, check=False, timeout=5).returncode == 0:
                return True
        except (OSError, subprocess.SubprocessError):
            continue
    return False


def replace_running():
    """The installer ran again while a tray is running (an update by the one-line installer): ask that tray to restart
    into the new files, as the in-app update does (the model that was running is loaded again). False when no
    tray answered."""
    import urllib.request
    from . import manager as mgr
    d = mgr.find_manager()
    if not d:
        return False
    try:
        req = urllib.request.Request(f"http://127.0.0.1:{d['port']}/api/restart_tray", b"{}",
                                     {"X-Onw-Token": d["token"], "Content-Type": "application/json"})
        return "error" not in json.loads(urllib.request.urlopen(req, timeout=10).read())
    except Exception:
        return False


def take_over():
    """A running tray that cannot restart itself (older than 0.1.3): stop it and its server, remember the model for the
    new tray to load again, and wait for its single-instance lock. True once the lock is ours."""
    from . import manager as mgr
    info = app.server_info()
    with open(os.path.join(app.config_dir(), "resume.json"), "w") as f:
        json.dump({"model": info["model"] if info else None}, f)
    app.stop_server()
    d = mgr.find_manager()
    if d:
        try:
            import psutil
            p = psutil.Process(d["pid"])
            procs = p.children(recursive=True) + [p]
            for q in procs:
                q.terminate()
            psutil.wait_procs(procs, timeout=10)
        except Exception:
            pass
    for _ in range(40):
        if app.single_instance("tray"):
            return True
        time.sleep(0.5)
    return False


def main(argv=None):
    argv = argv or []
    if not app.single_instance("tray"):                       # already in the tray (e.g. started again from the menu)
        if "--replace" in argv:                                # the installer ran again: the files are new
            if replace_running():                              # the old tray restarts itself on the new code
                return
            if not take_over():                                # an older tray: stop it and start here instead
                open_window()
                return
        else:
            open_window()
            return
    import pystray
    from . import manager as mgr
    app.clear_stale_server()                                   # a pid file from before a crash / reboot
    manager = mgr.Manager()
    mgr.serve(manager)                                         # the window's backend lives here
    from .i18n import T
    labels = {"stopped": "停止中", "loading": "ロード中…", "running": "動作中", "failed": "ロードに失敗"}
    state = {"s": app.server_status()}

    def cfg():
        return app.load_config()

    def status_text(_=None):
        info = app.server_info()
        name = (info or {}).get("model") or cfg().get("model") or T("モデル未選択")
        return f"onw AMD: {T(labels[state['s']])} — {name.split('/')[-1]}"

    def start(_=None, __=None):
        c = cfg()
        if not app.installed(c):
            open_window("models")
            return
        if c["model"] not in app.installed(c):
            c = app.update_config(model=next(iter(app.installed(c))))
        try:
            app.start_server(c)
        except Exception:                                      # e.g. the port in use: the window says why
            open_window("server")

    def quit_(icon, __=None):
        app.stop_server()
        try:                                                   # manager.json points at this process: a later
            os.remove(mgr.INFO())                              # `onw window` then starts its own manager
        except OSError:
            pass
        icon.stop()

    def restart_after_update(icon_=None):
        """Stop the server, start a fresh tray on the new code, reload the model that was running."""
        info = app.server_info()
        with open(os.path.join(app.config_dir(), "resume.json"), "w") as f:
            json.dump({"model": info["model"] if info else None}, f)
        app.stop_server()
        subprocess.Popen([sys.executable, "-m", "onw.relaunch", str(os.getpid())],
                         cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                         **({"creationflags": subprocess.CREATE_NO_WINDOW | subprocess.DETACHED_PROCESS} if os.name == "nt"
                            else {"start_new_session": True}))
        quit_(icon)

    def update_found(v):
        try:
            icon.notify(T("「更新」を押すと自動で入れ替えて再起動します。"), T("onw v{0} が使えます", v))
        except Exception:                                      # not every indicator backend can notify
            pass
        try:
            icon.update_menu()
        except Exception:
            pass

    def do_update(_=None, __=None):
        r = manager.run_update()
        if r.get("error"):
            open_window("server")

    manager.on_updated = lambda: restart_after_update()
    manager.on_restart = lambda: restart_after_update()        # the installer ran again (onw tray --replace)
    manager.on_update_found = update_found
    running = lambda _: state["s"] == "running"
    menu = pystray.Menu(
        pystray.MenuItem(status_text, None, enabled=False),
        pystray.MenuItem(lambda _: T("onw を開く（モデル・サーバー・設定）"), lambda *_: open_window(), default=True),
        pystray.MenuItem(lambda _: T("チャットを開く"), lambda *_: webbrowser.open(f"http://localhost:{cfg()['port']}/chat"), enabled=running),
        pystray.MenuItem(lambda _: T("検証ページを開く"), lambda *_: webbrowser.open(f"http://localhost:{cfg()['port']}/check"), enabled=running),
        pystray.MenuItem(lambda _: T("OpenAI API のURLをコピー"),
                         lambda *_: copy_text(app.api_url(cfg())) or open_window("server")),   # no clipboard tool: shown there
        pystray.MenuItem(lambda _: T("onw を更新（v{0}）", manager.upd["latest"]), do_update,
                         visible=lambda _: bool(manager.upd["latest"]) and manager.upd["state"] != "updating"),
        pystray.MenuItem(lambda _: T("モデルをロード"), start, visible=lambda _: state["s"] in ("stopped", "failed")),
        pystray.MenuItem(lambda _: T("アンロード（サーバー停止）"), lambda *_: app.stop_server(), visible=lambda _: state["s"] in ("loading", "running")),
        pystray.MenuItem(lambda _: T("終了（サーバーも停止）"), quit_))
    icon = pystray.Icon("onw-amd", icon_image(state["s"]), app.title(), menu)

    def refresh():
        last = None
        while True:
            state["s"] = app.server_status()
            if state["s"] != last:
                last = state["s"]
                icon.icon = icon_image(last)
                icon.title = status_text()
                try:
                    icon.update_menu()
                except Exception:
                    pass
            time.sleep(2)

    c = cfg()
    resume = os.path.join(app.config_dir(), "resume.json")
    resumed = None
    after_update = os.path.exists(resume)
    if after_update:                                 # just restarted by an update: reload that model
        try:
            resumed = json.load(open(resume)).get("model")
        except (OSError, ValueError):
            pass
        os.remove(resume)
    if not app.installed(c):
        open_window("models")                                  # first run: pick a model to download
    elif resumed and state["s"] == "stopped" and resumed in app.installed(c):
        try:
            app.start_server(app.update_config(model=resumed))
        except Exception:
            open_window("server")
    elif c["autostart_server"] and state["s"] == "stopped":
        start()
    if after_update and app.installed(c):                     # show the new version (an open window reopens itself
        open_window("server", updated=__version__)             # too; the single-instance lock keeps just one)
    threading.Thread(target=refresh, daemon=True).start()
    threading.Thread(target=app.refresh_shortcut, daemon=True).start()
    icon.run()