HarborForge.PlexumPlugin/internal/tools/kb.go

// kb.go — dynamic-kb-* tool implementations (DESIGN-DYNAMIC-BLOCK.md
// §3.3 / §4.4). Each tool returns a sdkplugin.ToolResult; errors are
// surfaced as IsError:true rather than RPC errors so the model sees a
// human-readable message.
//
// All 5 tools need the agent id (from ctx) + session id (looked up
// via Deps.SessionForAgent — main.go threads this from the per-turn
// RenderDynamicSubblock callback that updates the lookup table).

package tools

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

	sdkplugin "git.hangman-lab.top/hzhang/Plexum-sdk-go/plugin"

	"git.hangman-lab.top/zhi/HarborForge.PlexumPlugin/internal/kbblock"
	"git.hangman-lab.top/zhi/HarborForge.PlexumPlugin/internal/kbclient"
)

// KBDeps bundles the dependencies the KB tools need. Caller (main.go)
// constructs once and reuses across all tool calls.
type KBDeps struct {
	// Client is the fallback plugin-level KB client (auths with
	// Config.APIKey). When AgentClient resolves a per-agent client,
	// that takes priority.
	Client      *kbclient.Client
	ProfileRoot string                                          // <PLEXUM_PROFILE_ROOT>
	SessionFor  func(agentID string) string                     // returns sessionID or ""
	Turn        func(agentID string) int                        // current turn (best-effort, can return 0)
	HostCallTool func(ctx context.Context, agentID, toolName string, input json.RawMessage) (json.RawMessage, error)

	// AgentClient resolves a per-agent KBClient by reading the agent's
	// hf-token via HostAPI.GetAgentSecret (decision #20). Returns
	// (nil, nil) when no token is configured for the agent — caller
	// then falls back to Client. (nil, err) on RPC / decrypt failure.
	AgentClient func(ctx context.Context, agentID string) (*kbclient.Client, error)
}

// clientForCall picks the right KBClient: per-agent if AgentClient
// returns one, else falls back to the shared plugin-level Client.
// Returns (nil, "<reason>") only when neither path produces a usable
// client (no configured token + no plugin-level apiKey).
func (d KBDeps) clientForCall(ctx context.Context, agentID string) (*kbclient.Client, string) {
	if d.AgentClient != nil && agentID != "" {
		c, err := d.AgentClient(ctx, agentID)
		if err == nil && c != nil {
			return c, ""
		}
		if err != nil {
			return nil, "agent hf-token lookup failed: " + err.Error()
		}
		// no token for this agent — fall through to plugin-level
	}
	if d.Client != nil {
		return d.Client, ""
	}
	return nil, "HF KB backend unavailable (no per-agent hf-token + no plugin-level apiKey configured)"
}

// kb tool inputs

type listKBsInput struct {
	ProjectCode string `json:"project-code,omitempty"`
}

type listTopicsInput struct {
	KBCode string `json:"kb-code"`
}

type listFactsInput struct {
	KBCode   string `json:"kb-code"`
	TopicIDs []int  `json:"topic-ids"`
}

type cacheInput struct {
	KBCode            string `json:"kb-code"`
	FactIDs           []int  `json:"fact-ids"`
	PreviousDynamicID string `json:"previous-dynamic-id,omitempty"`
}

type evictInput struct {
	FactIDs []int `json:"fact-ids"`
}

type cacheResult struct {
	Added         []int  `json:"added,omitempty"`
	AlreadyCached []int  `json:"already_cached,omitempty"`
	Missing       []int  `json:"missing,omitempty"`
	Note          string `json:"note"`
}

type evictResult struct {
	Evicted   []int  `json:"evicted,omitempty"`
	NotCached []int  `json:"not_cached,omitempty"`
	Note      string `json:"note"`
}

// ToolKBListKBs implements dynamic-kb-list-kbs.
func ToolKBListKBs(ctx context.Context, deps KBDeps, raw json.RawMessage) (sdkplugin.ToolResult, error) {
	client, reason := deps.clientForCall(ctx, sdkplugin.AgentIDFromContext(ctx))
	if client == nil {
		return errResult("dynamic-kb-list-kbs: " + reason)
	}
	var in listKBsInput
	if len(raw) > 0 {
		_ = json.Unmarshal(raw, &in)
	}
	kbs, err := client.ListKBs(ctx, in.ProjectCode)
	if err != nil {
		return errResult("dynamic-kb-list-kbs: " + err.Error())
	}
	var sb strings.Builder
	if in.ProjectCode != "" {
		fmt.Fprintf(&sb, "project: %s\n", in.ProjectCode)
	}
	fmt.Fprintf(&sb, "kbs: %d\n\n", len(kbs))
	for _, k := range kbs {
		fmt.Fprintf(&sb, "[%s] %s\n  %s\n\n", k.KnowledgeBaseCode, k.Title, k.Description)
	}
	sb.WriteString("Call dynamic-kb-list-topics({kb-code: \"<code>\"}) to drill into a KB.\n")
	return okResult(sb.String())
}

// ToolKBListTopics implements dynamic-kb-list-topics.
func ToolKBListTopics(ctx context.Context, deps KBDeps, raw json.RawMessage) (sdkplugin.ToolResult, error) {
	client, reason := deps.clientForCall(ctx, sdkplugin.AgentIDFromContext(ctx))
	if client == nil {
		return errResult("dynamic-kb-list-topics: " + reason)
	}
	var in listTopicsInput
	if err := json.Unmarshal(raw, &in); err != nil {
		return errResult("dynamic-kb-list-topics: parse: " + err.Error())
	}
	if in.KBCode == "" {
		return errResult("dynamic-kb-list-topics: kb-code required")
	}
	topics, err := client.ListTopics(ctx, in.KBCode)
	if err != nil {
		return errResult("dynamic-kb-list-topics: " + err.Error())
	}
	var sb strings.Builder
	fmt.Fprintf(&sb, "kb: %s\ntopics: %d\n\n", in.KBCode, len(topics))
	for _, t := range topics {
		fmt.Fprintf(&sb, "[%d] %s\n  %s\n\n", t.ID, t.Topic, t.Description)
	}
	fmt.Fprintf(&sb, "Call dynamic-kb-list-facts({kb-code: %q, topic-ids: [<id>, ...]}) to drill into facts.\n",
		in.KBCode)
	return okResult(sb.String())
}

// ToolKBListFacts implements dynamic-kb-list-facts.
func ToolKBListFacts(ctx context.Context, deps KBDeps, raw json.RawMessage) (sdkplugin.ToolResult, error) {
	client, reason := deps.clientForCall(ctx, sdkplugin.AgentIDFromContext(ctx))
	if client == nil {
		return errResult("dynamic-kb-list-facts: " + reason)
	}
	var in listFactsInput
	if err := json.Unmarshal(raw, &in); err != nil {
		return errResult("dynamic-kb-list-facts: parse: " + err.Error())
	}
	if in.KBCode == "" || len(in.TopicIDs) == 0 {
		return errResult("dynamic-kb-list-facts: kb-code + topic-ids required")
	}
	facts, err := client.ListFacts(ctx, in.KBCode, in.TopicIDs)
	if err != nil {
		return errResult("dynamic-kb-list-facts: " + err.Error())
	}
	var sb strings.Builder
	fmt.Fprintf(&sb, "kb: %s\ntopics: %v\nfacts: %d\n\n", in.KBCode, in.TopicIDs, len(facts))
	for _, f := range facts {
		fmt.Fprintf(&sb, "[%d] (topic %d/%s) %s\n  %s\n\n",
			f.ID, f.TopicID, f.TopicSlug, f.Title, f.Snippet)
	}
	sb.WriteString(
		"Call dynamic-kb-cache({kb-code: \"" + in.KBCode +
			"\", fact-ids: [<id>, ...]}) to commit selected facts to your kb-block.\n",
	)
	return okResult(sb.String())
}

// ToolKBCache implements dynamic-kb-cache. Per-agent session lookup
// resolves the kb-block.json location.
func ToolKBCache(ctx context.Context, deps KBDeps, agentID string, raw json.RawMessage) (sdkplugin.ToolResult, error) {
	client, reason := deps.clientForCall(ctx, agentID)
	if client == nil {
		return errResult("dynamic-kb-cache: " + reason)
	}
	sessionID := ""
	if deps.SessionFor != nil {
		sessionID = deps.SessionFor(agentID)
	}
	if agentID == "" || sessionID == "" {
		return errResult("dynamic-kb-cache: no per-session context (agent/session unknown)")
	}
	var in cacheInput
	if err := json.Unmarshal(raw, &in); err != nil {
		return errResult("dynamic-kb-cache: parse: " + err.Error())
	}
	if in.KBCode == "" || len(in.FactIDs) == 0 {
		return errResult("dynamic-kb-cache: kb-code + fact-ids required")
	}
	block, err := kbblock.Open(deps.ProfileRoot, agentID, sessionID)
	if err != nil {
		return errResult("dynamic-kb-cache: open block: " + err.Error())
	}
	// Split into already-cached vs to-fetch.
	var toFetch []int
	var alreadyCached []int
	for _, id := range in.FactIDs {
		if block.Has(id) {
			alreadyCached = append(alreadyCached, id)
		} else {
			toFetch = append(toFetch, id)
		}
	}
	var fetched []kbclient.Fact
	if len(toFetch) > 0 {
		fetched, err = client.GetFacts(ctx, in.KBCode, toFetch)
		if err != nil {
			return errResult("dynamic-kb-cache: " + err.Error())
		}
	}
	turn := 0
	if deps.Turn != nil {
		turn = deps.Turn(agentID)
	}
	var added []int
	fetchedByID := map[int]kbclient.Fact{}
	for _, f := range fetched {
		fetchedByID[f.ID] = f
	}
	for _, id := range toFetch {
		f, ok := fetchedByID[id]
		if !ok {
			continue
		}
		if e := block.Add(f.ID, in.KBCode, f.TopicSlug, f.Content, turn); e != nil {
			added = append(added, f.ID)
		}
	}
	if err := block.Save(); err != nil {
		return errResult("dynamic-kb-cache: save: " + err.Error())
	}
	missing := diffInts(toFetch, added)
	sort.Ints(added)
	sort.Ints(alreadyCached)
	res := cacheResult{
		Added:         added,
		AlreadyCached: alreadyCached,
		Missing:       missing,
		Note:          "Newly cached facts are available starting your next turn.",
	}
	out, _ := json.Marshal(res)
	return okResult(string(out))
}

// ToolKBEvict implements dynamic-kb-evict.
func ToolKBEvict(ctx context.Context, deps KBDeps, agentID string, raw json.RawMessage) (sdkplugin.ToolResult, error) {
	sessionID := ""
	if deps.SessionFor != nil {
		sessionID = deps.SessionFor(agentID)
	}
	if agentID == "" || sessionID == "" {
		return errResult("dynamic-kb-evict: no per-session context")
	}
	var in evictInput
	if err := json.Unmarshal(raw, &in); err != nil {
		return errResult("dynamic-kb-evict: parse: " + err.Error())
	}
	if len(in.FactIDs) == 0 {
		return errResult("dynamic-kb-evict: fact-ids required")
	}
	block, err := kbblock.Open(deps.ProfileRoot, agentID, sessionID)
	if err != nil {
		return errResult("dynamic-kb-evict: open block: " + err.Error())
	}
	evicted := block.Remove(in.FactIDs...)
	if err := block.Save(); err != nil {
		return errResult("dynamic-kb-evict: save: " + err.Error())
	}
	notCached := diffInts(in.FactIDs, evicted)
	sort.Ints(evicted)
	res := evictResult{
		Evicted:   evicted,
		NotCached: notCached,
		Note:      "Evictions take effect starting your next turn.",
	}
	out, _ := json.Marshal(res)
	return okResult(string(out))
}

// ---- helpers ----

func diffInts(a, b []int) []int {
	if len(a) == 0 {
		return nil
	}
	want := make(map[int]bool, len(b))
	for _, x := range b {
		want[x] = true
	}
	out := make([]int, 0, len(a))
	for _, x := range a {
		if !want[x] {
			out = append(out, x)
		}
	}
	if len(out) == 0 {
		return nil
	}
	sort.Ints(out)
	return out
}