magpie/internal/gui/update_wait.go

package gui

import (
	"time"

	"github.com/yetone/magpie/internal/gateway"
)

// A restart to update ends the gateway this process serves, and with it
// whatever the agents have in flight through it (#577: an agent's reply
// streaming, a Claude Code turn waiting on its caller's tool results). So a
// restart asked for while the gateway is busy waits until it is idle:
// nothing in flight, no turn waiting on tools, and nothing ended in the
// last idleSettle — the gap of a tool chain, an agent running a tool
// between one request and the next. It waits waitMost at most; then the
// update stays downloaded, put in on quitting or at the next restart, and
// a newer version out meanwhile is the one restarted into. "Restart now"
// restarts at once, whatever is in flight.
var (
	idleSettle = 10 * time.Second
	waitMost   = time.Hour
	waitTick   = time.Second
)

// gatewayBusy is what the gateway this process serves has in flight;
// nothing when another magpie serves it, which a restart here leaves be.
var gatewayBusy = func() gateway.Busy {
	if gw := served.Load(); gw != nil {
		return gw.Busy()
	}
	return gateway.Busy{}
}

// waitRecheck asks the feed again before the restart (updater.recheck).
var waitRecheck = (*updater).recheck

// idleAt is whether b leaves nothing for a restart at now to cut short.
func idleAt(b gateway.Busy, now time.Time) bool {
	return !b.Any() && now.Sub(b.Last) >= idleSettle
}

// idle is whether the gateway is idle now.
func (u *updater) idle() bool { return idleAt(gatewayBusy(), time.Now()) }

// waitIdle has restart run once the gateway is idle; a wait under way
// takes the newer restart (the window as it is now).
func (u *updater) waitIdle(restart func() bool) {
	u.mu.Lock()
	u.whenIdle, u.gaveUp = restart, false
	if u.waiting {
		u.mu.Unlock()
		return
	}
	u.waiting, u.waitFrom = true, time.Now()
	u.waitGen++
	gen := u.waitGen
	u.mu.Unlock()
	u.told()
	u.loops.Add(1)
	go func() {
		defer u.loops.Done()
		u.waitFor(gen)
	}()
}

// cancelWait stops a restart waiting.
func (u *updater) cancelWait() {
	u.mu.Lock()
	was := u.waiting
	u.waiting, u.whenIdle = false, nil
	u.waitGen++
	u.mu.Unlock()
	if was {
		u.told()
	}
}

// waitingFor is whether a restart waits, and on what.
func (u *updater) waitingFor() (bool, gateway.Busy) {
	u.mu.Lock()
	w := u.waiting
	u.mu.Unlock()
	if !w {
		return false, gateway.Busy{}
	}
	return true, gatewayBusy()
}

func (u *updater) told() {
	u.mu.Lock()
	f := u.onWait
	u.mu.Unlock()
	if f != nil {
		go f() // off the caller's thread: the tray's relabel runs on the main one
	}
}

// waitFor restarts once the gateway is idle and the version downloaded is
// the latest, unless the wait is called off (gen no longer its) or runs out.
func (u *updater) waitFor(gen int) {
	var last gateway.Busy
	for {
		time.Sleep(waitTick)
		u.mu.Lock()
		if !u.waiting || u.waitGen != gen {
			u.mu.Unlock()
			return
		}
		if time.Since(u.waitFrom) >= waitMost {
			u.waiting, u.gaveUp, u.whenIdle = false, true, nil
			u.mu.Unlock()
			u.told()
			return
		}
		u.mu.Unlock()
		b := gatewayBusy()
		if b.Requests != last.Requests || b.Tools != last.Tools {
			last = b
			u.told()
		}
		if !idleAt(b, time.Now()) {
			continue
		}
		switch u.json().State {
		case "ready":
			// a newer version out since is downloaded first, and waited for
			waitRecheck(u)
			if u.json().State != "ready" {
				continue
			}
		case "checking", "downloading":
			continue
		default:
			// nothing left to restart into (a newer one failed to download)
			u.cancelWait()
			return
		}
		u.mu.Lock()
		if !u.waiting || u.waitGen != gen {
			u.mu.Unlock()
			return
		}
		restart := u.whenIdle
		u.waiting, u.whenIdle = false, nil
		u.mu.Unlock()
		u.told()
		if restart != nil {
			restart()
		}
		return
	}
}