refactor(cli): separate attach, launch, and upgrade
Keep restart approval in the CLI and make local discovery failures actionable. Cancellation leaves shared daemons running.
131 unmodified linesExpand all132133134135136136137138139202 unmodified linesExpand all131 unmodified linesExpand all capture_output=True, timeout=20, ) self.assertNotEqual(cleanup.returncode, 0) self.assertIn(b"storage is in use", cleanup.stderr) self.assertEqual(backup.read_bytes(), b"approved backup") self.assertEqual( json.loads(app.cli("storage", "--json")), json.loads(before) )202 unmodified linesExpand all1 unmodified line2345678910111213141511 unmodified lines27282930313233343536373839404142433544454647484950515253545556575859606162636465666768697071727374757677783779808182838485868788899091929394959641979899100101102103104105106107108109481101111 unmodified linepackage main
import ( "context" "errors" "fmt" "os" "os/signal" "path/filepath" "syscall" "time"
"albedo/cli/internal/app" "albedo/cli/internal/cli"11 unmodified lines os.Exit(1) }}func run(args []string) error { ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer cancel() streams := cli.Streams{In: os.Stdin, Out: os.Stdout, Err: os.Stderr} home := config.HomeDir() term := &terminal.Service{In: streams.In, Out: streams.Out, Err: streams.Err, Capable: terminal.IsTTY(streams.In, streams.Out), Home: home} if os.Getenv("ALBEDO_NO_BROWSER") == "" { term.OpenBrowser = config.OpenBrowser } options := daemon.LocalOptions{HomeDir: home, ProjectRoot: findProjectRoot()} rediscover := func(ctx context.Context) (daemon.ConnectionSnapshot, error) { return daemon.Rediscover(ctx, home) } var chosen *daemon.Connection application := &app.Service{ Connect: func(ctx context.Context) (*daemon.Connection, error) { return daemon.EnsureContext(ctx, home, findProjectRoot(), term.ReplaceStale) if err := ctx.Err(); err != nil { return nil, err } if chosen != nil { return chosen, nil } found, err := daemon.Discover(ctx, home) if err != nil { if _, compatibleFailure := errors.AsType[*daemon.CompatibilityError](err); !compatibleFailure || found.Kind != daemon.Running { return nil, err } } if found.Kind != daemon.Running { chosen, err = daemon.Launch(ctx, options) return chosen, err } // An unknown protocol cannot safely authorize the shutdown operation. if found.Health.Version != daemon.ProtocolVersion { return nil, daemon.CheckCompatible(found.Health) } selectedBuild := os.Getenv("ALBEDO_BUILD") if selectedBuild == "" { selectedBuild, _ = filepath.EvalSymlinks(os.Getenv("ALBEDO_DAEMON")) } restart, err := term.ConfirmRestart(ctx, found.Snapshot, selectedBuild) if err != nil { return nil, err } if restart { chosen, err = daemon.Upgrade(ctx, options, found.Snapshot) } else { chosen, err = daemon.Attach(ctx, found.Snapshot, rediscover) } return chosen, err }, Existing: func(ctx context.Context) (*daemon.Connection, error) { return daemon.ExistingContext(ctx, home) }, Existing: func(ctx context.Context) (*daemon.Connection, error) { found, err := daemon.Discover(ctx, home) if err != nil { if _, compatibleFailure := errors.AsType[*daemon.CompatibilityError](err); !compatibleFailure || found.Kind != daemon.Running { return nil, err } } if found.Kind != daemon.Running { return nil, nil } if found.Health.Version != daemon.ProtocolVersion { return nil, daemon.CheckCompatible(found.Health) } return daemon.NewConnection(found.Snapshot, nil), nil }, } store := &storage.Service{ Home: home, Now: time.Now, Running: func() (bool, error) { conn, err := daemon.Existing(home); return conn != nil, err }, Running: func() (bool, error) { found, err := daemon.Discover(ctx, home) if _, compatibilityFailure := errors.AsType[*daemon.CompatibilityError](err); compatibilityFailure && found.Kind == daemon.Running { return true, nil } return found.Kind == daemon.Running, err }, DeleteSessions: application.DeleteSessions, } workspace, err := os.Getwd() if err != nil { workspace = "." } return cli.Execute(context.Background(), args, cli.Dependencies{Application: application, Storage: store, Terminal: term, Workspace: workspace}, streams) return cli.Execute(ctx, args, cli.Dependencies{Application: application, Storage: store, Terminal: term, Workspace: workspace}, streams)}16 unmodified lines17181920212122232425262728292930313233113 unmodified linesExpand all16 unmodified lines "albedo/cli/internal/daemon")
func promptTestConnection(server *httptest.Server) *daemon.Connection { return daemon.NewConnection(daemon.ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, "") return daemon.NewConnection(daemon.ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil)}
func TestPromptRequiresSubmissionCancellationBeforeCreatingASession(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { if request.URL.Path != "/health" { t.Errorf("unsupported prompt issued %s", request.URL.Path) } _, _ = writer.Write([]byte(`{"capabilities":[]}`)) _, _ = writer.Write([]byte(`{"ok":true,"version":2,"capabilities":[]}`)) })) defer server.Close() service := Service{Connect: func(context.Context) (*daemon.Connection, error) { return promptTestConnection(server), nil }} _, err := service.RunPrompt(t.Context(), PromptOptions{Prompt: "hello"})113 unmodified linesExpand all41 unmodified lines4243444546464748495027 unmodified lines41 unmodified lines w.WriteHeader(404) } })) defer server.Close() conn := daemon.NewConnection(daemon.ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, "") conn := daemon.NewConnection(daemon.ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil) service := Service{Connect: func(context.Context) (*daemon.Connection, error) { return conn, nil }} for _, turns := range []int{3, 2} { cursors = nil result, err := service.Read(context.Background(), "s", turns)27 unmodified lines2 unmodified lines345677891012131112131427 unmodified lines4243444549465152535455565747486049505152535455567 unmodified lines2 unmodified linespackage app
import ( "context" "encoding/json" "errors" "net" "net/http" "net/http/httptest" "os" "path/filepath" "strconv" "sync/atomic" "testing"
27 unmodified lines port, err := strconv.Atoi(portText) if err != nil { t.Fatal(err) } home := t.TempDir() replacement := daemon.ConnectionSnapshot{Port: port, Token: "new", Version: 2} record, err := json.Marshal(replacement) if err != nil { t.Fatal(err) } if writeErr := os.WriteFile(filepath.Join(home, "daemon.json"), record, 0600); writeErr != nil { t.Fatal(writeErr) } original := replacement original.Token = "old" conn := daemon.NewConnection(original, home) conn := daemon.NewConnection(original, func(context.Context) (daemon.ConnectionSnapshot, error) { t.Error("shutdown must not rediscover a replacement") return replacement, nil }) t.Cleanup(conn.HTTPClient().CloseIdleConnections) service := Service{Existing: func(context.Context) (*daemon.Connection, error) { return conn, nil }} err = service.StopDaemon(context.Background()) apiError, ok := errors.AsType[*daemon.APIError](err)7 unmodified lines18 unmodified lines1920212223232425262717 unmodified lines454647484949505152534 unmodified lines18 unmodified lines "data: "+`{"events":[{"type":"overflow"}]}`+"\n\n"+ "data: "+`{"events":[{"type":"text","session":"s","text":"must stay hidden"}]}`+"\n\n") })) defer server.Close() conn := NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, "") conn := NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil) var delivered [][]AgentEvent err := StreamAgents(t.Context(), conn, func(events []AgentEvent) error { delivered = append(delivered, events); return nil }) if err != nil { t.Fatal(err)17 unmodified lines w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, "data: "+batch+"\n\n") })) defer server.Close() conn := NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, "") conn := NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil) err := StreamAgents(t.Context(), conn, func([]AgentEvent) error { t.Error("malformed batch partially delivered"); return nil }) failure, ok := errors.AsType[*StreamError](err) if !ok || failure.Kind != StreamProtocol { t.Fatalf("malformed batch classified incorrectly: %v", err)4 unmodified lines14 unmodified lines151617181919202122222324252627272829293031323334353637383940414243443345464736374849505141525354555610 unmodified lines14 unmodified lines "testing" "time")
func TestEnsureHonorsCancellationDuringDiscoveryAndLockWait(t *testing.T) {func TestDiscoveryAndLauncherWaitHonorCancellation(t *testing.T) { for _, health := range []bool{false, true} { t.Run(fmt.Sprintf("stalled_health_%t", health), func(t *testing.T) { t.Setenv("ALBEDO_DAEMON", "") home := t.TempDir() var operation func(context.Context) error if health { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() })) defer server.Close() writeDiscovery(t, home, ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port, Token: "t", Version: 2}) writeDiscovery(t, home, ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port, Pid: os.Getpid(), Token: "t", Version: ProtocolVersion}) operation = func(ctx context.Context) error { _, err := Discover(ctx, home); return err } } else { if err := os.WriteFile(filepath.Join(home, "starting.lock"), fmt.Append(nil, os.Getpid()), 0600); err != nil { if processAlive(0) { t.Skip("advisory locking requires Unix") } executable := filepath.Join(t.TempDir(), "unused-daemon") if err := os.WriteFile(executable, []byte("#!/bin/sh\nexit 99\n"), 0700); err != nil { t.Fatal(err) } t.Setenv("ALBEDO_DAEMON", executable) lock, err := acquireLauncher(t.Context(), home) if err != nil { t.Fatal(err) } defer lock.Close() operation = func(ctx context.Context) error { _, err := Launch(ctx, LocalOptions{HomeDir: home}); return err } } ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond) defer cancel() started := time.Now() _, err := EnsureContext(ctx, home, "", nil) if !errors.Is(err, context.DeadlineExceeded) { if err := operation(ctx); !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("deadline lost: %v", err) } if time.Since(started) > 300*time.Millisecond { t.Fatal("discovery or lock wait ignored caller deadline") t.Fatal("discovery or launcher wait ignored caller deadline") } }) }}10 unmodified lines47 unmodified lines4849505152525354555657 unmodified lines11411511611711811811912012112259 unmodified lines18218318418518618618718818919046 unmodified lines23723823924024124124224324424546 unmodified lines2922932942952962962972982993009 unmodified lines47 unmodified lines t.Errorf("unexpected automatic stream request: %s", request.URL) } })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session") var previews []string var completed []string consume := func(event StreamEvent) error { if event.Type == EventToolProgress && event.Progress != nil && event.Progress.Code != nil {57 unmodified lines t.Error("unexpected stream request") } })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session") consume := func(StreamEvent) error { return nil } for range 32 { if err := client.Stream(t.Context(), 0, consume); err != nil { t.Fatalf("accepted call rejected: %v", err)59 unmodified lines t.Error("unexpected stream request") } })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session") consume := func(StreamEvent) error { return nil } if err := client.Stream(t.Context(), 0, consume); err != nil { t.Fatalf("exact byte budget rejected: %v", err) }46 unmodified lines map[string]any{"type": "arguments_delta", "name": "python", "callId": "same", "text": arguments}, }))) })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session") ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) defer cancel() var err error if beforeHeaders {46 unmodified lines writer.Header().Set("Content-Type", "text/event-stream") _, _ = writer.Write([]byte(pages)) })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session") replayed := map[string]bool{} if err := client.Stream(t.Context(), 0, func(event StreamEvent) error { if event.Type == EventThinking { replayed[event.Text] = event.Replayed9 unmodified lines123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147package daemon
import ( "context" "encoding/json" "errors" "fmt" "net/http" "sync" "sync/atomic")
type ConnectionSnapshot struct { Token string `json:"token"` Build string `json:"build,omitempty"` Port int `json:"port"` Pid int `json:"pid"` Version int `json:"version"`}
type Connection struct { snapshot atomic.Pointer[ConnectionSnapshot] httpClient *http.Client refreshGate chan struct{} rediscover Rediscovery httpOnce sync.Once refreshOnce sync.Once}
// Rediscovery obtains an updated endpoint without owning API attachment.type Rediscovery func(context.Context) (ConnectionSnapshot, error)
// NewConnection creates an API client. A nil rediscover disables recovery.func NewConnection(snap ConnectionSnapshot, rediscover Rediscovery) *Connection { c := &Connection{rediscover: rediscover} c.snapshot.Store(&snap) return c}
func (c *Connection) Snapshot() ConnectionSnapshot { if c == nil { return ConnectionSnapshot{} } p := c.snapshot.Load() if p == nil { return ConnectionSnapshot{} } return *p}
func (c *Connection) Port() int { return c.Snapshot().Port}
func (c *Connection) Token() string { return c.Snapshot().Token}
func (c *Connection) Pid() int { return c.Snapshot().Pid}
func (c *Connection) Version() int { return c.Snapshot().Version}
func (c *Connection) Build() string { return c.Snapshot().Build}
func (c *Connection) BaseURL() string { port := c.Port() if port <= 0 { return "" } return fmt.Sprintf("http://127.0.0.1:%d", port)}
func (c *Connection) Update(other *Connection) { if c == nil || other == nil || c == other { return } snap := other.Snapshot() c.snapshot.Store(&snap)}
// Refresh serializes rediscovery and installs only a compatible live endpoint.func (c *Connection) Refresh(ctx context.Context) error { if c == nil { return errors.New("not connected to Albedo") } if ctx == nil { ctx = context.Background() } c.refreshOnce.Do(func() { c.refreshGate = make(chan struct{}, 1) }) select { case c.refreshGate <- struct{}{}: case <-ctx.Done(): return ctx.Err() } defer func() { <-c.refreshGate }() if err := ctx.Err(); err != nil { return err } if c.rediscover == nil { return errors.New("daemon rediscovery is unavailable") } snap, err := c.rediscover(ctx) if err != nil { return err } attached, err := Attach(ctx, snap, nil) if err != nil { return err } defer attached.HTTPClient().CloseIdleConnections() c.Update(attached) return nil}
// HTTPClient returns the reusable HTTP client owned by this Connection.// Address and token updates retain the same client and pool.func (c *Connection) HTTPClient() *http.Client { c.httpOnce.Do(func() { c.httpClient = newHTTPClient() }) return c.httpClient}
func (c *Connection) MarshalJSON() ([]byte, error) { snap := c.Snapshot() return json.Marshal(snap)}
func (c *Connection) UnmarshalJSON(data []byte) error { var snap ConnectionSnapshot if decodeErr := json.Unmarshal(data, &snap); decodeErr != nil { return decodeErr } c.snapshot.Store(&snap) return nil}
func ctxDone(ctx context.Context) <-chan struct{} { if ctx == nil { return nil } return ctx.Done()}1234556789101112131415161718192021//go:build !unix
package daemon
import "os/exec"import ( "errors" "os" "os/exec")
func detach(cmd *exec.Cmd) {}
func passFileLimit() {}
// processAlive cannot be checked portably here; assume the lock holder is live.func processAlive(int) bool { return true }
func tryLauncherLock(*os.File) (bool, error) { return false, errors.New("local daemon launch requires a platform with advisory file locks")}func endpointRefused(error) bool { return false }2 unmodified lines3456789101121 unmodified lines33343536373839404142434445462 unmodified linespackage daemon
import ( "errors" "os" "os/exec" "syscall")
21 unmodified lines } err := syscall.Kill(pid, 0) return err == nil || errors.Is(err, syscall.EPERM)}
func tryLauncherLock(file *os.File) (bool, error) { err := syscall.Flock(int(file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB) if errors.Is(err, syscall.EWOULDBLOCK) || errors.Is(err, syscall.EAGAIN) { return false, nil } return err == nil, err}
func endpointRefused(err error) bool { return errors.Is(err, syscall.ECONNREFUSED) }87 unmodified lines88899091929394959697989910010187 unmodified lines}
func (e *StreamError) Error() string { return e.Cause.Error() }func (e *StreamError) Unwrap() error { return e.Cause }
// CompatibilityError describes a live daemon that cannot serve this client.type CompatibilityError struct { Version int MissingCapabilities []string}
func (e *CompatibilityError) Error() string { return fmt.Sprintf("daemon protocol %d is incompatible with this client; required protocol is %d; missing capabilities: %v", e.Version, ProtocolVersion, e.MissingCapabilities)}30 unmodified lines3132333435353637383916 unmodified lines30 unmodified lines _, _ = io.WriteString(w, tc.body) })) defer server.Close() // This fixed peer has no discovery directory; authentication errors stay local. conn := NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, "") conn := NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil) _, requestErr := RequestOperation[any](context.Background(), conn, Operation{Name: "read failure", Method: http.MethodGet, Path: "/failure", Policy: ReadRecovery}) client := NewChatClient(conn, "session") _, chatErr := client.GetStatus(context.Background()) for _, err := range []error{requestErr, chatErr} {16 unmodified lines