From fded82ae5b1cbb14736274add19434e4bcb116af Mon Sep 17 00:00:00 2001
From: Brian McCallister <brianm@skife.org>
Date: Sat, 18 Apr 2026 15:28:55 -0700
Subject: [PATCH 1/3] feat: add remote agent auth tunneling via SSH agent
 protocol extensions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Implement proxying of auth requests back through SSH agent forwarding,
enabling users to SSH to a remote host and run `epithet agent fish` to
get certificates from an internal CA while auth plugins (FIDO2, browser)
run on the local machine.

Uses SSH agent protocol extensions (message type 27) with a registry
pattern for extensibility:
- epithet-hello@epithet.dev — probe if socket is an epithet agent
- epithet-auth@epithet.dev — request auth token from upstream

Key components:
- ProxyAgent: generic ExtendedAgent wrapper with extension handler registry
- ProxyListener: per-connection upstream dialing for agent protocol proxying
- Upstream client: ProbeUpstream() and RequestAuth() functions
- Broker: WithUpstream option and Authenticate() helper (try upstream, fall back to local)
- CLI: `epithet agent <shell>` wrapper mode mirroring ssh-agent's pattern

Multi-hop chaining works naturally — each layer proxies to its upstream.
No changes to CA or policy server required.
---
 cmd/epithet/agent.go             | 301 +++++++++++++++++++++++--------
 cmd/epithet/agent_test.go        |  51 ++++++
 pkg/agent/extensions.go          |  46 +++++
 pkg/agent/proxy.go               |  85 +++++++++
 pkg/agent/proxy_listener.go      | 122 +++++++++++++
 pkg/agent/proxy_listener_test.go | 296 ++++++++++++++++++++++++++++++
 pkg/agent/proxy_test.go          | 246 +++++++++++++++++++++++++
 pkg/agent/upstream.go            |  85 +++++++++
 pkg/agent/upstream_test.go       | 152 ++++++++++++++++
 pkg/broker/auth.go               |   9 +
 pkg/broker/broker.go             |  31 +++-
 pkg/broker/broker_test.go        | 154 ++++++++++++++++
 12 files changed, 1500 insertions(+), 78 deletions(-)
 create mode 100644 cmd/epithet/agent_test.go
 create mode 100644 pkg/agent/extensions.go
 create mode 100644 pkg/agent/proxy.go
 create mode 100644 pkg/agent/proxy_listener.go
 create mode 100644 pkg/agent/proxy_listener_test.go
 create mode 100644 pkg/agent/proxy_test.go
 create mode 100644 pkg/agent/upstream.go
 create mode 100644 pkg/agent/upstream_test.go

diff --git a/cmd/epithet/agent.go b/cmd/epithet/agent.go
index 7a55b00..d6b1f3b 100644
--- a/cmd/epithet/agent.go
+++ b/cmd/epithet/agent.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"log/slog"
 	"os"
+	"os/exec"
 	"os/signal"
 	"path/filepath"
 	"strconv"
@@ -14,6 +16,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/epithet-ssh/epithet/pkg/agent"
 	"github.com/epithet-ssh/epithet/pkg/broker"
 	"github.com/epithet-ssh/epithet/pkg/caclient"
 	"github.com/epithet-ssh/epithet/pkg/tlsconfig"
@@ -32,74 +35,243 @@ type AgentCLI struct {
 }
 
 // AgentStartCLI is the default subcommand that starts the agent/broker.
-type AgentStartCLI struct{}
+// When Shell is provided, it runs in wrapper mode: wrapping a shell with an
+// epithet-aware agent proxy (mirroring ssh-agent's "ssh-agent bash" pattern).
+type AgentStartCLI struct {
+	Shell string `arg:"" optional:"" help:"Shell to wrap with epithet agent (enables wrapper mode)"`
+}
 
 func (s *AgentStartCLI) Run(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error {
-	// Validate required fields for start
+	if s.Shell != "" {
+		return s.runWrapper(parent, logger, tlsCfg)
+	}
+	return s.runDaemon(parent, logger, tlsCfg)
+}
+
+// runDaemon runs the broker as a long-lived daemon (original behavior).
+func (s *AgentStartCLI) runDaemon(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error {
+	b, tempDir, brokerSock, agentDir, homeDir, err := setupBroker(parent, logger, tlsCfg)
+	if err != nil {
+		return err
+	}
+
+	// Clean up temp directory on exit.
+	defer func() {
+		if err := os.RemoveAll(tempDir); err != nil {
+			logger.Warn("failed to remove temp directory", "error", err, "path", tempDir)
+		} else {
+			logger.Debug("removed temp directory", "path", tempDir)
+		}
+	}()
+
+	// Set up context with cancellation on signals.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-sigChan
+		logger.Info("received shutdown signal")
+		cancel()
+	}()
+
+	// Generate SSH config.
+	sshConfigPath := filepath.Join(tempDir, "ssh-config.conf")
+	if err := parent.generateSSHConfig(sshConfigPath, agentDir, brokerSock, homeDir); err != nil {
+		logger.Warn("failed to generate SSH config", "error", err, "path", sshConfigPath)
+	} else {
+		includePattern := filepath.Join(homeDir, ".epithet", "run", "*", "ssh-config.conf")
+		if err := checkSSHConfigInclude(homeDir, includePattern, logger); err != nil {
+			logger.Warn(fmt.Sprintf("Add 'Include %s' to ~/.ssh/config", includePattern))
+		}
+		logger.Debug("generated SSH config", "path", sshConfigPath)
+	}
+
+	// Start broker.
+	logger.Info("starting broker", "socket", brokerSock)
+	err = b.Serve(ctx)
+	if err != nil && err != context.Canceled {
+		return fmt.Errorf("broker serve error: %w", err)
+	}
+
+	logger.Info("broker shutdown complete")
+	return nil
+}
+
+// runWrapper wraps a shell with an epithet-aware agent proxy. It probes the
+// current SSH_AUTH_SOCK for an upstream epithet agent, starts the broker, and
+// creates a proxy listener that handles epithet protocol extensions while
+// delegating standard SSH agent operations to the upstream agent.
+func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error {
+	upstreamSocket := os.Getenv("SSH_AUTH_SOCK")
+
+	// Probe for upstream epithet agent.
+	var depth int
+	var brokerOpts []broker.Option
+	if upstreamSocket != "" {
+		hello, err := agent.ProbeUpstream(upstreamSocket)
+		if err != nil {
+			logger.Warn("failed to probe upstream agent", "error", err)
+		} else if hello != nil {
+			depth = hello.ChainDepth + 1
+			brokerOpts = append(brokerOpts, broker.WithUpstream(upstreamSocket))
+			logger.Info("detected upstream epithet agent", "chain_depth", depth)
+		} else {
+			logger.Debug("upstream is vanilla ssh-agent")
+		}
+	} else {
+		logger.Debug("no SSH_AUTH_SOCK set")
+	}
+
+	b, tempDir, brokerSock, agentDir, homeDir, err := setupBrokerWithOptions(parent, logger, tlsCfg, brokerOpts...)
+	if err != nil {
+		return err
+	}
+
+	// Clean up temp directory on exit.
+	defer func() {
+		if err := os.RemoveAll(tempDir); err != nil {
+			logger.Warn("failed to remove temp directory", "error", err, "path", tempDir)
+		}
+	}()
+
+	// Set up context with cancellation.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Start broker in background.
+	brokerErr := make(chan error, 1)
+	go func() {
+		brokerErr <- b.Serve(ctx)
+	}()
+	<-b.Ready()
+
+	// Generate SSH config.
+	sshConfigPath := filepath.Join(tempDir, "ssh-config.conf")
+	if err := parent.generateSSHConfig(sshConfigPath, agentDir, brokerSock, homeDir); err != nil {
+		logger.Warn("failed to generate SSH config", "error", err, "path", sshConfigPath)
+	} else {
+		includePattern := filepath.Join(homeDir, ".epithet", "run", "*", "ssh-config.conf")
+		if err := checkSSHConfigInclude(homeDir, includePattern, logger); err != nil {
+			logger.Warn(fmt.Sprintf("Add 'Include %s' to ~/.ssh/config", includePattern))
+		}
+	}
+
+	// Create proxy listener if we have an upstream socket to proxy.
+	if upstreamSocket != "" {
+		proxySock := filepath.Join(tempDir, "proxy.sock")
+		setup := func(p *agent.ProxyAgent) {
+			p.RegisterExtension(agent.ExtensionHello, agent.HelloHandler(depth))
+			p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) {
+				return b.Authenticate(nil)
+			}))
+		}
+		proxyListener := agent.NewProxyListener(logger, proxySock, upstreamSocket, setup)
+		go func() {
+			if err := proxyListener.Serve(ctx); err != nil && err != context.Canceled {
+				logger.Error("proxy listener error", "error", err)
+			}
+		}()
+		// Wait briefly for the listener to start.
+		// TODO: add Ready() to ProxyListener like Broker has.
+		upstreamSocket = proxySock
+		logger.Info("proxy agent listening", "socket", proxySock)
+	}
+
+	// Build child environment with updated SSH_AUTH_SOCK.
+	childEnv := os.Environ()
+	if upstreamSocket != "" {
+		childEnv = replaceEnv(childEnv, "SSH_AUTH_SOCK", upstreamSocket)
+	}
+
+	// Run the shell as a child process.
+	cmd := exec.Command(s.Shell)
+	cmd.Env = childEnv
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	// Forward signals to child.
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		for sig := range sigChan {
+			if cmd.Process != nil {
+				cmd.Process.Signal(sig)
+			}
+		}
+	}()
+
+	logger.Info("starting shell", "shell", s.Shell)
+	if err := cmd.Run(); err != nil {
+		var exitErr *exec.ExitError
+		if errors.As(err, &exitErr) {
+			cancel()
+			signal.Stop(sigChan)
+			os.Exit(exitErr.ExitCode())
+		}
+		cancel()
+		return fmt.Errorf("shell failed: %w", err)
+	}
+
+	cancel()
+	signal.Stop(sigChan)
+	return nil
+}
+
+// setupBroker creates the broker and supporting directories with no extra options.
+func setupBroker(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) (*broker.Broker, string, string, string, string, error) {
+	return setupBrokerWithOptions(parent, logger, tlsCfg)
+}
+
+// setupBrokerWithOptions creates the broker, temp directories, and resolves auth.
+// Returns (broker, tempDir, brokerSock, agentDir, homeDir, error).
+func setupBrokerWithOptions(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config, opts ...broker.Option) (*broker.Broker, string, string, string, string, error) {
 	if len(parent.CaURL) == 0 {
-		return fmt.Errorf("--ca-url is required (at least one)")
+		return nil, "", "", "", "", fmt.Errorf("--ca-url is required (at least one)")
 	}
 
-	// Parse CA URLs into endpoints
 	caEndpoints, err := caclient.ParseCAURLs(parent.CaURL)
 	if err != nil {
-		return fmt.Errorf("invalid CA URL: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("invalid CA URL: %w", err)
 	}
 
-	logger.Debug("agent start command received", "ca_urls", parent.CaURL, "ca_timeout", parent.CaTimeout, "ca_cooldown", parent.CaCooldown)
+	logger.Debug("agent command received", "ca_urls", parent.CaURL, "ca_timeout", parent.CaTimeout, "ca_cooldown", parent.CaCooldown)
 
-	// Validate all CA URLs require TLS (unless --insecure)
 	for _, ep := range caEndpoints {
 		if err := tlsCfg.ValidateURL(ep.URL); err != nil {
-			return fmt.Errorf("CA URL %q: %w", ep.URL, err)
+			return nil, "", "", "", "", fmt.Errorf("CA URL %q: %w", ep.URL, err)
 		}
 	}
 
-	// Get home directory
 	homeDir, err := os.UserHomeDir()
 	if err != nil {
-		return fmt.Errorf("failed to get home directory: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("failed to get home directory: %w", err)
 	}
 
-	// Create a unique temporary directory for this broker instance
-	// Use a hash of the CA URLs to make it deterministic
 	instanceID := hashString(fmt.Sprintf("%v", parent.CaURL))
 	runDir := filepath.Join(homeDir, ".epithet", "run")
 	tempDir := filepath.Join(runDir, instanceID)
 
-	// Clean up stale run directories from dead processes
 	cleanupStaleRunDirs(runDir, logger)
 
-	// Clean up temp directory on exit
-	defer func() {
-		if err := os.RemoveAll(tempDir); err != nil {
-			logger.Warn("failed to remove temp directory", "error", err, "path", tempDir)
-		} else {
-			logger.Debug("removed temp directory", "path", tempDir)
-		}
-	}()
-
-	// Create temp directory
 	if err := os.MkdirAll(tempDir, 0700); err != nil {
-		return fmt.Errorf("failed to create temp directory: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("failed to create temp directory: %w", err)
 	}
 
-	// Write PID file for stale detection
 	pidFile := filepath.Join(tempDir, "broker.pid")
 	if err := os.WriteFile(pidFile, []byte(strconv.Itoa(os.Getpid())), 0600); err != nil {
-		return fmt.Errorf("failed to write PID file: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("failed to write PID file: %w", err)
 	}
 
-	// Define paths within temp directory
 	brokerSock := filepath.Join(tempDir, "broker.sock")
 	agentDir := filepath.Join(tempDir, "agent")
 
-	// Create agent directory
 	if err := os.MkdirAll(agentDir, 0700); err != nil {
-		return fmt.Errorf("failed to create agent directory: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("failed to create agent directory: %w", err)
 	}
 
-	// Create CA client
 	caClientOpts := []caclient.Option{
 		caclient.WithLogger(logger),
 		caclient.WithTimeout(parent.CaTimeout),
@@ -108,81 +280,60 @@ func (s *AgentStartCLI) Run(parent *AgentCLI, logger *slog.Logger, tlsCfg tlscon
 	}
 	caClient, err := caclient.New(caEndpoints, caClientOpts...)
 	if err != nil {
-		return fmt.Errorf("failed to create CA client: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("failed to create CA client: %w", err)
 	}
 
-	// Resolve auth command: use --auth if provided, otherwise discover from CA.
 	authCommand := parent.Auth
 	if authCommand == "" {
 		logger.Debug("no --auth provided, discovering from CA")
 
-		// Fetch public key (also caches discovery URL).
 		_, err := caClient.GetPublicKey(context.Background())
 		if err != nil {
-			return fmt.Errorf("failed to get CA public key: %w", err)
+			return nil, "", "", "", "", fmt.Errorf("failed to get CA public key: %w", err)
 		}
 
-		// Fetch unauthenticated discovery (auth config only, no token needed).
 		discovery, err := caClient.GetDiscovery(context.Background(), "")
 		if err != nil {
-			return fmt.Errorf("failed to get discovery config (try --auth to specify manually): %w", err)
+			return nil, "", "", "", "", fmt.Errorf("failed to get discovery config (try --auth to specify manually): %w", err)
 		}
 		if discovery == nil || discovery.Auth == nil {
-			return fmt.Errorf("CA did not return auth config in discovery (try --auth to specify manually)")
+			return nil, "", "", "", "", fmt.Errorf("CA did not return auth config in discovery (try --auth to specify manually)")
 		}
 
-		// Convert auth config to command.
 		authCommand, err = broker.AuthConfigToCommand(*discovery.Auth)
 		if err != nil {
-			return fmt.Errorf("failed to convert discovery auth config: %w", err)
+			return nil, "", "", "", "", fmt.Errorf("failed to convert discovery auth config: %w", err)
 		}
 
 		logger.Info("discovered auth config from CA", "type", discovery.Auth.Type, "issuer", discovery.Auth.Issuer)
 	}
 
-	// Create broker
-	b, err := broker.New(*logger, brokerSock, authCommand, caClient, agentDir)
+	b, err := broker.New(*logger, brokerSock, authCommand, caClient, agentDir, opts...)
 	if err != nil {
-		return fmt.Errorf("failed to create broker: %w", err)
+		return nil, "", "", "", "", fmt.Errorf("failed to create broker: %w", err)
 	}
 
-	// Set up context with cancellation on signals
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-	go func() {
-		<-sigChan
-		logger.Info("received shutdown signal")
-		cancel()
-	}()
-
-	// Generate SSH config file in the temp directory
-	sshConfigPath := filepath.Join(tempDir, "ssh-config.conf")
-
-	if err := parent.generateSSHConfig(sshConfigPath, agentDir, brokerSock, homeDir); err != nil {
-		logger.Warn("failed to generate SSH config", "error", err, "path", sshConfigPath)
-		// Don't fail startup, just warn
-	} else {
-		// Check if ~/.ssh/config has the Include directive
-		includePattern := filepath.Join(homeDir, ".epithet", "run", "*", "ssh-config.conf")
-		if err := checkSSHConfigInclude(homeDir, includePattern, logger); err != nil {
+	return b, tempDir, brokerSock, agentDir, homeDir, nil
+}
 
-			logger.Warn(fmt.Sprintf("Add 'Include %s' to ~/.ssh/config", includePattern))
+// replaceEnv returns a copy of environ with key set to value.
+// If key already exists, it is replaced; otherwise it is appended.
+func replaceEnv(environ []string, key, value string) []string {
+	prefix := key + "="
+	result := make([]string, 0, len(environ)+1)
+	found := false
+	for _, e := range environ {
+		if strings.HasPrefix(e, prefix) {
+			result = append(result, prefix+value)
+			found = true
+		} else {
+			result = append(result, e)
 		}
-		logger.Debug("generated SSH config", "path", sshConfigPath)
 	}
-
-	// Start broker
-	logger.Info("starting broker", "socket", brokerSock)
-	err = b.Serve(ctx)
-	if err != nil && err != context.Canceled {
-		return fmt.Errorf("broker serve error: %w", err)
+	if !found {
+		result = append(result, prefix+value)
 	}
-
-	logger.Info("broker shutdown complete")
-	return nil
+	return result
 }
 
 // generateSSHConfig writes an SSH config file for epithet
diff --git a/cmd/epithet/agent_test.go b/cmd/epithet/agent_test.go
new file mode 100644
index 0000000..d1cced1
--- /dev/null
+++ b/cmd/epithet/agent_test.go
@@ -0,0 +1,51 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestReplaceEnv_ExistingKey(t *testing.T) {
+	env := []string{"HOME=/home/user", "SSH_AUTH_SOCK=/old/path", "SHELL=/bin/fish"}
+	result := replaceEnv(env, "SSH_AUTH_SOCK", "/new/path")
+
+	found := false
+	for _, e := range result {
+		if e == "SSH_AUTH_SOCK=/new/path" {
+			found = true
+		}
+		if e == "SSH_AUTH_SOCK=/old/path" {
+			t.Error("old value still present")
+		}
+	}
+	if !found {
+		t.Error("new value not found")
+	}
+	if len(result) != 3 {
+		t.Errorf("expected 3 entries, got %d", len(result))
+	}
+}
+
+func TestReplaceEnv_NewKey(t *testing.T) {
+	env := []string{"HOME=/home/user", "SHELL=/bin/fish"}
+	result := replaceEnv(env, "SSH_AUTH_SOCK", "/new/path")
+
+	found := false
+	for _, e := range result {
+		if e == "SSH_AUTH_SOCK=/new/path" {
+			found = true
+		}
+	}
+	if !found {
+		t.Error("new key not appended")
+	}
+	if len(result) != 3 {
+		t.Errorf("expected 3 entries, got %d", len(result))
+	}
+}
+
+func TestReplaceEnv_EmptyEnv(t *testing.T) {
+	result := replaceEnv(nil, "KEY", "value")
+	if len(result) != 1 || result[0] != "KEY=value" {
+		t.Errorf("expected [KEY=value], got %v", result)
+	}
+}
diff --git a/pkg/agent/extensions.go b/pkg/agent/extensions.go
new file mode 100644
index 0000000..4f18a8b
--- /dev/null
+++ b/pkg/agent/extensions.go
@@ -0,0 +1,46 @@
+package agent
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+// Extension type constants follow the OpenSSH vendor extension naming convention.
+const (
+	ExtensionHello = "epithet-hello@epithet.dev"
+	ExtensionAuth  = "epithet-auth@epithet.dev"
+)
+
+// HelloResponse is returned by the epithet-hello extension.
+type HelloResponse struct {
+	ProtocolVersion int `json:"protocol_version"`
+	ChainDepth      int `json:"chain_depth"`
+}
+
+// AuthResponse is returned by the epithet-auth extension.
+type AuthResponse struct {
+	Token string `json:"token"`
+}
+
+// HelloHandler returns an ExtensionHandler for epithet-hello@epithet.dev.
+func HelloHandler(depth int) ExtensionHandler {
+	return func(_ []byte) ([]byte, error) {
+		return json.Marshal(HelloResponse{
+			ProtocolVersion: 1,
+			ChainDepth:      depth,
+		})
+	}
+}
+
+// AuthHandler returns an ExtensionHandler for epithet-auth@epithet.dev.
+// The authenticate function is called to obtain a token — typically wired
+// to the broker's auth flow.
+func AuthHandler(authenticate func() (string, error)) ExtensionHandler {
+	return func(_ []byte) ([]byte, error) {
+		token, err := authenticate()
+		if err != nil {
+			return nil, fmt.Errorf("auth failed: %w", err)
+		}
+		return json.Marshal(AuthResponse{Token: token})
+	}
+}
diff --git a/pkg/agent/proxy.go b/pkg/agent/proxy.go
new file mode 100644
index 0000000..dbf953c
--- /dev/null
+++ b/pkg/agent/proxy.go
@@ -0,0 +1,85 @@
+package agent
+
+import (
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
+)
+
+const agentSuccess = 6
+
+// ExtensionHandler handles a single epithet protocol extension.
+// contents is the raw request payload; returns raw response payload (without the type byte).
+type ExtensionHandler func(contents []byte) ([]byte, error)
+
+// ProxyAgent wraps an upstream ssh agent, proxying standard operations and
+// dispatching epithet protocol extensions to registered handlers.
+// Unregistered extensions are forwarded to the upstream agent.
+type ProxyAgent struct {
+	upstream   agent.ExtendedAgent
+	extensions map[string]ExtensionHandler
+}
+
+// NewProxyAgent creates a ProxyAgent that proxies to the given upstream agent.
+func NewProxyAgent(upstream agent.ExtendedAgent) *ProxyAgent {
+	return &ProxyAgent{
+		upstream:   upstream,
+		extensions: make(map[string]ExtensionHandler),
+	}
+}
+
+// RegisterExtension registers a handler for the given extension type.
+func (p *ProxyAgent) RegisterExtension(name string, handler ExtensionHandler) {
+	p.extensions[name] = handler
+}
+
+// Extension dispatches to a registered handler or forwards to upstream.
+func (p *ProxyAgent) Extension(extensionType string, contents []byte) ([]byte, error) {
+	if handler, ok := p.extensions[extensionType]; ok {
+		payload, err := handler(contents)
+		if err != nil {
+			return nil, err
+		}
+		// Prepend SSH_AGENT_SUCCESS byte — ServeAgent sends this as-is,
+		// and the client strips the type byte before returning.
+		return append([]byte{agentSuccess}, payload...), nil
+	}
+	return p.upstream.Extension(extensionType, contents)
+}
+
+// Standard agent methods — all delegate to upstream.
+
+func (p *ProxyAgent) List() ([]*agent.Key, error) {
+	return p.upstream.List()
+}
+
+func (p *ProxyAgent) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) {
+	return p.upstream.Sign(key, data)
+}
+
+func (p *ProxyAgent) Add(key agent.AddedKey) error {
+	return p.upstream.Add(key)
+}
+
+func (p *ProxyAgent) Remove(key ssh.PublicKey) error {
+	return p.upstream.Remove(key)
+}
+
+func (p *ProxyAgent) RemoveAll() error {
+	return p.upstream.RemoveAll()
+}
+
+func (p *ProxyAgent) Lock(passphrase []byte) error {
+	return p.upstream.Lock(passphrase)
+}
+
+func (p *ProxyAgent) Unlock(passphrase []byte) error {
+	return p.upstream.Unlock(passphrase)
+}
+
+func (p *ProxyAgent) Signers() ([]ssh.Signer, error) {
+	return p.upstream.Signers()
+}
+
+func (p *ProxyAgent) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
+	return p.upstream.SignWithFlags(key, data, flags)
+}
diff --git a/pkg/agent/proxy_listener.go b/pkg/agent/proxy_listener.go
new file mode 100644
index 0000000..5a67bf3
--- /dev/null
+++ b/pkg/agent/proxy_listener.go
@@ -0,0 +1,122 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"io"
+	"log/slog"
+	"net"
+	"os"
+	"sync"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// ProxyListener accepts connections on a Unix socket and serves each with a
+// ProxyAgent that proxies to an upstream agent. Each connection gets a fresh
+// upstream dial (the SSH agent protocol is connection-scoped).
+//
+// The setup callback is called for each new ProxyAgent, allowing the caller
+// to register extension handlers.
+type ProxyListener struct {
+	socketPath         string
+	upstreamSocketPath string
+	setup              func(*ProxyAgent)
+	log                *slog.Logger
+
+	listener  net.Listener
+	done      chan struct{}
+	closeOnce sync.Once
+}
+
+// NewProxyListener creates a ProxyListener. Call Serve() to begin accepting connections.
+func NewProxyListener(log *slog.Logger, socketPath, upstreamSocketPath string, setup func(*ProxyAgent)) *ProxyListener {
+	return &ProxyListener{
+		socketPath:         socketPath,
+		upstreamSocketPath: upstreamSocketPath,
+		setup:              setup,
+		log:                log,
+		done:               make(chan struct{}),
+	}
+}
+
+// Serve starts the listener and blocks until the context is cancelled.
+func (p *ProxyListener) Serve(ctx context.Context) error {
+	os.Remove(p.socketPath)
+	listener, err := net.Listen("unix", p.socketPath)
+	if err != nil {
+		return err
+	}
+	if err := os.Chmod(p.socketPath, 0600); err != nil {
+		listener.Close()
+		return err
+	}
+	p.listener = listener
+
+	go p.acceptLoop()
+
+	<-ctx.Done()
+	p.Close()
+	return ctx.Err()
+}
+
+func (p *ProxyListener) acceptLoop() {
+	for {
+		conn, err := p.listener.Accept()
+		if err != nil {
+			if errors.Is(err, net.ErrClosed) {
+				return
+			}
+			select {
+			case <-p.done:
+				return
+			default:
+				p.log.Warn("proxy listener accept error", "error", err)
+				continue
+			}
+		}
+		go p.serveConn(conn)
+	}
+}
+
+func (p *ProxyListener) serveConn(conn net.Conn) {
+	defer conn.Close()
+
+	// Dial a fresh connection to the upstream agent for this client.
+	upstreamConn, err := net.Dial("unix", p.upstreamSocketPath)
+	if err != nil {
+		p.log.Warn("failed to connect to upstream agent", "error", err)
+		return
+	}
+	defer upstreamConn.Close()
+
+	upstream := agent.NewClient(upstreamConn)
+	extUpstream, ok := upstream.(agent.ExtendedAgent)
+	if !ok {
+		p.log.Warn("upstream agent does not support extensions")
+		return
+	}
+
+	proxy := NewProxyAgent(extUpstream)
+	p.setup(proxy)
+
+	err = agent.ServeAgent(proxy, conn)
+	if err != nil && err != io.EOF {
+		p.log.Debug("proxy agent connection ended", "error", err)
+	}
+}
+
+// SocketPath returns the path to the proxy listener's Unix socket.
+func (p *ProxyListener) SocketPath() string {
+	return p.socketPath
+}
+
+// Close stops the listener and cleans up.
+func (p *ProxyListener) Close() {
+	p.closeOnce.Do(func() {
+		if p.listener != nil {
+			p.listener.Close()
+		}
+		close(p.done)
+	})
+}
diff --git a/pkg/agent/proxy_listener_test.go b/pkg/agent/proxy_listener_test.go
new file mode 100644
index 0000000..b55f38f
--- /dev/null
+++ b/pkg/agent/proxy_listener_test.go
@@ -0,0 +1,296 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"log/slog"
+	"net"
+	"os"
+	"testing"
+	"time"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// startVanillaAgent starts a plain ssh-agent keyring on a Unix socket.
+func startVanillaAgent(t *testing.T) (string, func()) {
+	t.Helper()
+
+	f, err := os.CreateTemp("", "vanilla-agent.*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	sockPath := f.Name()
+	f.Close()
+	os.Remove(sockPath)
+
+	listener, err := net.Listen("unix", sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	keyring := agent.NewKeyring()
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				agent.ServeAgent(keyring, c)
+			}(conn)
+		}
+	}()
+
+	return sockPath, func() {
+		listener.Close()
+		os.Remove(sockPath)
+	}
+}
+
+func TestProxyListenerEndToEnd(t *testing.T) {
+	// Set up a vanilla upstream agent.
+	upstreamSock, cleanupUpstream := startVanillaAgent(t)
+	defer cleanupUpstream()
+
+	// Create proxy listener that adds epithet extensions.
+	proxySock := tempSocketPath(t)
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+
+	authCalled := false
+	setup := func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(0))
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			authCalled = true
+			return "proxy-listener-token", nil
+		}))
+	}
+	proxy := NewProxyListener(logger, proxySock, upstreamSock, setup)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go proxy.Serve(ctx)
+	waitForSocket(t, proxySock)
+
+	// Connect to the proxy and test extensions.
+	conn, err := net.Dial("unix", proxySock)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	client := agent.NewClient(conn).(agent.ExtendedAgent)
+
+	// Test hello.
+	resp, err := client.Extension(ExtensionHello, nil)
+	if err != nil {
+		t.Fatalf("hello failed: %v", err)
+	}
+	var hello HelloResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil {
+		t.Fatalf("unmarshal hello: %v", err)
+	}
+	if hello.ProtocolVersion != 1 {
+		t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion)
+	}
+
+	// Test auth.
+	resp, err = client.Extension(ExtensionAuth, nil)
+	if err != nil {
+		t.Fatalf("auth failed: %v", err)
+	}
+	var authResp AuthResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &authResp); err != nil {
+		t.Fatalf("unmarshal auth: %v", err)
+	}
+	if authResp.Token != "proxy-listener-token" {
+		t.Errorf("expected 'proxy-listener-token', got %q", authResp.Token)
+	}
+	if !authCalled {
+		t.Error("auth handler was not called")
+	}
+
+	// Test standard agent ops pass through to upstream.
+	keys, err := client.List()
+	if err != nil {
+		t.Fatalf("List failed: %v", err)
+	}
+	if len(keys) != 0 {
+		t.Errorf("expected 0 keys, got %d", len(keys))
+	}
+}
+
+func TestProxyListenerMultipleConnections(t *testing.T) {
+	// Each connection should get its own upstream dial.
+	upstreamSock, cleanupUpstream := startVanillaAgent(t)
+	defer cleanupUpstream()
+
+	proxySock := tempSocketPath(t)
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+
+	callCount := 0
+	setup := func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(0))
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			callCount++
+			return "token", nil
+		}))
+	}
+	proxy := NewProxyListener(logger, proxySock, upstreamSock, setup)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go proxy.Serve(ctx)
+	waitForSocket(t, proxySock)
+
+	// Make two independent connections.
+	for i := range 3 {
+		conn, err := net.Dial("unix", proxySock)
+		if err != nil {
+			t.Fatalf("connection %d: %v", i, err)
+		}
+		client := agent.NewClient(conn).(agent.ExtendedAgent)
+		_, err = client.Extension(ExtensionAuth, nil)
+		if err != nil {
+			t.Fatalf("connection %d auth: %v", i, err)
+		}
+		conn.Close()
+	}
+
+	if callCount != 3 {
+		t.Errorf("expected 3 auth calls, got %d", callCount)
+	}
+}
+
+func TestMultiHopChain(t *testing.T) {
+	// Simulate: laptop → shell1 → shell2
+	// Each hop has its own proxy listener wrapping the previous.
+	// An auth request on shell2 should traverse back to the laptop.
+
+	// Laptop: vanilla agent + proxy with auth handler (the origin).
+	laptopUpstream, cleanupLaptop := startVanillaAgent(t)
+	defer cleanupLaptop()
+
+	laptopProxy := tempSocketPath(t)
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+
+	laptopAuthCalled := false
+	laptopSetup := func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(0))
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			laptopAuthCalled = true
+			return "laptop-origin-token", nil
+		}))
+	}
+	laptopListener := NewProxyListener(logger, laptopProxy, laptopUpstream, laptopSetup)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go laptopListener.Serve(ctx)
+	waitForSocket(t, laptopProxy)
+
+	// Verify laptop probe shows depth 0.
+	hello, err := ProbeUpstream(laptopProxy)
+	if err != nil {
+		t.Fatalf("probe laptop: %v", err)
+	}
+	if hello == nil {
+		t.Fatal("expected epithet agent at laptop, got nil")
+	}
+	if hello.ChainDepth != 0 {
+		t.Errorf("laptop depth: expected 0, got %d", hello.ChainDepth)
+	}
+
+	// Shell1: proxy wrapping laptop, with depth 1.
+	// Auth requests are forwarded to laptop via RequestAuth.
+	shell1Proxy := tempSocketPath(t)
+	shell1Setup := func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(1))
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			// Shell1 proxies auth to laptop.
+			return RequestAuth(laptopProxy)
+		}))
+	}
+	shell1Listener := NewProxyListener(logger, shell1Proxy, laptopProxy, shell1Setup)
+
+	go shell1Listener.Serve(ctx)
+	waitForSocket(t, shell1Proxy)
+
+	// Verify shell1 probe shows depth 1.
+	hello, err = ProbeUpstream(shell1Proxy)
+	if err != nil {
+		t.Fatalf("probe shell1: %v", err)
+	}
+	if hello.ChainDepth != 1 {
+		t.Errorf("shell1 depth: expected 1, got %d", hello.ChainDepth)
+	}
+
+	// Shell2: proxy wrapping shell1, with depth 2.
+	shell2Proxy := tempSocketPath(t)
+	shell2Setup := func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(2))
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			// Shell2 proxies auth to shell1.
+			return RequestAuth(shell1Proxy)
+		}))
+	}
+	shell2Listener := NewProxyListener(logger, shell2Proxy, shell1Proxy, shell2Setup)
+
+	go shell2Listener.Serve(ctx)
+	waitForSocket(t, shell2Proxy)
+
+	// Verify shell2 probe shows depth 2.
+	hello, err = ProbeUpstream(shell2Proxy)
+	if err != nil {
+		t.Fatalf("probe shell2: %v", err)
+	}
+	if hello.ChainDepth != 2 {
+		t.Errorf("shell2 depth: expected 2, got %d", hello.ChainDepth)
+	}
+
+	// Request auth from shell2 — should chain all the way back to laptop.
+	token, err := RequestAuth(shell2Proxy)
+	if err != nil {
+		t.Fatalf("multi-hop auth: %v", err)
+	}
+	if token != "laptop-origin-token" {
+		t.Errorf("expected 'laptop-origin-token', got %q", token)
+	}
+	if !laptopAuthCalled {
+		t.Error("laptop auth handler was never called — chain is broken")
+	}
+}
+
+// tempSocketPath returns a unique temporary socket path.
+func tempSocketPath(t *testing.T) string {
+	t.Helper()
+	// Use /tmp for short paths (macOS 104-byte socket path limit).
+	f, err := os.CreateTemp("/tmp", "ep-test.*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	path := f.Name()
+	f.Close()
+	os.Remove(path)
+	t.Cleanup(func() { os.Remove(path) })
+	return path
+}
+
+// waitForSocket polls until the Unix socket is accepting connections.
+func waitForSocket(t *testing.T, path string) {
+	t.Helper()
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		conn, err := net.Dial("unix", path)
+		if err == nil {
+			conn.Close()
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	t.Fatalf("socket %s never became available", path)
+}
diff --git a/pkg/agent/proxy_test.go b/pkg/agent/proxy_test.go
new file mode 100644
index 0000000..aadc7b1
--- /dev/null
+++ b/pkg/agent/proxy_test.go
@@ -0,0 +1,246 @@
+package agent
+
+import (
+	"encoding/json"
+	"errors"
+	"net"
+	"os"
+	"testing"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// startProxyServer creates a ProxyAgent over a real Unix socket, returning the
+// client-side ExtendedAgent and a cleanup function.
+func startProxyServer(t *testing.T, setup func(*ProxyAgent)) (agent.ExtendedAgent, func()) {
+	t.Helper()
+
+	upstream := agent.NewKeyring()
+
+	f, err := os.CreateTemp("", "proxy-test.*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	sockPath := f.Name()
+	f.Close()
+	os.Remove(sockPath)
+
+	listener, err := net.Listen("unix", sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				proxy := NewProxyAgent(upstream.(agent.ExtendedAgent))
+				setup(proxy)
+				agent.ServeAgent(proxy, c)
+			}(conn)
+		}
+	}()
+
+	conn, err := net.Dial("unix", sockPath)
+	if err != nil {
+		listener.Close()
+		os.Remove(sockPath)
+		t.Fatal(err)
+	}
+
+	client := agent.NewClient(conn)
+
+	cleanup := func() {
+		conn.Close()
+		listener.Close()
+		os.Remove(sockPath)
+	}
+
+	return client.(agent.ExtendedAgent), cleanup
+}
+
+// stripSuccessByte removes the leading SSH_AGENT_SUCCESS byte from an
+// extension response. The client returns the raw wire buffer which starts
+// with the type byte.
+func stripSuccessByte(t *testing.T, resp []byte) []byte {
+	t.Helper()
+	if len(resp) == 0 {
+		t.Fatal("empty response")
+	}
+	if resp[0] != agentSuccess {
+		t.Fatalf("expected first byte %d (SSH_AGENT_SUCCESS), got %d", agentSuccess, resp[0])
+	}
+	return resp[1:]
+}
+
+func TestHelloHandler(t *testing.T) {
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(0))
+	})
+	defer cleanup()
+
+	resp, err := client.Extension(ExtensionHello, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var hello HelloResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil {
+		t.Fatalf("failed to unmarshal hello response: %v", err)
+	}
+
+	if hello.ProtocolVersion != 1 {
+		t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion)
+	}
+	if hello.ChainDepth != 0 {
+		t.Errorf("expected chain_depth=0, got %d", hello.ChainDepth)
+	}
+}
+
+func TestHelloHandlerChainDepth(t *testing.T) {
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(3))
+	})
+	defer cleanup()
+
+	resp, err := client.Extension(ExtensionHello, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var hello HelloResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil {
+		t.Fatalf("failed to unmarshal: %v", err)
+	}
+
+	if hello.ChainDepth != 3 {
+		t.Errorf("expected chain_depth=3, got %d", hello.ChainDepth)
+	}
+}
+
+func TestAuthHandler(t *testing.T) {
+	called := false
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			called = true
+			return "test-token-123", nil
+		}))
+	})
+	defer cleanup()
+
+	resp, err := client.Extension(ExtensionAuth, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !called {
+		t.Error("authenticate function was not called")
+	}
+
+	var authResp AuthResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &authResp); err != nil {
+		t.Fatalf("failed to unmarshal: %v", err)
+	}
+
+	if authResp.Token != "test-token-123" {
+		t.Errorf("expected token 'test-token-123', got %q", authResp.Token)
+	}
+}
+
+func TestAuthHandlerError(t *testing.T) {
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			return "", errors.New("user cancelled")
+		}))
+	})
+	defer cleanup()
+
+	_, err := client.Extension(ExtensionAuth, nil)
+	if err == nil {
+		t.Fatal("expected error from failed auth, got nil")
+	}
+}
+
+func TestStandardOpsPassthrough(t *testing.T) {
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(0))
+	})
+	defer cleanup()
+
+	// List should work (empty keyring).
+	keys, err := client.List()
+	if err != nil {
+		t.Fatalf("List failed: %v", err)
+	}
+	if len(keys) != 0 {
+		t.Errorf("expected 0 keys, got %d", len(keys))
+	}
+}
+
+func TestUnregisteredExtensionForwardedToUpstream(t *testing.T) {
+	// The upstream keyring's Extension returns ErrExtensionUnsupported,
+	// which ServeAgent translates to SSH_AGENT_FAILURE.
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		// Register nothing — all extensions forward to upstream.
+	})
+	defer cleanup()
+
+	_, err := client.Extension("unknown@example.com", nil)
+	if err == nil {
+		t.Fatal("expected error for unknown extension, got nil")
+	}
+}
+
+func TestCustomExtensionHandler(t *testing.T) {
+	client, cleanup := startProxyServer(t, func(p *ProxyAgent) {
+		p.RegisterExtension("custom@example.com", func(contents []byte) ([]byte, error) {
+			return []byte("custom-response"), nil
+		})
+	})
+	defer cleanup()
+
+	resp, err := client.Extension("custom@example.com", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	payload := stripSuccessByte(t, resp)
+	if string(payload) != "custom-response" {
+		t.Errorf("expected 'custom-response', got %q", string(payload))
+	}
+}
+
+func TestDirectProxyAgentExtension(t *testing.T) {
+	// Test the ProxyAgent directly without going through the socket,
+	// verifying the raw wire format (includes success byte).
+	upstream := agent.NewKeyring()
+	proxy := NewProxyAgent(upstream.(agent.ExtendedAgent))
+	proxy.RegisterExtension(ExtensionHello, HelloHandler(0))
+
+	resp, err := proxy.Extension(ExtensionHello, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// First byte should be SSH_AGENT_SUCCESS (6).
+	if len(resp) == 0 {
+		t.Fatal("empty response")
+	}
+	if resp[0] != agentSuccess {
+		t.Errorf("expected first byte %d (SSH_AGENT_SUCCESS), got %d", agentSuccess, resp[0])
+	}
+
+	// Rest should be valid JSON.
+	var hello HelloResponse
+	if err := json.Unmarshal(resp[1:], &hello); err != nil {
+		t.Fatalf("failed to unmarshal payload: %v", err)
+	}
+	if hello.ProtocolVersion != 1 {
+		t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion)
+	}
+}
+
diff --git a/pkg/agent/upstream.go b/pkg/agent/upstream.go
new file mode 100644
index 0000000..2222e04
--- /dev/null
+++ b/pkg/agent/upstream.go
@@ -0,0 +1,85 @@
+package agent
+
+import (
+	"encoding/json"
+	"fmt"
+	"net"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// ProbeUpstream connects to the agent socket and sends epithet-hello.
+// Returns the hello response if the agent supports epithet extensions,
+// or nil if it's a vanilla ssh-agent.
+func ProbeUpstream(socketPath string) (*HelloResponse, error) {
+	conn, err := net.Dial("unix", socketPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to upstream agent: %w", err)
+	}
+	defer conn.Close()
+
+	client := agent.NewClient(conn)
+	extClient, ok := client.(agent.ExtendedAgent)
+	if !ok {
+		return nil, nil
+	}
+
+	resp, err := extClient.Extension(ExtensionHello, nil)
+	if err != nil {
+		// ErrExtensionUnsupported means it's a vanilla ssh-agent.
+		if err == agent.ErrExtensionUnsupported {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("hello extension failed: %w", err)
+	}
+
+	// Strip the SSH_AGENT_SUCCESS byte.
+	if len(resp) > 0 && resp[0] == agentSuccess {
+		resp = resp[1:]
+	}
+
+	var hello HelloResponse
+	if err := json.Unmarshal(resp, &hello); err != nil {
+		return nil, fmt.Errorf("failed to parse hello response: %w", err)
+	}
+
+	return &hello, nil
+}
+
+// RequestAuth sends an epithet-auth extension request to the upstream agent.
+// The upstream runs its own auth plugin with its own state — only the token
+// is returned.
+func RequestAuth(socketPath string) (string, error) {
+	conn, err := net.Dial("unix", socketPath)
+	if err != nil {
+		return "", fmt.Errorf("failed to connect to upstream agent: %w", err)
+	}
+	defer conn.Close()
+
+	client := agent.NewClient(conn)
+	extClient, ok := client.(agent.ExtendedAgent)
+	if !ok {
+		return "", fmt.Errorf("upstream agent does not support extensions")
+	}
+
+	resp, err := extClient.Extension(ExtensionAuth, nil)
+	if err != nil {
+		return "", fmt.Errorf("auth extension failed: %w", err)
+	}
+
+	// Strip the SSH_AGENT_SUCCESS byte.
+	if len(resp) > 0 && resp[0] == agentSuccess {
+		resp = resp[1:]
+	}
+
+	var authResp AuthResponse
+	if err := json.Unmarshal(resp, &authResp); err != nil {
+		return "", fmt.Errorf("failed to parse auth response: %w", err)
+	}
+
+	if authResp.Token == "" {
+		return "", fmt.Errorf("upstream returned empty token")
+	}
+
+	return authResp.Token, nil
+}
diff --git a/pkg/agent/upstream_test.go b/pkg/agent/upstream_test.go
new file mode 100644
index 0000000..fd54ae9
--- /dev/null
+++ b/pkg/agent/upstream_test.go
@@ -0,0 +1,152 @@
+package agent
+
+import (
+	"net"
+	"os"
+	"testing"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// startProxyServerSocket creates a ProxyAgent over a real Unix socket, returning
+// the socket path and a cleanup function. Used to test upstream client functions.
+func startProxyServerSocket(t *testing.T, setup func(*ProxyAgent)) (string, func()) {
+	t.Helper()
+
+	upstream := agent.NewKeyring()
+
+	f, err := os.CreateTemp("", "upstream-test.*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	sockPath := f.Name()
+	f.Close()
+	os.Remove(sockPath)
+
+	listener, err := net.Listen("unix", sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				proxy := NewProxyAgent(upstream.(agent.ExtendedAgent))
+				setup(proxy)
+				agent.ServeAgent(proxy, c)
+			}(conn)
+		}
+	}()
+
+	cleanup := func() {
+		listener.Close()
+		os.Remove(sockPath)
+	}
+
+	return sockPath, cleanup
+}
+
+func TestProbeUpstreamEpithet(t *testing.T) {
+	sockPath, cleanup := startProxyServerSocket(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(2))
+	})
+	defer cleanup()
+
+	hello, err := ProbeUpstream(sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if hello == nil {
+		t.Fatal("expected hello response, got nil")
+	}
+	if hello.ProtocolVersion != 1 {
+		t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion)
+	}
+	if hello.ChainDepth != 2 {
+		t.Errorf("expected chain_depth=2, got %d", hello.ChainDepth)
+	}
+}
+
+func TestProbeUpstreamVanillaAgent(t *testing.T) {
+	// Serve a plain keyring (no extensions registered) — behaves like a vanilla ssh-agent.
+	f, err := os.CreateTemp("", "vanilla-test.*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	sockPath := f.Name()
+	f.Close()
+	os.Remove(sockPath)
+
+	listener, err := net.Listen("unix", sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		listener.Close()
+		os.Remove(sockPath)
+	}()
+
+	keyring := agent.NewKeyring()
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				agent.ServeAgent(keyring, c)
+			}(conn)
+		}
+	}()
+
+	hello, err := ProbeUpstream(sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if hello != nil {
+		t.Errorf("expected nil for vanilla agent, got %+v", hello)
+	}
+}
+
+func TestProbeUpstreamBadSocket(t *testing.T) {
+	_, err := ProbeUpstream("/nonexistent/socket.sock")
+	if err == nil {
+		t.Fatal("expected error for bad socket path")
+	}
+}
+
+func TestRequestAuth(t *testing.T) {
+	sockPath, cleanup := startProxyServerSocket(t, func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			return "upstream-token-456", nil
+		}))
+	})
+	defer cleanup()
+
+	token, err := RequestAuth(sockPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if token != "upstream-token-456" {
+		t.Errorf("expected 'upstream-token-456', got %q", token)
+	}
+}
+
+func TestRequestAuthNoExtensionSupport(t *testing.T) {
+	// Serve a proxy with no auth handler registered — upstream returns ErrExtensionUnsupported.
+	sockPath, cleanup := startProxyServerSocket(t, func(p *ProxyAgent) {
+		// No extensions registered.
+	})
+	defer cleanup()
+
+	_, err := RequestAuth(sockPath)
+	if err == nil {
+		t.Fatal("expected error when auth extension not supported")
+	}
+}
diff --git a/pkg/broker/auth.go b/pkg/broker/auth.go
index cd69366..e0a8901 100644
--- a/pkg/broker/auth.go
+++ b/pkg/broker/auth.go
@@ -50,6 +50,15 @@ func (h *Auth) Token() string {
 	return h.token
 }
 
+// SetToken injects a token obtained from an upstream agent.
+// Used when the broker receives a token via the agent extension protocol
+// instead of running a local auth command.
+func (h *Auth) SetToken(token string) {
+	h.lock.Lock()
+	defer h.lock.Unlock()
+	h.token = token
+}
+
 // ClearToken clears the stored authentication token.
 // Keeps the state intact (refresh token may still be valid).
 // Used when receiving HTTP 401 from CA to force re-authentication.
diff --git a/pkg/broker/broker.go b/pkg/broker/broker.go
index 9d63a79..5da8225 100644
--- a/pkg/broker/broker.go
+++ b/pkg/broker/broker.go
@@ -69,6 +69,7 @@ type Broker struct {
 
 	caClient       *caclient.Client // Immutable after New()
 	agentSocketDir string           // Immutable after New()
+	upstreamSocket string           // Immutable after New(); empty if no upstream
 
 	// For graceful shutdown: track in-flight RPC connections
 	activeRPC       sync.WaitGroup
@@ -86,6 +87,15 @@ func (f optionFunc) apply(b *Broker) error {
 	return f(b)
 }
 
+// WithUpstream configures the broker to proxy auth requests to an upstream
+// epithet agent at the given socket path before falling back to local auth.
+func WithUpstream(socketPath string) Option {
+	return optionFunc(func(b *Broker) error {
+		b.upstreamSocket = socketPath
+		return nil
+	})
+}
+
 // New creates a new Broker instance. This does not start listening - call Serve() to begin accepting connections.
 func New(log slog.Logger, socketPath string, authCommand string, caClient *caclient.Client, agentSocketDir string, options ...Option) (*Broker, error) {
 	if caClient == nil {
@@ -276,7 +286,7 @@ func (b *Broker) MatchWithUserOutput(conn policy.Connection, userOutput io.Write
 		token := b.auth.Token()
 		if token == "" {
 			b.log.Debug("no auth token, authenticating")
-			token, err = b.auth.Run(nil, userOutput)
+			token, err = b.Authenticate(userOutput)
 			if err != nil {
 				// Auth command failed - don't retry automatically.
 				// User should fix the issue and retry the SSH connection.
@@ -454,6 +464,21 @@ func (b *Broker) ensureAgent(connectionHash policy.ConnectionHash, credential ag
 	return nil
 }
 
+// Authenticate obtains an auth token, trying upstream first if configured.
+// This is also used by the agent proxy's extension handler to run the local
+// auth flow when responding to downstream epithet-auth requests.
+func (b *Broker) Authenticate(userOutput io.Writer) (string, error) {
+	if b.upstreamSocket != "" {
+		token, err := agent.RequestAuth(b.upstreamSocket)
+		if err == nil {
+			b.auth.SetToken(token)
+			return token, nil
+		}
+		b.log.Warn("upstream auth failed, falling back to local", "error", err)
+	}
+	return b.auth.Run(nil, userOutput)
+}
+
 // shouldHandle checks if the given hostname matches discovery patterns.
 // Always fetches discovery patterns, authenticating if needed.
 // Returns true if epithet should handle this connection, false otherwise.
@@ -505,7 +530,7 @@ func (b *Broker) getDiscoveryPatterns(userOutput io.Writer) (*caclient.Discovery
 	if token == "" {
 		b.log.Debug("no auth token, authenticating to get discovery")
 		var err error
-		token, err = b.auth.Run(nil, userOutput)
+		token, err = b.Authenticate(userOutput)
 		if err != nil {
 			return nil, fmt.Errorf("authentication failed: %w", err)
 		}
@@ -520,7 +545,7 @@ func (b *Broker) getDiscoveryPatterns(userOutput io.Writer) (*caclient.Discovery
 		if errors.As(err, &invalidToken) {
 			b.log.Debug("token expired during Hello, re-authenticating")
 			b.auth.ClearToken()
-			token, err = b.auth.Run(nil, userOutput)
+			token, err = b.Authenticate(userOutput)
 			if err != nil {
 				return nil, fmt.Errorf("authentication failed: %w", err)
 			}
diff --git a/pkg/broker/broker_test.go b/pkg/broker/broker_test.go
index 5719cb2..4beb354 100644
--- a/pkg/broker/broker_test.go
+++ b/pkg/broker/broker_test.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -14,11 +15,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/epithet-ssh/epithet/pkg/agent"
 	pb "github.com/epithet-ssh/epithet/pkg/brokerv1"
 	"github.com/epithet-ssh/epithet/pkg/caclient"
 	"github.com/epithet-ssh/epithet/pkg/policy"
 	"github.com/lmittmann/tint"
 	"github.com/stretchr/testify/require"
+	sshagent "golang.org/x/crypto/ssh/agent"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 )
@@ -326,6 +329,157 @@ printf '%s' "test-token"
 	require.Empty(t, result2.Error)
 }
 
+func TestAuthenticate_WithUpstream(t *testing.T) {
+	t.Parallel()
+
+	// Start an upstream agent that serves epithet-auth.
+	upstreamSock := startUpstreamAuthAgent(t, "upstream-token-from-test")
+
+	tmpDir := shortTempDir(t)
+	socketPath := tmpDir + "/b.sock"
+	agentDir := tmpDir + "/a"
+
+	// Local auth should NOT be called — upstream provides the token.
+	localAuthScript := writeTestScript(t, `#!/bin/sh
+echo "ERROR: local auth should not have been called" >&2
+exit 1
+`)
+	b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir, WithUpstream(upstreamSock))
+	require.NoError(t, err)
+
+	token, err := b.Authenticate(nil)
+	require.NoError(t, err)
+	require.Equal(t, "upstream-token-from-test", token)
+
+	// Token should be cached in auth.
+	require.Equal(t, "upstream-token-from-test", b.auth.Token())
+}
+
+func TestAuthenticate_FallsBackToLocal(t *testing.T) {
+	t.Parallel()
+
+	// Start an upstream agent that does NOT serve epithet-auth (no extensions).
+	upstreamSock := startPlainAgent(t)
+
+	tmpDir := shortTempDir(t)
+	socketPath := tmpDir + "/b.sock"
+	agentDir := tmpDir + "/a"
+
+	// Local auth will be called as fallback.
+	localAuthScript := writeTestScript(t, `#!/bin/sh
+cat > /dev/null
+printf '%s' "local-fallback-token"
+`)
+	b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir, WithUpstream(upstreamSock))
+	require.NoError(t, err)
+
+	token, err := b.Authenticate(nil)
+	require.NoError(t, err)
+	// Local auth base64url-encodes the token.
+	require.NotEmpty(t, token)
+}
+
+func TestAuthenticate_NoUpstream(t *testing.T) {
+	t.Parallel()
+
+	tmpDir := shortTempDir(t)
+	socketPath := tmpDir + "/b.sock"
+	agentDir := tmpDir + "/a"
+
+	localAuthScript := writeTestScript(t, `#!/bin/sh
+cat > /dev/null
+printf '%s' "local-only-token"
+`)
+	// No WithUpstream — should go straight to local auth.
+	b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir)
+	require.NoError(t, err)
+
+	token, err := b.Authenticate(nil)
+	require.NoError(t, err)
+	require.NotEmpty(t, token)
+}
+
+// startUpstreamAuthAgent starts a proxy agent that handles epithet-auth and returns the given token.
+func startUpstreamAuthAgent(t *testing.T, token string) string {
+	t.Helper()
+
+	tmpDir := shortTempDir(t)
+	upstreamSock := tmpDir + "/upstream.sock"
+	vanillaSock := tmpDir + "/vanilla.sock"
+
+	// Start a vanilla agent as the upstream's upstream.
+	vanillaListener, err := net.Listen("unix", vanillaSock)
+	require.NoError(t, err)
+	t.Cleanup(func() { vanillaListener.Close() })
+
+	keyring := sshagent.NewKeyring()
+	go func() {
+		for {
+			conn, err := vanillaListener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				sshagent.ServeAgent(keyring, c)
+			}(conn)
+		}
+	}()
+
+	// Start the proxy listener with auth handler.
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+	setup := func(p *agent.ProxyAgent) {
+		p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) {
+			return token, nil
+		}))
+	}
+	proxy := agent.NewProxyListener(logger, upstreamSock, vanillaSock, setup)
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	go proxy.Serve(ctx)
+
+	// Wait for socket.
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		conn, err := net.Dial("unix", upstreamSock)
+		if err == nil {
+			conn.Close()
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	return upstreamSock
+}
+
+// startPlainAgent starts a vanilla ssh-agent (no epithet extensions).
+func startPlainAgent(t *testing.T) string {
+	t.Helper()
+
+	tmpDir := shortTempDir(t)
+	sockPath := tmpDir + "/plain.sock"
+
+	listener, err := net.Listen("unix", sockPath)
+	require.NoError(t, err)
+	t.Cleanup(func() { listener.Close() })
+
+	keyring := sshagent.NewKeyring()
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				sshagent.ServeAgent(keyring, c)
+			}(conn)
+		}
+	}()
+
+	return sockPath
+}
+
 func testLogger(t *testing.T) *slog.Logger {
 	logger := slog.New(tint.NewHandler(t.Output(), &tint.Options{
 		Level:      slog.LevelDebug,

From 309e833d246b44d5c6462766f2257c40e48741da Mon Sep 17 00:00:00 2001
From: Brian McCallister <brianm@skife.org>
Date: Sat, 18 Apr 2026 16:07:52 -0700
Subject: [PATCH 2/3] fix: address review findings for remote agent tunneling

- Add Ready() channel to ProxyListener with guaranteed close on all
  return paths, preventing caller hangs on listener startup failure
- Use per-process temp dirs in wrapper mode to avoid collisions between
  multiple wrapped shells against the same CA
- Only fall back to local auth on UpstreamUnavailableError; propagate
  genuine upstream auth failures (user cancelled, plugin error)
- Use errors.Is for ErrExtensionUnsupported sentinel comparison
- Fix test data races with atomic.Bool/Int32 for shared counters
- Document token forwarding trust model in extensions.go
- Drain brokerErr channel after shell exit in wrapper mode
---
 cmd/epithet/agent.go             | 64 +++++++++++++++++----------
 pkg/agent/extensions.go          |  7 +++
 pkg/agent/proxy_listener.go      | 12 +++++
 pkg/agent/proxy_listener_test.go | 46 +++++++------------
 pkg/agent/upstream.go            | 33 ++++++++++++--
 pkg/broker/broker.go             | 15 ++++++-
 pkg/broker/broker_test.go        | 76 +++++++++++++++++++++++++++-----
 7 files changed, 184 insertions(+), 69 deletions(-)

diff --git a/cmd/epithet/agent.go b/cmd/epithet/agent.go
index d6b1f3b..3c2be71 100644
--- a/cmd/epithet/agent.go
+++ b/cmd/epithet/agent.go
@@ -50,7 +50,7 @@ func (s *AgentStartCLI) Run(parent *AgentCLI, logger *slog.Logger, tlsCfg tlscon
 
 // runDaemon runs the broker as a long-lived daemon (original behavior).
 func (s *AgentStartCLI) runDaemon(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error {
-	b, tempDir, brokerSock, agentDir, homeDir, err := setupBroker(parent, logger, tlsCfg)
+	b, tempDir, brokerSock, agentDir, homeDir, err := setupBrokerWithOptions(parent, logger, tlsCfg, false)
 	if err != nil {
 		return err
 	}
@@ -124,7 +124,7 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 		logger.Debug("no SSH_AUTH_SOCK set")
 	}
 
-	b, tempDir, brokerSock, agentDir, homeDir, err := setupBrokerWithOptions(parent, logger, tlsCfg, brokerOpts...)
+	b, tempDir, brokerSock, agentDir, homeDir, err := setupBrokerWithOptions(parent, logger, tlsCfg, true, brokerOpts...)
 	if err != nil {
 		return err
 	}
@@ -173,8 +173,7 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 				logger.Error("proxy listener error", "error", err)
 			}
 		}()
-		// Wait briefly for the listener to start.
-		// TODO: add Ready() to ProxyListener like Broker has.
+		<-proxyListener.Ready()
 		upstreamSocket = proxySock
 		logger.Info("proxy agent listening", "socket", proxySock)
 	}
@@ -204,30 +203,32 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 	}()
 
 	logger.Info("starting shell", "shell", s.Shell)
-	if err := cmd.Run(); err != nil {
+	shellErr := cmd.Run()
+
+	// Shell exited — shut down broker and drain its error.
+	cancel()
+	signal.Stop(sigChan)
+	if err := <-brokerErr; err != nil && err != context.Canceled {
+		logger.Error("broker error", "error", err)
+	}
+
+	if shellErr != nil {
 		var exitErr *exec.ExitError
-		if errors.As(err, &exitErr) {
-			cancel()
-			signal.Stop(sigChan)
+		if errors.As(shellErr, &exitErr) {
 			os.Exit(exitErr.ExitCode())
 		}
-		cancel()
-		return fmt.Errorf("shell failed: %w", err)
+		return fmt.Errorf("shell failed: %w", shellErr)
 	}
 
-	cancel()
-	signal.Stop(sigChan)
 	return nil
 }
 
-// setupBroker creates the broker and supporting directories with no extra options.
-func setupBroker(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) (*broker.Broker, string, string, string, string, error) {
-	return setupBrokerWithOptions(parent, logger, tlsCfg)
-}
-
 // setupBrokerWithOptions creates the broker, temp directories, and resolves auth.
+// When wrapperMode is true, uses a per-process temp directory to avoid collisions
+// between multiple wrapped shells. When false (daemon mode), uses a deterministic
+// path derived from the CA URLs so SSH config includes can find it.
 // Returns (broker, tempDir, brokerSock, agentDir, homeDir, error).
-func setupBrokerWithOptions(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config, opts ...broker.Option) (*broker.Broker, string, string, string, string, error) {
+func setupBrokerWithOptions(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config, wrapperMode bool, opts ...broker.Option) (*broker.Broker, string, string, string, string, error) {
 	if len(parent.CaURL) == 0 {
 		return nil, "", "", "", "", fmt.Errorf("--ca-url is required (at least one)")
 	}
@@ -250,14 +251,31 @@ func setupBrokerWithOptions(parent *AgentCLI, logger *slog.Logger, tlsCfg tlscon
 		return nil, "", "", "", "", fmt.Errorf("failed to get home directory: %w", err)
 	}
 
-	instanceID := hashString(fmt.Sprintf("%v", parent.CaURL))
 	runDir := filepath.Join(homeDir, ".epithet", "run")
-	tempDir := filepath.Join(runDir, instanceID)
-
 	cleanupStaleRunDirs(runDir, logger)
 
-	if err := os.MkdirAll(tempDir, 0700); err != nil {
-		return nil, "", "", "", "", fmt.Errorf("failed to create temp directory: %w", err)
+	var tempDir string
+	if wrapperMode {
+		// Wrapper mode: each invocation gets its own directory to avoid
+		// collisions between multiple wrapped shells against the same CA.
+		tempDir, err = os.MkdirTemp(runDir, "wrap-")
+		if err != nil {
+			// Ensure runDir exists and retry.
+			if mkErr := os.MkdirAll(runDir, 0700); mkErr != nil {
+				return nil, "", "", "", "", fmt.Errorf("failed to create run directory: %w", mkErr)
+			}
+			tempDir, err = os.MkdirTemp(runDir, "wrap-")
+			if err != nil {
+				return nil, "", "", "", "", fmt.Errorf("failed to create temp directory: %w", err)
+			}
+		}
+	} else {
+		// Daemon mode: deterministic path so SSH config includes can find it.
+		instanceID := hashString(fmt.Sprintf("%v", parent.CaURL))
+		tempDir = filepath.Join(runDir, instanceID)
+		if err := os.MkdirAll(tempDir, 0700); err != nil {
+			return nil, "", "", "", "", fmt.Errorf("failed to create temp directory: %w", err)
+		}
 	}
 
 	pidFile := filepath.Join(tempDir, "broker.pid")
diff --git a/pkg/agent/extensions.go b/pkg/agent/extensions.go
index 4f18a8b..5d08bc4 100644
--- a/pkg/agent/extensions.go
+++ b/pkg/agent/extensions.go
@@ -35,6 +35,13 @@ func HelloHandler(depth int) ExtensionHandler {
 // AuthHandler returns an ExtensionHandler for epithet-auth@epithet.dev.
 // The authenticate function is called to obtain a token — typically wired
 // to the broker's auth flow.
+//
+// Security: the returned token is a bearer credential that transits the SSH
+// agent forwarding channel. Any host in the forwarding chain (i.e., the
+// remote sshd process) can observe the token. This is the same trust model
+// as standard SSH agent forwarding — the remote host can use the forwarded
+// agent to sign challenges on the user's behalf. Users should only forward
+// agents to hosts they trust, same as with ssh-agent -A.
 func AuthHandler(authenticate func() (string, error)) ExtensionHandler {
 	return func(_ []byte) ([]byte, error) {
 		token, err := authenticate()
diff --git a/pkg/agent/proxy_listener.go b/pkg/agent/proxy_listener.go
index 5a67bf3..888c0c4 100644
--- a/pkg/agent/proxy_listener.go
+++ b/pkg/agent/proxy_listener.go
@@ -25,6 +25,7 @@ type ProxyListener struct {
 	log                *slog.Logger
 
 	listener  net.Listener
+	ready     chan struct{} // Closed when listener is accepting connections.
 	done      chan struct{}
 	closeOnce sync.Once
 }
@@ -36,22 +37,28 @@ func NewProxyListener(log *slog.Logger, socketPath, upstreamSocketPath string, s
 		upstreamSocketPath: upstreamSocketPath,
 		setup:              setup,
 		log:                log,
+		ready:              make(chan struct{}),
 		done:               make(chan struct{}),
 	}
 }
 
 // Serve starts the listener and blocks until the context is cancelled.
+// The ready channel is always closed before Serve returns, whether
+// successfully or on error, so callers blocking on Ready() won't hang.
 func (p *ProxyListener) Serve(ctx context.Context) error {
 	os.Remove(p.socketPath)
 	listener, err := net.Listen("unix", p.socketPath)
 	if err != nil {
+		close(p.ready)
 		return err
 	}
 	if err := os.Chmod(p.socketPath, 0600); err != nil {
 		listener.Close()
+		close(p.ready)
 		return err
 	}
 	p.listener = listener
+	close(p.ready)
 
 	go p.acceptLoop()
 
@@ -106,6 +113,11 @@ func (p *ProxyListener) serveConn(conn net.Conn) {
 	}
 }
 
+// Ready returns a channel that is closed when the listener is accepting connections.
+func (p *ProxyListener) Ready() <-chan struct{} {
+	return p.ready
+}
+
 // SocketPath returns the path to the proxy listener's Unix socket.
 func (p *ProxyListener) SocketPath() string {
 	return p.socketPath
diff --git a/pkg/agent/proxy_listener_test.go b/pkg/agent/proxy_listener_test.go
index b55f38f..ed757f2 100644
--- a/pkg/agent/proxy_listener_test.go
+++ b/pkg/agent/proxy_listener_test.go
@@ -6,8 +6,8 @@ import (
 	"log/slog"
 	"net"
 	"os"
+	"sync/atomic"
 	"testing"
-	"time"
 
 	"golang.org/x/crypto/ssh/agent"
 )
@@ -58,11 +58,11 @@ func TestProxyListenerEndToEnd(t *testing.T) {
 	proxySock := tempSocketPath(t)
 	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
 
-	authCalled := false
+	var authCalled atomic.Bool
 	setup := func(p *ProxyAgent) {
 		p.RegisterExtension(ExtensionHello, HelloHandler(0))
 		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
-			authCalled = true
+			authCalled.Store(true)
 			return "proxy-listener-token", nil
 		}))
 	}
@@ -72,7 +72,7 @@ func TestProxyListenerEndToEnd(t *testing.T) {
 	defer cancel()
 
 	go proxy.Serve(ctx)
-	waitForSocket(t, proxySock)
+	<-proxy.Ready()
 
 	// Connect to the proxy and test extensions.
 	conn, err := net.Dial("unix", proxySock)
@@ -108,7 +108,7 @@ func TestProxyListenerEndToEnd(t *testing.T) {
 	if authResp.Token != "proxy-listener-token" {
 		t.Errorf("expected 'proxy-listener-token', got %q", authResp.Token)
 	}
-	if !authCalled {
+	if !authCalled.Load() {
 		t.Error("auth handler was not called")
 	}
 
@@ -130,11 +130,11 @@ func TestProxyListenerMultipleConnections(t *testing.T) {
 	proxySock := tempSocketPath(t)
 	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
 
-	callCount := 0
+	var callCount atomic.Int32
 	setup := func(p *ProxyAgent) {
 		p.RegisterExtension(ExtensionHello, HelloHandler(0))
 		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
-			callCount++
+			callCount.Add(1)
 			return "token", nil
 		}))
 	}
@@ -144,7 +144,7 @@ func TestProxyListenerMultipleConnections(t *testing.T) {
 	defer cancel()
 
 	go proxy.Serve(ctx)
-	waitForSocket(t, proxySock)
+	<-proxy.Ready()
 
 	// Make two independent connections.
 	for i := range 3 {
@@ -160,8 +160,8 @@ func TestProxyListenerMultipleConnections(t *testing.T) {
 		conn.Close()
 	}
 
-	if callCount != 3 {
-		t.Errorf("expected 3 auth calls, got %d", callCount)
+	if callCount.Load() != 3 {
+		t.Errorf("expected 3 auth calls, got %d", callCount.Load())
 	}
 }
 
@@ -177,11 +177,11 @@ func TestMultiHopChain(t *testing.T) {
 	laptopProxy := tempSocketPath(t)
 	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
 
-	laptopAuthCalled := false
+	var laptopAuthCalled atomic.Bool
 	laptopSetup := func(p *ProxyAgent) {
 		p.RegisterExtension(ExtensionHello, HelloHandler(0))
 		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
-			laptopAuthCalled = true
+			laptopAuthCalled.Store(true)
 			return "laptop-origin-token", nil
 		}))
 	}
@@ -191,7 +191,7 @@ func TestMultiHopChain(t *testing.T) {
 	defer cancel()
 
 	go laptopListener.Serve(ctx)
-	waitForSocket(t, laptopProxy)
+	<-laptopListener.Ready()
 
 	// Verify laptop probe shows depth 0.
 	hello, err := ProbeUpstream(laptopProxy)
@@ -218,7 +218,7 @@ func TestMultiHopChain(t *testing.T) {
 	shell1Listener := NewProxyListener(logger, shell1Proxy, laptopProxy, shell1Setup)
 
 	go shell1Listener.Serve(ctx)
-	waitForSocket(t, shell1Proxy)
+	<-shell1Listener.Ready()
 
 	// Verify shell1 probe shows depth 1.
 	hello, err = ProbeUpstream(shell1Proxy)
@@ -241,7 +241,7 @@ func TestMultiHopChain(t *testing.T) {
 	shell2Listener := NewProxyListener(logger, shell2Proxy, shell1Proxy, shell2Setup)
 
 	go shell2Listener.Serve(ctx)
-	waitForSocket(t, shell2Proxy)
+	<-shell2Listener.Ready()
 
 	// Verify shell2 probe shows depth 2.
 	hello, err = ProbeUpstream(shell2Proxy)
@@ -260,7 +260,7 @@ func TestMultiHopChain(t *testing.T) {
 	if token != "laptop-origin-token" {
 		t.Errorf("expected 'laptop-origin-token', got %q", token)
 	}
-	if !laptopAuthCalled {
+	if !laptopAuthCalled.Load() {
 		t.Error("laptop auth handler was never called — chain is broken")
 	}
 }
@@ -280,17 +280,3 @@ func tempSocketPath(t *testing.T) string {
 	return path
 }
 
-// waitForSocket polls until the Unix socket is accepting connections.
-func waitForSocket(t *testing.T, path string) {
-	t.Helper()
-	deadline := time.Now().Add(2 * time.Second)
-	for time.Now().Before(deadline) {
-		conn, err := net.Dial("unix", path)
-		if err == nil {
-			conn.Close()
-			return
-		}
-		time.Sleep(10 * time.Millisecond)
-	}
-	t.Fatalf("socket %s never became available", path)
-}
diff --git a/pkg/agent/upstream.go b/pkg/agent/upstream.go
index 2222e04..b6f28a5 100644
--- a/pkg/agent/upstream.go
+++ b/pkg/agent/upstream.go
@@ -2,6 +2,7 @@ package agent
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net"
 
@@ -27,7 +28,7 @@ func ProbeUpstream(socketPath string) (*HelloResponse, error) {
 	resp, err := extClient.Extension(ExtensionHello, nil)
 	if err != nil {
 		// ErrExtensionUnsupported means it's a vanilla ssh-agent.
-		if err == agent.ErrExtensionUnsupported {
+		if errors.Is(err, agent.ErrExtensionUnsupported) {
 			return nil, nil
 		}
 		return nil, fmt.Errorf("hello extension failed: %w", err)
@@ -46,25 +47,49 @@ func ProbeUpstream(socketPath string) (*HelloResponse, error) {
 	return &hello, nil
 }
 
+// UpstreamUnavailableError indicates that the upstream agent could not be
+// reached or does not support epithet extensions. This is the only class
+// of error where falling back to local auth is appropriate.
+type UpstreamUnavailableError struct {
+	Err error
+}
+
+func (e *UpstreamUnavailableError) Error() string {
+	return fmt.Sprintf("upstream unavailable: %v", e.Err)
+}
+
+func (e *UpstreamUnavailableError) Unwrap() error {
+	return e.Err
+}
+
 // RequestAuth sends an epithet-auth extension request to the upstream agent.
 // The upstream runs its own auth plugin with its own state — only the token
 // is returned.
+//
+// Returns UpstreamUnavailableError if the upstream can't be reached or doesn't
+// support the extension (safe to fall back to local auth). Other errors indicate
+// the upstream attempted auth and failed (should not fall back).
 func RequestAuth(socketPath string) (string, error) {
 	conn, err := net.Dial("unix", socketPath)
 	if err != nil {
-		return "", fmt.Errorf("failed to connect to upstream agent: %w", err)
+		return "", &UpstreamUnavailableError{Err: err}
 	}
 	defer conn.Close()
 
 	client := agent.NewClient(conn)
 	extClient, ok := client.(agent.ExtendedAgent)
 	if !ok {
-		return "", fmt.Errorf("upstream agent does not support extensions")
+		return "", &UpstreamUnavailableError{Err: fmt.Errorf("agent does not support extensions")}
 	}
 
 	resp, err := extClient.Extension(ExtensionAuth, nil)
 	if err != nil {
-		return "", fmt.Errorf("auth extension failed: %w", err)
+		if errors.Is(err, agent.ErrExtensionUnsupported) {
+			return "", &UpstreamUnavailableError{Err: err}
+		}
+		// Extension was dispatched but handler returned an error — this is a
+		// genuine auth failure (user cancelled, plugin error, etc.).
+		return "", fmt.Errorf("upstream auth failed: %w", err)
 	}
 
 	// Strip the SSH_AGENT_SUCCESS byte.
diff --git a/pkg/broker/broker.go b/pkg/broker/broker.go
index 5da8225..a381a73 100644
--- a/pkg/broker/broker.go
+++ b/pkg/broker/broker.go
@@ -467,6 +467,11 @@ func (b *Broker) ensureAgent(connectionHash policy.ConnectionHash, credential ag
 // Authenticate obtains an auth token, trying upstream first if configured.
 // This is also used by the agent proxy's extension handler to run the local
 // auth flow when responding to downstream epithet-auth requests.
+//
+// When an upstream is configured, only falls back to local auth if the upstream
+// is unreachable or doesn't support extensions (UpstreamUnavailableError).
+// Genuine upstream auth failures (user cancelled, plugin error) are not retried
+// locally — the remote host typically cannot complete browser/FIDO flows.
 func (b *Broker) Authenticate(userOutput io.Writer) (string, error) {
 	if b.upstreamSocket != "" {
 		token, err := agent.RequestAuth(b.upstreamSocket)
@@ -474,7 +479,15 @@ func (b *Broker) Authenticate(userOutput io.Writer) (string, error) {
 			b.auth.SetToken(token)
 			return token, nil
 		}
-		b.log.Warn("upstream auth failed, falling back to local", "error", err)
+
+		var unavailable *agent.UpstreamUnavailableError
+		if errors.As(err, &unavailable) {
+			b.log.Warn("upstream unavailable, falling back to local auth", "error", err)
+		} else {
+			// Upstream was reached and auth was attempted but failed.
+			// Don't fall back — the failure is authoritative.
+			return "", err
+		}
 	}
 	return b.auth.Run(nil, userOutput)
 }
diff --git a/pkg/broker/broker_test.go b/pkg/broker/broker_test.go
index 4beb354..cff3b3b 100644
--- a/pkg/broker/broker_test.go
+++ b/pkg/broker/broker_test.go
@@ -399,6 +399,70 @@ printf '%s' "local-only-token"
 	require.NotEmpty(t, token)
 }
 
+func TestAuthenticate_UpstreamAuthFailure_NoFallback(t *testing.T) {
+	t.Parallel()
+
+	// Start an upstream agent whose auth handler returns an error.
+	upstreamSock := startUpstreamFailingAuthAgent(t)
+
+	tmpDir := shortTempDir(t)
+	socketPath := tmpDir + "/b.sock"
+	agentDir := tmpDir + "/a"
+
+	// Local auth should NOT be called — upstream failure should propagate.
+	localAuthScript := writeTestScript(t, `#!/bin/sh
+echo "ERROR: local auth should not have been called" >&2
+exit 1
+`)
+	b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir, WithUpstream(upstreamSock))
+	require.NoError(t, err)
+
+	_, err = b.Authenticate(nil)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "upstream auth failed")
+}
+
+// startUpstreamFailingAuthAgent starts a proxy agent whose auth handler always fails.
+func startUpstreamFailingAuthAgent(t *testing.T) string {
+	t.Helper()
+
+	tmpDir := shortTempDir(t)
+	upstreamSock := tmpDir + "/upstream.sock"
+	vanillaSock := tmpDir + "/vanilla.sock"
+
+	vanillaListener, err := net.Listen("unix", vanillaSock)
+	require.NoError(t, err)
+	t.Cleanup(func() { vanillaListener.Close() })
+
+	keyring := sshagent.NewKeyring()
+	go func() {
+		for {
+			conn, err := vanillaListener.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				sshagent.ServeAgent(keyring, c)
+			}(conn)
+		}
+	}()
+
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+	setup := func(p *agent.ProxyAgent) {
+		p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) {
+			return "", fmt.Errorf("user cancelled authentication")
+		}))
+	}
+	proxy := agent.NewProxyListener(logger, upstreamSock, vanillaSock, setup)
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	go proxy.Serve(ctx)
+	<-proxy.Ready()
+
+	return upstreamSock
+}
+
 // startUpstreamAuthAgent starts a proxy agent that handles epithet-auth and returns the given token.
 func startUpstreamAuthAgent(t *testing.T, token string) string {
 	t.Helper()
@@ -437,17 +501,7 @@ func startUpstreamAuthAgent(t *testing.T, token string) string {
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 	go proxy.Serve(ctx)
-
-	// Wait for socket.
-	deadline := time.Now().Add(2 * time.Second)
-	for time.Now().Before(deadline) {
-		conn, err := net.Dial("unix", upstreamSock)
-		if err == nil {
-			conn.Close()
-			break
-		}
-		time.Sleep(10 * time.Millisecond)
-	}
+	<-proxy.Ready()
 
 	return upstreamSock
 }

From 3d67dc6a0f9f7ae01d3b872a495a38ae3b7cc5f3 Mon Sep 17 00:00:00 2001
From: Brian McCallister <brianm@skife.org>
Date: Sat, 18 Apr 2026 16:28:38 -0700
Subject: [PATCH 3/3] fix: harden wrapper agent startup and proxy fallback

---
 README.md                        |  2 +-
 cmd/epithet/agent.go             | 50 ++++++++++++++----------
 cmd/epithet/agent_test.go        | 48 +++++++++++++++++++++++
 pkg/agent/proxy_listener.go      | 29 +++++++-------
 pkg/agent/proxy_listener_test.go | 65 +++++++++++++++++++++++++++++++-
 5 files changed, 159 insertions(+), 35 deletions(-)

diff --git a/README.md b/README.md
index 24928cc..f499366 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
 
 [![Actions Status](https://github.com/epithet-ssh/epithet/workflows/build/badge.svg)](https://github.com/epithet-ssh/epithet/actions) [![Go Reportcard](https://goreportcard.com/badge/github.com/epithet-ssh/epithet)](https://goreportcard.com/report/github.com/epithet-ssh/epithet)
 
-Epithet is an SSH certificate authority that replaces static authorized_keys with short-lived certificates (2-10 minutes). It creates on-demand SSH agents for each outbound connection, enabling real-time policy enforcement without touching your target hosts.
+Epithet is an SSH agent and certificate authority that replaces static authorized_keys with short-lived certificates (2-10 minutes). It creates on-demand SSH agents for each outbound connection, enabling real-time policy enforcement without touching your target hosts.
 
 ## Quick start
 
diff --git a/cmd/epithet/agent.go b/cmd/epithet/agent.go
index 3c2be71..d8cba39 100644
--- a/cmd/epithet/agent.go
+++ b/cmd/epithet/agent.go
@@ -113,6 +113,7 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 		hello, err := agent.ProbeUpstream(upstreamSocket)
 		if err != nil {
 			logger.Warn("failed to probe upstream agent", "error", err)
+			upstreamSocket = ""
 		} else if hello != nil {
 			depth = hello.ChainDepth + 1
 			brokerOpts = append(brokerOpts, broker.WithUpstream(upstreamSocket))
@@ -145,7 +146,9 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 	go func() {
 		brokerErr <- b.Serve(ctx)
 	}()
-	<-b.Ready()
+	if err := waitForBrokerStartup(b, brokerErr); err != nil {
+		return err
+	}
 
 	// Generate SSH config.
 	sshConfigPath := filepath.Join(tempDir, "ssh-config.conf")
@@ -158,25 +161,22 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 		}
 	}
 
-	// Create proxy listener if we have an upstream socket to proxy.
-	if upstreamSocket != "" {
-		proxySock := filepath.Join(tempDir, "proxy.sock")
-		setup := func(p *agent.ProxyAgent) {
-			p.RegisterExtension(agent.ExtensionHello, agent.HelloHandler(depth))
-			p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) {
-				return b.Authenticate(nil)
-			}))
-		}
-		proxyListener := agent.NewProxyListener(logger, proxySock, upstreamSocket, setup)
-		go func() {
-			if err := proxyListener.Serve(ctx); err != nil && err != context.Canceled {
-				logger.Error("proxy listener error", "error", err)
-			}
-		}()
-		<-proxyListener.Ready()
-		upstreamSocket = proxySock
-		logger.Info("proxy agent listening", "socket", proxySock)
+	proxySock := filepath.Join(tempDir, "proxy.sock")
+	setup := func(p *agent.ProxyAgent) {
+		p.RegisterExtension(agent.ExtensionHello, agent.HelloHandler(depth))
+		p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) {
+			return b.Authenticate(nil)
+		}))
 	}
+	proxyListener := agent.NewProxyListener(logger, proxySock, upstreamSocket, setup)
+	go func() {
+		if err := proxyListener.Serve(ctx); err != nil && err != context.Canceled {
+			logger.Error("proxy listener error", "error", err)
+		}
+	}()
+	<-proxyListener.Ready()
+	upstreamSocket = proxySock
+	logger.Info("proxy agent listening", "socket", proxySock)
 
 	// Build child environment with updated SSH_AUTH_SOCK.
 	childEnv := os.Environ()
@@ -223,6 +223,18 @@ func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg
 	return nil
 }
 
+func waitForBrokerStartup(b *broker.Broker, brokerErr <-chan error) error {
+	select {
+	case <-b.Ready():
+		return nil
+	case err := <-brokerErr:
+		if err == nil || err == context.Canceled {
+			return fmt.Errorf("broker exited before becoming ready")
+		}
+		return fmt.Errorf("broker failed to start: %w", err)
+	}
+}
+
 // setupBrokerWithOptions creates the broker, temp directories, and resolves auth.
 // When wrapperMode is true, uses a per-process temp directory to avoid collisions
 // between multiple wrapped shells. When false (daemon mode), uses a deterministic
diff --git a/cmd/epithet/agent_test.go b/cmd/epithet/agent_test.go
index d1cced1..af1c863 100644
--- a/cmd/epithet/agent_test.go
+++ b/cmd/epithet/agent_test.go
@@ -1,7 +1,16 @@
 package main
 
 import (
+	"context"
+	"io"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"strings"
 	"testing"
+
+	"github.com/epithet-ssh/epithet/pkg/broker"
+	"github.com/epithet-ssh/epithet/pkg/caclient"
 )
 
 func TestReplaceEnv_ExistingKey(t *testing.T) {
@@ -49,3 +58,42 @@ func TestReplaceEnv_EmptyEnv(t *testing.T) {
 		t.Errorf("expected [KEY=value], got %v", result)
 	}
 }
+
+func TestWaitForBrokerStartup_ReturnsServeError(t *testing.T) {
+	t.Parallel()
+
+	tempDir := t.TempDir()
+	agentDir := filepath.Join(tempDir, "agent")
+	if err := os.MkdirAll(agentDir, 0700); err != nil {
+		t.Fatalf("mkdir agent dir: %v", err)
+	}
+
+	endpoints := []caclient.CAEndpoint{{URL: "http://127.0.0.1", Priority: caclient.DefaultPriority}}
+	caClient, err := caclient.New(endpoints)
+	if err != nil {
+		t.Fatalf("create CA client: %v", err)
+	}
+
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	socketPath := filepath.Join(tempDir, strings.Repeat("sock", 40))
+	b, err := broker.New(*logger, socketPath, "true", caClient, agentDir)
+	if err != nil {
+		t.Fatalf("create broker: %v", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	brokerErr := make(chan error, 1)
+	go func() {
+		brokerErr <- b.Serve(ctx)
+	}()
+
+	err = waitForBrokerStartup(b, brokerErr)
+	if err == nil {
+		t.Fatal("expected startup error, got nil")
+	}
+	if !strings.Contains(err.Error(), "broker failed to start") {
+		t.Fatalf("expected startup failure wrapper, got %v", err)
+	}
+}
diff --git a/pkg/agent/proxy_listener.go b/pkg/agent/proxy_listener.go
index 888c0c4..ef17303 100644
--- a/pkg/agent/proxy_listener.go
+++ b/pkg/agent/proxy_listener.go
@@ -89,25 +89,26 @@ func (p *ProxyListener) acceptLoop() {
 func (p *ProxyListener) serveConn(conn net.Conn) {
 	defer conn.Close()
 
-	// Dial a fresh connection to the upstream agent for this client.
-	upstreamConn, err := net.Dial("unix", p.upstreamSocketPath)
-	if err != nil {
-		p.log.Warn("failed to connect to upstream agent", "error", err)
-		return
-	}
-	defer upstreamConn.Close()
+	var upstream agent.ExtendedAgent
+	if p.upstreamSocketPath == "" {
+		// Preserve standard ssh-agent behavior even without an upstream socket.
+		upstream = agent.NewKeyring().(agent.ExtendedAgent)
+	} else {
+		// Dial a fresh connection to the upstream agent for this client.
+		upstreamConn, err := net.Dial("unix", p.upstreamSocketPath)
+		if err != nil {
+			p.log.Warn("failed to connect to upstream agent", "error", err)
+			return
+		}
+		defer upstreamConn.Close()
 
-	upstream := agent.NewClient(upstreamConn)
-	extUpstream, ok := upstream.(agent.ExtendedAgent)
-	if !ok {
-		p.log.Warn("upstream agent does not support extensions")
-		return
+		upstream = agent.NewClient(upstreamConn)
 	}
 
-	proxy := NewProxyAgent(extUpstream)
+	proxy := NewProxyAgent(upstream)
 	p.setup(proxy)
 
-	err = agent.ServeAgent(proxy, conn)
+	err := agent.ServeAgent(proxy, conn)
 	if err != nil && err != io.EOF {
 		p.log.Debug("proxy agent connection ended", "error", err)
 	}
diff --git a/pkg/agent/proxy_listener_test.go b/pkg/agent/proxy_listener_test.go
index ed757f2..2d18fd7 100644
--- a/pkg/agent/proxy_listener_test.go
+++ b/pkg/agent/proxy_listener_test.go
@@ -165,6 +165,70 @@ func TestProxyListenerMultipleConnections(t *testing.T) {
 	}
 }
 
+func TestProxyListenerWithoutUpstream(t *testing.T) {
+	proxySock := tempSocketPath(t)
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+
+	var authCalled atomic.Bool
+	setup := func(p *ProxyAgent) {
+		p.RegisterExtension(ExtensionHello, HelloHandler(0))
+		p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) {
+			authCalled.Store(true)
+			return "local-token", nil
+		}))
+	}
+	proxy := NewProxyListener(logger, proxySock, "", setup)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go proxy.Serve(ctx)
+	<-proxy.Ready()
+
+	conn, err := net.Dial("unix", proxySock)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	client := agent.NewClient(conn).(agent.ExtendedAgent)
+
+	resp, err := client.Extension(ExtensionHello, nil)
+	if err != nil {
+		t.Fatalf("hello failed: %v", err)
+	}
+	var hello HelloResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil {
+		t.Fatalf("unmarshal hello: %v", err)
+	}
+	if hello.ChainDepth != 0 {
+		t.Errorf("expected chain depth 0, got %d", hello.ChainDepth)
+	}
+
+	resp, err = client.Extension(ExtensionAuth, nil)
+	if err != nil {
+		t.Fatalf("auth failed: %v", err)
+	}
+	var authResp AuthResponse
+	if err := json.Unmarshal(stripSuccessByte(t, resp), &authResp); err != nil {
+		t.Fatalf("unmarshal auth: %v", err)
+	}
+	if authResp.Token != "local-token" {
+		t.Errorf("expected local token, got %q", authResp.Token)
+	}
+	if !authCalled.Load() {
+		t.Fatal("expected auth handler to be called")
+	}
+
+	keys, err := client.List()
+	if err != nil {
+		t.Fatalf("List failed: %v", err)
+	}
+	if len(keys) != 0 {
+		t.Errorf("expected empty local keyring, got %d keys", len(keys))
+	}
+}
+
 func TestMultiHopChain(t *testing.T) {
 	// Simulate: laptop → shell1 → shell2
 	// Each hop has its own proxy listener wrapping the previous.
@@ -279,4 +343,3 @@ func tempSocketPath(t *testing.T) string {
 	t.Cleanup(func() { os.Remove(path) })
 	return path
 }
-