ollie/cmd/olliesrv/internal/agent/runtime.go

package agent

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

	"ollie/cmd/olliesrv/internal/backend"
	"ollie/cmd/olliesrv/internal/prompts"
	toolclient "ollie/cmd/olliesrv/internal/toolclient"
	"ollie/toolsrv/protocol"
)

// --- 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 stable environment context for the preamble.
// Volatile values such as the current date do not belong in the provider-cached
// prefix and are supplied separately when needed.
func EnvironmentBlock(cwd, platform string, isGitRepo bool, sandbox string) string {
	git := "false"
	if isGitRepo {
		git = "true"
	}
	return fmt.Sprintf("# Environment\n\n- Working directory: %s\n- Platform: %s\n- Is git repo: %s\n- Sandbox: %s",
		cwd, platform, 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 []protocol.ToolInfo) string {
	infos = append([]protocol.ToolInfo(nil), infos...)
	sort.SliceStable(infos, func(i, j int) bool {
		if infos[i].Name == infos[j].Name {
			return infos[i].Server < infos[j].Server
		}
		return infos[i].Name < infos[j].Name
	})
	var sb strings.Builder
	for _, ti := range infos {
		if ti.Server != "" {
			continue
		}
		if ti.Description == "" {
			continue
		}
		if ti.Prompt != "" {
			// If prompt already starts with a markdown header, skip the generated one.
			if !strings.HasPrefix(ti.Prompt, "#") {
				sb.WriteString("## ")
				sb.WriteString(ti.Name)
				sb.WriteString("\n\n")
			}
			sb.WriteString(ti.Prompt)
		} else {
			sb.WriteString("## ")
			sb.WriteString(ti.Name)
			sb.WriteString("\n\n")
			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      *toolclient.ToolsrvConn
	Preamble        *Preamble
	Tools           []backend.Tool
	ToolMeta        map[string]protocol.ToolInfo
	ToolRevision    uint64
	Exec            toolExecutor
	GenParams       backend.GenerationParams
	CompactionModel string
	UserPrompt      string // resolved user prompts, prepended to each turn's input
}

// 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 *toolclient.ToolsrvConn, cwd string, env []string, systemPrompt, envBlock string) *Runtime {
	var allToolInfos []protocol.ToolInfo
	var allTools []backend.Tool

	if srv != nil {
		allToolInfos, _ = srv.ListTools()
		allTools = toolInfosToBackend(allToolInfos)
	}
	var toolRevision uint64
	if srv != nil {
		toolRevision = srv.ToolRegistryRevision()
	}

	var agentPrompt string
	var userPrompt string
	var genParams backend.GenerationParams
	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
		}
		if resolved, err := resolvePrompt(cfg.UserPrompts, cwd, env); err != nil {
			fmt.Fprintf(os.Stderr, "resolve user prompts: %v\n", err)
		} else {
			userPrompt = resolved
		}
		genParams = cfg.GenerationParams
	}
	genParams.CWD = cwd

	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]protocol.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,
		ToolRevision:    toolRevision,
		Exec:            exec,
		Preamble:        preamble,
		GenParams:       genParams,
		CompactionModel: compactionModel,
		UserPrompt:      userPrompt,
	}
}

// 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 []protocol.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:  injectDispatchFlags(t.InputSchema),
		}
	}
	return out
}

// dispatchFlags are optional parameters available on all tools.
// They are stripped before execution by the dispatch layer.
var dispatchFlags = map[string]json.RawMessage{
	"bypass":     json.RawMessage(`{"type":"boolean","description":"Run outside the sandbox via bypass broker."}`),
	"timeout":    json.RawMessage(`{"type":"integer","description":"Timeout in seconds (default: 30). Use 0 for no timeout."}`),
	"sandbox":    json.RawMessage(`{"type":"string","description":"Sandbox profile name (default: default)."}`),
	"background": json.RawMessage(`{"type":"boolean","description":"Run in background. Returns PID immediately; output is injected automatically."}`),
}

// injectDispatchFlags adds bypass/timeout/sandbox/background to a tool's schema
// if it has a properties object. Does not overwrite existing properties.
func injectDispatchFlags(schema json.RawMessage) json.RawMessage {
	if len(schema) == 0 {
		return schema
	}
	var obj map[string]json.RawMessage
	if json.Unmarshal(schema, &obj) != nil {
		return schema
	}
	propsRaw, ok := obj["properties"]
	if !ok {
		return schema
	}
	var props map[string]json.RawMessage
	if json.Unmarshal(propsRaw, &props) != nil {
		return schema
	}
	modified := false
	for k, v := range dispatchFlags {
		if _, exists := props[k]; !exists {
			props[k] = v
			modified = true
		}
	}
	if !modified {
		return schema
	}
	propsBytes, _ := json.Marshal(props)
	obj["properties"] = propsBytes
	out, _ := json.Marshal(obj)
	return out
}

// extractToolResult parses the tool server's JSON-RPC response into text
// and optional content blocks.
//
// See protocol.ToolResult for the wire format.
func extractToolResult(raw json.RawMessage) (text string, contentBlocks []backend.ContentBlock, isError bool) {
	var result protocol.ToolResult
	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
}