@pi-in-go/pigpen-ahp 0.1.0

package/extensions/ahp/internal/pi/restart_test.go

@@ -0,0 +1,185 @@
+package pi_test
+
+import (
+	"context"
+	"encoding/base64"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/microsoft/agent-host-protocol/clients/go/ahptypes"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/mapper"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/pi"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/pisession"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/testkit"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/twin"
+	"github.com/MichaelKinsy/pigpen/ahp/internal/wire"
+)
+
+// restartBackend answers every prompt with an empty settled run.
+type restartBackend struct {
+	mu        sync.Mutex
+	prompts   []string
+	listeners []func(mapper.Event)
+}
+
+func (b *restartBackend) Subscribe(l func(mapper.Event)) func() {
+	b.mu.Lock()
+	b.listeners = append(b.listeners, l)
+	b.mu.Unlock()
+	return func() {}
+}
+func (b *restartBackend) Prompt(_ context.Context, text string, _ []mapper.Image) error {
+	b.mu.Lock()
+	b.prompts = append(b.prompts, text)
+	ls := append([]func(mapper.Event){}, b.listeners...)
+	b.mu.Unlock()
+	for _, l := range ls {
+		l(mapper.Event{"type": "agent_start"})
+		l(mapper.Event{"type": "agent_settled"})
+	}
+	return nil
+}
+func (b *restartBackend) Steer(context.Context, string, []mapper.Image) error { return nil }
+func (b *restartBackend) Abort(context.Context) error                         { return nil }
+
+// Twin of upstream test/reconnect.test.ts "reconnect after host restart": the durable session
+// written by a previous host is restored when a client that knew the old host reconnects.
+func TestReconnectAfterHostRestart(t *testing.T) {
+	twin.Run(t, "reconnect", "repairs VS Code's channel-less reconnect and restores a durable session", func(t *testing.T) {
+		root, workspace := t.TempDir(), t.TempDir()
+		id := newID()
+		writer, err := pisession.Create(workspace, filepath.Join(root, "fixture"), pisession.Options{ID: id})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := writer.AppendMessage(obj{"role": "user", "content": "before restart", "timestamp": 0}); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := writer.AppendMessage(obj{"role": "assistant", "content": []any{obj{"type": "text", "text": "persisted reply"}}, "timestamp": 0}); err != nil {
+			t.Fatal(err)
+		}
+		backend := &restartBackend{}
+		h := startHarness(t, harnessOptions{sessionRoot: root, workingDir: workspace, createBackend: func(*pi.LiveSession) (pi.Backend, error) { return backend, nil }})
+
+		clientSession := "pi:/" + id
+		clientChat := "ahp-chat://default/" + base64.RawURLEncoding.EncodeToString([]byte(clientSession))
+		client := testkit.Connect(t, h.host)
+		t.Cleanup(client.Close)
+		// The reconnect names no channel, as VS Code's does.
+		var result struct {
+			Type      string
+			Snapshots []struct {
+				Resource string
+				State    map[string]any
+			}
+		}
+		client.Decode(client.Must("reconnect", obj{
+			"clientId": "vscode-from-previous-host", "lastSeenServerSeq": 42,
+			"subscriptions": []string{wire.RootChannel, clientSession}, "_meta": obj{"vscode.telemetryLevel": "off"},
+		}), &result)
+		if result.Type != "snapshot" {
+			t.Fatalf("type = %q", result.Type)
+		}
+		var resources []string
+		var sessionState map[string]any
+		for _, s := range result.Snapshots {
+			resources = append(resources, s.Resource)
+			if s.Resource == clientSession {
+				sessionState = s.State
+			}
+		}
+		if !reflect.DeepEqual(sortedCopy(resources), sortedCopy([]string{wire.RootChannel, clientSession})) {
+			t.Fatalf("snapshots = %v", resources)
+		}
+		if sessionState["defaultChat"] != clientChat {
+			t.Fatalf("defaultChat = %v, want %s", sessionState["defaultChat"], clientChat)
+		}
+
+		rpcErr := client.ExpectError("disposeSession", obj{"channel": clientSession}, wire.CodeInvalidRequest)
+		if !strings.Contains(rpcErr.Message, "VS Code provisional-session lifecycle bug") {
+			t.Fatalf("message = %q", rpcErr.Message)
+		}
+		if !h.host.Store().Has(wire.SessionURI(id)) {
+			t.Fatal("the durable session must survive the refused disposal")
+		}
+
+		chat := snapshotState[ahptypes.ChatState](t, client, clientChat)
+		turn := chat.Turns[0]
+		if turn.Message.Text != "before restart" || turn.State != ahptypes.TurnStateComplete {
+			t.Fatalf("restored turn = %+v", turn)
+		}
+		var reply string
+		for _, p := range turn.ResponseParts {
+			if md, ok := p.Value.(*ahptypes.MarkdownResponsePart); ok {
+				reply = md.Content
+			}
+		}
+		if reply != "persisted reply" {
+			t.Fatalf("reply = %q", reply)
+		}
+		if !h.host.Store().Has(wire.SessionURI(id)) || !h.host.Store().Has(wire.ChatURI(id)) {
+			t.Fatal("session and chat must both be restored")
+		}
+
+		client.Dispatch(clientChat, obj{"type": "chat/turnStarted", "turnId": "after-restart", "startedAt": "2025-01-01T00:00:00.000Z", "message": userMessage("continue after restart")})
+		testkit.Eventually(t, "the post-restart prompt to reach the backend", func() bool { backend.mu.Lock(); defer backend.mu.Unlock(); return len(backend.prompts) == 1 })
+		if !reflect.DeepEqual(backend.prompts, []string{"continue after restart"}) {
+			t.Fatalf("prompts = %v", backend.prompts)
+		}
+		testkit.Eventually(t, "the resumed turn to settle", func() bool {
+			s := h.host.Store().Chat(wire.ChatURI(id))
+			return s.ActiveTurn == nil && len(s.Turns) == 2
+		})
+		state := h.host.Store().Chat(wire.ChatURI(id))
+		if state.Turns[1].Message.Text != "continue after restart" || state.Turns[1].State != ahptypes.TurnStateComplete {
+			t.Fatalf("turns = %+v", turnShapes(state.Turns))
+		}
+	})
+}
+
+func sortedCopy(in []string) []string {
+	out := append([]string(nil), in...)
+	for i := range out {
+		for j := i + 1; j < len(out); j++ {
+			if out[j] < out[i] {
+				out[i], out[j] = out[j], out[i]
+			}
+		}
+	}
+	return out
+}
+
+// Twin of upstream test/session-storage.test.serial.ts.
+func TestDefaultSessionStorage(t *testing.T) {
+	twin.Run(t, "session-storage", "keeps the default catalogue and Pi session writer on the same isolated profile", func(t *testing.T) {
+		agentDir, workspace := t.TempDir(), t.TempDir()
+		catalogueRoot, create := pi.DefaultStorage(agentDir)
+		sessionID := newID()
+		store, err := create(workspace, sessionID)
+		if err != nil {
+			t.Fatal(err)
+		}
+		manager := store.(*pisession.Manager)
+		if _, err := manager.AppendMessage(obj{"role": "user", "content": "hello", "timestamp": 0}); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := manager.AppendMessage(obj{"role": "assistant", "content": []any{obj{"type": "text", "text": "hello"}}, "timestamp": 0}); err != nil {
+			t.Fatal(err)
+		}
+		file := store.File()
+		if file == "" {
+			t.Fatal("no session file")
+		}
+		found, err := pi.NewCatalogue(catalogueRoot).FindSessionFile(sessionID)
+		if err != nil || found != file {
+			t.Fatalf("catalogue found %q (%v), writer wrote %q", found, err, file)
+		}
+		if rel, err := filepath.Rel(filepath.Join(agentDir, "sessions"), file); err != nil || strings.HasPrefix(rel, "..") {
+			t.Fatalf("file %s is outside the profile (%q, %v)", file, rel, err)
+		}
+	})
+}

package/extensions/ahp/internal/pi/services.go

@@ -0,0 +1,70 @@
+package pi
+
+import (
+	"context"
+
+	"github.com/microsoft/agent-host-protocol/clients/go/ahptypes"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/host"
+)
+
+// ServicesOptions wire the session services together.
+type ServicesOptions struct {
+	Host *host.Host
+	// SessionRoot is where Pi's session files live (one sub-directory per working directory).
+	SessionRoot             string
+	DefaultWorkingDirectory string
+	CreateBackend           BackendFactory
+	CreateSessionManager    SessionManagerFactory
+	DefaultSelection        func() *ahptypes.ModelSelection
+	DeleteFile              func(path string) (SessionFileDeletionResult, error)
+	WrapStore               func(sessionID string, opened SessionStore) SessionStore
+	AdoptForeignTurns       bool
+	Log                     func(format string, args ...any)
+}
+
+// Services is the registry, catalogue and hydrator wired to one host.
+type Services struct {
+	Registry  *Registry
+	Catalogue *Catalogue
+	Hydrator  *Hydrator
+}
+
+// NewServices builds the session services.
+func NewServices(opts ServicesOptions) *Services {
+	catalogue := NewCatalogue(opts.SessionRoot)
+	registry := NewRegistry(RegistryOptions{
+		Host: opts.Host, DefaultWorkingDirectory: opts.DefaultWorkingDirectory, CreateBackend: opts.CreateBackend,
+		CreateSessionManager: opts.CreateSessionManager, DefaultSelection: opts.DefaultSelection, DeleteFile: opts.DeleteFile,
+		FindSessionFile: catalogue.FindSessionFile, AdoptForeignTurns: opts.AdoptForeignTurns, Log: opts.Log,
+	})
+	hydrator := NewHydrator(HydratorOptions{
+		Host: opts.Host, Catalogue: catalogue,
+		IsLive:            registry.Has,
+		IsDisposing:       registry.IsDisposing,
+		FallbackSelection: opts.DefaultSelection,
+		Adopt:             func(a AdoptedSession) { registry.Adopt(a) },
+		WrapStore:         opts.WrapStore,
+		Log:               opts.Log,
+	})
+	return &Services{Registry: registry, Catalogue: catalogue, Hydrator: hydrator}
+}
+
+type liveCatalogue struct{ s *Services }
+
+func (c liveCatalogue) List(ctx context.Context, limit *int64, cursor *string) (ahptypes.ListSessionsResult, error) {
+	return c.s.Catalogue.List(ctx, limit, cursor, c.s.Registry.CatalogueOverrides)
+}
+
+type turnPaging struct{ r *Registry }
+
+func (p turnPaging) FetchTurns(ctx context.Context, params ahptypes.FetchTurnsParams) error {
+	return p.r.FetchTurns(ctx, params)
+}
+
+// Capabilities are what the host serves through these services.
+func (s *Services) Capabilities() host.Capabilities {
+	return host.Capabilities{
+		Catalogue: liveCatalogue{s}, Sessions: s.Registry, Hydrator: s.Hydrator, TurnPaging: turnPaging{s.Registry},
+	}
+}

package/extensions/ahp/internal/pi/sessionconfig.go

@@ -0,0 +1,64 @@
+package pi
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/microsoft/agent-host-protocol/clients/go/ahptypes"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/wire"
+)
+
+// ProjectTrustKey is the pre-creation config property that reports project-resource trust.
+const ProjectTrustKey = "projectResources"
+
+// SessionConfigOptions configure a [SessionConfigService].
+type SessionConfigOptions struct {
+	DefaultWorkingDirectory string
+	ProjectTrustPolicy      ProjectTrustPolicy
+	AgentDir                string
+}
+
+// SessionConfigService answers resolveSessionConfig: what a session would be like before it is
+// created. The exchange is iterative: the answer changes as the user picks a directory, which is
+// the whole reason the command exists (port of src/pi/session-config.ts).
+type SessionConfigService struct{ opts SessionConfigOptions }
+
+// NewSessionConfigService creates the service.
+func NewSessionConfigService(opts SessionConfigOptions) *SessionConfigService {
+	return &SessionConfigService{opts: opts}
+}
+
+// Resolve implements host.SessionConfigHandler.
+func (s *SessionConfigService) Resolve(_ context.Context, params ahptypes.ResolveSessionConfigParams) (ahptypes.ResolveSessionConfigResult, error) {
+	workingDirectory := s.opts.DefaultWorkingDirectory
+	if params.WorkingDirectory != nil && *params.WorkingDirectory != "" {
+		path, err := wire.FileURIToPath(*params.WorkingDirectory)
+		if err != nil {
+			return ahptypes.ResolveSessionConfigResult{}, wire.InvalidParams(err.Error())
+		}
+		workingDirectory = path
+	}
+	trust, err := ResolveProjectTrust(workingDirectory, s.opts.ProjectTrustPolicy, s.opts.AgentDir)
+	if err != nil {
+		return ahptypes.ResolveSessionConfigResult{}, err
+	}
+	description := "Trust-gated project resources (" + trust.Reason + ")."
+	if trust.Reason == "no-project-resources" {
+		description = "This folder has no trust-gated project resources."
+	}
+	readOnly := true
+	value, _ := json.Marshal(trust.Trusted)
+	return ahptypes.ResolveSessionConfigResult{
+		Schema: ahptypes.SessionConfigSchema{Type: "object", Properties: map[string]ahptypes.SessionConfigPropertySchema{
+			// Host policy, not a per-session choice.
+			ProjectTrustKey: {Type: "boolean", Title: "Load project resources", Description: &description, ReadOnly: &readOnly},
+		}},
+		Values: map[string]json.RawMessage{ProjectTrustKey: value},
+	}, nil
+}
+
+// Completions implements host.SessionConfigHandler: this host offers no dynamic completions.
+func (s *SessionConfigService) Completions(context.Context, ahptypes.SessionConfigCompletionsParams) (ahptypes.SessionConfigCompletionsResult, error) {
+	return ahptypes.SessionConfigCompletionsResult{Items: []ahptypes.SessionConfigValueItem{}}, nil
+}

package/extensions/ahp/internal/pi/sessionfiles_test.go

@@ -0,0 +1,74 @@
+package pi_test
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+)
+
+// fixtureSessionDirectory mirrors Pi's default cwd partition solely for synthetic JSONL fixtures
+// (upstream test/support/session-files.ts).
+func fixtureSessionDirectory(t testing.TB, root, cwd string) string {
+	t.Helper()
+	abs, err := filepath.Abs(cwd)
+	if err != nil {
+		t.Fatal(err)
+	}
+	abs = strings.TrimLeft(abs, `/\`)
+	partition := "--" + strings.NewReplacer("/", "-", `\`, "-", ":", "-").Replace(abs) + "--"
+	dir := filepath.Join(root, partition)
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+	return dir
+}
+
+type fakeSession struct {
+	cwd              string
+	firstUserMessage string
+	name             string
+	mtimeSeconds     int64 // controls catalogue ordering
+}
+
+var entryCounter int
+
+// writeFakeSession writes a minimal but genuine Pi session file: a session header followed by
+// entries linked through id/parentId (upstream writeFakeSession).
+func writeFakeSession(t testing.TB, root, id string, o fakeSession) string {
+	t.Helper()
+	dir := fixtureSessionDirectory(t, root, o.cwd)
+	timestamp := time.Unix(o.mtimeSeconds, 0).UTC().Format("2006-01-02T15:04:05.000Z")
+	line := func(v map[string]any) string {
+		raw, _ := json.Marshal(v)
+		return string(raw)
+	}
+	lines := []string{line(map[string]any{"type": "session", "id": id, "parentId": nil, "timestamp": timestamp, "version": 3, "cwd": o.cwd})}
+	var parent any
+	push := func(entry map[string]any) {
+		entryCounter++
+		entryID := fmt.Sprintf("entry-%d", entryCounter)
+		entry["id"], entry["parentId"], entry["timestamp"] = entryID, parent, timestamp
+		lines = append(lines, line(entry))
+		parent = entryID
+	}
+	if o.firstUserMessage != "" {
+		push(map[string]any{"type": "message", "message": map[string]any{"role": "user", "content": o.firstUserMessage, "timestamp": 0}})
+	}
+	if o.name != "" {
+		push(map[string]any{"type": "session_info", "name": o.name})
+	}
+	path := filepath.Join(dir, strconv.FormatInt(o.mtimeSeconds, 10)+"_"+id+".jsonl")
+	if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")+"\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	at := time.Unix(o.mtimeSeconds, 0)
+	if err := os.Chtimes(path, at, at); err != nil {
+		t.Fatal(err)
+	}
+	return path
+}

package/extensions/ahp/internal/pi/sessionstore.go

@@ -0,0 +1,83 @@
+package pi
+
+import (
+	"path/filepath"
+	"strings"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/pisession"
+)
+
+// SessionStore is the slice of Pi's session manager the host uses. *pisession.Manager satisfies
+// it; the extension adapter provides one over the live PiG session.
+type SessionStore interface {
+	// File is the session file, "" when the session is not persisted; it may not exist yet.
+	File() string
+	LeafEntry() pisession.Entry
+	SessionName() string
+	AppendSessionInfo(name string) (string, error)
+	// Branch is the path from the root to the leaf.
+	Branch(fromID string) []pisession.Entry
+	// ContextEntries is the compaction-aware entry list along the current branch.
+	ContextEntries() []pisession.Entry
+}
+
+// SessionManagerFactory creates the storage for a new session.
+type SessionManagerFactory func(workingDirectory, sessionID string) (SessionStore, error)
+
+// InMemoryStorage is a factory for sessions that never touch disk.
+func InMemoryStorage(workingDirectory, sessionID string) (SessionStore, error) {
+	return pisession.InMemory(workingDirectory, pisession.Options{ID: sessionID})
+}
+
+// PersistentStorage is a factory for sessions written under dir.
+func PersistentStorage(dir string) SessionManagerFactory {
+	return func(workingDirectory, sessionID string) (SessionStore, error) {
+		return pisession.Create(workingDirectory, dir, pisession.Options{ID: sessionID})
+	}
+}
+
+// HistorySource is what history serving needs from a live session.
+type HistorySource interface {
+	Store() SessionStore
+	ID() string
+	// Anchor is the recorded truncation anchor of a turn.
+	Anchor(turnID string) (string, bool)
+}
+
+// TruncationAnchor is the session entry a truncation to turnID would move the agent to ("" turnID
+// clears everything). ok is false when no entry matches: accepting an impossible truncation would
+// shorten the client's view while Pi kept using context the user believes is gone.
+func TruncationAnchor(source HistorySource, turnID *string) (string, bool) {
+	key := ClearAllAnchor
+	if turnID != nil {
+		key = *turnID
+	}
+	if known, ok := source.Anchor(key); ok && known != "" {
+		return known, true
+	}
+	// A turn this host never ran: recompute from the file, which is where a hydrated session's ids
+	// came from in the first place.
+	a, ok := RebuildHistoryFromSession(source.Store(), RebuildOptions{TurnIDPrefix: source.ID()}).Anchors[key]
+	return a, ok && a != ""
+}
+
+// DefaultSessionDir is where Pi keeps the sessions of a working directory: one sub-directory of
+// <agentDir>/sessions per cwd, named after the encoded path (Pi's getDefaultSessionDirPath).
+func DefaultSessionDir(agentDir, cwd string) string {
+	abs, err := filepath.Abs(cwd)
+	if err != nil {
+		abs = cwd
+	}
+	abs = strings.TrimLeft(abs, `/\`)
+	safe := "--" + strings.NewReplacer("/", "-", `\`, "-", ":", "-").Replace(abs) + "--"
+	return filepath.Join(agentDir, "sessions", safe)
+}
+
+// DefaultStorage keeps new sessions where Pi itself would, under <agentDir>/sessions, and pairs the
+// catalogue root with the same profile so the catalogue and the writer never disagree (port of
+// createDefaultPiSessionStorage).
+func DefaultStorage(agentDir string) (catalogueRoot string, create SessionManagerFactory) {
+	return filepath.Join(agentDir, "sessions"), func(workingDirectory, sessionID string) (SessionStore, error) {
+		return pisession.Create(workingDirectory, DefaultSessionDir(agentDir, workingDirectory), pisession.Options{ID: sessionID})
+	}
+}