ollie/cmd/olliesrv/internal/session/setup.go

package session

import (
	"context"
	"fmt"

	"ollie/cmd/olliesrv/internal/agent"
	"ollie/cmd/olliesrv/internal/prompts"
	"ollie/cmd/olliesrv/internal/toolclient"
	toolsrvclient "ollie/toolsrv/client"
	"ollie/util"
)

// SessionInfra holds the tool server infrastructure for a session.
// Created by SetupToolServer and used by both CreateAgent and restoreSession.
type SessionInfra struct {
	Proc          *toolclient.Process
	Keeper        *toolclient.ProcessKeeper
	ToolsConn     *toolsrvclient.Conn
	NewToolServer func() *toolsrvclient.Conn
	Platform      string
	IsGitRepo     bool
}

// ToolServerConfig configures tool server spawning.
type ToolServerConfig struct {
	Ctx          context.Context
	CWD          string
	RemoteTarget string
	SessionID    string // session ID for tool registry
	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      *toolclient.Process
	Keeper    *toolclient.ProcessKeeper
	ToolsConn *toolsrvclient.Conn
}

// 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() *toolsrvclient.Conn
		if r.Keeper != nil {
			newToolServer = func() *toolsrvclient.Conn {
				conn, err := r.Keeper.Dial()
				if err != nil {
					return nil
				}
				return conn
			}
		} else {
			newToolServer = func() *toolsrvclient.Conn {
				return r.ToolsConn
			}
		}
		// Update CWD on the reused connection.
		if r.ToolsConn != nil {
			r.ToolsConn.SetCWD(cfg.CWD)
		}
		return &SessionInfra{
			Proc:          r.Proc,
			Keeper:        r.Keeper,
			ToolsConn:     newToolServer(),
			NewToolServer: newToolServer,
			Platform:      "linux",
			IsGitRepo:     util.IsGitRepo(cfg.CWD),
		}, nil
	}
	var proc *toolclient.Process
	var keeper *toolclient.ProcessKeeper
	var newToolServer func() *toolsrvclient.Conn
	var err error
	var platform string
	var isGitRepo bool

	if cfg.RemoteTarget != "" {
		proc, err = toolclient.SpawnRemote(cfg.Ctx, toolclient.RemoteConfig{
			SSHTarget: cfg.RemoteTarget,
			CWD:       cfg.CWD,
			SessionID: cfg.SessionID,
			Yolo:      cfg.Yolo,
		})
		if err != nil {
			return nil, fmt.Errorf("remote spawn: %w", err)
		}
		keeper = toolclient.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolclient.Process, error) {
			return toolclient.SpawnRemote(ctx, toolclient.RemoteConfig{
				SSHTarget: cfg.RemoteTarget,
				CWD:       cfg.CWD,
				SessionID: cfg.SessionID,
				Yolo:      cfg.Yolo,
			})
		})
		newToolServer = func() *toolsrvclient.Conn {
			conn, err := keeper.Dial()
			if err != nil {
				return nil
			}
			return conn
		}
		platform = proc.Info.Platform
		isGitRepo = proc.Info.IsGitRepo
	} else {
		var dialOpts []toolclient.Option
		if cfg.Yolo {
			dialOpts = append(dialOpts, toolclient.WithYolo())
		}
		if cfg.SessionID != "" {
			dialOpts = append(dialOpts, toolclient.WithSessionID(cfg.SessionID))
		}
		proc, err = toolclient.Spawn(cfg.Ctx, cfg.CWD, dialOpts...)
		if err != nil {
			return nil, fmt.Errorf("local spawn: %w", err)
		}
		keeper = toolclient.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolclient.Process, error) {
			return toolclient.Spawn(ctx, cfg.CWD, dialOpts...)
		})
		newToolServer = func() *toolsrvclient.Conn {
			conn, err := keeper.Dial()
			if err != nil {
				return nil
			}
			return conn
		}
		platform = "linux"
		isGitRepo = util.IsGitRepo(cfg.CWD)
	}

	toolsConn := newToolServer()

	return &SessionInfra{
		Proc:          proc,
		Keeper:        keeper,
		ToolsConn:     toolsConn,
		NewToolServer: newToolServer,
		Platform:      platform,
		IsGitRepo:     isGitRepo,
	}, 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, platform string, isGitRepo bool, systemPromptOverride string) PromptLayers {
	spOverride := systemPromptOverride
	if spOverride == "" && cfg != nil {
		spOverride = cfg.SystemPrompt
	}
	sysPrompt := prompts.ResolveSystemPrompt(spOverride)
	envBlock := agent.EnvironmentBlock(cwd, platform, isGitRepo, "")

	return PromptLayers{
		SystemPrompt: sysPrompt,
		EnvBlock:     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, conn *toolsrvclient.Conn, sessID, uname string, logError func(string, ...any)) {
	if cfg == nil || conn == nil {
		return
	}

	// Send session ID to the tool server first — ollie-remote uses this
	// to attach its tool registry via the OnEnvSet callback.
	if sessID != "" {
		conn.SetEnv("OLLIE_SESSION_ID", sessID)
	}
	if uname != "" {
		conn.SetAgentID(uname)
	}

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