fix(chat): register permission prompts before slow broker work

internal/chat/transport.go

@@ -273,23 +273,39 @@ type Remote struct {
 }
 
 func Dial(ctx context.Context, dataDir string) (*Remote, error) {
-	path := SocketPath(dataDir)
-	c, err := dialPath(ctx, path)
+	c, err := DialHook(ctx, dataDir)
+	if err != nil {
+		return nil, err
+	}
+	if err := c.call(ctx, "GET", "/health", nil, nil); err != nil {
+		c.Close()
+		return nil, err
+	}
+	return c, nil
+}
+
+// DialHook prepares an authenticated hook client without a health preflight.
+// Registration must be the first request: a slow health response must not
+// consume the hook's short connection budget and silently lose the prompt.
+// The request itself verifies the token and bounds the actual socket dial.
+func DialHook(ctx context.Context, dataDir string) (*Remote, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	c, err := dialPath(ctx, SocketPath(dataDir))
 	if err != nil {
 		return nil, err
 	}
 	c.keyPath = filepath.Join(dataDir, "chat.key")
 	key, err := os.ReadFile(c.keyPath)
-	if err == nil {
-		c.token = string(key)
-		err = c.call(ctx, "GET", "/health", nil, nil)
-	}
 	if err != nil {
 		c.Close()
 		return nil, err
 	}
+	c.token = string(key)
 	return c, nil
 }
+
 func dialPath(ctx context.Context, path string) (*Remote, error) {
 	if err := fleet.CheckSocketDir(path); err != nil {
 		return nil, err

internal/chat/transport_test.go

@@ -2,13 +2,17 @@ package chat
 
 import (
 	"context"
+	"fmt"
 	"github.com/labstack/driga/internal/fleet"
 	"github.com/labstack/driga/internal/secfile"
+	"net"
+	"net/http"
 	"os"
 	"path/filepath"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestSocketAndFallback(t *testing.T) {
@@ -156,3 +160,82 @@ func TestChatRefusesUnsafeFallbackDirectory(t *testing.T) {
 		})
 	}
 }
+
+// A loaded collector can take longer than the hook's old 200 ms health
+// budget even though it is listening and able to register the permission.
+func TestHookRegistersWithSlowBroker(t *testing.T) {
+	b, _, _ := fixture(t, true)
+	dir := t.TempDir()
+	if err := secfile.MkdirPrivate(dir); err != nil {
+		t.Fatal(err)
+	}
+	if err := secfile.Protect(dir); err != nil {
+		t.Fatal(err)
+	}
+	path := SocketPath(dir)
+	if err := fleet.EnsureSocketDir(path); err != nil {
+		t.Fatal(err)
+	}
+	l, err := net.Listen("unix", path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(path)
+	if runtime.GOOS != "windows" {
+		if err := os.Chmod(path, 0600); err != nil {
+			t.Fatal(err)
+		}
+	}
+	b.clientToken = "test-hook-token"
+	if err := os.WriteFile(filepath.Join(dir, "chat.key"), []byte(b.clientToken), 0600); err != nil {
+		t.Fatal(err)
+	}
+	release := make(chan struct{})
+	defer close(release)
+	h := handler(b)
+	server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/health" {
+			// Never finish a preflight before its deadline. No timing race or
+			// platform-specific scheduling is needed to reproduce the lost prompt.
+			select {
+			case <-release:
+			case <-r.Context().Done():
+			}
+			return
+		}
+		if r.URL.Path == "/hook-prompt" {
+			select {
+			case <-time.After(300 * time.Millisecond):
+			case <-r.Context().Done():
+				return
+			}
+		}
+		h.ServeHTTP(w, r)
+	})}
+	defer server.Close()
+	go server.Serve(l)
+	dialCtx, cancelDial := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	c, err := DialHook(dialCtx, dir)
+	cancelDial()
+	if err != nil {
+		t.Fatalf("hook did not register prompt: %v", err)
+	}
+	defer c.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	done := make(chan error, 1)
+	go func() {
+		out, err := c.WaitPrompt(ctx, testPrompt("slow"))
+		if err == nil && out.Behavior != "allow" {
+			err = fmt.Errorf("decision: %+v", out)
+		}
+		done <- err
+	}()
+	waitPrompt(t, b, "slow")
+	if err := b.AnswerPrompt(ctx, "slow", true); err != nil {
+		t.Fatal(err)
+	}
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+}

internal/cli/hook_prompt.go

@@ -20,7 +20,7 @@ type permissionPayload struct {
 func promptBroker() (*chat.Remote, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
 	defer cancel()
-	return chat.Dial(ctx, config.DataDir())
+	return chat.DialHook(ctx, config.DataDir())
 }
 func endPromptSession(session, cwd string) {
 	c, err := promptBroker()