Merge pull request #1056 from rumpl/runtime-session-store

Add the session store to the runtime

Author: dgageot

cmd/root/run.go

@@ -6,13 +6,15 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"path/filepath"
 
 	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
 	"go.opentelemetry.io/otel"
 
 	"github.com/docker/cagent/pkg/cli"
 	"github.com/docker/cagent/pkg/config"
+	"github.com/docker/cagent/pkg/paths"
 	"github.com/docker/cagent/pkg/runtime"
 	"github.com/docker/cagent/pkg/session"
 	"github.com/docker/cagent/pkg/team"
@@ -28,6 +30,7 @@ type runExecFlags struct {
 	modelOverrides []string
 	dryRun         bool
 	runConfig      config.RuntimeConfig
+	sessionDB      string
 
 	// Exec only
 	hideToolCalls bool
@@ -64,6 +67,7 @@ func addRunOrExecFlags(cmd *cobra.Command, flags *runExecFlags) {
 	cmd.PersistentFlags().StringArrayVar(&flags.modelOverrides, "model", nil, "Override agent model: [agent=]provider/model (repeatable)")
 	cmd.PersistentFlags().BoolVar(&flags.dryRun, "dry-run", false, "Initialize the agent without executing anything")
 	cmd.PersistentFlags().StringVar(&flags.remoteAddress, "remote", "", "Use remote runtime with specified address")
+	cmd.PersistentFlags().StringVarP(&flags.sessionDB, "session-db", "s", filepath.Join(paths.GetHomeDir(), ".cagent", "session.db"), "Path to the session database")
 }
 
 func (f *runExecFlags) runRunCommand(cmd *cobra.Command, args []string) error {
@@ -105,7 +109,7 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s
 			return err
 		}
 
-		rt, sess, err = f.createLocalRuntimeAndSession(t)
+		rt, sess, err = f.createLocalRuntimeAndSession(ctx, t)
 		if err != nil {
 			return err
 		}
@@ -166,26 +170,35 @@ func (f *runExecFlags) createRemoteRuntimeAndSession(ctx context.Context, origin
 	return remoteRt, sess, nil
 }
 
-func (f *runExecFlags) createLocalRuntimeAndSession(t *team.Team) (runtime.Runtime, *session.Session, error) {
+func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, t *team.Team) (runtime.Runtime, *session.Session, error) {
 	agent, err := t.Agent(f.agentName)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	sess := session.New(
-		session.WithMaxIterations(agent.MaxIterations()),
-		session.WithToolsApproved(f.autoApprove),
-	)
+	sessStore, err := session.NewSQLiteSessionStore(f.sessionDB)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create session store: %w", err)
+	}
 
 	localRt, err := runtime.New(t,
+		runtime.WithSessionStore(sessStore),
 		runtime.WithCurrentAgent(f.agentName),
 		runtime.WithTracer(otel.Tracer(AppName)),
-		runtime.WithRootSessionID(sess.ID),
 	)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create runtime: %w", err)
 	}
 
+	sess := session.New(
+		session.WithMaxIterations(agent.MaxIterations()),
+		session.WithToolsApproved(f.autoApprove),
+	)
+
+	if err := sessStore.AddSession(ctx, sess); err != nil {
+		return nil, nil, err
+	}
+
 	slog.Debug("Using local runtime", "agent", f.agentName)
 	return localRt, sess, nil
 }

e2e/cagent_exec_test.go

@@ -104,7 +104,7 @@ func cagentExec(t *testing.T, moreArgs ...string) string {
 
 	// Start a recording AI proxy to record and replay traffic.
 	svr, _ := startRecordingAIProxy(t)
-	args = append(args, "--models-gateway", svr.URL)
+	args = append(args, "--models-gateway", svr.URL, "--session-db", "/tmp/session.db")
 
 	// Run cagent exec
 	var stdout bytes.Buffer

pkg/mcp/server.go

@@ -167,7 +167,6 @@ func CreateToolHandler(t *team.Team, agentName string) func(context.Context, *mc
 
 		rt, err := runtime.New(t,
 			runtime.WithCurrentAgent(agentName),
-			runtime.WithRootSessionID(sess.ID),
 		)
 		if err != nil {
 			return nil, ToolOutput{}, fmt.Errorf("failed to create runtime: %w", err)

pkg/runtime/runtime.go

@@ -32,6 +32,10 @@ import (
 	mcptools "github.com/docker/cagent/pkg/tools/mcp"
 )
 
+type SessionStore interface {
+	UpdateSession(ctx context.Context, sess *session.Session) error
+}
+
 // UnwrapMCPToolset extracts an MCP toolset from a potentially wrapped StartableToolSet.
 // Returns the MCP toolset if found, or nil if the toolset is not an MCP toolset.
 func UnwrapMCPToolset(toolset tools.ToolSet) *mcptools.Toolset {
@@ -94,16 +98,18 @@ type Runtime interface {
 	CurrentWelcomeMessage(ctx context.Context) string
 	// EmitStartupInfo emits initial agent, team, and toolset information for immediate display
 	EmitStartupInfo(ctx context.Context, events chan Event)
+
 	// RunStream starts the agent's interaction loop and returns a channel of events
 	RunStream(ctx context.Context, sess *session.Session) <-chan Event
 	// Run starts the agent's interaction loop and returns the final messages
 	Run(ctx context.Context, sess *session.Session) ([]session.Message, error)
 	// Resume allows resuming execution after user confirmation
 	Resume(ctx context.Context, confirmationType ResumeType)
-	// Summarize generates a summary for the session
-	Summarize(ctx context.Context, sess *session.Session, events chan Event)
 	// ResumeElicitation sends an elicitation response back to a waiting elicitation request
 	ResumeElicitation(_ context.Context, action tools.ElicitationAction, content map[string]any) error
+
+	// Summarize generates a summary for the session
+	Summarize(ctx context.Context, sess *session.Session, events chan Event)
 }
 
 type ModelStore interface {
@@ -121,7 +127,6 @@ type LocalRuntime struct {
 	toolMap                     map[string]ToolHandler
 	team                        *team.Team
 	currentAgent                string
-	rootSessionID               string // Root session ID for OAuth state encoding (preserved across sub-sessions)
 	resumeChan                  chan ResumeType
 	tracer                      trace.Tracer
 	modelsStore                 ModelStore
@@ -133,6 +138,7 @@ type LocalRuntime struct {
 	elicitationEventsChannelMux sync.RWMutex           // Protects elicitationEventsChannel
 	ragInitialized              atomic.Bool
 	titleGen                    *titleGenerator
+	sessionStore                SessionStore
 }
 
 type streamResult struct {
@@ -158,12 +164,6 @@ func WithManagedOAuth(managed bool) Opt {
 	}
 }
 
-func WithRootSessionID(sessionID string) Opt {
-	return func(r *LocalRuntime) {
-		r.rootSessionID = sessionID
-	}
-}
-
 // WithTracer sets a custom OpenTelemetry tracer; if not provided, tracing is disabled (no-op).
 func WithTracer(t trace.Tracer) Opt {
 	return func(r *LocalRuntime) {
@@ -183,6 +183,12 @@ func WithModelStore(store ModelStore) Opt {
 	}
 }
 
+func WithSessionStore(store SessionStore) Opt {
+	return func(r *LocalRuntime) {
+		r.sessionStore = store
+	}
+}
+
 // New creates a new runtime for an agent and its team
 func New(agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
 	modelsStore, err := modelsdev.NewStore()
@@ -199,6 +205,7 @@ func New(agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
 		modelsStore:          modelsStore,
 		sessionCompaction:    true,
 		managedOAuth:         true,
+		sessionStore:         session.NewInMemorySessionStore(),
 	}
 
 	for _, opt := range opts {
@@ -610,6 +617,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 							CreatedAt: time.Now().Format(time.RFC3339),
 						}
 						sess.AddMessage(session.NewAgentMessage(a, &assistantMessage))
+						_ = r.sessionStore.UpdateSession(ctx, sess)
 						return
 					}
 				case <-ctx.Done():
@@ -707,6 +715,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 				}
 
 				sess.AddMessage(session.NewAgentMessage(a, &assistantMessage))
+				_ = r.sessionStore.UpdateSession(ctx, sess)
 				slog.Debug("Added assistant message to session", "agent", a.Name(), "total_messages", len(sess.GetAllMessages()))
 			} else {
 				slog.Debug("Skipping empty assistant message (no content and no tool calls)", "agent", a.Name())
@@ -1023,14 +1032,14 @@ func (r *LocalRuntime) processToolCalls(ctx context.Context, sess *session.Sessi
 						r.runAgentTool(callCtx, def.handler, sess, toolCall, def.tool, events, a)
 					case ResumeTypeReject:
 						slog.Debug("Resume signal received, rejecting tool handler", "tool", toolCall.Function.Name, "session_id", sess.ID)
-						r.addToolRejectedResponse(sess, toolCall, def.tool, events)
+						r.addToolRejectedResponse(ctx, sess, toolCall, def.tool, events)
 					}
 				case <-callCtx.Done():
 					slog.Debug("Context cancelled while waiting for resume", "tool", toolCall.Function.Name, "session_id", sess.ID)
 					// Synthesize cancellation responses for the current and any remaining tool calls
-					r.addToolCancelledResponse(sess, toolCall, def.tool, events)
+					r.addToolCancelledResponse(ctx, sess, toolCall, def.tool, events)
 					for j := i + 1; j < len(calls); j++ {
-						r.addToolCancelledResponse(sess, calls[j], def.tool, events)
+						r.addToolCancelledResponse(ctx, sess, calls[j], def.tool, events)
 					}
 					callSpan.SetStatus(codes.Ok, "tool call canceled by user")
 					return
@@ -1066,17 +1075,17 @@ func (r *LocalRuntime) processToolCalls(ctx context.Context, sess *session.Sessi
 						r.runTool(callCtx, tool, toolCall, events, sess, a)
 					case ResumeTypeReject:
 						slog.Debug("Resume signal received, rejecting tool handler", "tool", toolCall.Function.Name, "session_id", sess.ID)
-						r.addToolRejectedResponse(sess, toolCall, tool, events)
+						r.addToolRejectedResponse(ctx, sess, toolCall, tool, events)
 					}
 
 					slog.Debug("Added tool response to session", "tool", toolCall.Function.Name, "session_id", sess.ID, "total_messages", len(sess.GetAllMessages()))
 					break toolLoop
 				case <-callCtx.Done():
 					slog.Debug("Context cancelled while waiting for resume", "tool", toolCall.Function.Name, "session_id", sess.ID)
 					// Synthesize cancellation responses for the current and any remaining tool calls
-					r.addToolCancelledResponse(sess, toolCall, tool, events)
+					r.addToolCancelledResponse(ctx, sess, toolCall, tool, events)
 					for j := i + 1; j < len(calls); j++ {
-						r.addToolCancelledResponse(sess, calls[j], tool, events)
+						r.addToolCancelledResponse(ctx, sess, calls[j], tool, events)
 					}
 					callSpan.SetStatus(codes.Ok, "tool call canceled by user")
 					return
@@ -1146,6 +1155,7 @@ func (r *LocalRuntime) runTool(ctx context.Context, tool tools.Tool, toolCall to
 		CreatedAt:  time.Now().Format(time.RFC3339),
 	}
 	sess.AddMessage(session.NewAgentMessage(a, &toolResponseMsg))
+	_ = r.sessionStore.UpdateSession(ctx, sess)
 }
 
 func (r *LocalRuntime) runAgentTool(ctx context.Context, handler ToolHandlerFunc, sess *session.Session, toolCall tools.ToolCall, tool tools.Tool, events chan Event, a *agent.Agent) {
@@ -1199,9 +1209,10 @@ func (r *LocalRuntime) runAgentTool(ctx context.Context, handler ToolHandlerFunc
 		CreatedAt:  time.Now().Format(time.RFC3339),
 	}
 	sess.AddMessage(session.NewAgentMessage(a, &toolResponseMsg))
+	_ = r.sessionStore.UpdateSession(ctx, sess)
 }
 
-func (r *LocalRuntime) addToolRejectedResponse(sess *session.Session, toolCall tools.ToolCall, tool tools.Tool, events chan Event) {
+func (r *LocalRuntime) addToolRejectedResponse(ctx context.Context, sess *session.Session, toolCall tools.ToolCall, tool tools.Tool, events chan Event) {
 	a := r.CurrentAgent()
 
 	result := "The user rejected the tool call."
@@ -1215,9 +1226,10 @@ func (r *LocalRuntime) addToolRejectedResponse(sess *session.Session, toolCall t
 		CreatedAt:  time.Now().Format(time.RFC3339),
 	}
 	sess.AddMessage(session.NewAgentMessage(a, &toolResponseMsg))
+	_ = r.sessionStore.UpdateSession(ctx, sess)
 }
 
-func (r *LocalRuntime) addToolCancelledResponse(sess *session.Session, toolCall tools.ToolCall, tool tools.Tool, events chan Event) {
+func (r *LocalRuntime) addToolCancelledResponse(ctx context.Context, sess *session.Session, toolCall tools.ToolCall, tool tools.Tool, events chan Event) {
 	a := r.CurrentAgent()
 
 	result := "The tool call was canceled by the user."
@@ -1231,6 +1243,7 @@ func (r *LocalRuntime) addToolCancelledResponse(sess *session.Session, toolCall
 		CreatedAt:  time.Now().Format(time.RFC3339),
 	}
 	sess.AddMessage(session.NewAgentMessage(a, &toolResponseMsg))
+	_ = r.sessionStore.UpdateSession(ctx, sess)
 }
 
 // startSpan wraps tracer.Start, returning a no-op span if the tracer is nil.
@@ -1418,6 +1431,7 @@ func (r *LocalRuntime) Summarize(ctx context.Context, sess *session.Session, eve
 	}
 	// Add the summary to the session as a summary item
 	sess.Messages = append(sess.Messages, session.Item{Summary: summary})
+	_ = r.sessionStore.UpdateSession(ctx, sess)
 	slog.Debug("Generated session summary", "session_id", sess.ID, "summary_length", len(summary))
 	events <- SessionSummary(sess.ID, summary, r.currentAgent)
 }

pkg/server/session_manager.go

@@ -132,6 +132,7 @@ func (sm *sessionManager) RunSession(ctx context.Context, sessionID, agentFilena
 	for _, msg := range messages {
 		sess.AddMessage(session.UserMessage(msg.Content, msg.MultiContent...))
 	}
+
 	if err := sm.sessionStore.UpdateSession(ctx, sess); err != nil {
 		return nil, err
 	}
@@ -228,7 +229,7 @@ func (sm *sessionManager) runtimeForSession(ctx context.Context, sess *session.S
 	opts := []runtime.Opt{
 		runtime.WithCurrentAgent(currentAgent),
 		runtime.WithManagedOAuth(false),
-		runtime.WithRootSessionID(sess.ID),
+		runtime.WithSessionStore(sm.sessionStore),
 	}
 	run, err := runtime.New(t, opts...)
 	if err != nil {

pkg/session/store.go

@@ -9,6 +9,8 @@ import (
 	"time"
 
 	_ "modernc.org/sqlite"
+
+	"github.com/docker/cagent/pkg/concurrent"
 )
 
 var (
@@ -34,6 +36,68 @@ type Store interface {
 	UpdateSession(ctx context.Context, session *Session) error
 }
 
+type InMemorySessionStore struct {
+	sessions *concurrent.Map[string, *Session]
+}
+
+func NewInMemorySessionStore() Store {
+	return &InMemorySessionStore{
+		sessions: concurrent.NewMap[string, *Session](),
+	}
+}
+
+func (s *InMemorySessionStore) AddSession(ctx context.Context, session *Session) error {
+	if session.ID == "" {
+		return ErrEmptyID
+	}
+	s.sessions.Store(session.ID, session)
+	return nil
+}
+
+func (s *InMemorySessionStore) GetSession(ctx context.Context, id string) (*Session, error) {
+	if id == "" {
+		return nil, ErrEmptyID
+	}
+	session, exists := s.sessions.Load(id)
+	if !exists {
+		return nil, ErrNotFound
+	}
+	return session, nil
+}
+
+func (s *InMemorySessionStore) GetSessions(ctx context.Context) ([]*Session, error) {
+	sessions := make([]*Session, 0, s.sessions.Length())
+	s.sessions.Range(func(key string, value *Session) bool {
+		sessions = append(sessions, value)
+		return true
+	})
+	return sessions, nil
+}
+
+func (s *InMemorySessionStore) DeleteSession(ctx context.Context, id string) error {
+	if id == "" {
+		return ErrEmptyID
+	}
+	_, exists := s.sessions.Load(id)
+	if !exists {
+		return ErrNotFound
+	}
+	s.sessions.Delete(id)
+	return nil
+}
+
+func (s *InMemorySessionStore) UpdateSession(ctx context.Context, session *Session) error {
+	if session.ID == "" {
+		return ErrEmptyID
+	}
+	_, exists := s.sessions.Load(session.ID)
+	if !exists {
+		return ErrNotFound
+	}
+	s.sessions.Store(session.ID, session)
+	return nil
+}
+
 // SQLiteSessionStore implements Store using SQLite
 type SQLiteSessionStore struct {
 	db *sql.DB