okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffauffix(stream): restore unfinished responses on session attachmenta229862

main
branchesmainall refs →
2.4 KiB63 linesraw
1// Shutdown recovery must be checked through the Service caller: shared-daemon
2// E2E cannot safely refuse shutdown while advertising replacement credentials.
3package app
4
5import (
6	"context"
7	"errors"
8	"net"
9	"net/http"
10	"net/http/httptest"
11	"strconv"
12	"sync/atomic"
13	"testing"
14
15	"albedo/cli/internal/daemon"
16)
17
18func TestStopDaemonDoesNotRecoverAuthenticationRefusal(t *testing.T) {
19	var shutdowns, healthProbes atomic.Int32
20	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
21		switch request.URL.Path {
22		case "/server/shutdown":
23			shutdowns.Add(1)
24			if request.Method != http.MethodPost || request.Header.Get("Authorization") != "Bearer old" {
25				t.Errorf("unexpected shutdown attempt: %s %q", request.Method, request.Header.Get("Authorization"))
26			}
27			writer.WriteHeader(http.StatusForbidden)
28			_, _ = writer.Write([]byte(`{"code":"authentication_required","detail":"invalid bearer"}`))
29		case "/server":
30			healthProbes.Add(1)
31			_, _ = writer.Write([]byte(`{"instance_id":"replacement","protocol":3,"state":"ready","capabilities":{"durable_inputs":1,"session_replay":2,"collection_invalidation":1,"tool_progress":1},"build":null,"digest":null,"extensions":[],"quota":[],"notices":[]}`))
32		default:
33			t.Errorf("unexpected request %s", request.URL.Path)
34			writer.WriteHeader(http.StatusNotFound)
35		}
36	}))
37	defer server.Close()
38	_, portText, err := net.SplitHostPort(server.Listener.Addr().String())
39	if err != nil {
40		t.Fatal(err)
41	}
42	port, err := strconv.Atoi(portText)
43	if err != nil {
44		t.Fatal(err)
45	}
46	replacement := daemon.ConnectionSnapshot{Port: port, Token: "new", Version: daemon.ProtocolVersion, InstanceID: "instance-a"}
47	original := replacement
48	original.Token = "old"
49	conn := daemon.NewConnection(original, func(context.Context) (daemon.ConnectionSnapshot, error) {
50		t.Error("shutdown must not rediscover a replacement")
51		return replacement, nil
52	})
53	t.Cleanup(conn.HTTPClient().CloseIdleConnections)
54	service := Service{Existing: func(context.Context) (*daemon.Connection, error) { return conn, nil }}
55	err = service.StopDaemon(context.Background())
56	apiError, ok := errors.AsType[*daemon.APIError](err)
57	if !ok || apiError.StatusCode != http.StatusForbidden {
58		t.Fatalf("want original authentication rejection, got %v", err)
59	}
60	if shutdowns.Load() != 1 || healthProbes.Load() != 0 {
61		t.Fatalf("shutdown was recovered: attempts=%d health probes=%d", shutdowns.Load(), healthProbes.Load())
62	}
63}