ollie/agent/runtime.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"time"

	"ollie/backend"
	"ollie/prompts"
	"ollie/toolsrv"
)

// --- Preamble sections ---

// Preamble is an ordered list of named text sections. The final system prompt
// is produced by joining all non-empty sections with newlines. Sections are
// identified by name so individual parts can be updated without reparsing.
type Preamble struct {
	sections []section
}

type section struct {
	name Section
	text string
}

// Set replaces (or appends) the named section.
func (p *Preamble) Set(name Section, text string) {
	for i := range p.sections {
		if p.sections[i].name == name {
			p.sections[i].text = text
			return
		}
	}
	p.sections = append(p.sections, section{name, text})
}

// Get returns the text of a named section.
func (p *Preamble) Get(name Section) string {
	for _, s := range p.sections {
		if s.name == name {
			return s.text
		}
	}
	return ""
}

// String joins all non-empty sections with a newline separator.
func (p *Preamble) String() string {
	var sb strings.Builder
	for _, s := range p.sections {
		if s.text == "" {
			continue
		}
		if sb.Len() > 0 {
			sb.WriteByte('\n')
		}
		sb.WriteString(s.text)
	}
	return sb.String()
}

// Section is a named preamble section identifier.
type Section string

const (
	SectionSystem Section = "system"
	SectionEnv    Section = "env"
	SectionAgent  Section = "agent"
	SectionTools  Section = "tools"
)

// --- Environment block ---

// EnvironmentBlock renders the environment context section for the preamble.
func EnvironmentBlock(cwd, platform string, isGitRepo bool, sandbox string) string {
	date := time.Now().Format("2006-01-02")
	git := "false"
	if isGitRepo {
		git = "true"
	}
	return fmt.Sprintf("# Environment\n\n- Working directory: %s\n- Platform: %s\n- Current date: %s\n- Is git repo: %s\n- Sandbox: %s",
		cwd, platform, date, git, sandbox)
}

// --- Tool rendering ---

// renderTools produces a single combined section: a one-line summary per tool
// followed immediately by its documentation (if any). This keeps the listing
// and usage instructions adjacent so models don't lose context between them.
func renderTools(infos []toolsrv.ToolInfo) string {
	var sb strings.Builder
	for _, ti := range infos {
		if ti.Server != "" {
			continue
		}
		if ti.Description == "" {
			continue
		}
		sb.WriteString("## ")
		sb.WriteString(ti.Name)
		sb.WriteString("\n\n")
		if ti.Prompt != "" {
			sb.WriteString(ti.Prompt)
		} else {
			sb.WriteString(ti.Description)
		}
		sb.WriteString("\n\n")
	}
	if sb.Len() == 0 {
		return ""
	}
	return "# Tools\n\n" + sb.String()
}

// --- Runtime ---

// Runtime holds the swappable per-agent configuration. It contains everything
// that changes on an /agent switch but is stable across turns within the same
// agent. The agent struct stores a pointer to the active Runtime; switching
// agents replaces it atomically.
type Runtime struct {
	Backend    backend.Backend
	ToolServer *toolsrv.Conn
	Preamble   *Preamble
	Tools      []backend.Tool
	ToolMeta   map[string]toolsrv.ToolInfo
	Exec       toolExecutor
	GenParams  backend.GenerationParams
	MaxSteps   int
	// toolRegistryRevision tracks the last-seen tool registry revision.
	toolRegistryRevision uint64
	CompactionModel      string
	Messages             []string
}

// PreambleString returns the fully rendered system prompt.
func (rt *Runtime) PreambleString() string {
	if rt.Preamble == nil {
		return ""
	}
	return rt.Preamble.String()
}

// BuildRuntime constructs a Runtime from a pre-configured tool server and
// optional agent config.
func BuildRuntime(cfg *AgentConfig, srv *toolsrv.Conn, cwd string, env []string, systemPrompt, envBlock string) *Runtime {
	var messages []string

	var allToolInfos []toolsrv.ToolInfo
	var allTools []backend.Tool

	if srv != nil && (cfg == nil || cfg.ToolsEnabled()) {
		var listErr error
		allToolInfos, listErr = srv.ListTools()
		if listErr != nil {
			messages = append(messages, fmt.Sprintf("list tools: %v", listErr))
		}
		allTools = toolInfosToBackend(allToolInfos)
	}

	var agentPrompt string
	var genParams backend.GenerationParams
	var maxSteps int
	if cfg != nil {
		if resolved, err := resolvePrompt(cfg.Prompt, cwd, env); err != nil {
			fmt.Fprintf(os.Stderr, "resolve prompt: %v\n", err)
		} else {
			agentPrompt = resolved
		}
		genParams = cfg.GenerationParams
		maxSteps = cfg.MaxSteps
	}

	exec := func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) {
		if srv == nil {
			return "", nil, fmt.Errorf("no tool server available")
		}
		raw, err := srv.CallTool(ctx, name, args)
		if err != nil {
			return "", nil, err
		}
		text, blocks, isErr := extractToolResult(raw)
		if isErr {
			return "", nil, fmt.Errorf("%s", text)
		}
		return text, blocks, nil
	}

	toolMeta := make(map[string]toolsrv.ToolInfo, len(allToolInfos))
	for _, ti := range allToolInfos {
		toolMeta[ti.Name] = ti
	}

	var compactionModel string
	if cfg != nil {
		compactionModel = cfg.CompactionModel
	}

	// Build preamble from sections
	preamble := &Preamble{}
	preamble.Set(SectionSystem, systemPrompt)
	preamble.Set(SectionEnv, envBlock)
	preamble.Set(SectionAgent, agentPrompt)
	preamble.Set(SectionTools, renderTools(allToolInfos))

	return &Runtime{
		ToolServer:      srv,
		Tools:           allTools,
		ToolMeta:        toolMeta,
		Exec:            exec,
		Preamble:        preamble,
		GenParams:       genParams,
		MaxSteps:        maxSteps,
		CompactionModel: compactionModel,
		Messages:        messages,
	}
}

// ResolveSystemPrompt resolves the system prompt — either from the agent config
// override path or the embedded default.
func ResolveSystemPrompt(cfg *AgentConfig) string {
	if cfg != nil && cfg.SystemPrompt != "" {
		return prompts.ResolveSystemPrompt(cfg.SystemPrompt)
	}
	return prompts.DefaultSystemPrompt
}

func toolInfosToBackend(infos []toolsrv.ToolInfo) []backend.Tool {
	out := make([]backend.Tool, len(infos))
	for i, t := range infos {
		out[i] = backend.Tool{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.InputSchema,
		}
	}
	return out
}

// extractToolResult parses the tool server's JSON-RPC response into text
// and optional content blocks.
//
// Protocol: tool scripts write JSON to stdout matching this schema:
//
//	{
//	  "isError": false,
//	  "content": [
//	    {"type": "text", "text": "..."},
//	    {"type": "image", "media_type": "image/png", "data": "<base64>"}
//	  ]
//	}
//
// If the response is not valid JSON (e.g. a plain string from a simple script),
// it is returned as-is with isError=false. This fallback allows tools to emit
// raw text without wrapping it in the content-block envelope.
func extractToolResult(raw json.RawMessage) (text string, contentBlocks []backend.ContentBlock, isError bool) {
	var result struct {
		IsError bool `json:"isError"`
		Content []struct {
			Type      string `json:"type"`
			Text      string `json:"text"`
			MediaType string `json:"media_type"`
			Data      string `json:"data"`
		} `json:"content"`
	}
	if err := json.Unmarshal(raw, &result); err != nil {
		return string(raw), nil, false
	}
	var parts []string
	for _, c := range result.Content {
		switch c.Type {
		case "text":
			parts = append(parts, c.Text)
		case "image":
			contentBlocks = append(contentBlocks, backend.ContentBlock{
				Type: "image",
				ImageSource: &backend.ImageSource{
					Type:      "base64",
					MediaType: c.MediaType,
					Data:      c.Data,
				},
			})
		}
	}
	return strings.Join(parts, "\n"), contentBlocks, result.IsError
}