From 189fc943e45f7546f1ec76f286609b969df33fa9 Mon Sep 17 00:00:00 2001 From: ZanzyTHEbar Date: Sun, 22 Feb 2026 22:49:10 +0000 Subject: [PATCH] feat(tools): add focus_history tool, expand focus and search - 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 --- pkg/tools/focus.go | 246 ++++++++++++++++++++++++++++++-- pkg/tools/focus_history.go | 183 ++++++++++++++++++++++++ pkg/tools/focus_history_test.go | 159 +++++++++++++++++++++ pkg/tools/focus_test.go | 110 +++++++++++++- pkg/tools/search.go | 71 ++++++++- pkg/tools/search_test.go | 35 +++++ 6 files changed, 789 insertions(+), 15 deletions(-) create mode 100644 pkg/tools/focus_history.go create mode 100644 pkg/tools/focus_history_test.go diff --git a/pkg/tools/focus.go b/pkg/tools/focus.go index a13402415..bdb36fc77 100644 --- a/pkg/tools/focus.go +++ b/pkg/tools/focus.go @@ -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 +} diff --git a/pkg/tools/focus_history.go b/pkg/tools/focus_history.go new file mode 100644 index 000000000..9f7de7744 --- /dev/null +++ b/pkg/tools/focus_history.go @@ -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 +} diff --git a/pkg/tools/focus_history_test.go b/pkg/tools/focus_history_test.go new file mode 100644 index 000000000..5e40d7a23 --- /dev/null +++ b/pkg/tools/focus_history_test.go @@ -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) +} diff --git a/pkg/tools/focus_test.go b/pkg/tools/focus_test.go index b995418f0..17847a6d8 100644 --- a/pkg/tools/focus_test.go +++ b/pkg/tools/focus_test.go @@ -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:") +} diff --git a/pkg/tools/search.go b/pkg/tools/search.go index 6a69f3d9c..d168ca355 100644 --- a/pkg/tools/search.go +++ b/pkg/tools/search.go @@ -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) { diff --git a/pkg/tools/search_test.go b/pkg/tools/search_test.go index d1850a388..faf7dab4e 100644 --- a/pkg/tools/search_test.go +++ b/pkg/tools/search_test.go @@ -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()