ollie/fs/session/create.go

package fs

import (
	"context"
	"fmt"
	"os"
	"strings"

	"ollie/agent"
	"ollie/backend"
	"ollie/paths"
	"ollie/prompts"
	coresession "ollie/session"
	"ollie/toolsrv"
)

// Create creates a new empty session from key=value args and registers it.
// Returns the session name. The session has no agent (Core is nil) until
// CreateAgent is called.
func Create(rs *rootState, args []string) (string, error) {
	name := ""
	for _, arg := range args {
		k, v, ok := strings.Cut(arg, "=")
		if !ok {
			return "", fmt.Errorf("invalid option %q (expected key=value)", arg)
		}
		if v == "" {
			continue
		}
		switch k {
		case "name":
			name = v
		default:
			return "", fmt.Errorf("unknown option %q (valid: name)", k)
		}
	}

	sessID := coresession.NewSessionID()
	if name == "" {
		name = sessID
		if idx := strings.IndexByte(sessID, '-'); idx > 0 {
			name = sessID[:idx]
		}
	}

	rs.mu.RLock()
	_, exists := rs.sessions[name]
	rs.mu.RUnlock()
	if exists {
		return "", fmt.Errorf("session already exists: %s", name)
	}

	sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
	sess := NewSession(sessID, nil, sessionCtx, sessionCancel)
	sess.SetName(name)

	rs.mu.Lock()
	rs.sessions[name] = sess
	rs.mu.Unlock()

	rs.cfg.Log.Info("created empty session %s (id=%s)", name, sessID)
	broadcastEvents()
	return name, nil
}

// CreateAgent creates and configures an agent within a new or existing session.
// If the session already exists (has a Core), it only replaces the agent.
// If the session doesn't exist yet, it creates the full session with agent.
func CreateAgent(rs *rootState, sessName string, args []string) error {
	backendOverride := ""
	modelOverride := ""
	agentName := ""
	agentAlias := ""
	cwd := ""
	remoteTarget := ""
	systemPromptOverride := ""
	name := ""
	for _, arg := range args {
		k, v, ok := strings.Cut(arg, "=")
		if !ok {
			return fmt.Errorf("invalid option %q (expected key=value)", arg)
		}
		if v == "" {
			continue
		}
		switch k {
		case "name":
			name = v
		case "backend":
			backendOverride = v
		case "model":
			modelOverride = v
		case "agent":
			agentName = v
		case "agent_alias":
			agentAlias = v
		case "cwd":
			cwd = v
		case "remote":
			remoteTarget = v
		case "system_prompt":
			systemPromptOverride = v
		default:
			return fmt.Errorf("unknown option %q (valid: name, backend, model, agent, agent_alias, cwd, remote, system_prompt)", k)
		}
	}

	// If name is provided but no agent_alias, use name as alias.
	if name != "" && agentAlias == "" {
		agentAlias = name
	}

	rs.mu.RLock()
	sess, ok := rs.sessions[sessName]
	rs.mu.RUnlock()
	if !ok {
		// Session doesn't exist yet — create it later after Core is set
		sessID := coresession.NewSessionID()
		sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
		sess = NewSession(sessID, nil, sessionCtx, sessionCancel)
		sess.SetName(sessName)
	}

	if agentName == "" {
		if v := os.Getenv("OLLIE_DEFAULT_AGENT"); v != "" {
			agentName = v
		} else {
			agentName = "default"
		}
	}

	cwd = paths.ExpandHome(os.ExpandEnv(cwd))
	if cwd == "" {
		return fmt.Errorf("cwd is required (set cwd=<path> when creating an agent)")
	}
	if remoteTarget == "" {
		if info, err := os.Stat(cwd); err != nil {
			return fmt.Errorf("cwd %q does not exist", cwd)
		} else if !info.IsDir() {
			return fmt.Errorf("cwd %q is not a directory", cwd)
		}
	}

	uname := rs.nextUname()
	sessID := sess.id

	cfg := LoadAgentConfig(rs.cfg.AgentsDir, agentName, nil)
	if cfg != nil {
		if backendOverride == "" && cfg.Backend != "" {
			backendOverride = cfg.Backend
		}
		if modelOverride == "" && cfg.Model != "" {
			modelOverride = cfg.Model
		}
	}

	be, err := backend.NewWithName(backendOverride)
	if err != nil {
		return fmt.Errorf("backend: %w", err)
	}

	if modelOverride == "" {
		modelOverride = os.Getenv("OLLIE_MODEL")
	}
	if modelOverride != "" {
		be.SetModel(modelOverride)
	}

	if err := rs.cfg.MkdirAll(rs.cfg.SessionsDir, PermMkdirPrivate); err != nil {
		return fmt.Errorf("sessions dir: %w", err)
	}

	var dialOpts []toolsrv.Option
	if rs.cfg.Strict {
		dialOpts = append(dialOpts, toolsrv.WithStrict())
	}
	if rs.cfg.Yolo {
		dialOpts = append(dialOpts, toolsrv.WithYolo())
	}

	var proc *toolsrv.Process
	var newToolServer func() toolsrv.Runner
	var promptEnv []string

	if remoteTarget != "" {
		proc, err = toolsrv.SpawnRemote(rs.cfg.Ctx, toolsrv.RemoteConfig{
			SSHTarget: remoteTarget,
			CWD:       cwd,
		})
		if err != nil {
			return fmt.Errorf("remote spawn: %w", err)
		}
		newToolServer = func() toolsrv.Runner {
			conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket))
			return conn
		}
		promptEnv = []string{
			"PRIME_CWD=" + cwd,
			"PRIME_PLATFORM=" + proc.Info.Platform,
			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo),
		}
	} else {
		proc, err = toolsrv.Spawn(rs.cfg.Ctx, cwd, dialOpts...)
		if err != nil {
			return fmt.Errorf("local spawn: %w", err)
		}
		newToolServer = func() toolsrv.Runner {
			conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket))
			return conn
		}
		promptEnv = agent.PromptEnv(cwd)
	}

	env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname}
	env = append(env, promptEnv...)

	spOverride := systemPromptOverride
	if spOverride == "" && cfg != nil {
		spOverride = cfg.SystemPrompt
	}
	sysPrompt := prompts.ResolveSystemPrompt(spOverride)

	envMap := make(map[string]string)
	for _, e := range env {
		if k, v, ok := strings.Cut(e, "="); ok {
			envMap[k] = v
		}
	}
	opModel := prompts.OperationalModel(rs.cfg.Enable9P, false, envMap)

	platform := "linux"
	isGitRepo := false
	for _, e := range promptEnv {
		if k, v, ok := strings.Cut(e, "="); ok {
			switch k {
			case "PRIME_PLATFORM":
				platform = v
			case "PRIME_IS_GIT_REPO":
				isGitRepo = v == "true"
			}
		}
	}
	envBlock := prompts.Environment(cwd, platform, isGitRepo, "")

	toolSrv := newToolServer()
	rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, sysPrompt, opModel, envBlock)

	al := NewAgentLog(remoteTarget)

	if sess.Core == nil {
		// No core session yet — create one with the first agent.
		core := coresession.New(coresession.Config{
			Backend:        be,
			AgentName:      agentName,
			AgentsDir:      rs.cfg.AgentsDir,
			SessionsDir:    rs.cfg.SessionsDir,
			SessionID:      sessID,
			AgentID:        uname,
			CWD:            cwd,
			Remote:         remoteTarget,
			Runtime:        rt,
			NewToolServer:  newToolServer,
			PromptEnvExtra: promptEnv,
			BaseLayers:     []string{sysPrompt, opModel, envBlock},
			Log:            rs.cfg.Sink.NewLogger("core"),
			Output:         NewEventHandler(al),
			ReadPlanStep: func() string {
				return ""
			},
		})

		sess.mu.Lock()
		sess.Core = core
		sess.uname = uname
		sess.setAgentLogLocked(uname, al)
		sess.proc = proc
		sess.mu.Unlock()
	} else {
		// Core session exists — add another agent to it.
		saveFn, flushFn := sess.Core.SaveFuncs()
		ag := agent.NewAgent(agent.AgentCfg{
			History:        nil,
			Runtime:        rt,
			AgentName:      agentName,
			AgentsDir:      rs.cfg.AgentsDir,
			AgentID:        uname,
			CWD:            cwd,
			BaseLayers:     []string{sysPrompt, opModel, envBlock},
			PromptEnvExtra: promptEnv,
			NewToolServer:  newToolServer,
			NewBackend:     backend.NewWithName,
			Output:         NewEventHandler(al),
			Log:            rs.cfg.Sink.NewLogger("core"),
			AuditLog:       rs.cfg.Sink.NewLogger("core").Sub("audit"),
			SessionID:      sessID,
			ReadPlanStep: func() string {
				return ""
			},
			SaveSession: saveFn,
			FlushSave:   flushFn,
		})
		sess.Core.AddAgent(ag)

		sess.mu.Lock()
		sess.setAgentLogLocked(uname, al)
		sess.mu.Unlock()
	}

	// Register session in the map if not already there
	rs.mu.Lock()
	if _, exists := rs.sessions[sessName]; !exists {
		rs.sessions[sessName] = sess
	}
	rs.mu.Unlock()

	ag := sess.Core.Agent()
	ag.Submit(sess.SessionCtx, "/agent "+agentName)
	if agentAlias != "" {
		ag.SetName(agentAlias)
	} else {
		ag.SetName(uname)
	}

	rs.cfg.Log.Info("created agent in session %s (backend=%s model=%s agent=%s)",
		sessName, ag.BackendName(), ag.ModelName(), ag.Name())
	broadcastEvents()
	return nil
}

// CreateFromRoot is the public API for callers that have a *Tree.
func CreateFromRoot(root *Tree, args []string) (string, error) {
	return Create(state(root), args)
}

// CreateAgentFromRoot is the public API for callers that have a *Tree.
// Returns the agent name on success.
func CreateAgentFromRoot(root *Tree, sessName string, args []string) (string, error) {
	rs := state(root)
	if err := CreateAgent(rs, sessName, args); err != nil {
		return "", err
	}
	rs.mu.RLock()
	sess := rs.sessions[sessName]
	rs.mu.RUnlock()
	if sess == nil || sess.Core == nil {
		return "", fmt.Errorf("agent not created")
	}
	return sess.Core.Agent().Name(), nil
}