Merge pull request #972 from rumpl/feat-server

Session manager for the server

Author: dgageot

cmd/root/api.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -77,12 +78,11 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return fmt.Errorf("failed to resolve agent sources: %w", err)
 	}
-	s, err := server.New(sessionStore, &f.runConfig, sources)
+
+	s, err := server.New(ctx, sessionStore, &f.runConfig, time.Duration(f.pullIntervalMins)*time.Minute, sources)
 	if err != nil {
 		return fmt.Errorf("failed to create server: %w", err)
 	}
 
-	// TODO(rumpl): implement pull interval
-
 	return s.Serve(ctx, ln)
 }

pkg/server/server.go

@@ -22,43 +22,35 @@ import (
 	"github.com/docker/cagent/pkg/config"
 	"github.com/docker/cagent/pkg/runtime"
 	"github.com/docker/cagent/pkg/session"
-	"github.com/docker/cagent/pkg/team"
-	"github.com/docker/cagent/pkg/teamloader"
 	"github.com/docker/cagent/pkg/tools"
 )
 
 type Server struct {
 	e              *echo.Echo
-	runtimes       *concurrent.Map[string, runtime.Runtime]
 	runtimeCancels *concurrent.Map[string, context.CancelFunc]
 	sessionStore   session.Store
 	runConfig      *config.RuntimeConfig
-	agentSources   config.Sources
+	sm             *sessionManager
 }
 
-type Opt func(*Server) error
-
-func New(sessionStore session.Store, runConfig *config.RuntimeConfig, agentSources config.Sources) (*Server, error) {
+func New(ctx context.Context, sessionStore session.Store, runConfig *config.RuntimeConfig, refreshInterval time.Duration, agentSources config.Sources) (*Server, error) {
 	e := echo.New()
 	e.Use(middleware.CORS())
 	e.Use(middleware.Logger())
 
 	s := &Server{
 		e:              e,
-		runtimes:       concurrent.NewMap[string, runtime.Runtime](),
 		runtimeCancels: concurrent.NewMap[string, context.CancelFunc](),
 		sessionStore:   sessionStore,
 		runConfig:      runConfig,
-		agentSources:   agentSources,
+		sm:             newSessionManager(ctx, agentSources, refreshInterval),
 	}
 
 	group := e.Group("/api")
 
 	// List all available agents
 	group.GET("/agents", s.getAgents)
 
-	// SESSIONS
-
 	// List all sessions
 	group.GET("/sessions", s.getSessions)
 	// Get a session by id
@@ -95,11 +87,9 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
 	return nil
 }
 
-// API handlers
-
 func (s *Server) getAgents(c echo.Context) error {
 	agents := []api.Agent{}
-	for k, agentSource := range s.agentSources {
+	for k, agentSource := range s.sm.sources {
 		slog.Debug("API source", "source", agentSource.Name())
 
 		c, err := config.Load(c.Request().Context(), agentSource)
@@ -252,7 +242,7 @@ func (s *Server) resumeSession(c echo.Context) error {
 		return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid request body: %v", err))
 	}
 
-	rt, exists := s.runtimes.Load(sessionID)
+	rt, exists := s.sm.runtimes.Load(sessionID)
 	if !exists {
 		return echo.NewHTTPError(http.StatusNotFound, fmt.Sprintf("runtime not found: %s", sessionID))
 	}
@@ -273,9 +263,9 @@ func (s *Server) deleteSession(c echo.Context) error {
 	}
 
 	// Clean up the runtime
-	if _, exists := s.runtimes.Load(sessionID); exists {
+	if _, exists := s.sm.runtimes.Load(sessionID); exists {
 		slog.Debug("Removing runtime for session", "session_id", sessionID)
-		s.runtimes.Delete(sessionID)
+		s.sm.runtimes.Delete(sessionID)
 	}
 
 	// Delete the session from storage
@@ -307,7 +297,7 @@ func (s *Server) runAgent(c echo.Context) error {
 	rc := s.runConfig.Clone()
 	rc.WorkingDir = sess.WorkingDir
 
-	rt, err := s.runtimeForSession(c.Request().Context(), sess, agentFilename, currentAgent, rc)
+	rt, err := s.sm.runtimeForSession(c.Request().Context(), sess, agentFilename, currentAgent, rc)
 	if err != nil {
 		return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to get runtime for session: %v", err))
 	}
@@ -332,7 +322,6 @@ func (s *Server) runAgent(c echo.Context) error {
 	c.Response().Header().Set("Connection", "keep-alive")
 	c.Response().WriteHeader(http.StatusOK)
 
-	// Create a cancellable context for this stream
 	streamCtx, cancel := context.WithCancel(c.Request().Context())
 	s.runtimeCancels.Store(sess.ID, cancel)
 	defer func() {
@@ -356,56 +345,14 @@ func (s *Server) runAgent(c echo.Context) error {
 	return nil
 }
 
-func (s *Server) runtimeForSession(ctx context.Context, sess *session.Session, agentFilename, currentAgent string, runConfig *config.RuntimeConfig) (runtime.Runtime, error) {
-	rt, exists := s.runtimes.Load(sess.ID)
-	if exists {
-		return rt, nil
-	}
-
-	t, err := s.loadTeam(ctx, agentFilename, runConfig)
-	if err != nil {
-		return nil, err
-	}
-
-	agent, err := t.Agent(currentAgent)
-	if err != nil {
-		return nil, echo.NewHTTPError(http.StatusNotFound, fmt.Sprintf("agent not found: %v", err))
-	}
-	sess.MaxIterations = agent.MaxIterations()
-
-	opts := []runtime.Opt{
-		runtime.WithCurrentAgent(currentAgent),
-		runtime.WithManagedOAuth(false),
-		runtime.WithRootSessionID(sess.ID),
-	}
-	rt, err = runtime.New(t, opts...)
-	if err != nil {
-		slog.Error("Failed to create runtime", "error", err)
-		return nil, echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to create runtime: %v", err))
-	}
-	s.runtimes.Store(sess.ID, rt)
-	slog.Debug("Runtime created for session", "session_id", sess.ID)
-
-	return rt, nil
-}
-
-func (s *Server) loadTeam(ctx context.Context, agentFilename string, runConfig *config.RuntimeConfig) (*team.Team, error) {
-	agentSource, found := s.agentSources[agentFilename]
-	if !found {
-		return nil, fmt.Errorf("agent not found: %s", agentFilename)
-	}
-
-	return teamloader.Load(ctx, agentSource, runConfig)
-}
-
 func (s *Server) elicitation(c echo.Context) error {
 	sessionID := c.Param("id")
 	var req api.ResumeElicitationRequest
 	if err := c.Bind(&req); err != nil {
 		return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid request body: %v", err))
 	}
 
-	rt, exists := s.runtimes.Load(sessionID)
+	rt, exists := s.sm.runtimes.Load(sessionID)
 	if !exists {
 		return c.JSON(http.StatusNotFound, map[string]string{"error": fmt.Sprintf("runtime not found: %s", sessionID)})
 	}

pkg/server/server_test.go

@@ -95,7 +95,7 @@ func startServer(t *testing.T, ctx context.Context, agentsDir string) string {
 
 	sources, err := config.ResolveSources(agentsDir)
 	require.NoError(t, err)
-	srv, err := New(store, &runConfig, sources)
+	srv, err := New(ctx, store, &runConfig, 0, sources)
 	require.NoError(t, err)
 
 	socketPath := "unix://" + filepath.Join(t.TempDir(), "sock")

pkg/server/session_manager.go

@@ -0,0 +1,82 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"time"
+
+	"github.com/labstack/echo/v4"
+
+	"github.com/docker/cagent/pkg/concurrent"
+	"github.com/docker/cagent/pkg/config"
+	"github.com/docker/cagent/pkg/runtime"
+	"github.com/docker/cagent/pkg/session"
+	"github.com/docker/cagent/pkg/team"
+	"github.com/docker/cagent/pkg/teamloader"
+)
+
+type sessionManager struct {
+	runtimes *concurrent.Map[string, runtime.Runtime]
+	sources  config.Sources
+
+	refreshInterval time.Duration
+}
+
+func newSessionManager(ctx context.Context, sources config.Sources, refreshInterval time.Duration) *sessionManager {
+	loaders := make(config.Sources)
+	for name, source := range sources {
+		loaders[name] = newSourceLoader(ctx, source, refreshInterval)
+	}
+
+	sm := &sessionManager{
+		runtimes:        concurrent.NewMap[string, runtime.Runtime](),
+		sources:         loaders,
+		refreshInterval: refreshInterval,
+	}
+
+	return sm
+}
+
+func (sm *sessionManager) runtimeForSession(ctx context.Context, sess *session.Session, agentFilename, currentAgent string, rc *config.RuntimeConfig) (runtime.Runtime, error) {
+	rt, exists := sm.runtimes.Load(sess.ID)
+	if exists {
+		return rt, nil
+	}
+
+	t, err := sm.loadTeam(ctx, agentFilename, rc)
+	if err != nil {
+		return nil, err
+	}
+
+	agent, err := t.Agent(currentAgent)
+	if err != nil {
+		return nil, echo.NewHTTPError(http.StatusNotFound, fmt.Sprintf("agent not found: %v", err))
+	}
+	sess.MaxIterations = agent.MaxIterations()
+
+	opts := []runtime.Opt{
+		runtime.WithCurrentAgent(currentAgent),
+		runtime.WithManagedOAuth(false),
+		runtime.WithRootSessionID(sess.ID),
+	}
+	rt, err = runtime.New(t, opts...)
+	if err != nil {
+		slog.Error("Failed to create runtime", "error", err)
+		return nil, echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to create runtime: %v", err))
+	}
+	sm.runtimes.Store(sess.ID, rt)
+	slog.Debug("Runtime created for session", "session_id", sess.ID)
+
+	return rt, nil
+}
+
+func (sm *sessionManager) loadTeam(ctx context.Context, agentFilename string, runConfig *config.RuntimeConfig) (*team.Team, error) {
+	agentSource, found := sm.sources[agentFilename]
+	if !found {
+		return nil, fmt.Errorf("agent not found: %s", agentFilename)
+	}
+
+	return teamloader.Load(ctx, agentSource, runConfig)
+}

pkg/server/source_loader.go

@@ -0,0 +1,83 @@
+package server
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/docker/cagent/pkg/config"
+)
+
+type sourceLoader struct {
+	inner           config.Source
+	refreshInterval time.Duration
+
+	mu   sync.RWMutex
+	data []byte
+	err  error
+}
+
+func newSourceLoader(ctx context.Context, inner config.Source, refreshInterval time.Duration) *sourceLoader {
+	sl := &sourceLoader{
+		inner:           inner,
+		refreshInterval: refreshInterval,
+	}
+
+	sl.load(ctx)
+
+	if refreshInterval > 0 {
+		go sl.refreshLoop(ctx)
+	}
+
+	return sl
+}
+
+func (sl *sourceLoader) Name() string {
+	return sl.inner.Name()
+}
+
+func (sl *sourceLoader) ParentDir() string {
+	return sl.inner.ParentDir()
+}
+
+func (sl *sourceLoader) Read(_ context.Context) ([]byte, error) {
+	sl.mu.RLock()
+	defer sl.mu.RUnlock()
+	return sl.data, sl.err
+}
+
+func (sl *sourceLoader) load(ctx context.Context) {
+	data, err := sl.inner.Read(ctx)
+
+	sl.mu.Lock()
+	defer sl.mu.Unlock()
+
+	if err != nil {
+		// Only log errors, keep previous data if available
+		slog.Warn("Failed to refresh source",
+			"source", sl.inner.Name(),
+			"error", err)
+		// Only update error if we don't have data yet
+		if len(sl.data) == 0 {
+			sl.err = err
+		}
+	} else {
+		sl.data = data
+		sl.err = nil
+	}
+}
+
+func (sl *sourceLoader) refreshLoop(ctx context.Context) {
+	ticker := time.NewTicker(sl.refreshInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			sl.load(ctx)
+		}
+	}
+}