opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(tools): add focus_history tool, expand focus and search

Browse source 
- pkg/tools/focus_history.go: FocusHistoryTool for querying completed
  focus entries (goal, topic, steps, outcome, summary)
- focus.go, focus_test.go: focus tool enhancements
- search.go, search_test.go: search tool updates
This commit is contained in:
ZanzyTHEbar 2026-02-22 22:49:10 +00:00
parent 2dd54a8269
commit 189fc943e4
6 changed files with 789 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

246
pkg/tools/focus.go
View file

@ -30,9 +30,15 @@ var focusAgentID = pkg.NAME
// FocusState tracks an active focus investigation.
type FocusState struct {
Topic string `json:"topic"`
CheckpointIndex int `json:"checkpoint_index"`
StartedAt time.Time `json:"started_at"`
Topic string `json:"topic"`
Goal string `json:"goal"`
Steps []string `json:"steps"`
Deadline string `json:"deadline"`
CheckpointIndex int `json:"checkpoint_index"`
StartedAt time.Time `json:"started_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"`
Outcome string `json:"outcome,omitempty"`
Status string `json:"status,omitempty"`
}
// KnowledgeBlock stores accumulated summaries from completed focus sessions.
@ -42,9 +48,52 @@ type KnowledgeBlock struct {
// KnowledgeEntry is a single completed focus summary.
type KnowledgeEntry struct {
Topic string `json:"topic"`
Summary string `json:"summary"`
CreatedAt time.Time `json:"created_at"`
Topic string `json:"topic"`
Goal string `json:"goal"`
Steps []string `json:"steps,omitempty"`
Deadline string `json:"deadline,omitempty"`
Outcome string `json:"outcome,omitempty"`
Summary string `json:"summary"`
CreatedAt time.Time `json:"created_at"`
CompletedAt time.Time `json:"completed_at,omitempty"`
}
// FocusText returns the canonical topic/goal label for this focus state.
func (fs FocusState) FocusText() string {
if strings.TrimSpace(fs.Goal) != "" {
return fs.Goal
}
return fs.Topic
}
// FormatBlock renders active focus metadata for system-prompt injection.
func (fs *FocusState) FormatBlock() string {
if fs == nil || (strings.TrimSpace(fs.Topic) == "" && strings.TrimSpace(fs.Goal) == "") {
return ""
}
lines := make([]string, 0, 6)
lines = append(lines, "# Focus")
lines = append(lines, "")
lines = append(lines, "## "+fs.FocusText())
if len(fs.Steps) > 0 {
lines = append(lines, "### Planned Steps")
for _, step := range fs.Steps {
step = strings.TrimSpace(step)
if step == "" {
continue
}
lines = append(lines, "- "+step)
}
}
if strings.TrimSpace(fs.Deadline) != "" {
lines = append(lines, "Deadline: "+fs.Deadline)
}
if strings.TrimSpace(fs.Status) != "" {
lines = append(lines, "Status: "+fs.Status)
}
return strings.Join(lines, "\n") + "\n"
}
// FormatBlock renders the knowledge block for system prompt injection.
@ -55,7 +104,33 @@ func (kb *KnowledgeBlock) FormatBlock() string {
var sb strings.Builder
sb.WriteString("# Knowledge\n\n")
for _, e := range kb.Entries {
sb.WriteString(fmt.Sprintf("## %s\n%s\n\n", e.Topic, e.Summary))
title := e.Topic
if strings.TrimSpace(title) == "" {
title = e.Goal
}
sb.WriteString(fmt.Sprintf("## %s\n", title))
if strings.TrimSpace(e.Outcome) != "" {
sb.WriteString(e.Outcome + "\n")
} else if strings.TrimSpace(e.Summary) != "" {
sb.WriteString(e.Summary + "\n")
}
if len(e.Steps) > 0 {
sb.WriteString("\n### Steps\n")
for _, step := range e.Steps {
step = strings.TrimSpace(step)
if step == "" {
continue
}
sb.WriteString("- " + step + "\n")
}
}
if strings.TrimSpace(e.Deadline) != "" {
sb.WriteString("\nDeadline: " + e.Deadline + "\n")
}
if !e.CompletedAt.IsZero() {
sb.WriteString(fmt.Sprintf("Completed at: %s\n", e.CompletedAt.Format(time.RFC3339)))
}
sb.WriteString("\n")
}
return sb.String()
}
@ -65,6 +140,7 @@ type StartFocusTool struct {
delegate KVStore
sessions *session.SessionManager
sessionKey func() string
OnChange func() // called after focus state changes; used for cache invalidation
}
func NewStartFocusTool(delegate KVStore, sessions *session.SessionManager, sessionKeyFn func() string) *StartFocusTool {
@ -89,6 +165,21 @@ func (t *StartFocusTool) Parameters() map[string]interface{} {
"type": "string",
"description": "What you are investigating or working on (e.g., 'debug authentication flow', 'implement caching layer')",
},
"goal": map[string]interface{}{
"type": "string",
"description": "Optional target outcome for this focus. If set, supersedes topic as the canonical label.",
},
"steps": map[string]interface{}{
"type": "array",
"description": "Optional ordered list of expected investigation steps.",
"items": map[string]interface{}{
"type": "string",
},
},
"deadline": map[string]interface{}{
"type": "string",
"description": "Optional ISO-8601 deadline or scheduling hint (e.g., '2026-02-28T17:00:00Z').",
},
},
"required": []string{"topic"},
}
@ -99,6 +190,18 @@ func (t *StartFocusTool) Execute(ctx context.Context, args map[string]interface{
if topic == "" {
return ErrorResult("topic is required")
}
goal, err := parseOptionalStringArg(args, "goal")
if err != nil {
return ErrorResult(fmt.Sprintf("goal %v", err))
}
steps, err := parseStringSliceArg(args, "steps")
if err != nil {
return ErrorResult(fmt.Sprintf("steps %v", err))
}
deadline, err := parseOptionalStringArg(args, "deadline")
if err != nil {
return ErrorResult(fmt.Sprintf("deadline %v", err))
}
sk := t.sessionKey()
if sk == "" {
@ -108,6 +211,9 @@ func (t *StartFocusTool) Execute(ctx context.Context, args map[string]interface{
history := t.sessions.GetHistory(sk)
state := FocusState{
Topic: topic,
Goal: goal,
Steps: steps,
Deadline: deadline,
CheckpointIndex: len(history),
StartedAt: time.Now(),
}
@ -125,11 +231,23 @@ func (t *StartFocusTool) Execute(ctx context.Context, args map[string]interface{
logger.InfoCF("focus", "Focus started",
map[string]interface{}{
"topic": topic,
"goal": goal,
"steps": len(steps),
"deadline": deadline,
"checkpoint": state.CheckpointIndex,
"session": sk,
})
return SilentResult(fmt.Sprintf("Focus started on: %s\nCheckpoint at message %d. Explore freely, then call complete_focus when done.", topic, state.CheckpointIndex))
goalLine := ""
if strings.TrimSpace(goal) != "" {
goalLine = fmt.Sprintf("Focus label: %s\n", state.FocusText())
}
if t.OnChange != nil {
t.OnChange()
}
return SilentResult(fmt.Sprintf("Focus started on: %s\n%sCheckpoint at message %d. Explore freely, then call complete_focus when done.", topic, goalLine, state.CheckpointIndex))
}
// CompleteFocusTool summarizes findings, persists to knowledge, and prunes context.
@ -137,6 +255,7 @@ type CompleteFocusTool struct {
delegate KVStore
sessions *session.SessionManager
sessionKey func() string
OnChange func() // called after focus state + knowledge changes; used for cache invalidation
}
func NewCompleteFocusTool(delegate KVStore, sessions *session.SessionManager, sessionKeyFn func() string) *CompleteFocusTool {
@ -161,6 +280,21 @@ func (t *CompleteFocusTool) Parameters() map[string]interface{} {
"type": "string",
"description": "Concise summary of the investigation: what was attempted, what was learned, and the outcome",
},
"outcome": map[string]interface{}{
"type": "string",
"description": "Optional explicit outcome statement for this focus.",
},
"steps": map[string]interface{}{
"type": "array",
"description": "Optional ordered list of steps taken during the focus session.",
"items": map[string]interface{}{
"type": "string",
},
},
"status": map[string]interface{}{
"type": "string",
"description": "Optional completion status (default: completed).",
},
},
"required": []string{"summary"},
}
@ -188,8 +322,33 @@ func (t *CompleteFocusTool) Execute(ctx context.Context, args map[string]interfa
return ErrorResult(fmt.Sprintf("corrupt focus state: %v", err))
}
outcome, err := parseOptionalStringArg(args, "outcome")
if err != nil {
return ErrorResult(fmt.Sprintf("outcome %v", err))
}
steps, err := parseStringSliceArg(args, "steps")
if err != nil {
return ErrorResult(fmt.Sprintf("steps %v", err))
}
status, err := parseOptionalStringArg(args, "status")
if err != nil {
return ErrorResult(fmt.Sprintf("status %v", err))
}
if status == "" {
status = "completed"
}
if len(steps) > 0 {
state.Steps = steps
}
state.Status = status
completedAt := time.Now()
state.CompletedAt = &completedAt
if outcome != "" {
state.Outcome = outcome
}
// Append to knowledge block
if err := t.appendKnowledge(ctx, sk, state.Topic, summary); err != nil {
if err := t.appendKnowledge(ctx, sk, state, summary); err != nil {
return ErrorResult(fmt.Sprintf("save knowledge: %v", err))
}
@ -214,6 +373,10 @@ func (t *CompleteFocusTool) Execute(ctx context.Context, args map[string]interfa
"session": sk,
})
if t.OnChange != nil {
t.OnChange()
}
return SilentResult(fmt.Sprintf(
"Focus completed: '%s'\nSummary saved to Knowledge block.\nPruned %d messages (%d → %d). Context is now lean.",
state.Topic, msgsBefore-msgsAfter, msgsBefore, msgsAfter,
@ -249,7 +412,7 @@ func (t *CompleteFocusTool) pruneHistory(history []messages.Message, checkpointI
return result
}
func (t *CompleteFocusTool) appendKnowledge(ctx context.Context, sessionKey, topic, summary string) error {
func (t *CompleteFocusTool) appendKnowledge(ctx context.Context, sessionKey string, state FocusState, summary string) error {
kvKey := knowledgeKVPrefix + sessionKey
kb := &KnowledgeBlock{}
@ -261,9 +424,14 @@ func (t *CompleteFocusTool) appendKnowledge(ctx context.Context, sessionKey, top
}
kb.Entries = append(kb.Entries, KnowledgeEntry{
Topic: topic,
Summary: summary,
CreatedAt: time.Now(),
Topic: state.Topic,
Goal: state.Goal,
Steps: state.Steps,
Deadline: state.Deadline,
Outcome: state.Outcome,
Summary: summary,
CreatedAt: time.Now(),
CompletedAt: *state.CompletedAt,
})
data, err := jsonv2.Marshal(kb)
@ -274,6 +442,39 @@ func (t *CompleteFocusTool) appendKnowledge(ctx context.Context, sessionKey, top
return t.delegate.UpsertKV(ctx, focusAgentID, kvKey, string(data))
}
func parseStringSliceArg(args map[string]interface{}, key string) ([]string, error) {
raw, ok := args[key]
if !ok {
return nil, nil
}
var values []string
switch v := raw.(type) {
case []string:
values = append(values, v...)
case []interface{}:
for i, item := range v {
s, ok := item.(string)
if !ok {
return nil, fmt.Errorf("%s[%d] must be a string", key, i)
}
values = append(values, strings.TrimSpace(s))
}
default:
return nil, fmt.Errorf("%s must be an array of strings", key)
}
normalized := values[:0]
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" {
continue
}
normalized = append(normalized, value)
}
return normalized, nil
}
// LoadKnowledgeBlock loads the persistent knowledge block for a session.
func LoadKnowledgeBlock(ctx context.Context, delegate KVStore, sessionKey string) string {
if delegate == nil {
@ -292,3 +493,22 @@ func LoadKnowledgeBlock(ctx context.Context, delegate KVStore, sessionKey string
return kb.FormatBlock()
}
// LoadFocusState loads an active focus state for the given session.
func LoadFocusState(ctx context.Context, delegate KVStore, sessionKey string) (FocusState, bool) {
if delegate == nil {
return FocusState{}, false
}
key := focusKVPrefix + sessionKey
raw, err := delegate.GetKV(ctx, focusAgentID, key)
if err != nil || raw == "" {
return FocusState{}, false
}
var state FocusState
if err := jsonv2.Unmarshal([]byte(raw), &state); err != nil {
logger.WarnCF("focus", "failed to parse active focus state", map[string]interface{}{"error": err, "session": sessionKey})
return FocusState{}, false
}
return state, true
}

183
pkg/tools/focus_history.go Normal file
View file

@ -0,0 +1,183 @@
package tools
import (
"context"
"fmt"
"sort"
"strings"
jsonv2 "github.com/go-json-experiment/json"
)
const defaultFocusHistoryLimit = 5
const maxFocusHistoryLimit = 25
// FocusHistoryTool lets the agent query completed focus entries.
type FocusHistoryTool struct {
delegate KVStore
sessionKey func() string
}
func NewFocusHistoryTool(delegate KVStore, sessionKeyFn func() string) *FocusHistoryTool {
return &FocusHistoryTool{
delegate: delegate,
sessionKey: sessionKeyFn,
}
}
func (t *FocusHistoryTool) Name() string { return "focus_history" }
func (t *FocusHistoryTool) Description() string {
return "Search and retrieve completed focus session history entries for reflection, handoff, and continuity."
}
func (t *FocusHistoryTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
"description": "Optional substring to search across goal/topic/steps/outcome/summary.",
},
"limit": map[string]interface{}{
"type": "number",
"description": "Maximum number of history items to return. Defaults to 5, max 25.",
},
},
}
}
func (t *FocusHistoryTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
sessionKey := ""
if t.sessionKey != nil {
sessionKey = t.sessionKey()
}
if strings.TrimSpace(sessionKey) == "" {
return ErrorResult("no active session")
}
encountered, err := t.loadCompletedFocus(ctx, sessionKey)
if err != nil {
return ErrorResult(err.Error())
}
if len(encountered) == 0 {
return &ToolResult{ForLLM: "No completed focus history found."}
}
query, _ := args["query"].(string)
query = strings.ToLower(strings.TrimSpace(query))
limit := defaultFocusHistoryLimit
if raw, ok := args["limit"]; ok {
switch l := raw.(type) {
case int:
limit = l
case int64:
limit = int(l)
case float64:
limit = int(l)
case string:
if strings.TrimSpace(l) == "" {
break
}
if parsed, parseErr := parseIntArg(l); parseErr != nil {
return ErrorResult("invalid limit")
} else {
limit = parsed
}
default:
return ErrorResult("invalid limit")
}
}
if limit <= 0 {
limit = defaultFocusHistoryLimit
}
if limit > maxFocusHistoryLimit {
limit = maxFocusHistoryLimit
}
filtered := filterFocusHistory(encountered, query)
if len(filtered) == 0 {
return &ToolResult{ForLLM: "No focus history matches the query."}
}
if limit > len(filtered) {
limit = len(filtered)
}
filtered = filtered[:limit]
response := struct {
Query string `json:"query"`
Count int `json:"count"`
Total int `json:"total"`
Items []KnowledgeEntry `json:"items"`
}{query, len(filtered), len(filtered), filtered}
payload, err := jsonv2.Marshal(response)
if err != nil {
return ErrorResult("failed to marshal focus history")
}
return &ToolResult{ForLLM: string(payload)}
}
func parseIntArg(value string) (int, error) {
var parsed int
if _, err := fmt.Sscanf(value, "%d", &parsed); err != nil {
return 0, err
}
return parsed, nil
}
func (t *FocusHistoryTool) loadCompletedFocus(ctx context.Context, sessionKey string) ([]KnowledgeEntry, error) {
knowledge, err := t.delegate.GetKV(ctx, focusAgentID, knowledgeKVPrefix+sessionKey)
if err != nil {
return nil, err
}
if strings.TrimSpace(knowledge) == "" {
return nil, nil
}
var block KnowledgeBlock
if err := jsonv2.Unmarshal([]byte(knowledge), &block); err != nil {
return nil, err
}
entries := append([]KnowledgeEntry(nil), block.Entries...)
sort.Slice(entries, func(i, j int) bool {
return entries[i].CompletedAt.After(entries[j].CompletedAt)
})
return entries, nil
}
func filterFocusHistory(entries []KnowledgeEntry, query string) []KnowledgeEntry {
if query == "" {
return entries
}
filtered := make([]KnowledgeEntry, 0, len(entries))
for _, entry := range entries {
if matchesFocusHistoryEntry(entry, query) {
filtered = append(filtered, entry)
}
}
return filtered
}
func matchesFocusHistoryEntry(entry KnowledgeEntry, query string) bool {
fields := []string{
entry.Topic,
entry.Goal,
entry.Outcome,
entry.Summary,
strings.Join(entry.Steps, " "),
entry.Deadline,
}
query = strings.ToLower(strings.TrimSpace(query))
for _, field := range fields {
if strings.Contains(strings.ToLower(field), query) {
return true
}
}
return false
}

159
pkg/tools/focus_history_test.go Normal file
View file

@ -0,0 +1,159 @@
package tools
import (
"testing"
"time"
jsonv2 "github.com/go-json-experiment/json"
"github.com/stretchr/testify/require"
)
func TestFocusHistoryTool_NoHistory(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
tool := NewFocusHistoryTool(delegate, func() string { return "session-no-history" })
result := tool.Execute(t.Context(), map[string]interface{}{})
require.Contains(t, result.ForLLM, "No completed focus history found")
}
func TestFocusHistoryTool_NoSession(t *testing.T) {
t.Parallel()
tool := NewFocusHistoryTool(newMockFocusDelegate(), func() string { return "" })
result := tool.Execute(t.Context(), map[string]interface{}{})
require.Contains(t, result.ForLLM, "no active session")
}
func TestFocusHistoryTool_FiltersByQueryAndOrdersByCompletion(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
ctx := t.Context()
sk := "session-1"
now := time.Now()
block := KnowledgeBlock{
Entries: []KnowledgeEntry{
{
Topic: "reduce API latency",
Goal: "improve api latency",
CompletedAt: now.Add(-2 * time.Hour),
Summary: "Added cache",
Outcome: "done",
},
{
Topic: "fix auth regression",
Goal: "fix auth regression",
CompletedAt: now.Add(-1 * time.Hour),
Summary: "Patched token parsing",
Outcome: "resolved",
},
},
}
raw, err := jsonv2.Marshal(block)
require.NoError(t, err)
err = delegate.UpsertKV(ctx, focusAgentID, knowledgeKVPrefix+sk, string(raw))
require.NoError(t, err)
tool := NewFocusHistoryTool(delegate, func() string { return sk })
result := tool.Execute(ctx, map[string]interface{}{"query": "auth", "limit": 10.0})
require.NotNil(t, result)
require.NotNil(t, result)
var response struct {
Items []KnowledgeEntry `json:"items"`
}
require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &response))
require.Len(t, response.Items, 1)
require.Equal(t, "fix auth regression", response.Items[0].Topic)
}
func TestFocusHistoryTool_InvalidLimit(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
sk := "session-2"
block := KnowledgeBlock{
Entries: []KnowledgeEntry{
{
Topic: "first focus",
CompletedAt: time.Now(),
Summary: "seed",
},
},
}
raw, err := jsonv2.Marshal(block)
require.NoError(t, err)
err = delegate.UpsertKV(t.Context(), focusAgentID, knowledgeKVPrefix+sk, string(raw))
require.NoError(t, err)
tool := NewFocusHistoryTool(delegate, func() string { return sk })
result := tool.Execute(t.Context(), map[string]interface{}{"limit": "x"})
require.Contains(t, result.ForLLM, "invalid limit")
}
func TestFocusHistoryTool_IsolatesSessionData(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
ctx := t.Context()
skA := "session-a"
skB := "session-b"
blockA := KnowledgeBlock{
Entries: []KnowledgeEntry{
{Topic: "session a item", CompletedAt: time.Now(), Summary: "A"},
},
}
blockB := KnowledgeBlock{
Entries: []KnowledgeEntry{
{Topic: "session b item", CompletedAt: time.Now().Add(10 * time.Second), Summary: "B"},
},
}
rawA, err := jsonv2.Marshal(blockA)
require.NoError(t, err)
rawB, err := jsonv2.Marshal(blockB)
require.NoError(t, err)
require.NoError(t, delegate.UpsertKV(ctx, focusAgentID, knowledgeKVPrefix+skA, string(rawA)))
require.NoError(t, delegate.UpsertKV(ctx, focusAgentID, knowledgeKVPrefix+skB, string(rawB)))
tool := NewFocusHistoryTool(delegate, func() string { return skB })
result := tool.Execute(ctx, map[string]interface{}{})
var response struct {
Items []KnowledgeEntry `json:"items"`
}
require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &response))
require.Len(t, response.Items, 1)
require.Equal(t, "session b item", response.Items[0].Topic)
}
func TestFocusHistoryTool_SortsByCompletedAtDescending(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
ctx := t.Context()
sk := "session-sort"
oldest := time.Now().Add(-2 * time.Hour)
newest := time.Now()
middle := time.Now().Add(-1 * time.Hour)
block := KnowledgeBlock{
Entries: []KnowledgeEntry{
{Topic: "old", CompletedAt: oldest},
{Topic: "new", CompletedAt: newest},
{Topic: "mid", CompletedAt: middle},
},
}
raw, err := jsonv2.Marshal(block)
require.NoError(t, err)
require.NoError(t, delegate.UpsertKV(ctx, focusAgentID, knowledgeKVPrefix+sk, string(raw)))
tool := NewFocusHistoryTool(delegate, func() string { return sk })
result := tool.Execute(ctx, map[string]interface{}{"limit": 3.0})
var response struct {
Items []KnowledgeEntry `json:"items"`
}
require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &response))
require.Len(t, response.Items, 3)
require.Equal(t, "new", response.Items[0].Topic)
require.Equal(t, "mid", response.Items[1].Topic)
require.Equal(t, "old", response.Items[2].Topic)
}

110
pkg/tools/focus_test.go
View file

@ -3,6 +3,7 @@ package tools
import (
"context"
"testing"
"time"
jsonv2 "github.com/go-json-experiment/json"
"github.com/google/go-cmp/cmp"
@ -67,7 +68,10 @@ func TestStartFocus(t *testing.T) {
ctx := t.Context()
result := tool.Execute(ctx, map[string]interface{}{
"topic": "investigate auth bug",
"topic": "investigate auth bug",
"goal": "Resolve login failures",
"steps": []interface{}{"collect logs", "reproduce issue", "", "run tests"},
"deadline": "2026-02-21T17:00:00Z",
})
require.NotNil(t, result)
@ -80,9 +84,28 @@ func TestStartFocus(t *testing.T) {
var state FocusState
require.NoError(t, jsonv2.Unmarshal([]byte(raw), &state))
assert.Empty(t, cmp.Diff("investigate auth bug", state.Topic))
assert.Empty(t, cmp.Diff("Resolve login failures", state.Goal))
assert.Empty(t, cmp.Diff([]string{"collect logs", "reproduce issue", "run tests"}, state.Steps))
assert.Empty(t, cmp.Diff("2026-02-21T17:00:00Z", state.Deadline))
assert.Empty(t, cmp.Diff(2, state.CheckpointIndex))
}
func TestStartFocus_InvalidStepsPayload(t *testing.T) {
t.Parallel()
sm := session.NewSessionManager("")
sk := "test-session"
sm.GetOrCreate(sk)
delegate := newMockFocusDelegate()
tool := NewStartFocusTool(delegate, sm, func() string { return sk })
ctx := t.Context()
result := tool.Execute(ctx, map[string]interface{}{
"topic": "investigate auth bug",
"steps": "bad-input",
})
assert.Contains(t, result.ForLLM, "steps must be an array of strings")
}
func TestStartFocus_MissingTopic(t *testing.T) {
t.Parallel()
sm := session.NewSessionManager("")
@ -121,6 +144,9 @@ func TestCompleteFocus(t *testing.T) {
completeTool := NewCompleteFocusTool(delegate, sm, func() string { return sk })
result := completeTool.Execute(ctx, map[string]interface{}{
"summary": "Found auth bug in token validation. Fixed by adding expiry check.",
"outcome": "Token validation now enforces expiry",
"status": "completed",
"steps": []interface{}{"collect logs", "inspect token", "add expiry check"},
})
require.NotNil(t, result)
@ -142,7 +168,10 @@ func TestCompleteFocus(t *testing.T) {
require.NoError(t, jsonv2.Unmarshal([]byte(knowledgeRaw), &kb))
require.Len(t, kb.Entries, 1)
assert.Empty(t, cmp.Diff("debug auth", kb.Entries[0].Topic))
assert.Empty(t, cmp.Diff("Token validation now enforces expiry", kb.Entries[0].Outcome))
assert.Empty(t, cmp.Diff([]string{"collect logs", "inspect token", "add expiry check"}, kb.Entries[0].Steps))
assert.Contains(t, kb.Entries[0].Summary, "token validation")
assert.False(t, kb.Entries[0].CompletedAt.IsZero())
// Focus state should be cleaned up
focusRaw, _ := delegate.GetKV(ctx, focusAgentID, focusKVPrefix+sk)
@ -164,6 +193,25 @@ func TestCompleteFocus_NoActiveFocus(t *testing.T) {
assert.Contains(t, result.ForLLM, "no active focus session")
}
func TestCompleteFocus_InvalidStepsPayload(t *testing.T) {
t.Parallel()
sm := session.NewSessionManager("")
sk := "test-session"
sm.GetOrCreate(sk)
sm.AddMessage(sk, "user", "hello")
delegate := newMockFocusDelegate()
startTool := NewStartFocusTool(delegate, sm, func() string { return sk })
completeTool := NewCompleteFocusTool(delegate, sm, func() string { return sk })
startTool.Execute(t.Context(), map[string]interface{}{"topic": "debug"})
result := completeTool.Execute(t.Context(), map[string]interface{}{
"summary": "done",
"steps": "bad-step-list",
})
assert.Contains(t, result.ForLLM, "steps must be an array of strings")
}
func TestCompleteFocus_MultipleKnowledgeEntries(t *testing.T) {
t.Parallel()
sm := session.NewSessionManager("")
@ -305,3 +353,63 @@ func TestLoadKnowledgeBlock(t *testing.T) {
assert.Contains(t, block, "## Test")
assert.Contains(t, block, "Test summary")
}
func TestLoadFocusState(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
ctx := t.Context()
sk := "focus-session"
raw := `{"topic":"legacy topic","checkpoint_index":2,"started_at":"2026-01-01T00:00:00Z"}`
_ = delegate.UpsertKV(ctx, focusAgentID, focusKVPrefix+sk, raw)
state, ok := LoadFocusState(ctx, delegate, sk)
require.True(t, ok)
assert.Empty(t, cmp.Diff("legacy topic", state.FocusText()))
assert.Empty(t, cmp.Diff(2, state.CheckpointIndex))
}
func TestLoadFocusStateInvalidJSON(t *testing.T) {
t.Parallel()
delegate := newMockFocusDelegate()
ctx := t.Context()
sk := "bad-focus-session"
_ = delegate.UpsertKV(ctx, focusAgentID, focusKVPrefix+sk, "{bad-json")
_, ok := LoadFocusState(ctx, delegate, sk)
assert.False(t, ok)
}
func TestFocusStateFormatBlock(t *testing.T) {
t.Parallel()
fs := &FocusState{
Goal: "Investigate auth flow",
Steps: []string{"collect logs", "trace middleware"},
Deadline: "2026-02-21T17:00:00Z",
Status: "active",
}
block := fs.FormatBlock()
assert.Contains(t, block, "# Focus")
assert.Contains(t, block, "## Investigate auth flow")
assert.Contains(t, block, "Planned Steps")
assert.Contains(t, block, "- collect logs")
assert.Contains(t, block, "- trace middleware")
assert.Contains(t, block, "Deadline: 2026-02-21T17:00:00Z")
}
func TestKnowledgeBlockFormat_UsesGoalAndOutcome(t *testing.T) {
t.Parallel()
kb := &KnowledgeBlock{
Entries: []KnowledgeEntry{
{Goal: "Debug auth", Outcome: "Found token expiry bug", Deadline: "2026-02-25T10:00:00Z", CompletedAt: time.Date(2026, 2, 20, 15, 0, 0, 0, time.UTC)},
},
}
formatted := kb.FormatBlock()
assert.Contains(t, formatted, "# Knowledge")
assert.Contains(t, formatted, "## Debug auth")
assert.Contains(t, formatted, "Found token expiry bug")
assert.Contains(t, formatted, "Deadline: 2026-02-25T10:00:00Z")
assert.Contains(t, formatted, "Completed at:")
}

71
pkg/tools/search.go
View file

@ -16,6 +16,8 @@ import (
type ToolSearchTool struct {
registry *ToolRegistry
skillsLoader *skills.SkillsLoader
focusStore KVStore
sessionKeyFn func() string
}
// NewToolSearchTool creates a tool that searches the registry.
@ -28,6 +30,11 @@ func (t *ToolSearchTool) SetSkillsLoader(sl *skills.SkillsLoader) {
t.skillsLoader = sl
}
func (t *ToolSearchTool) SetFocusContext(delegate KVStore, sessionKeyFn func() string) {
t.focusStore = delegate
t.sessionKeyFn = sessionKeyFn
}
func (t *ToolSearchTool) Name() string { return "tool_search" }
func (t *ToolSearchTool) Description() string {
@ -58,7 +65,7 @@ type toolSearchResult struct {
Required []string `json:"required,omitempty"`
}
func (t *ToolSearchTool) Execute(_ context.Context, args map[string]interface{}) *ToolResult {
func (t *ToolSearchTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
query, _ := args["query"].(string)
if query == "" {
return t.listAll()
@ -66,6 +73,7 @@ func (t *ToolSearchTool) Execute(_ context.Context, args map[string]interface{})
queryLower := strings.ToLower(query)
queryTerms := strings.Fields(queryLower)
focusTerms := t.focusTerms(ctx)
var results []toolSearchResult
@ -76,6 +84,7 @@ func (t *ToolSearchTool) Execute(_ context.Context, args map[string]interface{})
continue
}
score := fuzzyScore(tool.Name(), tool.Description(), queryTerms)
score += focusBiasScore(tool.Name(), tool.Description(), focusTerms)
if score > 0 {
params, required := extractSchemaFields(tool.Parameters())
results = append(results, toolSearchResult{
@ -135,6 +144,66 @@ func (t *ToolSearchTool) Execute(_ context.Context, args map[string]interface{})
return &ToolResult{ForLLM: string(b)}
}
func (t *ToolSearchTool) focusTerms(ctx context.Context) []string {
if t.focusStore == nil || t.sessionKeyFn == nil {
return nil
}
sessionKey := t.sessionKeyFn()
if strings.TrimSpace(sessionKey) == "" {
return nil
}
state, ok := LoadFocusState(ctx, t.focusStore, sessionKey)
if !ok {
return nil
}
return extractFocusTerms(state)
}
func extractFocusTerms(state FocusState) []string {
rawTerms := []string{
state.FocusText(),
state.Deadline,
strings.Join(state.Steps, " "),
state.Status,
state.Outcome,
}
stopWords := map[string]struct{}{
"and": {}, "the": {}, "to": {}, "for": {}, "a": {}, "an": {}, "on": {}, "in": {}, "of": {}, "is": {}, "it": {},
}
combined := strings.Join(rawTerms, " ")
normalized := strings.Fields(strings.ToLower(strings.NewReplacer(";", " ", ",", " ", ".", " ", "(", " ", ")", " ", "-", " ", "_", " ", "/", " ", "\\", " ").Replace(combined)))
terms := make([]string, 0, len(normalized))
seen := make(map[string]struct{}, len(normalized))
for _, term := range normalized {
term = strings.TrimSpace(term)
if len(term) < 2 {
continue
}
if _, stop := stopWords[term]; stop {
continue
}
if _, exists := seen[term]; exists {
continue
}
seen[term] = struct{}{}
terms = append(terms, term)
}
return terms
}
func focusBiasScore(name, description string, focusTerms []string) int {
if len(focusTerms) == 0 {
return 0
}
return fuzzyScore(name, description, focusTerms)
}
// extractSchemaFields pulls the properties map and required list from a
// tool's full JSON Schema parameters object.
func extractSchemaFields(params map[string]interface{}) (map[string]interface{}, []string) {

35
pkg/tools/search_test.go
View file

@ -4,9 +4,11 @@ import (
"context"
"os"
"testing"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/skills"
jsonv2 "github.com/go-json-experiment/json"
"github.com/stretchr/testify/require"
)
func TestToolSearchTool_Name(t *testing.T) {
@ -193,6 +195,39 @@ func TestToolSearchTool_NoMatch(t *testing.T) {
}
}
func TestToolSearchTool_FocusAwareBias(t *testing.T) {
t.Parallel()
r := NewToolRegistry()
r.Register(&stubTool{name: "auth_reliability_probe", desc: "Collect auth traces and track reliability signals"})
r.Register(&stubTool{name: "auth_reader", desc: "Read authentication logs"})
delegate := newMockFocusDelegate()
ctx := t.Context()
focusState := FocusState{
Topic: "investigate service reliability",
Goal: "improve auth reliability",
Steps: []string{"check reliability", "validate auth metrics"},
StartedAt: time.Now(),
}
raw, err := jsonv2.Marshal(focusState)
require.NoError(t, err)
err = delegate.UpsertKV(ctx, focusAgentID, focusKVPrefix+"agent-1", string(raw))
require.NoError(t, err)
s := NewToolSearchTool(r)
s.SetFocusContext(delegate, func() string { return "agent-1" })
result := s.Execute(ctx, map[string]interface{}{"query": "auth"})
var results []toolSearchResult
require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &results))
if len(results) < 2 {
t.Fatalf("expected at least 2 results, got %d", len(results))
}
if results[0].Name != "auth_reliability_probe" {
t.Fatalf("expected focused tool first, got %s", results[0].Name)
}
}
func TestToolSearchTool_ExcludesMetaTools(t *testing.T) {
t.Parallel()
r := NewToolRegistry()