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
	}

	// Build ToolMeta lookup from the tool listing — single source of truth
	// for all per-tool metadata (tier, readOnly, outputFormat, resetsCounter).
	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
	}

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

	return &Runtime{
		ToolServer:      srv,
		Tools:           allTools,
		ToolMeta:        toolMeta,
		Exec:            exec,
		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
}