Merge pull request #11 from hackertron/feat/gateway-server

feat: add HTTP/WebSocket gateway server with REST API and streaming

Author: hackertron

cmd/yantra/main.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"syscall"
 
+	"github.com/hackertron/Yantra/internal/gateway"
 	"github.com/hackertron/Yantra/internal/memory"
 	"github.com/hackertron/Yantra/internal/provider"
 	"github.com/hackertron/Yantra/internal/runtime"
@@ -307,7 +308,74 @@ func runTUI(cmd *cobra.Command, args []string) error {
 }
 
 func runServe(cmd *cobra.Command, args []string) error {
-	fmt.Println("Starting Yantra API server...")
-	// TODO: implement API server
-	return fmt.Errorf("not yet implemented")
+	ctx := cmd.Context()
+	logger := slog.Default()
+
+	cfg, err := types.LoadConfig(configPath)
+	if err != nil {
+		return fmt.Errorf("loading config: %w", err)
+	}
+
+	p, err := provider.BuildFromConfig(cfg)
+	if err != nil {
+		return fmt.Errorf("building provider: %w", err)
+	}
+	p = provider.NewReliable(p, provider.DefaultReliableConfig())
+
+	absWorkspace, err := filepath.Abs(".")
+	if err != nil {
+		return fmt.Errorf("resolving workspace: %w", err)
+	}
+
+	// Set up memory if enabled.
+	var mem types.MemoryRetrieval
+	var memDB *memory.DB
+	var sessStore types.SessionStore
+
+	if cfg.Memory.Enabled {
+		dbPath := cfg.Memory.DBPath
+		if dbPath == "" {
+			dbPath = ".yantra/memory.db"
+		}
+		if !filepath.IsAbs(dbPath) {
+			dbPath = filepath.Join(absWorkspace, dbPath)
+		}
+
+		memDB, err = memory.OpenDB(dbPath)
+		if err != nil {
+			slog.Warn("failed to open memory DB, continuing without memory", "error", err)
+		} else {
+			embedder, err := memory.NewEmbeddingBackend(cfg.Memory)
+			if err != nil {
+				slog.Warn("failed to create embedding backend, continuing without embeddings", "error", err)
+			}
+			mem = memory.NewStore(memDB, embedder, cfg.Memory.Retrieval)
+			sessStore = memory.NewSessionStore(memDB)
+		}
+	}
+	if memDB != nil {
+		defer memDB.Close()
+	}
+
+	// Sessions are required for the gateway even when memory is disabled.
+	// Fall back to an in-memory SQLite DB for session tracking only.
+	if sessStore == nil {
+		sessDB, err := memory.OpenDB(":memory:")
+		if err != nil {
+			return fmt.Errorf("opening session DB: %w", err)
+		}
+		defer sessDB.Close()
+		sessStore = memory.NewSessionStore(sessDB)
+	}
+
+	policy := tool.NewWorkspacePolicy(cfg.Tools.Shell)
+	reg := tool.NewRegistry(policy)
+	if err := tool.RegisterBuiltins(reg, cfg.Tools, mem); err != nil {
+		return fmt.Errorf("registering tools: %w", err)
+	}
+
+	srv := gateway.NewServer(cfg.Gateway, cfg, p, reg, mem, sessStore, absWorkspace, logger)
+
+	logger.Info("starting Yantra API server", "listen", cfg.Gateway.Listen)
+	return srv.ListenAndServe(ctx)
 }

go.mod

@@ -4,6 +4,7 @@ go 1.26
 
 require (
 	github.com/anthropics/anthropic-sdk-go v1.26.0
+	github.com/gorilla/websocket v1.5.3
 	github.com/knadh/koanf/parsers/toml v0.1.0
 	github.com/knadh/koanf/providers/env v1.1.0
 	github.com/knadh/koanf/providers/file v1.2.1
@@ -28,7 +29,6 @@ require (
 	github.com/google/s2a-go v0.1.8 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
-	github.com/gorilla/websocket v1.5.3 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/knadh/koanf/maps v0.1.2 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect

internal/gateway/auth.go

@@ -0,0 +1,36 @@
+package gateway
+
+import (
+	"net/http"
+	"strings"
+)
+
+// authMiddleware rejects requests without a valid Bearer token.
+// If the server's APIKey is empty (dev mode), all requests pass through.
+func (s *Server) authMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if s.cfg.APIKey == "" {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		auth := r.Header.Get("Authorization")
+		if auth == "" {
+			http.Error(w, `{"error":"missing authorization header"}`, http.StatusUnauthorized)
+			return
+		}
+
+		token := strings.TrimPrefix(auth, "Bearer ")
+		if token == auth || !s.validateAPIKey(token) {
+			http.Error(w, `{"error":"invalid api key"}`, http.StatusUnauthorized)
+			return
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}
+
+// validateAPIKey checks whether key matches the configured API key.
+func (s *Server) validateAPIKey(key string) bool {
+	return s.cfg.APIKey != "" && key == s.cfg.APIKey
+}

internal/gateway/auth_test.go

@@ -0,0 +1,75 @@
+package gateway
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/hackertron/Yantra/internal/types"
+)
+
+func TestAuthMiddleware_EmptyKey_PassesAll(t *testing.T) {
+	s := &Server{cfg: types.GatewayConfig{APIKey: ""}}
+
+	handler := s.authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/api/v1/sessions", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Errorf("expected 200, got %d", rec.Code)
+	}
+}
+
+func TestAuthMiddleware_CorrectToken(t *testing.T) {
+	s := &Server{cfg: types.GatewayConfig{APIKey: "secret-key"}}
+
+	handler := s.authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/api/v1/sessions", nil)
+	req.Header.Set("Authorization", "Bearer secret-key")
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Errorf("expected 200, got %d", rec.Code)
+	}
+}
+
+func TestAuthMiddleware_WrongToken(t *testing.T) {
+	s := &Server{cfg: types.GatewayConfig{APIKey: "secret-key"}}
+
+	handler := s.authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/api/v1/sessions", nil)
+	req.Header.Set("Authorization", "Bearer wrong-key")
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusUnauthorized {
+		t.Errorf("expected 401, got %d", rec.Code)
+	}
+}
+
+func TestAuthMiddleware_MissingHeader(t *testing.T) {
+	s := &Server{cfg: types.GatewayConfig{APIKey: "secret-key"}}
+
+	handler := s.authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/api/v1/sessions", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusUnauthorized {
+		t.Errorf("expected 401, got %d", rec.Code)
+	}
+}

internal/gateway/health.go

@@ -0,0 +1,21 @@
+package gateway
+
+import (
+	"encoding/json"
+	"net/http"
+)
+
+func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(map[string]string{
+		"status":  "ok",
+		"version": "0.1.0",
+	})
+}
+
+func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(map[string]bool{
+		"ready": true,
+	})
+}

internal/gateway/health_test.go

@@ -0,0 +1,53 @@
+package gateway
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/hackertron/Yantra/internal/types"
+)
+
+func TestHandleHealth(t *testing.T) {
+	s := &Server{cfg: types.GatewayConfig{}}
+
+	req := httptest.NewRequest("GET", "/health", nil)
+	rec := httptest.NewRecorder()
+	s.handleHealth(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("expected 200, got %d", rec.Code)
+	}
+
+	var body map[string]string
+	if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
+		t.Fatalf("decode error: %v", err)
+	}
+	if body["status"] != "ok" {
+		t.Errorf("expected status=ok, got %q", body["status"])
+	}
+	if body["version"] != "0.1.0" {
+		t.Errorf("expected version=0.1.0, got %q", body["version"])
+	}
+}
+
+func TestHandleReady(t *testing.T) {
+	s := &Server{cfg: types.GatewayConfig{}}
+
+	req := httptest.NewRequest("GET", "/ready", nil)
+	rec := httptest.NewRecorder()
+	s.handleReady(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("expected 200, got %d", rec.Code)
+	}
+
+	var body map[string]bool
+	if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
+		t.Fatalf("decode error: %v", err)
+	}
+	if !body["ready"] {
+		t.Error("expected ready=true")
+	}
+}