@pi-in-go/pigpen-ahp 0.1.0

package/extensions/ahp/runtime.go

@@ -0,0 +1,311 @@
+package ahp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/url"
+	"os"
+	"path/filepath"
+	"strconv"
+	"sync"
+
+	sdk "github.com/MichaelKinsy/PiG/extensions/sdk"
+	ahptypes "github.com/microsoft/agent-host-protocol/clients/go/ahptypes"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/compose"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/live"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/pi"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/settings"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/wire"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/ws"
+)
+
+// Version is reported to clients as the server version.
+const Version = "0.1.0"
+
+// runtime owns the (optional) listener of one PiG process.
+type runtime struct {
+	live *live.Session
+
+	mu      sync.Mutex
+	server  *ws.Server
+	built   *compose.Built
+	addr    string
+	ctx     sdk.Context
+	haveCtx bool
+}
+
+func newRuntime() *runtime { return &runtime{live: live.New()} }
+
+func (rt *runtime) status() string {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+	if rt.server == nil {
+		return "AHP listener is off. /ahp start opens it (settings: " + settingsPath(nil) + ")"
+	}
+	return rt.addr
+}
+
+// listenNotice is the message printed when the listener is up. The first address is the form VS Code takes
+// (ws://<host>:<port>?tkn=<token>) and is followed by a space: VS Code's SSH launcher (chat.sshRemoteAgentHostCommand) and
+// "Agents: Add Remote Agent Host..." both read the first ws:// URL in the text, and the launcher's token runs to the next
+// space or '&'. The generic form (/?token=) follows for other AHP clients; the listener takes either. hostport is the
+// listener's address. A listener without a token has one address and no token form.
+func listenNotice(hostport, token string) string {
+	if token == "" {
+		return "AHP listening on ws://" + hostport
+	}
+	escaped := url.QueryEscape(token)
+	return "AHP listening on ws://" + hostport + "?tkn=" + escaped +
+		" - for VS Code, paste it into \"Agents: Add Remote Agent Host...\"; other AHP clients can also use ws://" + hostport + "/?token=" + escaped
+}
+
+// settingsPath is the settings file: the --ahp-settings flag, else <pig home>/ahp/settings.json.
+func settingsPath(ctx *sdk.Context) string {
+	if ctx != nil {
+		if v, err := ctx.GetFlag(flagSettings); err == nil {
+			if path, ok := v.(string); ok && path != "" {
+				return path
+			}
+		}
+	}
+	if home := os.Getenv("PIG_HOME"); home != "" {
+		return filepath.Join(home, "ahp", "settings.json")
+	}
+	if dir, err := os.UserHomeDir(); err == nil {
+		return filepath.Join(dir, ".pig", "ahp", "settings.json")
+	}
+	return filepath.Join(".pig", "ahp", "settings.json")
+}
+
+// onSessionStart opens the listener when --ahp was given. PiG emits session_shutdown before a
+// replacement session starts (and re-runs the extension factory), so a listener opened with
+// /ahp start has already been closed here; only --ahp brings it back for the new session.
+func (rt *runtime) onSessionStart(ctx sdk.Context, id, cwd, reason string) {
+	rt.mu.Lock()
+	rt.ctx, rt.haveCtx = ctx, true
+	rt.mu.Unlock()
+	requested := false
+	if v, err := ctx.GetFlag(flagStart); err == nil {
+		requested, _ = v.(bool)
+	}
+	if requested {
+		if addr, err := rt.start(ctx); err != nil {
+			ctx.Notify("AHP: "+err.Error(), "error")
+		} else {
+			ctx.Notify(addr, "info")
+		}
+	}
+}
+
+func (rt *runtime) onSessionShutdown(ctx sdk.Context, reason string) {
+	restarts := false
+	if v, err := ctx.GetFlag(flagStart); err == nil {
+		restarts, _ = v.(bool)
+	}
+	if notice := replacementNotice(reason, rt.stop(), restarts); notice != "" {
+		ctx.Notify(notice, "warning")
+	}
+}
+
+func (rt *runtime) onModelChanged() {
+	rt.mu.Lock()
+	built := rt.built
+	rt.mu.Unlock()
+	if built == nil {
+		return
+	}
+	built.Host.DispatchServerAction(wire.RootChannel, ahptypes.StateAction{Value: &ahptypes.RootAgentsChangedAction{
+		Type: ahptypes.ActionTypeRootAgentsChanged, Agents: []ahptypes.AgentInfo{pi.BuildAgentInfo(rt.live.Models())},
+	}})
+}
+
+func isLoopback(host string) bool {
+	if host == "localhost" {
+		return true
+	}
+	ip := net.ParseIP(host)
+	return ip != nil && ip.IsLoopback()
+}
+
+// start opens the listener and returns the notice to show (listenNotice). It is only ever called on the user's explicit request.
+func (rt *runtime) start(ctx sdk.Context) (string, error) {
+	rt.mu.Lock()
+	if rt.server != nil {
+		addr := rt.addr
+		rt.mu.Unlock()
+		return addr, nil
+	}
+	rt.mu.Unlock()
+	if rt.live.ID() == "" {
+		return "", errors.New("no PiG session is running yet")
+	}
+	path := settingsPath(&ctx)
+	direct, err := settings.LoadDirect(path)
+	if err != nil {
+		return "", err
+	}
+	access, err := settings.LoadAccess(path)
+	if err != nil {
+		return "", err
+	}
+	token := ""
+	if direct.Token != nil {
+		token = *direct.Token
+	}
+	if err := requireTokenOffLoopback(path, direct.Host, token); err != nil {
+		return "", err
+	}
+
+	cwd := rt.live.WorkingDirectory()
+	sessionRoot := sessionRootOf(ctx)
+	opts := compose.Options{
+		ServerInfo:       &ahptypes.Implementation{Name: "pigpen-ahp", Version: strPtr(Version)},
+		WorkingDirectory: cwd, SessionRoot: sessionRoot,
+		Models:               rt.live.Models,
+		CreateBackend:        rt.live.BackendFactory,
+		CreateSessionManager: rt.createStore,
+		DefaultSelection:     rt.live.CurrentSelection,
+		DeleteFile:           rt.deleteFile(access),
+		WrapStore:            rt.wrapStore,
+		AdoptForeignTurns:    true,
+		Log:                  func(format string, args ...any) { fmt.Fprintf(os.Stderr, "ahp: "+format+"\n", args...) },
+	}
+	opts.Filesystem, opts.Terminals = accessOptions(access, cwd)
+	built, err := compose.Build(opts)
+	if err != nil {
+		return "", err
+	}
+	if err := rt.registerLiveSession(built, cwd); err != nil {
+		built.Close()
+		return "", err
+	}
+	server, err := ws.Serve(built.Host, ws.Options{
+		Addr: net.JoinHostPort(direct.Host, strconv.Itoa(direct.Port)), Token: token, AllowedOrigins: access.AllowedOrigins,
+	})
+	if err != nil {
+		built.Close()
+		return "", err
+	}
+	addr := listenNotice(server.Addr().String(), token)
+	rt.mu.Lock()
+	rt.server, rt.built, rt.addr = server, built, addr
+	rt.mu.Unlock()
+	return addr, nil
+}
+
+// stop closes the listener and the services behind it; false when nothing was running.
+func (rt *runtime) stop() bool {
+	rt.mu.Lock()
+	server, built := rt.server, rt.built
+	rt.server, rt.built, rt.addr = nil, nil, ""
+	rt.mu.Unlock()
+	if server == nil {
+		return false
+	}
+	_ = server.Close()
+	built.Close()
+	return true
+}
+
+func sessionRootOf(ctx sdk.Context) string {
+	if f, err := ctx.GetSessionFile(); err == nil && f != nil {
+		// <root>/<encoded cwd>/<file>.jsonl
+		return filepath.Dir(filepath.Dir(*f))
+	}
+	return ""
+}
+
+func strPtr(s string) *string { return &s }
+
+// createStore is the storage for sessions created through the host: only the running session
+// exists, so a client cannot make new ones.
+func (rt *runtime) createStore(cwd, id string) (pi.SessionStore, error) {
+	if id == rt.live.ID() {
+		return rt.live.NewStore(), nil
+	}
+	return nil, errors.New("this host serves the running PiG session only; a client cannot create sessions")
+}
+
+func (rt *runtime) wrapStore(sessionID string, opened pi.SessionStore) pi.SessionStore {
+	if sessionID == rt.live.ID() {
+		return rt.live.NewStore()
+	}
+	return opened
+}
+
+// deleteFile refuses to delete anything unless the settings allow it, and never the running
+// session's file.
+func (rt *runtime) deleteFile(access settings.Access) func(string) (pi.SessionFileDeletionResult, error) {
+	return guardedDelete(access.AllowSessionDeletion, func() string { return rt.live.NewStore().File() }, pi.FileDeleter)
+}
+
+func guardedDelete(allow bool, liveFile func() string, remove func(string) (pi.SessionFileDeletionResult, error)) func(string) (pi.SessionFileDeletionResult, error) {
+	return func(path string) (pi.SessionFileDeletionResult, error) {
+		if !allow {
+			return pi.SessionFileDeletionResult{Error: "deleting sessions from a client is disabled (settings: allowSessionDeletion)"}, nil
+		}
+		if live := liveFile(); live != "" && filepath.Clean(live) == filepath.Clean(path) {
+			return pi.SessionFileDeletionResult{Error: "the running PiG session cannot be deleted from a client"}, nil
+		}
+		return remove(path)
+	}
+}
+
+// accessOptions turns the capability settings into the optional services of the composition:
+// nil (not exposed) unless the settings enable them. A filesystem without roots is confined to
+// the working directory; only an explicit unrestricted lifts that.
+func accessOptions(access settings.Access, cwd string) (*compose.FilesystemOptions, *compose.TerminalOptions) {
+	var fs *compose.FilesystemOptions
+	var terminals *compose.TerminalOptions
+	if access.Filesystem.Enabled {
+		roots := access.Filesystem.Roots
+		if len(roots) == 0 && !access.Filesystem.Unrestricted {
+			roots = []string{cwd}
+		}
+		fs = &compose.FilesystemOptions{Roots: roots, Unrestricted: access.Filesystem.Unrestricted}
+	}
+	if access.Terminals.Enabled {
+		terminals = &compose.TerminalOptions{}
+	}
+	return fs, terminals
+}
+
+// replacementNotice is what the user is told when a running listener closes with its session
+// (session_shutdown for new, resume, fork or reload). With --ahp the next session opens it
+// again, so there is nothing to say; after /ahp start it stays closed until asked again.
+func replacementNotice(reason string, running, restarts bool) string {
+	if !running || restarts || reason == "quit" {
+		return ""
+	}
+	return "AHP listener stopped: the session was replaced (" + reason + "). /ahp start serves the new session."
+}
+
+// requireTokenOffLoopback refuses to expose an unauthenticated listener beyond this machine.
+func requireTokenOffLoopback(path, host, token string) error {
+	if !isLoopback(host) && token == "" {
+		return fmt.Errorf("%s: host %q is not loopback, so a token is required", path, host)
+	}
+	return nil
+}
+
+// registerLiveSession makes the running session available to clients: hydrated from its file
+// when it has history on disk, created empty otherwise, then its agent attached so events are
+// mirrored from the start.
+func (rt *runtime) registerLiveSession(built *compose.Built, cwd string) error {
+	id := rt.live.ID()
+	session := wire.SessionURI(id)
+	ctx := context.Background()
+	hydrated, err := built.Sessions.Hydrator.Hydrate(ctx, session)
+	if err != nil {
+		return err
+	}
+	if !hydrated {
+		if err := built.Sessions.Registry.Create(ctx, ahptypes.CreateSessionParams{Channel: session, WorkingDirectories: []ahptypes.URI{wire.PathToFileURI(cwd)}}); err != nil {
+			return err
+		}
+	}
+	return built.Sessions.Registry.EnsureBackend(session)
+}

package/extensions/ahp/runtime_test.go

@@ -0,0 +1,110 @@
+package ahp
+
+import (
+	"net/url"
+	"regexp"
+	"strings"
+	"testing"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/pi"
+)
+
+func TestLoopbackHosts(t *testing.T) {
+	for host, want := range map[string]bool{
+		"127.0.0.1": true, "::1": true, "localhost": true, "127.1.2.3": true,
+		"0.0.0.0": false, "::": false, "192.168.1.5": false, "example.com": false, "": false,
+	} {
+		if got := isLoopback(host); got != want {
+			t.Errorf("isLoopback(%q) = %v, want %v", host, got, want)
+		}
+	}
+}
+
+func TestRemoteHostNeedsAToken(t *testing.T) {
+	if err := requireTokenOffLoopback("s.json", "0.0.0.0", ""); err == nil || !strings.Contains(err.Error(), "token is required") {
+		t.Fatalf("a public host without a token was accepted: %v", err)
+	}
+	for _, c := range []struct{ host, token string }{{"0.0.0.0", "secret"}, {"127.0.0.1", ""}, {"localhost", ""}} {
+		if err := requireTokenOffLoopback("s.json", c.host, c.token); err != nil {
+			t.Errorf("%+v refused: %v", c, err)
+		}
+	}
+}
+
+func TestSessionDeletionIsGuarded(t *testing.T) {
+	var removed []string
+	remove := func(p string) (pi.SessionFileDeletionResult, error) {
+		removed = append(removed, p)
+		return pi.SessionFileDeletionResult{OK: true}, nil
+	}
+	live := func() string { return "/sessions/./live.jsonl" }
+
+	off := guardedDelete(false, live, remove)
+	if r, _ := off("/sessions/other.jsonl"); r.OK || r.Error == "" {
+		t.Fatalf("deletion is off by default, got %+v", r)
+	}
+	on := guardedDelete(true, live, remove)
+	if r, _ := on("/sessions/live.jsonl"); r.OK || !strings.Contains(r.Error, "running") {
+		t.Fatalf("the running session's file must never be deleted, got %+v", r)
+	}
+	if len(removed) != 0 {
+		t.Fatalf("removed %v before any allowed deletion", removed)
+	}
+	if r, _ := on("/sessions/other.jsonl"); !r.OK {
+		t.Fatalf("an allowed deletion failed: %+v", r)
+	}
+	if len(removed) != 1 || removed[0] != "/sessions/other.jsonl" {
+		t.Fatalf("removed = %v", removed)
+	}
+	if r, _ := guardedDelete(true, func() string { return "" }, remove)("/sessions/x.jsonl"); !r.OK {
+		t.Fatalf("with no live file everything else may go: %+v", r)
+	}
+}
+
+// VS Code's "Agents: Add Remote Agent Host..." takes ws://<host>:<port>?tkn=<token>; the notice prints that form first and the
+// generic one after it, and prints no token form for a listener without a token.
+func TestListenNoticePrintsTheVSCodeAddress(t *testing.T) {
+	got := listenNotice("127.0.0.1:7117", "s3cret")
+	for _, want := range []string{"AHP listening on ws://127.0.0.1:7117?tkn=s3cret ", "VS Code", "Agents: Add Remote Agent Host...", "ws://127.0.0.1:7117/?token=s3cret"} {
+		if !strings.Contains(got, want) {
+			t.Errorf("notice %q lacks %q", got, want)
+		}
+	}
+	if got := listenNotice("127.0.0.1:7117", ""); got != "AHP listening on ws://127.0.0.1:7117" {
+		t.Errorf("a tokenless notice = %q", got)
+	}
+}
+
+func TestListenNoticeEscapesAToken(t *testing.T) {
+	got := listenNotice("127.0.0.1:1", "a b&c")
+	if !strings.Contains(got, "?token=a+b%26c") || !strings.Contains(got, "?tkn=a+b%26c") {
+		t.Errorf("a token with reserved characters must be query-escaped: %q", got)
+	}
+}
+
+// VS Code's SSH launcher (chat.sshRemoteAgentHostCommand) scrapes the command's output with this expression and takes the
+// first match: microsoft/vscode src/vs/platform/agentHost/node/sshRemoteAgentHostHelpers.ts (AGENT_HOST_WS_URL_RE,
+// extractAgentHostWebSocketURL). The token it captures is used as is.
+var vsCodeSSHListeningURL = regexp.MustCompile(`ws://(?:127\.0\.0\.1|localhost):(\d+)(?:\?tkn=([^\s&]+))?`)
+
+func TestListenNoticeIsReadByVSCodeSSHLauncher(t *testing.T) {
+	m := vsCodeSSHListeningURL.FindStringSubmatch(listenNotice("127.0.0.1:7117", "s3cret"))
+	if m == nil || m[1] != "7117" || m[2] != "s3cret" {
+		t.Fatalf("VS Code's SSH launcher reads port and token %q from the notice, want 7117 and s3cret", m)
+	}
+}
+
+// "Agents: Add Remote Agent Host..." takes the first ws:// URL in what is pasted, drops trailing ),.;] and moves `tkn` out
+// of the query into the connection token: parseRemoteAgentHostInput in microsoft/vscode
+// src/vs/platform/agentHost/common/remoteAgentHostService.ts. Pasting the whole notice must give the right token.
+func TestListenNoticeIsReadByAddRemoteAgentHost(t *testing.T) {
+	notice := listenNotice("127.0.0.1:7117", "s3cret")
+	candidate := strings.TrimRight(regexp.MustCompile(`(?i)(?:https?|wss?)://\S+`).FindString(notice), "),.;]")
+	u, err := url.Parse(candidate)
+	if err != nil {
+		t.Fatalf("%q: %v", candidate, err)
+	}
+	if got := u.Query().Get("tkn"); got != "s3cret" || u.Host != "127.0.0.1:7117" {
+		t.Fatalf("pasting the notice gives host %q and tkn %q, want 127.0.0.1:7117 and s3cret (from %q)", u.Host, got, candidate)
+	}
+}