enkly-ai/services/agent/internal/infrastructure/agentrun/effort.go

// effort.go adds three related capabilities to the round-trip loop in
// agent.go, all modeled on Claude's agentic-coding feature set:
//
//   - Effort control: a caller picks one of pkg/effort's five levels
//     for a turn (EffortThinkPolicy, and Run reading a.Effort to size
//     MaxRounds/ContextBudgetChars — see Agent.effectiveLimits).
//   - Adaptive thinking: instead of (or layered on top of) a fixed
//     per-level floor, the model's *own* per-round signal — did the
//     last tool call fail? is this a late round in a long turn? — can
//     push a round's thinking effort up, so token spend tracks actual
//     difficulty rather than one static setting for the whole turn
//     (AdaptiveThinkPolicy, RoundSignal).
//   - System entries: a caller can hand Run() additional instruction/
//     permission/budget text to fold into the system prompt for this
//     run, without touching SystemPrompt itself — see WithSystemEntries
//     and its doc comment for exactly what "without breaking cache"
//     can and can't mean against a vLLM-served backend.
package agentrun

import (
	"context"
	"fmt"
	"strings"

	"enkly/pkg/effort"
)

// ---------------------------------------------------------------------
// Effort control
// ---------------------------------------------------------------------

// effortKey is the context.Context key WithEffort stores a level under
// — same unexported-struct-key idiom as preferredLanguageKey.
type effortKey struct{}

// WithEffort attaches the caller's chosen reasoning-depth level to ctx
// for Run to read (see Agent.effectiveLimits and EffortThinkPolicy). An
// unset context (effortFrom returns "", false) means "use the Agent's
// own configured Effort field" — see New's default of effort.Medium —
// so a caller that never touches this gets exactly the behavior the
// agent had before effort control existed.
func WithEffort(ctx context.Context, level effort.Level) context.Context {
	return context.WithValue(ctx, effortKey{}, level)
}

func effortFrom(ctx context.Context) (effort.Level, bool) {
	l, ok := ctx.Value(effortKey{}).(effort.Level)
	return l, ok
}

// EffortFromContext returns the effort level in force on ctx (see
// WithEffort), or effort.Default when none was set. Exported — unlike
// effortFrom — for callers outside this package that need to read the
// resolved level without going through an *Agent (e.g. toolexec's
// dispatch_workflow orchestrator sizing its own parallelism from
// effort.Level.Resolve().MaxParallelSubagents).
func EffortFromContext(ctx context.Context) effort.Level {
	if l, ok := effortFrom(ctx); ok && l.Valid() {
		return l
	}
	return effort.Default
}

// resolveEffort picks the level actually in force for one Run call:
// an explicit WithEffort on ctx wins, otherwise the Agent's own
// configured default (a.Effort), otherwise effort.Default.
func (a *Agent) resolveEffort(ctx context.Context) effort.Level {
	if l, ok := effortFrom(ctx); ok && l.Valid() {
		return l
	}
	if a.Effort.Valid() {
		return a.Effort
	}
	return effort.Default
}

// EffortThinkPolicy returns a ThinkPolicy that applies profile.Think as
// a floor on every round — never asking for *less* thinking than the
// level calls for — while still preserving DefaultThinkPolicy's two
// situational overrides: round 0 still gets "high" (planning benefits
// from deep thinking regardless of level, unless the level's own floor
// is already "high" or above) and a same-round retry still asks for
// "low" (a fast, simple re-attempt is always what a retry wants, even
// at Max effort — spending more reasoning on a request that just broke
// the chat template isn't going to fix it).
func EffortThinkPolicy(profile effort.Profile) ThinkPolicy {
	return func(round int, retrying bool) string {
		if retrying {
			return "low"
		}
		if round == 0 && !thinkAtLeast(profile.Think, "high") {
			return "high"
		}
		return profile.Think
	}
}

// thinkRank orders the agent's think settings from least to most
// effort, for thinkAtLeast's comparison. Bare "true" (thinking on, no
// explicit level) is treated as equivalent to "medium" — vLLM/Qwen3
// has no server-side notion of graded reasoning depth (see
// vllmkit.Client.Think's doc comment: every non-empty, non-"false"
// setting just toggles chat_template_kwargs.enable_thinking on), so
// this ranking exists purely for agentrun's own floor/escalation
// bookkeeping, not because the wire protocol distinguishes the levels.
var thinkRank = map[string]int{"": 0, "false": 0, "low": 1, "true": 2, "medium": 2, "high": 3}

func thinkAtLeast(setting, floor string) bool {
	return thinkRank[strings.ToLower(strings.TrimSpace(setting))] >= thinkRank[floor]
}

// ---------------------------------------------------------------------
// Adaptive thinking
// ---------------------------------------------------------------------

// RoundSignal summarizes what just happened in the previous round of
// the same turn, for AdaptiveThinkPolicy to read when deciding this
// round's effort. Computed by Run from the round it just finished —
// see adaptiveSignalFrom.
type RoundSignal struct {
	// Round is the round about to run (0-indexed) — i.e. len(previous
	// rounds), not the round the signal describes.
	Round int
	// PriorRoundFailed is true when the immediately preceding round's
	// tool call(s) included at least one IsError result — a concrete
	// sign the model is dealing with something that didn't go as
	// planned and may benefit from more deliberate reasoning about
	// what to try next, rather than a quick pattern-matched retry.
	PriorRoundFailed bool
	// PriorToolResultBytes is the size of the largest tool result the
	// previous round produced. A very large result (a big file, a wide
	// grep) is a signal the model now has a lot to digest before its
	// next move is obvious.
	PriorToolResultBytes int
	// ToolCallsSoFar is the running count of tool calls issued earlier
	// in this same turn.
	ToolCallsSoFar int
}

// adaptiveDeepRoundBytes is the PriorToolResultBytes threshold above
// which AdaptiveThinkPolicy treats "a lot just happened" as a reason to
// think harder about it — calibrated against compactedHeadChars +
// compactedTailChars (800): a result bigger than what compaction would
// eventually trim is, by definition, a substantial one.
const adaptiveDeepRoundBytes = compactedHeadChars + compactedTailChars

// adaptiveLateRoundThreshold is how many tool calls into a turn
// "this is getting long, don't lose the thread" starts to apply.
const adaptiveLateRoundThreshold = 8

// AdaptiveThinkPolicy wraps a floor policy (normally
// EffortThinkPolicy(profile)) so a round whose RoundSignal looks
// difficult gets escalated to "high" even if the floor would have
// allowed less — "the model automatically determines whether to use
// turn-based deep thinking" from a concrete, legible signal rather than
// the model guessing about its own difficulty, since vLLM's
// OpenAI-compatible chat API gives no hook for the model to request
// more effort mid-turn on its own. A round that looks ordinary just gets whatever floor already
// says, so this never spends *less* than the configured level, only
// sometimes more — "optimizing token usage" by reserving the escalation
// for rounds that actually show a difficulty signal, instead of every
// round paying a high-effort tax regardless of how routine it is.
func AdaptiveThinkPolicy(floor ThinkPolicy) func(round int, retrying bool, sig RoundSignal) string {
	return func(round int, retrying bool, sig RoundSignal) string {
		base := ""
		if floor != nil {
			base = floor(round, retrying)
		}
		if retrying {
			return base // never escalate a retry — see EffortThinkPolicy's own reasoning
		}
		if sig.PriorRoundFailed || sig.PriorToolResultBytes > adaptiveDeepRoundBytes || sig.ToolCallsSoFar > adaptiveLateRoundThreshold {
			if !thinkAtLeast(base, "high") {
				return "high"
			}
		}
		return base
	}
}

// ---------------------------------------------------------------------
// System entries: mid-task instruction/permission/budget updates
// ---------------------------------------------------------------------

// SystemEntryKind labels what a SystemEntry is updating, purely for the
// heading it renders under in the system prompt — the model doesn't
// enforce any of these as hard constraints beyond what the prose itself
// says (there's no separate permissions engine here — see toolexec for
// where real enforcement, e.g. the dispatch_subagent read-only
// whitelist, actually lives).
type SystemEntryKind string

const (
	// SystemEntryInstruction is a standing instruction added mid-turn
	// — e.g. "also update the changelog for every file you touch from
	// now on."
	SystemEntryInstruction SystemEntryKind = "instruction"
	// SystemEntryPermission narrows or widens what the model should
	// consider itself allowed to do — e.g. "do not run_bash_command
	// against anything under /prod" or "you may now run destructive
	// migrations in this workspace, it's a throwaway clone."
	SystemEntryPermission SystemEntryKind = "permission"
	// SystemEntryBudget communicates a resource constraint — e.g. "you
	// have 4 tool calls left before this session's budget runs out,
	// wrap up."
	SystemEntryBudget SystemEntryKind = "budget"
)

func (k SystemEntryKind) heading() string {
	switch k {
	case SystemEntryPermission:
		return "Permission update"
	case SystemEntryBudget:
		return "Budget update"
	default:
		return "Instruction update"
	}
}

// SystemEntry is one mid-conversation update a caller can inject into a
// run via WithSystemEntries.
type SystemEntry struct {
	Kind    SystemEntryKind
	Content string
}

// systemEntriesKey is WithSystemEntries' context key.
type systemEntriesKey struct{}

// WithSystemEntries attaches additional system-level entries to ctx for
// Run to fold into the system prompt, appended after everything else
// buildSystemPrompt already assembles (base SystemPrompt, project
// memory, language directive) — see buildSystemPrompt. This is enkly's
// answer to "update instructions, permissions or token budgets in the
// middle of a conversation... without breaking cache": be honest about
// what that can mean here. vLLM has no equivalent of the Anthropic
// API's explicit prompt-cache breakpoints; what it actually does is
// automatic prefix caching (--enable-prefix-caching on the chat
// deployment) — reusing cached KV-cache blocks for however much of the
// new request's token prefix is byte-identical, block-for-block, to a
// previously served request's, across any request that lands on the
// same vLLM replica. The guarantee this function and
// buildSystemPrompt together provide is the one piece of that which is
// actually in enkly's control: every entry is appended, never inserted
// or edited in place, and always after the fixed SystemPrompt + memory
// + language block — so a conversation's prefix up to "## Live session
// updates" stays byte-identical turn over turn whether or not new
// entries have been added, maximizing how much of it a KV-cache
// implementation *can* reuse, rather than scattering updates throughout
// the prompt in a way that would invalidate everything after the first
// change on every single turn.
func WithSystemEntries(ctx context.Context, entries []SystemEntry) context.Context {
	if len(entries) == 0 {
		return ctx
	}
	return context.WithValue(ctx, systemEntriesKey{}, entries)
}

func systemEntriesFrom(ctx context.Context) []SystemEntry {
	entries, _ := ctx.Value(systemEntriesKey{}).([]SystemEntry)
	return entries
}

// maxSystemEntryChars caps one entry's rendered length, for the same
// reason memoryPromptChars caps ENKLY.md: a caller-supplied string
// shouldn't be able to blow out the context budget of every single
// round for the rest of the turn.
const maxSystemEntryChars = 2_000

// renderSystemEntries turns entries into the "## Live session updates"
// section buildSystemPrompt appends. Returns "" for no entries, so
// callers can unconditionally concatenate it.
func renderSystemEntries(entries []SystemEntry) string {
	if len(entries) == 0 {
		return ""
	}
	var b strings.Builder
	b.WriteString("\n\n## Live session updates\nThe following were added to this session after it started; they take " +
		"precedence over anything above that they conflict with, for the remainder of this conversation.\n")
	for _, e := range entries {
		content := strings.TrimSpace(e.Content)
		if content == "" {
			continue
		}
		if len(content) > maxSystemEntryChars {
			content = safeHeadCut(content, maxSystemEntryChars) + " […truncated]"
		}
		fmt.Fprintf(&b, "\n### %s\n%s\n", e.Kind.heading(), content)
	}
	return b.String()
}