byteclaw/internal/agent/agent.go

// Package agent implements the LLM tool-calling loop, cancellation and
// summarisation of long sessions.
package agent

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"

	"github.com/crypto/byteclaw/internal/llm"
	"github.com/crypto/byteclaw/internal/store"
	"github.com/crypto/byteclaw/internal/tools"
)

// ErrCancelled is returned by Run when the user stopped the answer.
var ErrCancelled = errors.New("cancelled by user")

// Event is streamed to clients while the agent works.
type Event struct {
	Type    string `json:"type"` // session | delta | tool_call | tool_result | summarized | done | cancelled | error
	Text    string `json:"text,omitempty"`
	Tool    string `json:"tool,omitempty"`
	Args    string `json:"args,omitempty"`
	CallID  string `json:"call_id,omitempty"`
	Session string `json:"session_id,omitempty"`
}

// SummarizeConfig controls automatic summarisation.
type SummarizeConfig struct {
	Enabled     bool
	MaxMessages int
	MaxTokens   int
	KeepRecent  int
}

type Agent struct {
	// LLM is used when Resolve is nil.
	LLM llm.Provider
	// Resolve picks the provider for a session (e.g. per-session override).
	Resolve       func(ctx context.Context, sessionID string) (llm.Provider, error)
	Tools         *tools.Registry
	Store         *store.Store
	SystemPrompt  string
	MaxIterations int
	HistoryLimit  int
	Summarize     SummarizeConfig

	mu    sync.Mutex
	locks map[string]*sync.Mutex
	runs  map[string]context.CancelCauseFunc
}

func (a *Agent) sessionLock(id string) *sync.Mutex {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.locks == nil {
		a.locks = map[string]*sync.Mutex{}
	}
	l, ok := a.locks[id]
	if !ok {
		l = &sync.Mutex{}
		a.locks[id] = l
	}
	return l
}

func (a *Agent) provider(ctx context.Context, sessionID string) (llm.Provider, error) {
	if a.Resolve != nil {
		return a.Resolve(ctx, sessionID)
	}
	if a.LLM == nil {
		return nil, errors.New("no LLM provider configured")
	}
	return a.LLM, nil
}

// Cancel stops the running answer of a session. Returns false if none is running.
func (a *Agent) Cancel(sessionID string) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	if c, ok := a.runs[sessionID]; ok {
		c(ErrCancelled)
		return true
	}
	return false
}

// Running reports whether an answer is in progress for the session.
func (a *Agent) Running(sessionID string) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	_, ok := a.runs[sessionID]
	return ok
}

func (a *Agent) register(sessionID string, c context.CancelCauseFunc) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.runs == nil {
		a.runs = map[string]context.CancelCauseFunc{}
	}
	a.runs[sessionID] = c
}

func (a *Agent) unregister(sessionID string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	delete(a.runs, sessionID)
}

// Run appends the user message to the session and loops until the model
// answers without tool calls. Returns the final assistant text. If the run is
// cancelled via Cancel, the partial answer is stored (marked cancelled) and
// ErrCancelled is returned together with the partial text.
func (a *Agent) Run(parent context.Context, sessionID, userText string, emit func(Event)) (string, error) {
	if emit == nil {
		emit = func(Event) {}
	}
	l := a.sessionLock(sessionID)
	l.Lock()
	defer l.Unlock()

	ctx, cancel := context.WithCancelCause(parent)
	defer cancel(nil)
	a.register(sessionID, cancel)
	defer a.unregister(sessionID)
	// Writes after a cancellation must still succeed.
	saveCtx := context.WithoutCancel(ctx)

	if _, err := a.Store.AppendMessage(saveCtx, sessionID, llm.Message{Role: llm.RoleUser, Content: userText}); err != nil {
		return "", err
	}
	prov, err := a.provider(ctx, sessionID)
	if err != nil {
		return "", err
	}
	if a.Summarize.Enabled {
		if sm, err := a.summarizeLocked(ctx, prov, sessionID, false); err == nil && sm != nil {
			emit(Event{Type: "summarized", Text: sm.Content})
		}
	}
	var specs []llm.ToolSpec
	if a.Tools != nil {
		specs = a.Tools.Specs()
	}
	cancelled := func() bool { return errors.Is(context.Cause(ctx), ErrCancelled) }
	for i := 0; i < a.MaxIterations; i++ {
		if cancelled() {
			return "", ErrCancelled
		}
		msgs, err := a.history(saveCtx, sessionID)
		if err != nil {
			return "", err
		}
		var partial strings.Builder
		resp, err := prov.Chat(ctx, llm.Request{Messages: msgs, Tools: specs}, func(t string) {
			partial.WriteString(t)
			emit(Event{Type: "delta", Text: t})
		})
		if err != nil {
			if cancelled() {
				text := partial.String()
				if text != "" {
					if _, serr := a.Store.AppendMessage(saveCtx, sessionID, llm.Message{Role: llm.RoleAssistant, Content: text, Cancelled: true}); serr != nil {
						return text, serr
					}
				}
				return text, ErrCancelled
			}
			return "", err
		}
		asst := llm.Message{Role: llm.RoleAssistant, Content: resp.Content, ToolCalls: resp.ToolCalls}
		if _, err := a.Store.AppendMessage(saveCtx, sessionID, asst); err != nil {
			return "", err
		}
		if len(resp.ToolCalls) == 0 {
			return resp.Content, nil
		}
		for _, tc := range resp.ToolCalls {
			emit(Event{Type: "tool_call", Tool: tc.Name, Args: tc.Arguments, CallID: tc.ID})
			out := "error: cancelled by user"
			if !cancelled() {
				out = a.Tools.Call(ctx, tc.Name, tc.Arguments)
			}
			emit(Event{Type: "tool_result", Tool: tc.Name, Text: clip(out, 2000), CallID: tc.ID})
			// Every tool call gets a result so the history stays valid for the provider.
			if _, err := a.Store.AppendMessage(saveCtx, sessionID, llm.Message{Role: llm.RoleTool, ToolCallID: tc.ID, Name: tc.Name, Content: out}); err != nil {
				return "", err
			}
		}
	}
	return "", fmt.Errorf("stopped after %d iterations without a final answer", a.MaxIterations)
}

func (a *Agent) history(ctx context.Context, sessionID string) ([]llm.Message, error) {
	var after int64
	summary := ""
	if sm, err := a.Store.LatestSummary(ctx, sessionID); err == nil {
		after, summary = sm.UptoMessageID, sm.Content
	} else if !errors.Is(err, store.ErrNotFound) {
		return nil, err
	}
	stored, err := a.Store.MessagesAfter(ctx, sessionID, after, a.HistoryLimit)
	if err != nil {
		return nil, err
	}
	// A truncated window must not start with orphaned tool results.
	for len(stored) > 0 && stored[0].Role == llm.RoleTool {
		stored = stored[1:]
	}
	msgs := make([]llm.Message, 0, len(stored)+2)
	if a.SystemPrompt != "" {
		msgs = append(msgs, llm.Message{Role: llm.RoleSystem, Content: a.SystemPrompt})
	}
	if summary != "" {
		msgs = append(msgs, llm.Message{Role: llm.RoleSystem, Content: "Summary of the earlier conversation:\n" + summary})
	}
	for _, m := range stored {
		mm := m.Message
		mm.Cancelled = false
		msgs = append(msgs, mm)
	}
	return msgs, nil
}

func clip(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}