okami.mom/albedo

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

Aly Raffauffeat(api)!: migrate daemon and clients to protocol 361c43ca2026-10-03

// Hanging refusal bodies and failures after stream delivery require controlled
// peers; real daemon events cannot reliably force these transport boundaries.
package daemon

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"
)

type streamCaller struct {
	start func(context.Context, *Connection, func() error) error
	name  string
}

func streamCallers() []streamCaller {
	return []streamCaller{
		{name: "agents", start: func(ctx context.Context, conn *Connection, deliver func() error) error {
			return StreamAgents(ctx, conn, func(_ []AgentEvent) error { return deliver() })
		}},
		{name: "chat", start: func(ctx context.Context, conn *Connection, deliver func() error) error {
			return NewChatClient(conn, "test").Stream(ctx, 0, func(event StreamEvent) error {
				if event.Type == EventText {
					return deliver()
				}
				return nil
			})
		}},
	}
}

func TestStreamAuthenticationRecoveryBeforeDelivery(t *testing.T) {
	for _, caller := range streamCallers() {
		t.Run(caller.name, func(t *testing.T) {
			for _, test := range []struct {
				name, token string
				marked      bool
				requests    int32
			}{
				{name: "changed credentials", token: "new", marked: true, requests: 2},
				{name: "unchanged credentials", token: "old", marked: true, requests: 1},
				{name: "unmarked refusal", token: "new", requests: 1},
			} {
				t.Run(test.name, func(t *testing.T) {
					var requests atomic.Int32
					var deliveries atomic.Int32
					server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
						if request.URL.Path == "/server" {
							_, _ = writer.Write([]byte(readyServer))
							return
						}
						if requests.Add(1) == 1 {
							if test.marked {
								writer.Header().Set("WWW-Authenticate", "Bearer")
								writer.WriteHeader(http.StatusUnauthorized)
							} else {
								writer.WriteHeader(http.StatusForbidden)
							}
							if caller.name == "chat" {
								if test.marked {
									_, _ = writer.Write([]byte(`{"code":"authentication_required"}`))
								} else {
									_, _ = writer.Write([]byte(`{"detail":"denied"}`))
								}
								return
							}
							// The refusal marker allows closing the unfinished
							// body must cancel this handler without waiting for more bytes.
							writer.(http.Flusher).Flush()
							<-request.Context().Done()
							return
						}
						if request.Header.Get("Authorization") != "Bearer new" {
							t.Error("stream reused old credentials")
						}
						writer.Header().Set("Content-Type", "text/event-stream")
						writeInitialStream(writer, caller.name)
					}))
					defer server.Close()

					snapshot := ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port, Token: test.token, Version: ProtocolVersion}
					replacement := snapshot
					snapshot.Token = "old"
					conn := NewConnection(snapshot, func(context.Context) (ConnectionSnapshot, error) { return replacement, nil })
					defer conn.HTTPClient().CloseIdleConnections()
					ctx, cancel := context.WithTimeout(context.Background(), time.Second)
					defer cancel()
					streamErr := caller.start(ctx, conn, func() error { deliveries.Add(1); return nil })
					if errors.Is(streamErr, context.DeadlineExceeded) {
						t.Fatalf("stream waited for unfinished refusal body: %v", streamErr)
					}
					if requests.Load() != test.requests {
						t.Fatalf("got %d requests, expected %d", requests.Load(), test.requests)
					}
					if test.requests == 2 {
						if streamErr != nil || deliveries.Load() != 1 {
							t.Fatalf("stream failed recovery: %v deliveries=%d", streamErr, deliveries.Load())
						}
					} else {
						failure, ok := errors.AsType[*APIError](streamErr)
						expectedStatus := http.StatusForbidden
						if test.marked {
							expectedStatus = http.StatusUnauthorized
						}
						if !ok || failure.StatusCode != expectedStatus || deliveries.Load() != 0 {
							t.Fatalf("refusal changed: %v deliveries=%d", streamErr, deliveries.Load())
						}
					}
				})
			}
		})
	}
}

func TestAcceptedStreamIsNeverReplayed(t *testing.T) {
	for _, caller := range streamCallers() {
		t.Run(caller.name, func(t *testing.T) {
			for _, failure := range []string{"truncated body", "callback error"} {
				t.Run(failure, func(t *testing.T) {
					var requests atomic.Int32
					var deliveries atomic.Int32
					callbackFailure := errors.New("consumer stopped")
					server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
						if request.URL.Path == "/server" {
							_, _ = writer.Write([]byte(readyServer))
							return
						}
						requests.Add(1)
						if failure == "truncated body" {
							writer.Header().Set("Content-Length", "1000000")
						}
						writer.Header().Set("Content-Type", "text/event-stream")
						writeInitialStream(writer, caller.name)
					}))
					defer server.Close()
					snapshot := ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port, Token: "token", Version: ProtocolVersion}

					replacement := snapshot
					conn := NewConnection(snapshot, func(context.Context) (ConnectionSnapshot, error) { return replacement, nil })
					defer conn.HTTPClient().CloseIdleConnections()
					ctx, cancel := context.WithTimeout(context.Background(), time.Second)
					defer cancel()
					streamErr := caller.start(ctx, conn, func() error {
						deliveries.Add(1)
						if failure == "callback error" {
							return callbackFailure
						}
						return nil
					})
					if failure == "callback error" && !errors.Is(streamErr, callbackFailure) {
						t.Fatalf("callback failure lost: %v", streamErr)
					}
					if failure == "truncated body" && !errors.Is(streamErr, io.ErrUnexpectedEOF) {
						t.Fatalf("stream truncation lost: %v", streamErr)
					}
					if requests.Load() != 1 || deliveries.Load() != 1 {
						t.Fatalf("accepted stream replayed: requests=%d deliveries=%d", requests.Load(), deliveries.Load())
					}
				})
			}
		})
	}
}