ollie/agent/build_runtime.go

package agent

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

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

// BuildRuntime constructs a Runtime from a pre-configured Dispatcher and
// optional agent config. cwd sets the working directory reported in the
// system prompt; if empty, the process working directory is used.
// env provides additional environment variables injected into prompt resolution
// subprocesses (e.g. OLLIE_SESSION_ID=xxx).
// The caller is responsible for registering all servers on d before calling this.
func BuildRuntime(cfg *AgentConfig, srv toolsrv.Runner, 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))
		}
		// Only loaded tools become backend callable tools.
		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.GenParams()
		maxSteps = cfg.MaxSteps
		if len(cfg.AllowTools) > 0 {
			if rs, ok := srv.(interface{ SetAllowTools([]string) }); ok {
				rs.SetAllowTools(cfg.AllowTools)
			}
		}
	}

	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
	}

	var classify toolClassifier
	if pc, ok := srv.(toolsrv.ParallelClassifier); ok {
		classify = pc.IsParallelRead
	}

	// Capture output formats from the same tool listing used to build the
	// backend tool definitions. Remote runners expose ToolInfo through
	// ListTools, but do not necessarily implement OutputFormatClassifier.
	outputFormats := make(map[string]string, len(allToolInfos))
	for _, ti := range allToolInfos {
		outputFormats[ti.Name] = ti.OutputFormat
	}
	outputFormat := func(name string) string {
		return outputFormats[name]
	}
	if oc, ok := srv.(toolsrv.OutputFormatClassifier); ok {
		outputFormat = oc.OutputFormat
	}

	var tierFn func(string, json.RawMessage) ResultTier
	if tc, ok := srv.(toolsrv.TierClassifier); ok {
		tierFn = func(name string, args json.RawMessage) ResultTier {
			switch tc.ResultTierArgs(name, args) {
			case "cold":
				return TierCold
			case "warm":
				return TierWarm
			default:
				return TierHot
			}
		}
	}

	// Build resetsCounter lookup from tool infos.
	resetsCounterMap := make(map[string]bool, len(allToolInfos))
	for _, ti := range allToolInfos {
		if ti.ResetsCounter {
			resetsCounterMap[ti.Name] = true
		}
	}
	var resetsCounterFn func(string) bool
	if len(resetsCounterMap) > 0 {
		resetsCounterFn = func(name string) bool {
			return resetsCounterMap[name]
		}
	}

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

	preamble := BuildPreamble(systemPrompt, envBlock, agentPrompt, allToolInfos)

	return &Runtime{
		ToolServer:      srv,
		Tools:           allTools,
		Exec:            exec,
		ClassifyTool:    classify,
		OutputFormat:    outputFormat,
		ClassifyTier:    tierFn,
		ResetsCounter:   resetsCounterFn,
		Preamble:        preamble,
		GenParams:       genParams,
		MaxSteps:        maxSteps,
		CompactionModel: compactionModel,
		Messages:        messages,
	}
}

// BuildPreamble constructs the complete system prompt from its components.
// The preamble evolves when:
// - Working directory changes (envBlock updates)
// - Tools are loaded/unloaded (allToolInfos changes)
// - Agent profile switches (agentPrompt changes)
// This function should be called whenever any component changes.
func BuildPreamble(systemPrompt, envBlock, agentPrompt string, allToolInfos []toolsrv.ToolInfo) string {
	var sb strings.Builder

	// 1. Core system prompt (identity, rules, operating principles)
	sb.WriteString(systemPrompt)
	sb.WriteByte('\n')

	// 2. Environment context (working directory, platform, date, etc.)
	sb.WriteString(envBlock)
	sb.WriteByte('\n')

	// 3. Agent-specific prompt (role, task instructions from config)
	sb.WriteString(agentPrompt)
	sb.WriteByte('\n')

	// 4. Compact tool surface listing (name + description)
	var toolListing strings.Builder
	for _, ti := range allToolInfos {
		if ti.Description != "" && ti.Server == "" {
			fmt.Fprintf(&toolListing, "- **%s** — %s\n", ti.Name, ti.Description)
		}
	}

	// 5. Full tool documentation (prompts, schemas, examples)
	var toolDocs strings.Builder
	for _, ti := range allToolInfos {
		if ti.Prompt != "" && ti.Server == "" {
			toolDocs.WriteString(ti.Prompt)
			toolDocs.WriteString("\n\n")
			if len(ti.InputSchema) > 0 && string(ti.InputSchema) != "{}" {
				fmt.Fprintf(&toolDocs, "**Schema**:\n```json\n%s\n```\n\n", ti.InputSchema)
			}
		}
	}

	// 6. Render tool sections
	if toolListing.Len() > 0 || toolDocs.Len() > 0 {
		toolSections := prompts.ToolSections(toolListing.String(), toolDocs.String())
		if toolSections != "" {
			sb.WriteString("\n")
			sb.WriteString(toolSections)
		}
	}

	// 6. Tool discovery instructions (9P filesystem)
	sb.WriteString("\n")
	sb.WriteString(prompts.ToolDiscoveryPrompt)

	return sb.String()
}

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
}

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
}