ollie/session/setup.go

package session

import (
	"context"
	"fmt"
	"strings"

	"ollie/agent"
	"ollie/prompts"
	"ollie/toolsrv"
)

// SessionInfra holds the tool server infrastructure for a session.
// Created by SetupToolServer and used by both CreateAgent and restoreSession.
type SessionInfra struct {
	Proc          *toolsrv.Process
	Keeper        *toolsrv.ProcessKeeper
	ToolsConn     toolsrv.Runner
	NewToolServer func() toolsrv.Runner
	PromptEnv     []string
}

// ToolServerConfig configures tool server spawning.
type ToolServerConfig struct {
	Ctx          context.Context
	CWD          string
	RemoteTarget string
	Yolo         bool

	// ReuseFrom, if non-nil, reuses existing infrastructure instead of spawning.
	// Used when a session already has a tool server from Create().
	ReuseFrom *InfraConfig
}

// InfraConfig describes existing infrastructure to reuse.
type InfraConfig struct {
	Proc      *toolsrv.Process
	Keeper    *toolsrv.ProcessKeeper
	ToolsConn toolsrv.Runner
}

// SetupToolServer spawns the tool server process and creates the infrastructure.
// For remote targets, it spawns via SSH. For local, it spawns a local process.
// If ReuseInfra is set, reuses existing infrastructure instead of spawning.
// Returns the infrastructure or an error.
func SetupToolServer(cfg ToolServerConfig) (*SessionInfra, error) {
	// Reuse existing infrastructure if provided.
	if cfg.ReuseFrom != nil {
		r := cfg.ReuseFrom
		var newToolServer func() toolsrv.Runner
		if r.Keeper != nil {
			newToolServer = func() toolsrv.Runner {
				conn, err := r.Keeper.Dial()
				if err != nil {
					return nil
				}
				return conn
			}
		} else {
			newToolServer = func() toolsrv.Runner {
				return r.ToolsConn
			}
		}
		// Update CWD if the connection supports it.
		if cwc, ok := r.ToolsConn.(interface{ SetCWD(string) }); ok {
			cwc.SetCWD(cfg.CWD)
		}
		return &SessionInfra{
			Proc:          r.Proc,
			Keeper:        r.Keeper,
			ToolsConn:     newToolServer(),
			NewToolServer: newToolServer,
			PromptEnv:     agent.PromptEnv(cfg.CWD),
		}, nil
	}
	var proc *toolsrv.Process
	var keeper *toolsrv.ProcessKeeper
	var newToolServer func() toolsrv.Runner
	var promptEnv []string
	var err error

	if cfg.RemoteTarget != "" {
		proc, err = toolsrv.SpawnRemote(cfg.Ctx, toolsrv.RemoteConfig{
			SSHTarget: cfg.RemoteTarget,
			CWD:       cfg.CWD,
		})
		if err != nil {
			return nil, fmt.Errorf("remote spawn: %w", err)
		}
		keeper = toolsrv.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) {
			return toolsrv.SpawnRemote(ctx, toolsrv.RemoteConfig{
				SSHTarget: cfg.RemoteTarget,
				CWD:       cfg.CWD,
			})
		})
		newToolServer = func() toolsrv.Runner {
			conn, err := keeper.Dial()
			if err != nil {
				return nil
			}
			return conn
		}
		promptEnv = []string{
			"PRIME_CWD=" + cfg.CWD,
			"PRIME_PLATFORM=" + proc.Info.Platform,
			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo),
		}
	} else {
		var dialOpts []toolsrv.Option
		if cfg.Yolo {
			dialOpts = append(dialOpts, toolsrv.WithYolo())
		}
		proc, err = toolsrv.Spawn(cfg.Ctx, cfg.CWD, dialOpts...)
		if err != nil {
			return nil, fmt.Errorf("local spawn: %w", err)
		}
		keeper = toolsrv.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) {
			return toolsrv.Spawn(ctx, cfg.CWD, dialOpts...)
		})
		newToolServer = func() toolsrv.Runner {
			conn, err := keeper.Dial()
			if err != nil {
				return nil
			}
			return conn
		}
		promptEnv = agent.PromptEnv(cfg.CWD)
	}

	toolsConn := newToolServer()

	return &SessionInfra{
		Proc:          proc,
		Keeper:        keeper,
		ToolsConn:     toolsConn,
		NewToolServer: newToolServer,
		PromptEnv:     promptEnv,
	}, nil
}

// PromptLayers holds the resolved prompt components for an agent.
type PromptLayers struct {
	SystemPrompt string
	EnvBlock     string
}

// BuildPromptLayers resolves system prompt, operational model, and environment block.
func BuildPromptLayers(cfg *agent.AgentConfig, cwd, sessID, uname string, promptEnv []string, systemPromptOverride string) PromptLayers {
	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)

	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, "")

	return PromptLayers{
		SystemPrompt: sysPrompt,
		EnvBlock:     envBlock,
	}
}

// BaseLayers returns the prompt layers as a slice for agent configuration.
func (p PromptLayers) BaseLayers() []string {
	return []string{p.SystemPrompt, p.EnvBlock}
}

// LoadAutoLoadTools loads all tools from the agent config's autoLoad list.
// It first sends the session ID and uname to the tool server so that
// ollie-remote can attach its tool registry before processing tool_load RPCs.
func LoadAutoLoadTools(cfg *agent.AgentConfig, runner toolsrv.Runner, sessID, uname string, logError func(string, ...any)) {
	if cfg == nil || runner == nil {
		return
	}

	// Send session ID to the tool server first — ollie-remote uses this
	// to attach its tool registry via the OnEnvSet callback.
	if es, ok := runner.(interface{ SetEnv(string, string) }); ok {
		if sessID != "" {
			es.SetEnv("OLLIE_SESSION_ID", sessID)
		}
		if uname != "" {
			es.SetEnv("OLLIE_UNAME", uname)
		}
	}

	for _, tl := range cfg.AutoLoad {
		if err := LoadToolOnRunner(runner, tl); err != nil {
			if logError != nil {
				logError("autoLoad tool %q: %v", tl, err)
			}
		}
	}
}