package mcp import ( "context" "fmt" "log/slog" "os" "path/filepath" "strings" "time" "github.com/containerd/containerd/v2/pkg/oci" "github.com/containerd/errdefs" "github.com/jackc/pgx/v5/pgxpool" "github.com/opencontainers/runtime-spec/specs-go" "github.com/memohai/memoh/internal/config" ctr "github.com/memohai/memoh/internal/containerd" dbsqlc "github.com/memohai/memoh/internal/db/sqlc" "github.com/memohai/memoh/internal/identity" ) const ( BotLabelKey = "mcp.bot_id" ContainerPrefix = "mcp-" ) type ExecRequest struct { BotID string Command []string Env []string WorkDir string Terminal bool UseStdio bool } type ExecResult struct { ExitCode uint32 } type Manager struct { service ctr.Service cfg config.MCPConfig containerID func(string) string db *pgxpool.Pool queries *dbsqlc.Queries logger *slog.Logger } func NewManager(log *slog.Logger, service ctr.Service, cfg config.MCPConfig) *Manager { return &Manager{ service: service, cfg: cfg, logger: log.With(slog.String("component", "mcp")), containerID: func(botID string) string { return ContainerPrefix + botID }, } } func (m *Manager) WithDB(db *pgxpool.Pool) *Manager { m.db = db m.queries = dbsqlc.New(db) return m } func (m *Manager) Init(ctx context.Context) error { image := m.cfg.BusyboxImage if image == "" { image = config.DefaultBusyboxImg } _, err := m.service.PullImage(ctx, image, &ctr.PullImageOptions{ Unpack: true, Snapshotter: m.cfg.Snapshotter, }) return err } // EnsureBot creates the MCP container for a bot if it does not exist. func (m *Manager) EnsureBot(ctx context.Context, botID string) error { if err := validateBotID(botID); err != nil { return err } dataDir, err := m.ensureBotDir(botID) if err != nil { return err } dataMount := m.cfg.DataMount if dataMount == "" { dataMount = config.DefaultDataMount } image := m.cfg.BusyboxImage if image == "" { image = config.DefaultBusyboxImg } specOpts := []oci.SpecOpts{ oci.WithMounts([]specs.Mount{ { Destination: dataMount, Type: "bind", Source: dataDir, Options: []string{"rbind", "rw"}, }, { Destination: "/app", Type: "bind", Source: dataDir, Options: []string{"rbind", "rw"}, }, }), } _, err = m.service.CreateContainer(ctx, ctr.CreateContainerRequest{ ID: m.containerID(botID), ImageRef: image, Snapshotter: m.cfg.Snapshotter, Labels: map[string]string{ BotLabelKey: botID, }, SpecOpts: specOpts, }) if err == nil { return nil } if !errdefs.IsAlreadyExists(err) { return err } return nil } // ListBots returns the bot IDs that have MCP containers. func (m *Manager) ListBots(ctx context.Context) ([]string, error) { containers, err := m.service.ListContainers(ctx) if err != nil { return nil, err } botIDs := make([]string, 0, len(containers)) for _, container := range containers { info, err := container.Info(ctx) if err != nil { return nil, err } if strings.HasPrefix(info.ID, ContainerPrefix) { if botID, ok := info.Labels[BotLabelKey]; ok { botIDs = append(botIDs, botID) } } } return botIDs, nil } func (m *Manager) Start(ctx context.Context, botID string) error { if err := m.EnsureBot(ctx, botID); err != nil { return err } _, err := m.service.StartTask(ctx, m.containerID(botID), &ctr.StartTaskOptions{ UseStdio: false, }) return err } func (m *Manager) Stop(ctx context.Context, botID string, timeout time.Duration) error { if err := validateBotID(botID); err != nil { return err } return m.service.StopTask(ctx, m.containerID(botID), &ctr.StopTaskOptions{ Timeout: timeout, Force: true, }) } func (m *Manager) Delete(ctx context.Context, botID string) error { if err := validateBotID(botID); err != nil { return err } _ = m.service.DeleteTask(ctx, m.containerID(botID), &ctr.DeleteTaskOptions{Force: true}) return m.service.DeleteContainer(ctx, m.containerID(botID), &ctr.DeleteContainerOptions{ CleanupSnapshot: true, }) } func (m *Manager) Exec(ctx context.Context, req ExecRequest) (*ExecResult, error) { if err := validateBotID(req.BotID); err != nil { return nil, err } if len(req.Command) == 0 { return nil, fmt.Errorf("%w: empty command", ctr.ErrInvalidArgument) } if m.queries == nil { return nil, fmt.Errorf("db is not configured") } startedAt := time.Now() if _, err := m.CreateVersion(ctx, req.BotID); err != nil { return nil, err } result, err := m.service.ExecTask(ctx, m.containerID(req.BotID), ctr.ExecTaskRequest{ Args: req.Command, Env: req.Env, WorkDir: req.WorkDir, Terminal: req.Terminal, UseStdio: req.UseStdio, }) if err != nil { return nil, err } if err := m.insertEvent(ctx, m.containerID(req.BotID), "exec", map[string]any{ "bot_id": req.BotID, "command": req.Command, "work_dir": req.WorkDir, "exit_code": result.ExitCode, "duration": time.Since(startedAt).String(), }); err != nil { return nil, err } return &ExecResult{ExitCode: result.ExitCode}, nil } // DataDir returns the host data directory for a bot. func (m *Manager) DataDir(botID string) (string, error) { if err := validateBotID(botID); err != nil { return "", err } root := m.cfg.DataRoot if root == "" { root = config.DefaultDataRoot } return filepath.Join(root, "bots", botID), nil } func (m *Manager) ensureBotDir(botID string) (string, error) { root := m.cfg.DataRoot if root == "" { root = config.DefaultDataRoot } dir := filepath.Join(root, "bots", botID) if err := os.MkdirAll(dir, 0o755); err != nil { return "", err } return dir, nil } func validateBotID(botID string) error { return identity.ValidateUserID(botID) }