zzet--gortex/cmd/gortex/daemon_streamable.go

package main

import (
	"context"
	"net/http"
	"sync"

	"github.com/zzet/gortex/internal/daemon"
	gortexmcp "github.com/zzet/gortex/internal/mcp"
	"github.com/zzet/gortex/internal/mcp/streamable"
	"github.com/zzet/gortex/internal/server"
	"go.uber.org/zap"
)

// composeDaemonHTTPHandler combines the daemon's two HTTP surfaces onto a
// single listener: the MCP Streamable transport (streamH, serving /mcp
// and /healthz with its own auth) under the catch-all, and the /v1 REST
// API (v1, the former `gortex server`) under /v1/ behind a bearer-auth
// wrapper. CORS, when an origin is set, wraps the whole mux so browser
// clients (the web UI) can reach either surface cross-origin.
func composeDaemonHTTPHandler(streamH, v1 http.Handler, tokenFn func() string, corsOrigin string) http.Handler {
	top := http.NewServeMux()
	top.Handle("/", streamH)
	top.Handle("/v1/", server.WithAuthFunc(v1, tokenFn))
	if corsOrigin != "" {
		return server.WithCORS(top, server.CORSOptions{AllowOrigins: []string{corsOrigin}})
	}
	return top
}

// daemonStreamableDispatcher bridges streamable.Transport (which has
// its own session model keyed by `Mcp-Session-Id`) to the daemon's
// existing MCPDispatcher (which wants a *daemon.Session). For every
// streamable session ID it lazily registers a detached daemon
// session, so existing daemon-level features — status visibility,
// per-session MCP state on the in-process server, router proxying —
// see HTTP-arriving traffic exactly the same way they see unix-socket
// traffic.
//
// The bridge is owned by cmd/gortex (not the daemon package) so it
// can reach across to the streamable package without introducing a
// daemon→streamable import edge. Symmetric with the existing
// mcpDispatcher pattern.
type daemonStreamableDispatcher struct {
	inner    daemon.MCPDispatcher
	registry *daemon.SessionRegistry
	logger   *zap.Logger

	mu      sync.Mutex
	bridged map[string]*daemon.Session // streamableSessionID → synthetic daemon.Session
}

func newDaemonStreamableDispatcher(inner daemon.MCPDispatcher, reg *daemon.SessionRegistry, logger *zap.Logger) *daemonStreamableDispatcher {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &daemonStreamableDispatcher{
		inner:    inner,
		registry: reg,
		logger:   logger,
		bridged:  make(map[string]*daemon.Session),
	}
}

// Dispatch implements streamable.Dispatcher. It pulls the streamable
// session ID out of the context (set by the transport's
// localDispatch), pairs it with a synthetic daemon.Session, and hands
// the frame down to the wrapped dispatcher. The synthetic session is
// created on first use; subsequent calls in the same streamable
// session reuse it so per-session in-process state (savings counters,
// recent activity, frecency feedback) accumulates across requests.
func (d *daemonStreamableDispatcher) Dispatch(ctx context.Context, frame []byte) ([]byte, error) {
	if d.inner == nil {
		return nil, nil
	}
	sid := gortexmcp.SessionIDFromContext(ctx)
	cwd := gortexmcp.SessionCWDFromContext(ctx)
	sess := d.acquire(sid, cwd)
	if sess == nil {
		// Stateless mode (no Mcp-Session-Id header). Run the
		// frame through a one-shot session so the dispatcher sees
		// a valid pointer; tear it down after the call so memory
		// doesn't grow unbounded.
		sess = d.registry.RegisterDetached("", daemon.Handshake{
			Mode: daemon.ModeMCP,
			CWD:  cwd,
		})
		defer d.unbridge(sess.ID)
		defer d.registry.RemoveByID(sess.ID)
	}
	return d.inner.Dispatch(ctx, sess, frame)
}

// acquire returns the synthetic daemon.Session for a streamable
// session ID, creating it on demand. Returns nil when sid == "" so
// the caller can fall back to one-shot semantics.
func (d *daemonStreamableDispatcher) acquire(sid, cwd string) *daemon.Session {
	if sid == "" {
		return nil
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	if sess, ok := d.bridged[sid]; ok {
		// Refresh the cwd if the client moved between requests.
		if cwd != "" && sess.CWD != cwd {
			sess.CWD = cwd
		}
		return sess
	}
	// Use the streamable session id verbatim — that way the daemon
	// status command's MCP-sessions block shows the same id the
	// client saw on its first response header.
	sess := d.registry.RegisterDetached(sid, daemon.Handshake{
		Mode:       daemon.ModeMCP,
		CWD:        cwd,
		ClientName: "http",
	})
	d.bridged[sid] = sess
	return sess
}

func (d *daemonStreamableDispatcher) unbridge(sid string) {
	if sid == "" {
		return
	}
	d.mu.Lock()
	delete(d.bridged, sid)
	d.mu.Unlock()
}

// onSessionEnded mirrors the unix-socket transport's defer cleanup:
// when a streamable session is dropped (DELETE /mcp or store
// eviction) the daemon-side bookkeeping must let go too. Wired into
// the streamable transport via the on-delete cleanup path defined
// below.
func (d *daemonStreamableDispatcher) onSessionEnded(sid string) {
	sess := d.registry.RemoveByID(sid)
	d.unbridge(sid)
	if sess != nil {
		if hook, ok := d.inner.(daemon.SessionEndedHook); ok {
			hook.SessionEnded(sess)
		}
	}
}

// observeStoreCleanup wraps a streamable.SessionStore so the daemon
// gets a callback every time a session is dropped (explicit DELETE or
// TTL eviction). The wrapper preserves the interface contract while
// letting the daemon-side dispatcher purge its bridged map.
type observingStore struct {
	inner streamable.SessionStore
	on    func(string)
}

func wrapStreamableStoreWithCleanup(inner streamable.SessionStore, on func(string)) streamable.SessionStore {
	if inner == nil || on == nil {
		return inner
	}
	return &observingStore{inner: inner, on: on}
}

func (s *observingStore) Create(state streamable.SessionState) (string, error) {
	return s.inner.Create(state)
}

func (s *observingStore) Get(id string) (streamable.SessionState, bool) {
	return s.inner.Get(id)
}

func (s *observingStore) Update(state streamable.SessionState) error {
	return s.inner.Update(state)
}

func (s *observingStore) Delete(id string) {
	s.inner.Delete(id)
	if id != "" {
		s.on(id)
	}
}

func (s *observingStore) Len() int { return s.inner.Len() }

// buildDaemonStreamableHandler stitches the transport, store, and
// session bridge into a single http.Handler the daemon can mount on
// /mcp. Pulled out so cmd/gortex/daemon.go stays terse: one call
// returns the handler ready to assign to daemon.Server.HTTPHandler.
func buildDaemonStreamableHandler(disp daemon.MCPDispatcher, reg *daemon.SessionRegistry, router *daemon.Router, logger *zap.Logger, tokenFn func() string) http.Handler {
	bridge := newDaemonStreamableDispatcher(disp, reg, logger)
	store := streamable.NewMemoryStore(daemon.DefaultOverlayIdleTTL)
	wrapped := wrapStreamableStoreWithCleanup(store, bridge.onSessionEnded)
	transport := streamable.New(streamable.Config{
		Dispatcher: bridge,
		Store:      wrapped,
		Logger:     logger,
		Router:     router,
		InitializeHook: func(_ context.Context, state *streamable.SessionState) {
			// Bridge the new streamable session into the daemon's
			// session registry up-front so the daemon-status block
			// shows it before the first tool call lands.
			bridge.acquire(state.ID, state.CWD)
		},
	})
	mux := http.NewServeMux()
	mux.Handle("POST /mcp", transport)
	mux.Handle("GET /mcp", transport)
	mux.Handle("DELETE /mcp", transport)
	mux.Handle("OPTIONS /mcp", transport)
	// Health probe so operators (and `gortex daemon http-status`
	// scripts) can verify the listener is up without dispatching an
	// MCP frame.
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"status":"ok","transport":"streamable-http","spec":"mcp-2026-03-26"}`))
	})
	// Always wrap: the middleware resolves the token per request, so a
	// request with no configured token is served unauthenticated and the
	// token can be rotated (added/changed/removed) without a restart.
	return bearerAuthMiddleware(mux, tokenFn)
}

// bearerAuthMiddleware gates every request behind a bearer token resolved
// per request via tokenFn (so the token can rotate without a restart).
// When tokenFn returns "", the request is served unauthenticated.
// /healthz is exempt so liveness probes don't need a token.
func bearerAuthMiddleware(next http.Handler, tokenFn func() string) http.Handler {
	if tokenFn == nil {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/healthz" {
			next.ServeHTTP(w, r)
			return
		}
		token := tokenFn()
		if token == "" {
			next.ServeHTTP(w, r)
			return
		}
		if r.Header.Get("Authorization") != "Bearer "+token {
			w.Header().Set("WWW-Authenticate", `Bearer realm="gortex"`)
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// Compile-time check: the dispatcher satisfies streamable.Dispatcher.
var _ streamable.Dispatcher = (*daemonStreamableDispatcher)(nil)

// Compile-time check: the observing store satisfies SessionStore.
var _ streamable.SessionStore = (*observingStore)(nil)