ollie-9p/session/create.go

package session

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

	"ollie/agent"
	"ollie/backend"
	coresession "ollie/session"
	"ollie/toolsrv"
	"ollie/tools"
	"ollie/paths"
	"olliesrv/fs"
	"olliesrv/prompts"
)

// Create creates a new session from key=value args. Returns the session ID.
func Create(rs *rootState, args []string) (string, error) {
	name := ""
	backendOverride := ""
	modelOverride := ""
	agentName := ""
	cwd := ""
	remoteTarget := ""
	systemPromptOverride := ""
	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 "cwd":
			cwd = v
		case "remote":
			remoteTarget = v
		case "system_prompt":
			systemPromptOverride = v
		default:
			return "", fmt.Errorf("unknown option %q (valid: name, backend, model, agent, cwd, remote, system_prompt)", k)
		}
	}
	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 (e.g. new cwd=/path/to/project)")
	}

	sessID := name
	if sessID == "" {
		sessID = coresession.NewSessionID()
	}

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

	var core *coresession.Session
	var sessPtr *Session
	uname := rs.nextUname()
	if rs.cfg.NewCore != nil {
		var err error
		core, err = rs.cfg.NewCore(sessID, agentName, cwd)
		if err != nil {
			return "", err
		}
	} else {
		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, fs.PermMkdirPrivate); err != nil {
			return "", fmt.Errorf("sessions dir: %w", err)
		}

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

		if cfg != nil && len(cfg.AllowTools) > 0 {
			execOpts = append(execOpts, toolsrv.WithAllowTools(cfg.AllowTools))
		}
		if rs.cfg.ToolRegistry != nil {
			execOpts = append(execOpts, toolsrv.WithToolRegistry(rs.cfg.ToolRegistry, sessID))
		}
		if rs.cfg.SkillsRegistry != nil {
			execOpts = append(execOpts, toolsrv.WithSkillsRegistry(rs.cfg.SkillsRegistry))
			execOpts = append(execOpts, toolsrv.WithBuiltins(tools.Builtins()))
		}
		var newToolServer func() toolsrv.Runner
		var remoteEnv []string
		if remoteTarget != "" {
			rsrv, dialErr := toolsrv.RemoteDial(rs.cfg.Ctx, toolsrv.RemoteConfig{
				SSHTarget: remoteTarget,
				CWD:       cwd,
			})
			if dialErr != nil {
				return "", fmt.Errorf("remote dial: %w", dialErr)
			}
			newToolServer = func() toolsrv.Runner { return rsrv }
			remoteEnv = []string{
				"PRIME_CWD=" + cwd,
				"PRIME_PLATFORM=" + rsrv.Info.Platform,
				"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo),
			}
		} else {
			newToolServer = func() toolsrv.Runner {
				srv, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(cwd), execOpts...)
				return srv
			}
		}

		var promptEnv []string
		if len(remoteEnv) > 0 {
			promptEnv = remoteEnv
		} else {
			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, rs.cfg.EnableDBus, 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)

		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"),
			ReadPlanStep: func() string {
				if sessPtr == nil {
					return ""
				}
				sessPtr.mu.RLock()
				data := make([]byte, len(sessPtr.plan))
				copy(data, sessPtr.plan)
				sessPtr.mu.RUnlock()
				return coresession.NextUncheckedStep(data)
			},
		})
	}

	sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
	sess := NewSession(sessID, core, sessionCtx, sessionCancel)
	sessPtr = sess
	sess.remote = remoteTarget

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

	// Load the agent now that the session is registered.
	core.Agent().Submit(sessionCtx, "/agent "+agentName)

	rs.cfg.Log.Info("new session %s (backend=%s model=%s agent=%s)",
		sessID, core.Agent().BackendName(), core.Agent().ModelName(), core.Agent().Name())
	if rs.cfg.OnSessionCreated != nil {
		rs.cfg.OnSessionCreated(sessID, sess)
	}
	return sessID, nil
}

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