amcs/internal/tools/world_model.go

package tools

import (
	"context"
	"strings"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"

	"git.warky.dev/wdevs/amcs/internal/session"
	"git.warky.dev/wdevs/amcs/internal/store"
	ext "git.warky.dev/wdevs/amcs/internal/types"
)

type WorldModelTool struct {
	store    *store.DB
	sessions *session.ActiveProjects
}

func NewWorldModelTool(db *store.DB, sessions *session.ActiveProjects) *WorldModelTool {
	return &WorldModelTool{store: db, sessions: sessions}
}

type BootstrapWorldModelInput struct {
	Project      string `json:"project" jsonschema:"project name or id; must be explicit for deterministic startup"`
	ContextLimit int    `json:"context_limit,omitempty" jsonschema:"recent project thoughts to include (default 10, maximum 50)"`
}

type EffectiveSkill struct {
	ID          int64    `json:"id"`
	Name        string   `json:"name"`
	Description string   `json:"description,omitempty"`
	Content     string   `json:"content,omitempty"`
	Tags        []string `json:"tags"`
	Source      string   `json:"source"`
	Override    bool     `json:"override"`
}

type WorldModelContextItem struct {
	ID        int64               `json:"id"`
	Content   string              `json:"content"`
	Metadata  ext.ThoughtMetadata `json:"metadata"`
	CreatedAt time.Time           `json:"created_at"`
}

type BootstrapWorldModelOutput struct {
	Version          int                     `json:"version"`
	GeneratedAt      time.Time               `json:"generated_at"`
	Project          ext.Project             `json:"project"`
	Skills           []EffectiveSkill        `json:"skills"`
	Guardrails       []ext.AgentGuardrail    `json:"guardrails"`
	Personas         []ext.ProjectPersona    `json:"personas"`
	PersonaManifests []ext.PersonaManifest   `json:"persona_manifests"`
	DefaultPersona   *ext.PersonaFull        `json:"default_persona,omitempty"`
	Context          []WorldModelContextItem `json:"context"`
}

func (t *WorldModelTool) Bootstrap(ctx context.Context, req *mcp.CallToolRequest, in BootstrapWorldModelInput) (*mcp.CallToolResult, BootstrapWorldModelOutput, error) {
	project, err := resolveProject(ctx, t.store, t.sessions, req, in.Project, true)
	if err != nil {
		return nil, BootstrapWorldModelOutput{}, err
	}
	if t.sessions != nil && req != nil && req.Session != nil {
		t.sessions.Set(req.Session.ID(), project.ID)
	}

	projectSkills, err := t.store.ListProjectSkills(ctx, project.NumericID)
	if err != nil {
		return nil, BootstrapWorldModelOutput{}, err
	}
	guardrails, err := t.store.ListProjectGuardrails(ctx, project.NumericID)
	if err != nil {
		return nil, BootstrapWorldModelOutput{}, err
	}
	personas, err := t.store.ListProjectPersonas(ctx, project.NumericID)
	if err != nil {
		return nil, BootstrapWorldModelOutput{}, err
	}

	manifests := make([]ext.PersonaManifest, 0, len(personas))
	for _, linked := range personas {
		manifest, err := t.store.GetPersonaManifest(ctx, linked.Persona.Name)
		if err != nil {
			return nil, BootstrapWorldModelOutput{}, err
		}
		manifests = append(manifests, manifest)
	}

	var defaultPersona *ext.PersonaFull
	for _, linked := range personas {
		if !linked.IsDefault {
			continue
		}
		loaded, err := t.store.GetPersona(ctx, linked.Persona.Name, true, nil)
		if err != nil {
			return nil, BootstrapWorldModelOutput{}, err
		}
		defaultPersona = &loaded
		break
	}

	skills := mergeWorldModelSkills(projectSkills, defaultPersona)
	guardrails = mergeWorldModelGuardrails(guardrails, defaultPersona)
	limit := in.ContextLimit
	if limit <= 0 {
		limit = 10
	}
	if limit > 50 {
		limit = 50
	}
	recent, err := t.store.RecentThoughts(ctx, &project.NumericID, limit, 0)
	if err != nil {
		return nil, BootstrapWorldModelOutput{}, err
	}
	contextItems := make([]WorldModelContextItem, 0, len(recent))
	for _, thought := range recent {
		contextItems = append(contextItems, WorldModelContextItem{ID: thought.ID, Content: thought.Content, Metadata: thought.Metadata, CreatedAt: thought.CreatedAt})
	}
	if guardrails == nil {
		guardrails = []ext.AgentGuardrail{}
	}
	if personas == nil {
		personas = []ext.ProjectPersona{}
	}
	_ = t.store.TouchProject(ctx, project.NumericID)

	return nil, BootstrapWorldModelOutput{
		Version: 1, GeneratedAt: time.Now().UTC(), Project: *project, Skills: skills,
		Guardrails: guardrails, Personas: personas, PersonaManifests: manifests,
		DefaultPersona: defaultPersona, Context: contextItems,
	}, nil
}

func mergeWorldModelSkills(projectSkills []ext.AgentSkill, persona *ext.PersonaFull) []EffectiveSkill {
	result := make([]EffectiveSkill, 0, len(projectSkills))
	index := make(map[string]int, len(projectSkills))
	for _, skill := range projectSkills {
		key := strings.ToLower(skill.Name)
		index[key] = len(result)
		result = append(result, EffectiveSkill{ID: skill.ID, Name: skill.Name, Description: skill.Description, Content: skill.Content, Tags: skill.Tags, Source: "project", Override: skill.Override})
	}
	if persona == nil {
		return result
	}
	for _, skill := range persona.Skills {
		key := strings.ToLower(skill.Name)
		entry := EffectiveSkill{ID: skill.ID, Name: skill.Name, Description: skill.Description, Content: skill.Content, Tags: skill.Tags, Source: "persona", Override: skill.Override}
		if i, ok := index[key]; ok {
			if skill.Override {
				result[i] = entry
			}
			continue
		}
		index[key] = len(result)
		result = append(result, entry)
	}
	return result
}

func mergeWorldModelGuardrails(project []ext.AgentGuardrail, persona *ext.PersonaFull) []ext.AgentGuardrail {
	result := append([]ext.AgentGuardrail(nil), project...)
	index := make(map[string]int, len(result))
	for i, guardrail := range result {
		index[strings.ToLower(guardrail.Name)] = i
	}
	if persona == nil {
		return result
	}
	severity := map[string]int{"low": 1, "medium": 2, "high": 3, "critical": 4}
	for _, guardrail := range persona.Guardrails {
		key := strings.ToLower(guardrail.Name)
		if i, ok := index[key]; ok {
			if severity[guardrail.Severity] > severity[result[i].Severity] {
				result[i].Severity = guardrail.Severity
				result[i].Content = guardrail.Content
			}
			continue
		}
		index[key] = len(result)
		result = append(result, ext.AgentGuardrail{ID: guardrail.ID, Name: guardrail.Name, Description: guardrail.Description, Content: guardrail.Content, Severity: guardrail.Severity, Tags: guardrail.Tags})
	}
	return result
}