lvmh/daemon/docker.go

package main

// docker.go — worker container spawner (clone → ensure image → run).

import (
	"archive/tar"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api/types/build"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/image"
	"github.com/docker/docker/api/types/volume"
	"github.com/docker/docker/client"
)

// Spawn job states and docker wiring constants.
const (
	stateCloning  string = "cloning"
	stateBuilding string = "building"
	stateCreating string = "creating"
	stateRunning  string = "running"
	stateError    string = "error"

	imageRefWorker    string = "lvmh-worker:latest"
	labelSession      string = "lvmh.session"
	volumeRepoPrefix  string = "lvmh-repo-"
	volumeSessions    string = "lvmh-sessions"
	workspaceMount    string = "/workspace"
	sessionsMount     string = "/pi-sessions"
	modelsMountTarget string = "/root/.pi/agent/models.json"

	envWorkerDockerfile   string = "LVMH_WORKER_DOCKERFILE"
	defaultDockerfile     string = "/app/build/docker/worker.Dockerfile"
	envWorkerContext      string = "LVMH_WORKER_CONTEXT"
	envRepoDir            string = "LVMH_REPO_DIR"
	defaultRepoDir        string = "/data/repos"
	envContainerLVMHURL   string = "LVMH_CONTAINER_LVMH_URL"
	defaultContainerURL   string = "ws://lvmh:8686/agent/ws"
	envContainerNetwork   string = "LVMH_CONTAINER_NETWORK"
	defaultNetwork        string = "lvmh-net"
	envWorkerModels       string = "LVMH_WORKER_MODELS"
	defaultWorkerModels   string = "/app/build/docker/worker-models.json"
	envProviderAPIKey     string = "ZAI_RENAUD_API_KEY"
	envLVMHSessionID      string = "LVMH_SESSION_ID"
	envLVMHAgent          string = "LVMH_AGENT"
	envLVMHRepo           string = "LVMH_REPO"
	stopTimeoutSeconds    int    = 10
	buildContextReadLimit int64  = 64 << 20
)

var errNoContainer = errors.New("no container for session")

var repoPathRe = regexp.MustCompile(`^[A-Za-z0-9_.-]+(/[A-Za-z0-9_.-]+)+$`)

// validRepoPath accepts "group/project" style paths (at least two segments).
func validRepoPath(repo string) bool { return repoPathRe.MatchString(repo) }

func repoSlug(repo string) string { return strings.ReplaceAll(repo, "/", "-") }

// newUUID returns a random RFC 4122 v4 UUID string.
func newUUID() string {
	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		// crypto/rand failure is fatal-grade; never expected in practice.
		panic("newUUID: crypto/rand unavailable: " + err.Error())
	}
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	dst := make([]byte, 36)
	hex.Encode(dst[0:8], b[0:4])
	dst[8] = '-'
	hex.Encode(dst[9:13], b[4:6])
	dst[13] = '-'
	hex.Encode(dst[14:18], b[6:8])
	dst[18] = '-'
	hex.Encode(dst[19:23], b[8:10])
	dst[23] = '-'
	hex.Encode(dst[24:36], b[10:16])
	return string(dst)
}

// SpawnJob is one async spawn pipeline run, exposed via /api/spawn/status.
type SpawnJob struct {
	Repo        string `json:"repo"`
	State       string `json:"state"`
	ContainerID string `json:"containerId"`
	SessionID   string `json:"sessionId,omitempty"`
	Message     string `json:"message,omitempty"`
	UpdatedAt   int64  `json:"-"`
}

// SpawnResult is the immediate POST /api/spawn response.
type SpawnResult struct {
	SessionID   string `json:"sessionId"`
	ContainerID string `json:"containerId"`
}

// Spawner owns the docker client, per-repo clone serialization and job state.
type Spawner struct {
	store   *Store
	hub     *Hub
	cli     *client.Client
	baseURL string // gitlab base, for clone URLs

	reposDir     string
	dockerfile   string
	buildContext string
	containerURL string
	network      string
	modelsPath   string

	mu        sync.Mutex
	jobs      map[string]*SpawnJob // keyed by sessionId
	slugLocks map[string]*sync.Mutex
}

func NewSpawner(store *Store, hub *Hub, baseURL string) (*Spawner, error) {
	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
	if err != nil {
		return nil, fmt.Errorf("docker client: %w", err)
	}
	dockerfile := envOr(envWorkerDockerfile, defaultDockerfile)
	return &Spawner{
		store:        store,
		hub:          hub,
		cli:          cli,
		baseURL:      strings.TrimRight(baseURL, "/"),
		reposDir:     envOr(envRepoDir, defaultRepoDir),
		dockerfile:   dockerfile,
		buildContext: envOr(envWorkerContext, filepath.Dir(filepath.Dir(dockerfile))),
		containerURL: envOr(envContainerLVMHURL, defaultContainerURL),
		network:      envOr(envContainerNetwork, defaultNetwork),
		modelsPath:   envOr(envWorkerModels, defaultWorkerModels),
		jobs:         make(map[string]*SpawnJob),
		slugLocks:    make(map[string]*sync.Mutex),
	}, nil
}

func envOr(key, def string) string {
	if v := strings.TrimSpace(os.Getenv(key)); v != "" {
		return v
	}
	return def
}

// JobsSnapshot returns current spawn jobs for /api/spawn/status and /ws pushes.
func (s *Spawner) JobsSnapshot() []SpawnJob {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]SpawnJob, 0, len(s.jobs))
	for _, j := range s.jobs {
		out = append(out, *j)
	}
	return out
}

func (s *Spawner) setJob(sessionID, repo, state, containerID, message string) {
	s.mu.Lock()
	j, ok := s.jobs[sessionID]
	if !ok {
		j = &SpawnJob{SessionID: sessionID, Repo: repo}
		s.jobs[sessionID] = j
	}
	j.Repo = repo
	j.State = state
	j.ContainerID = containerID
	j.Message = message
	j.UpdatedAt = time.Now().UnixMilli()
	s.mu.Unlock()
	s.hub.BroadcastSpawnStatus()
}

// Start launches the async spawn pipeline and returns the new sessionId.
func (s *Spawner) Start(ctx context.Context, repo, branch string) (SpawnResult, error) {
	if _, err := s.imageExists(ctx); err != nil {
		return SpawnResult{}, fmt.Errorf("docker unavailable: %w", err)
	}
	exists, err := s.imageExists(ctx)
	if err == nil && !exists {
		if _, statErr := os.Stat(s.dockerfile); statErr != nil {
			return SpawnResult{}, fmt.Errorf(
				"image %s not found and no worker Dockerfile at %s (set %s or run `make worker-image`)",
				imageRefWorker, s.dockerfile, envWorkerDockerfile)
		}
	}
	sessionID := newUUID()
	s.setJob(sessionID, repo, stateCloning, "", "")
	go s.runJob(repo, branch, sessionID)
	return SpawnResult{SessionID: sessionID}, nil
}

func (s *Spawner) imageExists(ctx context.Context) (bool, error) {
	args := filters.NewArgs(filters.Arg("reference", imageRefWorker))
	summaries, err := s.cli.ImageList(ctx, image.ListOptions{Filters: args})
	if err != nil {
		return false, err
	}
	return len(summaries) > 0, nil
}

func (s *Spawner) slugLock(slug string) *sync.Mutex {
	s.mu.Lock()
	defer s.mu.Unlock()
	l, ok := s.slugLocks[slug]
	if !ok {
		l = &sync.Mutex{}
		s.slugLocks[slug] = l
	}
	return l
}

// runJob is the async clone→build→create→start pipeline.
func (s *Spawner) runJob(repo, branch, sessionID string) {
	slug := repoSlug(repo)
	lock := s.slugLock(slug)
	lock.Lock()
	defer lock.Unlock()

	if err := s.cloneOrUpdate(repo, branch, slug); err != nil {
		s.setJob(sessionID, repo, stateError, "", err.Error())
		return
	}
	s.setJob(sessionID, repo, stateBuilding, "", "")
	if err := s.ensureImage(context.Background()); err != nil {
		s.setJob(sessionID, repo, stateError, "", err.Error())
		return
	}
	s.setJob(sessionID, repo, stateCreating, "", "")
	containerID, err := s.createAndStart(context.Background(), repo, slug, sessionID)
	if err != nil {
		s.setJob(sessionID, repo, stateError, "", err.Error())
		return
	}
	if err := s.store.UpsertContainer(sessionID, containerID, repo); err != nil {
		s.setJob(sessionID, repo, stateError, containerID, "container started but not persisted: "+err.Error())
		return
	}
	s.setJob(sessionID, repo, stateRunning, containerID, "")
}

// cloneOrUpdate clones the repo into reposDir/<slug> or fast-forwards it.
func (s *Spawner) cloneOrUpdate(repo, branch, slug string) error {
	dir := filepath.Join(s.reposDir, slug)
	cloneURL, err := s.cloneURL(repo)
	if err != nil {
		return err
	}
	if st, err := os.Stat(filepath.Join(dir, ".git")); err == nil && st.IsDir() {
		if err := gitRun(dir, "pull", "--ff-only"); err != nil {
			return fmt.Errorf("git pull %s: %w", repo, err)
		}
		return nil
	}
	if err := os.MkdirAll(s.reposDir, 0o755); err != nil {
		return err
	}
	args := []string{"clone"}
	if branch != "" {
		args = append(args, "--branch", branch)
	}
	args = append(args, "--", cloneURL, dir)
	if err := gitRun("", args...); err != nil {
		return fmt.Errorf("git clone %s: %w", repo, err)
	}
	return nil
}

// cloneURL builds an authenticated https clone URL when a PAT is stored.
func (s *Spawner) cloneURL(repo string) (string, error) {
	u, err := url.Parse(s.baseURL + "/" + repo + ".git")
	if err != nil {
		return "", err
	}
	if token, ok, _ := s.store.GetSetting(settingGitLabToken); ok && token != "" && u.User == nil {
		u.User = url.UserPassword("oauth2", token)
	}
	return u.String(), nil
}

func gitRun(dir string, args ...string) error {
	cmd := exec.Command("git", args...)
	if dir != "" {
		cmd.Dir = dir
	}
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		msg := strings.TrimSpace(out.String())
		if len(msg) > 500 {
			msg = msg[len(msg)-500:]
		}
		if msg != "" {
			return fmt.Errorf("%w: %s", err, msg)
		}
		return err
	}
	return nil
}

// ensureImage builds lvmh-worker:latest when missing.
func (s *Spawner) ensureImage(ctx context.Context) error {
	exists, err := s.imageExists(ctx)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	if _, err := os.Stat(s.dockerfile); err != nil {
		return fmt.Errorf("worker Dockerfile missing at %s: %w", s.dockerfile, err)
	}
	relDockerfile, err := filepath.Rel(s.buildContext, s.dockerfile)
	if err != nil {
		return err
	}
	var buf bytes.Buffer
	if err := tarDir(&buf, s.buildContext); err != nil {
		return fmt.Errorf("build context %s: %w", s.buildContext, err)
	}
	resp, err := s.cli.ImageBuild(ctx, &buf, build.ImageBuildOptions{
		Tags:       []string{imageRefWorker},
		Dockerfile: relDockerfile,
		Remove:     true,
	})
	if err != nil {
		return fmt.Errorf("docker build: %w", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, buildContextReadLimit))
	if msg := extractBuildError(body); msg != "" {
		return fmt.Errorf("docker build failed: %s", msg)
	}
	return nil
}

// extractBuildError scans the build output stream for an "error" field.
func extractBuildError(body []byte) string {
	for _, line := range bytes.Split(body, []byte("\n")) {
		var msg struct {
			Error       string `json:"error"`
			ErrorDetail struct {
				Message string `json:"message"`
			} `json:"errorDetail"`
		}
		if json.Unmarshal(line, &msg) != nil {
			continue
		}
		if msg.Error != "" || msg.ErrorDetail.Message != "" {
			if msg.ErrorDetail.Message != "" {
				return msg.ErrorDetail.Message
			}
			return msg.Error
		}
	}
	return ""
}

// createAndStart provisions volumes, creates and starts the worker container.
func (s *Spawner) createAndStart(ctx context.Context, repo, slug, sessionID string) (string, error) {
	repoVolume := volumeRepoPrefix + slug
	fresh, err := s.ensureRepoVolume(ctx, repoVolume)
	if err != nil {
		return "", err
	}
	if fresh {
		if err := s.seedVolume(ctx, slug, repoVolume); err != nil {
			return "", fmt.Errorf("seed %s: %w", repoVolume, err)
		}
	}
	if _, err := s.cli.VolumeCreate(ctx, volume.CreateOptions{Name: volumeSessions}); err != nil {
		return "", fmt.Errorf("volume %s: %w", volumeSessions, err)
	}

	binds := []string{repoVolume + ":" + workspaceMount, volumeSessions + ":" + sessionsMount}
	if _, err := os.Stat(s.modelsPath); err == nil {
		binds = append(binds, s.modelsPath+":"+modelsMountTarget+":ro")
	}
	cfg := &container.Config{
		Image: imageRefWorker,
		Env: []string{
			envProviderAPIKey + "=" + os.Getenv(envProviderAPIKey),
			envToken + "=" + daemonToken,
			"LVMH_URL=" + s.containerURL,
			envLVMHSessionID + "=" + sessionID,
			envLVMHAgent + "=1",
			envLVMHRepo + "=" + repo,
		},
		Labels: map[string]string{labelSession: sessionID},
	}
	hostCfg := &container.HostConfig{
		Binds:       binds,
		NetworkMode: container.NetworkMode(s.network),
		AutoRemove:  false,
	}
	name := "lvmh-agent-" + strings.ReplaceAll(sessionID, "-", "")[:12]
	created, err := s.cli.ContainerCreate(ctx, cfg, hostCfg, nil, nil, name)
	if err != nil {
		return "", fmt.Errorf("docker create: %w", err)
	}
	if err := s.cli.ContainerStart(ctx, created.ID, container.StartOptions{}); err != nil {
		s.removeContainer(context.Background(), created.ID)
		return "", fmt.Errorf("docker start: %w", err)
	}
	return created.ID, nil
}

// ensureRepoVolume creates the per-repo volume; reports whether it is fresh.
func (s *Spawner) ensureRepoVolume(ctx context.Context, name string) (bool, error) {
	if _, err := s.cli.VolumeInspect(ctx, name); err == nil {
		return false, nil
	}
	if _, err := s.cli.VolumeCreate(ctx, volume.CreateOptions{Name: name}); err != nil {
		return false, fmt.Errorf("volume %s: %w", name, err)
	}
	return true, nil
}

// seedVolume copies the cloned repo into the fresh named volume via a
// one-shot container from the worker image (no extra images needed).
func (s *Spawner) seedVolume(ctx context.Context, slug, repoVolume string) error {
	repoDir := filepath.Join(s.reposDir, slug)
	cfg := &container.Config{
		Image: imageRefWorker,
		Cmd:   []string{"sh", "-c", "cp -a /src/. " + workspaceMount + "/"},
	}
	hostCfg := &container.HostConfig{
		Binds: []string{
			repoVolume + ":" + workspaceMount,
			repoDir + ":/src:ro",
		},
	}
	created, err := s.cli.ContainerCreate(ctx, cfg, hostCfg, nil, nil, "lvmh-seed-"+slug)
	if err != nil {
		return err
	}
	defer s.removeContainer(context.Background(), created.ID)
	if err := s.cli.ContainerStart(ctx, created.ID, container.StartOptions{}); err != nil {
		return err
	}
	waitCh, errCh := s.cli.ContainerWait(ctx, created.ID, container.WaitConditionNotRunning)
	select {
	case <-waitCh:
		return nil
	case err := <-errCh:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *Spawner) removeContainer(ctx context.Context, id string) {
	_ = s.cli.ContainerRemove(ctx, id, container.RemoveOptions{Force: true})
}

// RemoveSession stops and removes the container spawned for a session.
func (s *Spawner) RemoveSession(ctx context.Context, sessionID string) error {
	row, ok, err := s.store.GetContainer(sessionID)
	if err != nil {
		return err
	}
	if !ok {
		return errNoContainer
	}
	stopCtx, cancel := context.WithTimeout(context.Background(), time.Duration(stopTimeoutSeconds)*time.Second)
	if err := s.cli.ContainerStop(stopCtx, row.ContainerID, container.StopOptions{}); err != nil {
		cancel()
		return fmt.Errorf("docker stop: %w", err)
	}
	cancel()
	s.removeContainer(ctx, row.ContainerID)
	if err := s.store.DeleteContainer(sessionID); err != nil {
		return err
	}
	s.mu.Lock()
	if j, ok := s.jobs[sessionID]; ok {
		j.State = stateError
		j.Message = "container removed"
		j.UpdatedAt = time.Now().UnixMilli()
	}
	s.mu.Unlock()
	s.hub.BroadcastSpawnStatus()
	return nil
}

// tarDir writes a tar archive of dir (excluding .git-heavy junk) into w.
func tarDir(w io.Writer, dir string) error {
	tw := tar.NewWriter(w)
	defer tw.Close()
	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		rel, err := filepath.Rel(dir, path)
		if err != nil {
			return err
		}
		if rel == "." {
			return nil
		}
		name := filepath.ToSlash(rel)
		hdr, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}
		hdr.Name = name
		if info.IsDir() {
			hdr.Name += "/"
		}
		hdr.Uname = "root"
		hdr.Gname = "root"
		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		if info.Mode().IsRegular() {
			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer f.Close()
			if _, err := io.Copy(tw, f); err != nil {
				return err
			}
		}
		return nil
	})
}