BlackBean/pigo/internal/dream/apply.go

package dream

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/smallnest/pigo/internal/agentcore"
	"github.com/smallnest/pigo/internal/provider"
)

// This file wires the llmConsolidator to a real provider (the main-session
// model, SPEC Q3) and holds the MEMORY.md index cleanup the Runner runs after
// writeback. The provider plumbing mirrors internal/cli/run.SetupEnv: resolve a
// Provider for the model/base-url/protocol/provider tuple, resolve the API key
// through a CredentialStore (--api-key override → env → config), and run a
// single StreamCompletion, draining the event stream to the final text.

// NewLLMConsolidator builds the production Consolidator backed by the given
// model configuration — the same tuple cmd/pigo resolves for the main session
// (CLI flags overlaid with config.toml). It resolves the Provider once so every
// Consolidate call reuses it. A resolution failure (bad model / missing
// provider) is returned so the caller can decide whether to fall back to the
// no-op Consolidator or fail the run.
func NewLLMConsolidator(model, baseURL, protocol, providerName, apiKey string, thinking agentcore.ThinkingLevel) (Consolidator, error) {
	complete, err := newModelCompleter(model, baseURL, protocol, providerName, apiKey, thinking)
	if err != nil {
		return nil, err
	}
	return &llmConsolidator{complete: complete}, nil
}

// newModelCompleter resolves the provider and returns a completeFn that performs
// one non-streaming-consuming completion: it sends the system+user prompt as a
// single user turn (no tools — the dream agent only reasons and replies) and
// returns the concatenated assistant text. A hard "cannot build the stream"
// error is returned directly; a runtime failure rides the stream as a terminal
// error event whose message we convert to an error (so the Runner marks the run
// failed rather than silently deleting nothing, SPEC §5.5).
func newModelCompleter(model, baseURL, protocol, providerName, apiKey string, thinking agentcore.ThinkingLevel) (completeFn, error) {
	prov, resolvedName, err := provider.ResolveProvider(model, baseURL, protocol, providerName, os.Getenv)
	if err != nil {
		return nil, fmt.Errorf("dream: resolve provider: %w", err)
	}
	creds := provider.NewCredentialStore(nil)
	creds.SetOverride(resolvedName, apiKey)

	return func(ctx context.Context, systemPrompt, userPrompt string) (string, error) {
		key := creds.GetAPIKey(ctx, resolvedName)
		llm := provider.LlmContext{
			SystemPrompt: systemPrompt,
			Messages: agentcore.MessageList{
				agentcore.UserMessage{
					RoleField: agentcore.RoleUser,
					Content:   agentcore.ContentList{agentcore.NewTextContent(userPrompt)},
				},
			},
		}
		stream, err := prov.StreamCompletion(ctx, provider.CompletionRequest{
			Model:   model,
			Context: llm,
			Config:  provider.StreamConfig{APIKey: key, ThinkingLevel: thinking},
		})
		if err != nil {
			return "", err
		}
		final, err := drainToMessage(ctx, stream)
		if err != nil {
			return "", err
		}
		if final.StopReason == agentcore.StopReasonError {
			if final.ErrorMessage != "" {
				return "", fmt.Errorf("model error: %s", final.ErrorMessage)
			}
			return "", fmt.Errorf("model returned an error response")
		}
		return agentcore.ContentToText(final.Content), nil
	}, nil
}

// drainToMessage consumes the provider event stream to completion and returns
// the terminal assistant message. It mirrors the loop's stream-drain contract:
// the done/error event carries the final message; if the stream closes without
// one, it falls back to the stream Result. Draining is required because the
// producer blocks on the event channel until consumed.
func drainToMessage(ctx context.Context, stream *provider.AssistantMessageEventStream) (agentcore.AssistantMessage, error) {
	for ev := range stream.Events() {
		switch e := ev.(type) {
		case provider.StreamDoneEvent:
			return e.Message, nil
		case provider.StreamErrorEvent:
			return e.Message, nil
		}
	}
	final, err := stream.Result(ctx)
	if err != nil {
		return agentcore.AssistantMessage{}, err
	}
	return final, nil
}

// updateScopeIndexes rewrites each affected scope's MEMORY.md to drop any line
// that references a now-deleted memory file, keeping the index consistent with
// the entries on disk and free of dangling links (PRD US-003). It is safe to
// call when no MEMORY.md exists (no-op) and when deleted is empty. Each rewrite
// is atomic (temp+rename) and guarded by withinScope, so it cannot escape the
// memory store.
func updateScopeIndexes(memoryRoot, projectDir string, deleted map[string]struct{}) error {
	if len(deleted) == 0 {
		return nil
	}
	scopes := []string{filepath.Join(memoryRoot, "global")}
	if projectDir != "" {
		scopes = append(scopes, filepath.Join(memoryRoot, "projects", projectID(projectDir)))
	}
	for _, scope := range scopes {
		index := filepath.Join(scope, "MEMORY.md")
		if _, err := os.Stat(index); err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return err
		}
		if !withinScope(memoryRoot, projectDir, index) {
			continue
		}
		tokens := indexRefTokens(memoryRoot, scope, deleted)
		if len(tokens) == 0 {
			continue
		}
		raw, err := os.ReadFile(index)
		if err != nil {
			return err
		}
		newBody, changed := stripDanglingIndexLines(string(raw), tokens)
		if changed {
			if err := atomicWrite(index, []byte(newBody)); err != nil {
				return err
			}
		}
	}
	return nil
}

// indexRefTokens is the set of substrings that identify a deleted file inside a
// MEMORY.md index line: its absolute path, its path relative to the memory root
// and to the scope root, and its bare basename. A line containing any of these
// is treated as a link/reference to the removed entry. MEMORY.md itself is never
// a token (it is never a consolidation deletion target).
func indexRefTokens(memoryRoot, scope string, deleted map[string]struct{}) map[string]struct{} {
	tokens := make(map[string]struct{})
	for p := range deleted {
		clean := filepath.Clean(p)
		add := func(s string) {
			if s != "" && s != "." {
				tokens[filepath.ToSlash(s)] = struct{}{}
			}
		}
		add(clean)
		if rel, err := filepath.Rel(memoryRoot, clean); err == nil && !strings.HasPrefix(rel, "..") {
			add(rel)
		}
		if rel, err := filepath.Rel(scope, clean); err == nil && !strings.HasPrefix(rel, "..") {
			add(rel)
		}
		add(filepath.Base(clean))
	}
	return tokens
}

// stripDanglingIndexLines removes every line of body that references any of the
// reference tokens, returning the rewritten body and whether anything changed.
// It matches on the forward-slash form of each line so Windows-style separators
// in the index still match the slash tokens. Matching is boundary-aware: a token
// (e.g. the basename "b.md") only matches when it is not embedded inside a longer
// filename token (so "club.md" or "b.mdx" is not mistaken for "b.md"), avoiding
// dropping unrelated index lines.
func stripDanglingIndexLines(body string, tokens map[string]struct{}) (string, bool) {
	lines := strings.Split(body, "\n")
	kept := make([]string, 0, len(lines))
	changed := false
	for _, line := range lines {
		probe := filepath.ToSlash(line)
		drop := false
		for tok := range tokens {
			if containsRefToken(probe, tok) {
				drop = true
				break
			}
		}
		if drop {
			changed = true
			continue
		}
		kept = append(kept, line)
	}
	if !changed {
		return body, false
	}
	return strings.Join(kept, "\n"), true
}

// containsRefToken reports whether tok occurs in line at a filename boundary:
// the characters immediately before and after the match must not be filename
// continuation characters ([A-Za-z0-9_-]). This lets "b.md" match "user/b.md",
// "(b.md)" and "- b.md" while rejecting "club.md" and "b.mdx".
func containsRefToken(line, tok string) bool {
	if tok == "" {
		return false
	}
	from := 0
	for {
		i := strings.Index(line[from:], tok)
		if i < 0 {
			return false
		}
		start := from + i
		end := start + len(tok)
		if !isFilenameChar(byteAt(line, start-1)) && !isFilenameChar(byteAt(line, end)) {
			return true
		}
		from = start + 1
	}
}

// byteAt returns line[i], or 0 when i is out of range (treated as a boundary).
func byteAt(line string, i int) byte {
	if i < 0 || i >= len(line) {
		return 0
	}
	return line[i]
}

// isFilenameChar reports whether b can appear inside a bare filename token, used
// to detect whether a reference-token match is embedded in a longer name.
func isFilenameChar(b byte) bool {
	switch {
	case b >= 'a' && b <= 'z', b >= 'A' && b <= 'Z', b >= '0' && b <= '9':
		return true
	case b == '_' || b == '-':
		return true
	}
	return false
}