mirror of
https://github.com/docker/docker-agent.git
synced 2026-11-04 10:02:15 +00:00
2200 lines
85 KiB
Go
2200 lines
85 KiB
Go
package runtime
|
|
|
|
import (
|
|
"cmp"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"maps"
|
|
"os"
|
|
"slices"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
"go.opentelemetry.io/otel/trace"
|
|
|
|
"github.com/docker/docker-agent/pkg/agent"
|
|
"github.com/docker/docker-agent/pkg/chat"
|
|
"github.com/docker/docker-agent/pkg/config/latest"
|
|
"github.com/docker/docker-agent/pkg/config/types"
|
|
"github.com/docker/docker-agent/pkg/effort"
|
|
"github.com/docker/docker-agent/pkg/evaluator"
|
|
"github.com/docker/docker-agent/pkg/harness"
|
|
"github.com/docker/docker-agent/pkg/hooks"
|
|
"github.com/docker/docker-agent/pkg/hooks/builtins"
|
|
"github.com/docker/docker-agent/pkg/httpclient"
|
|
"github.com/docker/docker-agent/pkg/model/provider"
|
|
"github.com/docker/docker-agent/pkg/model/provider/dmr/dmrmodels"
|
|
"github.com/docker/docker-agent/pkg/modelsdev"
|
|
"github.com/docker/docker-agent/pkg/session"
|
|
"github.com/docker/docker-agent/pkg/sessiontitle"
|
|
"github.com/docker/docker-agent/pkg/team"
|
|
"github.com/docker/docker-agent/pkg/telemetry/genai"
|
|
"github.com/docker/docker-agent/pkg/tools"
|
|
agenttool "github.com/docker/docker-agent/pkg/tools/builtin/agent"
|
|
"github.com/docker/docker-agent/pkg/tools/builtin/skills"
|
|
"github.com/docker/docker-agent/pkg/tools/lifecycle"
|
|
)
|
|
|
|
// ToolHandlerFunc handles a runtime-managed tool call. rt is the dispatcher's
|
|
// per-call handle, used by handlers that need to talk back to the in-flight
|
|
// call (streaming output, asking the user to approve an action).
|
|
type ToolHandlerFunc func(ctx context.Context, sess *session.Session, toolCall tools.ToolCall, events EventSink, rt tools.Runtime) (*tools.ToolCallResult, error)
|
|
|
|
// Runtime defines the contract for runtime execution
|
|
type Runtime interface {
|
|
// CurrentAgentInfo returns information about the currently active agent
|
|
CurrentAgentInfo(ctx context.Context) CurrentAgentInfo
|
|
// CurrentAgentName returns the name of the currently active agent
|
|
CurrentAgentName(ctx context.Context) string
|
|
// SetCurrentAgent sets the currently active agent for subsequent user messages
|
|
SetCurrentAgent(ctx context.Context, agentName string) error
|
|
// CurrentAgentTools returns the tools for the active agent
|
|
CurrentAgentTools(ctx context.Context) ([]tools.Tool, error)
|
|
// CurrentAgentToolsetStatuses returns lifecycle status for each toolset of
|
|
// the active agent (name, kind, state, last error, restart count).
|
|
// Used by the /tools dialog. Best-effort: toolsets that don't expose
|
|
// state appear with State == StateStopped/Ready as appropriate.
|
|
CurrentAgentToolsetStatuses() []tools.ToolsetStatus
|
|
|
|
// RestartToolset finds the named toolset on the active agent and asks
|
|
// its supervisor to drop the current session and reconnect. Returns
|
|
// an error if no toolset matches name or the toolset does not
|
|
// implement tools.Restartable.
|
|
RestartToolset(ctx context.Context, name string) error
|
|
// EmitStartupInfo emits initial agent, team, and toolset information for immediate display.
|
|
// When sess is non-nil and contains token data, a TokenUsageEvent is also emitted
|
|
// so the UI can display context usage percentage on session restore.
|
|
EmitStartupInfo(ctx context.Context, sess *session.Session, events EventSink)
|
|
// EmitAgentInfo emits up-to-date agent and team info (model, thinking
|
|
// level, description) without re-running heavier startup work such as
|
|
// tool discovery. Used to refresh the UI after a lightweight change like
|
|
// cycling the thinking level, avoiding the tool-count flicker that
|
|
// re-emitting full startup info would cause.
|
|
EmitAgentInfo(ctx context.Context, events EventSink)
|
|
// ResetStartupInfo resets the startup info emission flag, allowing re-emission
|
|
ResetStartupInfo()
|
|
// RunStream starts the agent's interaction loop and returns a channel of events.
|
|
// Consumers must drain it to closure before releasing turn ownership, even
|
|
// after cancellation. Cancel the supplied context before draining on errors.
|
|
RunStream(ctx context.Context, sess *session.Session) <-chan Event
|
|
// Run starts the agent's interaction loop and returns the final messages
|
|
Run(ctx context.Context, sess *session.Session) ([]session.Message, error)
|
|
// Resume allows resuming execution after user confirmation.
|
|
// The ResumeRequest carries the decision type and an optional reason (for rejections).
|
|
Resume(ctx context.Context, req ResumeRequest)
|
|
// ResumeElicitation sends an elicitation response back to a waiting
|
|
// elicitation request. elicitationID correlates the response with a
|
|
// specific concurrent request (see ElicitationRequestEvent.ElicitationID);
|
|
// omit it (or pass "") to fall back to resolving the sole pending
|
|
// request, for backward compatibility with older clients. The parameter
|
|
// is variadic — rather than a required 4th positional argument — purely
|
|
// so pre-#3584 callers of this interface method keep compiling unchanged;
|
|
// implementations should treat more than one value as a caller error and
|
|
// use only the first.
|
|
ResumeElicitation(_ context.Context, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error
|
|
// SessionStore returns the session store for browsing/loading past sessions.
|
|
// Returns nil if no persistent session store is configured.
|
|
SessionStore() session.Store
|
|
|
|
// Summarize generates a summary for the session
|
|
Summarize(ctx context.Context, sess *session.Session, additionalPrompt string, events EventSink)
|
|
|
|
// PermissionsInfo returns the team-level permission patterns (allow/ask/deny).
|
|
// Returns nil if no permissions are configured.
|
|
PermissionsInfo() *PermissionsInfo
|
|
|
|
// ReadSkillContent loads a skill for sess without executing embedded commands.
|
|
ReadSkillContent(ctx context.Context, sess *session.Session, name string) (string, error)
|
|
|
|
// CurrentAgentSkillsToolset returns the skills toolset for the current agent, or nil if skills are not enabled.
|
|
CurrentAgentSkillsToolset() *skills.ToolSet
|
|
|
|
// RunSkillFork executes a `context: fork` skill as an isolated
|
|
// sub-session of sess. Used by the run_skill tool and the App's
|
|
// slash-command path. Returns a ToolCallResult carrying the sub-agent's
|
|
// last assistant message, or a user-facing error reason; err is non-nil
|
|
// only for unexpected runtime failures.
|
|
RunSkillFork(ctx context.Context, sess *session.Session, args skills.RunSkillArgs, events EventSink) (*tools.ToolCallResult, error)
|
|
|
|
// CurrentMCPPrompts returns MCP prompts available from the current agent's toolsets.
|
|
// Returns an empty map if no MCP prompts are available.
|
|
CurrentMCPPrompts(ctx context.Context) map[string]tools.PromptInfo
|
|
|
|
// ExecuteMCPPrompt executes a named MCP prompt with the given arguments.
|
|
ExecuteMCPPrompt(ctx context.Context, promptName string, arguments map[string]string) (string, error)
|
|
|
|
// UpdateSessionTitle persists a new title for the current session.
|
|
UpdateSessionTitle(ctx context.Context, sess *session.Session, title string) error
|
|
|
|
// TitleGenerator returns a generator for automatic session titles, or nil
|
|
// if the runtime does not support local title generation (e.g. remote
|
|
// runtimes) or no configured model is a usable title candidate.
|
|
TitleGenerator(ctx context.Context) *sessiontitle.Generator
|
|
|
|
// Steer enqueues a user message for urgent mid-turn injection into the
|
|
// running agent loop. Returns an error if the queue is full or steering
|
|
// is not available.
|
|
Steer(ctx context.Context, msg QueuedMessage) error
|
|
// FollowUp enqueues a message for end-of-turn processing. Each follow-up
|
|
// gets a full undivided agent turn. Returns an error if the queue is full.
|
|
FollowUp(ctx context.Context, msg QueuedMessage) error
|
|
|
|
// SetAgentModel sets a model override for the named agent.
|
|
// modelRef can be:
|
|
// - "" (empty) to clear the override and use the agent's default model
|
|
// - A model name from the config (e.g., "my_fast_model")
|
|
// - An inline spec (e.g., "openai/gpt-4o")
|
|
// Returns [ErrUnsupported] for runtimes that don't expose model
|
|
// switching (e.g. remote runtimes, where the server owns the choice).
|
|
SetAgentModel(ctx context.Context, agentName, modelRef string) error
|
|
|
|
// CycleAgentThinkingLevel advances the named agent's thinking-effort
|
|
// level to the next value in its provider's cycle (wrapping around),
|
|
// applies it as a runtime model override, and returns the newly
|
|
// selected level. Returns [ErrUnsupported] for runtimes that can't
|
|
// switch models (e.g. remote runtimes) or for models that don't
|
|
// support thinking-effort selection.
|
|
CycleAgentThinkingLevel(ctx context.Context, agentName string) (effort.Level, error)
|
|
|
|
// SetAgentThinkingLevel sets the named agent's thinking-effort level to
|
|
// the requested value, applies it as a runtime model override, and
|
|
// returns the applied level. The level must be one the current model
|
|
// supports. Returns [ErrUnsupported] for runtimes that can't switch
|
|
// models (e.g. remote runtimes) or for models that don't support
|
|
// thinking-effort selection.
|
|
SetAgentThinkingLevel(ctx context.Context, agentName string, level effort.Level) (effort.Level, error)
|
|
|
|
// AvailableModels returns the models the user can pick from in the
|
|
// /model picker. Returns nil for runtimes that don't expose model
|
|
// switching; see SupportsModelSwitching for a cheap pre-check.
|
|
AvailableModels(ctx context.Context) []ModelChoice
|
|
|
|
// SupportsModelSwitching reports whether SetAgentModel and
|
|
// AvailableModels are wired for this runtime. Use it to gate UI
|
|
// affordances (e.g. show /model in the menu) without paying the
|
|
// cost of AvailableModels.
|
|
SupportsModelSwitching() bool
|
|
|
|
// OnToolsChanged registers a handler invoked outside of any RunStream
|
|
// when a toolset reports a tool list change (e.g. after an MCP
|
|
// ToolListChanged notification). A nil handler unregisters. Runtimes
|
|
// that don't emit such events can implement this as a no-op.
|
|
OnToolsChanged(handler func(Event))
|
|
|
|
// OnBackgroundEvent registers a handler invoked outside of any RunStream
|
|
// for events surfaced from detached background work (e.g. token usage
|
|
// from background agent tasks started via run_background_agent), which
|
|
// has no live event channel of its own. Runtimes that don't run local
|
|
// background work can implement this as a no-op.
|
|
OnBackgroundEvent(handler func(Event))
|
|
|
|
// OnElicitationRequest registers a handler invoked whenever an MCP
|
|
// toolset raises an elicitation request, regardless of which stream
|
|
// (foreground or a detached background job) is currently "active" for
|
|
// the runtime's internal elicitation bridge. This is the reliable
|
|
// delivery route for background-job elicitations (#3584); runtimes
|
|
// that don't run local background work can implement this as a no-op.
|
|
OnElicitationRequest(handler func(Event))
|
|
|
|
// QueueStatus returns the current depth and capacity of message queues
|
|
QueueStatus() QueueStatus
|
|
|
|
// TogglePause toggles whether the run loop is paused at iteration
|
|
// boundaries. Returns the new state (true if now paused). Returns
|
|
// [ErrUnsupported] for runtimes that don't expose pause control.
|
|
TogglePause(ctx context.Context) (paused bool, err error)
|
|
|
|
// Close releases resources held by the runtime (e.g., background
|
|
// agents and pending lifecycles). The session store is *not* closed
|
|
// here: it is supplied by the embedder via WithSessionStore and may
|
|
// be shared with other runtimes (e.g. the TUI's spawner reuses the
|
|
// same store across spawned sessions). Embedders that own the store
|
|
// must close it themselves.
|
|
Close() error
|
|
}
|
|
|
|
// PermissionsInfo contains the allow, ask, and deny patterns for tool permissions.
|
|
type PermissionsInfo struct {
|
|
Allow []string
|
|
Ask []string
|
|
Deny []string
|
|
}
|
|
|
|
type CurrentAgentInfo struct {
|
|
Name string
|
|
Description string
|
|
Commands types.Commands
|
|
}
|
|
|
|
type ModelStore interface {
|
|
GetModel(ctx context.Context, id modelsdev.ID) (*modelsdev.Model, error)
|
|
GetDatabase(ctx context.Context) (*modelsdev.Database, error)
|
|
}
|
|
|
|
// LocalRuntime manages the execution of agents
|
|
type LocalRuntime struct {
|
|
harnessFactory *harness.Factory
|
|
commandEvaluator *CommandEvaluatorFactory
|
|
|
|
ctx func() context.Context
|
|
toolMap map[string]ToolHandlerFunc
|
|
toolDeferrals tools.DeferralTracker
|
|
team *team.Team
|
|
agents *agentRouter
|
|
resumeChan chan ResumeRequest
|
|
tracer trace.Tracer
|
|
modelsStore ModelStore
|
|
sessionCompaction bool
|
|
managedOAuth bool
|
|
unmanagedOAuthRedirectURI string
|
|
nonInteractive bool
|
|
directElicitation tools.ElicitationHandler
|
|
startupInfoEmitted atomic.Bool // Track if startup info has been emitted to avoid unnecessary duplication
|
|
elicitation elicitationBridge // Owns the per-stream events channel for outbound elicitation requests
|
|
elicitationWaiters elicitationWaiters // Routes elicitation responses to the request awaiting them, keyed by ID (#3584)
|
|
elicitationDeclines elicitationDeclineNotes
|
|
|
|
// elicitationSinkMu guards onElicitationRequest; elicitation requests can
|
|
// arrive from background-job goroutines (runCollecting) concurrently with
|
|
// the sink being (re)registered.
|
|
elicitationSinkMu sync.RWMutex
|
|
onElicitationRequest func(Event)
|
|
onElicitationContext func(context.Context, Event)
|
|
sessionStore session.Store
|
|
// generatedFiles caches per-owner-session workspace roots and manifest
|
|
// records for [LocalRuntime.ResolveGeneratedFile]. Seeded by
|
|
// materialization, filled lazily for restored sessions.
|
|
generatedFiles generatedFileCache
|
|
workingDir string // Working directory for hooks execution
|
|
env []string // Environment variables for hooks execution
|
|
modelSwitcherCfg *ModelSwitcherConfig
|
|
providerRegistry *provider.Registry
|
|
gatewayModels gatewayModelsCache
|
|
dmrModels dmrModelsCache
|
|
modelContexts modelContextCache
|
|
|
|
// hooksRegistry is the runtime-private hooks.Registry used to build
|
|
// every Executor. It carries the runtime-owned builtin hooks
|
|
// (add_date, add_environment_info) registered once during
|
|
// NewLocalRuntime, so they're available to every agent without
|
|
// touching any process-wide state.
|
|
hooksRegistry *hooks.Registry
|
|
|
|
// autoInjectors run on every per-agent hook config during
|
|
// [buildHooksExecutors] so embedders can plug in builtins (today
|
|
// snapshot via [builtins.SnapshotController]) without the runtime
|
|
// hard-coding their wiring. Set via [WithAutoInjector].
|
|
autoInjectors []builtins.AutoInjector
|
|
|
|
// hooksExecByAgent holds the per-agent [hooks.Executor], keyed by
|
|
// agent name. Built once in [NewLocalRuntime.buildHooksExecutors]
|
|
// after team and runtime config are finalized; agents with no hooks
|
|
// have no entry, so [hooksExec] returns nil for them. Read-only after
|
|
// construction, so no locking is needed.
|
|
hooksExecByAgent map[string]*hooks.Executor
|
|
|
|
// transforms is the runtime's [MessageTransform] chain, applied to
|
|
// every LLM call in registration order. Populated by
|
|
// [NewLocalRuntime] (for the runtime-shipped strip transform) and by
|
|
// [WithMessageTransform] (for embedder-supplied transforms).
|
|
// Read-only after construction.
|
|
transforms []registeredTransform
|
|
|
|
fallback *fallbackExecutor
|
|
|
|
// observers receive every event the runtime produces, in
|
|
// registration order. Built up via [WithEventObserver] during
|
|
// construction; read-only afterwards. Always contains at least one
|
|
// entry: the auto-registered [PersistenceObserver] for the
|
|
// configured session store. See [EventObserver] for the contract.
|
|
observers []EventObserver
|
|
|
|
// fallback owns the model-fallback chain (primary + configured
|
|
// fallbacks), per-attempt retry/backoff for transient errors, and
|
|
// the per-agent "sticky" cooldown after a fallback succeeds. It
|
|
// holds the cooldownManager and rate-limit retry flag so that state
|
|
// stays out of LocalRuntime. See [fallbackExecutor].
|
|
|
|
// steerQueue stores urgent mid-turn messages. The agent loop drains
|
|
// ALL pending messages after tool execution, before the stop check.
|
|
steerQueue MessageQueue
|
|
|
|
// followUpQueue stores end-of-turn messages. The agent loop pops
|
|
// exactly ONE message after the model stops and stop-hooks have run.
|
|
followUpQueue MessageQueue
|
|
|
|
// budgetCfg is the run-wide budget (nil when unset), budgetsCfg the
|
|
// named budget definitions, and agentBudgets the names each agent
|
|
// declared. All three are immutable templates; budget below is the
|
|
// live accumulator built from them.
|
|
budgetCfg *latest.BudgetConfig
|
|
budgetsCfg map[string]latest.BudgetConfig
|
|
agentBudgets map[string][]string
|
|
|
|
// budgetMu guards budget and budgetStarted. The first root stream or
|
|
// manual compaction installs the trackers; subsequent runs reuse them.
|
|
// Delegated and compaction work share the same wallets.
|
|
budgetMu sync.Mutex
|
|
budget *budgetSet
|
|
// budgetStarted distinguishes "not built yet" from "built, and there
|
|
// was nothing to budget" — without it a nil budget would be rebuilt on
|
|
// every message, and the reset bug would come back for unbudgeted runs.
|
|
budgetStarted bool
|
|
|
|
// activeRootStreams tracks top-level RunStream loops. Recalls use this to
|
|
// preserve mid-turn steering while still waking the embedder once the parent
|
|
// stream has gone idle.
|
|
activeRootStreams atomic.Int32
|
|
|
|
// liveSessions is the registry of sessions with an active RunStream,
|
|
// keyed by session ID. It backs the /context team view (LiveSessions)
|
|
// and targeted manual compaction (CompactLiveSession). Guarded by
|
|
// liveSessionsMu; see live_sessions.go for the enqueue/drain contract.
|
|
liveSessionsMu sync.Mutex
|
|
liveSessions map[string]*liveSessionEntry
|
|
|
|
// recallHandler wakes embedders (TUI/App) when a tool recall arrives after
|
|
// the active RunStream has gone idle. Protected by recallMu because tool
|
|
// callbacks can fire from background goroutines.
|
|
recallMu sync.RWMutex
|
|
recallHandler RecallHandler
|
|
|
|
// onToolsChanged is called when an MCP toolset reports a tool list
|
|
// change. Protected by toolsChangedMu because MCP change-notification
|
|
// goroutines call emitToolsChanged concurrently, mirroring
|
|
// onBackgroundEvent/backgroundEventMu below. toolsChangedUnsubs holds
|
|
// this runtime's subscriptions on ChangeSubscriber toolsets, released
|
|
// on re-registration and at Close.
|
|
toolsChangedMu sync.RWMutex
|
|
onToolsChanged func(Event)
|
|
toolsChangedUnsubs []func()
|
|
|
|
// onBackgroundEvent is called for events surfaced from detached
|
|
// background work (e.g. background agent tasks). Protected by
|
|
// backgroundEventMu because background tasks read it from their own
|
|
// goroutines.
|
|
backgroundEventMu sync.RWMutex
|
|
onBackgroundEvent func(Event)
|
|
onBackgroundContext func(context.Context, Event)
|
|
|
|
bgAgents *agenttool.Handler
|
|
|
|
// dmrModelLister lists the models pulled locally in Docker Model Runner,
|
|
// used to populate DMR entries in the model picker. Defaults to
|
|
// dmrmodels.ListModelsWithMetadata in NewLocalRuntime; left nil by runtimes built directly
|
|
// (e.g. tests) so DMR discovery stays opt-in. Tests inject a stub here.
|
|
dmrModelLister func(ctx context.Context) ([]dmrmodels.Model, error)
|
|
|
|
// now is the runtime's clock. Defaults to time.Now and can be replaced
|
|
// in tests via WithClock to make timestamps and cooldown windows
|
|
// deterministic. Every time-dependent call inside the runtime (message
|
|
// CreatedAt, fallback cooldown windows, tool-call latency) goes through
|
|
// this hook so a single fake clock controls them all.
|
|
now func() time.Time
|
|
|
|
// telemetry receives the runtime's observability events (session
|
|
// start/end, tool calls, token usage, errors). Defaults to
|
|
// defaultTelemetry which forwards to pkg/telemetry. Tests can inject
|
|
// a recorder via WithTelemetry to assert the lifecycle without
|
|
// standing up an OTel pipeline.
|
|
telemetry Telemetry
|
|
|
|
// maxOverflowCompactions caps the number of consecutive context-
|
|
// overflow auto-compactions the run loop attempts before surfacing the
|
|
// error. Defaults to defaultMaxOverflowCompactions; tests use
|
|
// WithMaxOverflowCompactions to exercise both the "compaction
|
|
// succeeded" and "compaction exhausted" branches.
|
|
maxOverflowCompactions int
|
|
|
|
// toolListTimeout bounds how long EmitStartupInfo waits for a single
|
|
// toolset to enumerate its tools before skipping it, so one hung toolset
|
|
// cannot stall the sidebar's "Loading tools…" state forever. Defaults to
|
|
// defaultToolListTimeout; overridden via WithToolListTimeout.
|
|
toolListTimeout time.Duration
|
|
|
|
// toolStartTimeout bounds how long EmitStartupInfo waits for a single
|
|
// toolset to start before skipping it for the startup sidebar pass.
|
|
// Defaults to defaultToolStartTimeout; overridden via WithToolStartTimeout.
|
|
toolStartTimeout time.Duration
|
|
|
|
// streamStoppedDeliveryTimeout bounds how long finalizeEventChannel
|
|
// blocks trying to deliver StreamStopped when the event buffer is full.
|
|
// Defaults to defaultStreamStoppedDeliveryTimeout; tests shrink it to
|
|
// exercise the drop-when-abandoned path without a real-time wait. Zero
|
|
// falls back to the default (see streamStoppedTimeout), so runtimes built
|
|
// directly via a struct literal in tests still get bounded delivery.
|
|
streamStoppedDeliveryTimeout time.Duration
|
|
|
|
// pauseMu guards pauseCh.
|
|
pauseMu sync.Mutex
|
|
// pauseCh is non-nil and open while /pause has paused the run loop;
|
|
// nil otherwise. See TogglePause and waitIfPaused.
|
|
pauseCh chan struct{}
|
|
}
|
|
|
|
type Opt func(*LocalRuntime)
|
|
|
|
func WithCurrentAgent(agentName string) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.agents.Set(agentName)
|
|
}
|
|
}
|
|
|
|
func WithManagedOAuth(managed bool) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.managedOAuth = managed
|
|
}
|
|
}
|
|
|
|
// WithUnmanagedOAuthRedirectURI configures the redirect_uri the runtime
|
|
// advertises when running MCP server OAuth flows in unmanaged mode (i.e.
|
|
// when WithManagedOAuth(false) is set). When set, docker-agent generates
|
|
// state + PKCE + DCR in-process and emits an elicitation carrying the
|
|
// `authorize_url` + `state`; the client returns `{code, state}` via
|
|
// ResumeElicitation and docker-agent does the token exchange itself.
|
|
// When empty, the runtime falls back to the legacy unmanaged contract
|
|
// where the client performs the OAuth flow and returns an access token.
|
|
func WithUnmanagedOAuthRedirectURI(uri string) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.unmanagedOAuthRedirectURI = uri
|
|
}
|
|
}
|
|
|
|
// WithElicitationHandler handles requests synchronously instead of emitting elicitation events.
|
|
// Non-interactive runtimes still decline; the handler must honor the request context.
|
|
func WithElicitationHandler(handler tools.ElicitationHandler) Opt {
|
|
return func(r *LocalRuntime) { r.directElicitation = handler }
|
|
}
|
|
|
|
// WithNonInteractive marks the runtime as headless (e.g., MCP serve mode).
|
|
// When set, blocking operations like elicitation requests are automatically
|
|
// declined instead of waiting for user interaction that will never come.
|
|
//
|
|
// Note: this complements session.WithNonInteractive, which controls per-session
|
|
// loop behavior (e.g., auto-stop on max iterations). Both should be set for
|
|
// fully headless operation - the runtime flag prevents elicitation hangs, while
|
|
// the session flag adjusts iteration behavior. In MCP serve mode, both are set
|
|
// by the server code in pkg/mcp/server.go.
|
|
func WithNonInteractive(nonInteractive bool) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.nonInteractive = nonInteractive
|
|
}
|
|
}
|
|
|
|
// WithTracer sets a custom OpenTelemetry tracer; if not provided, tracing is disabled (no-op).
|
|
func WithTracer(t trace.Tracer) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.tracer = t
|
|
}
|
|
}
|
|
|
|
// WithSteerQueue sets a custom MessageQueue for mid-turn message injection.
|
|
// If not provided, an in-memory buffered queue is used.
|
|
func WithSteerQueue(q MessageQueue) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.steerQueue = q
|
|
}
|
|
}
|
|
|
|
// WithFollowUpQueue sets a custom MessageQueue for end-of-turn follow-up
|
|
// messages. If not provided, an in-memory buffered queue is used.
|
|
func WithFollowUpQueue(q MessageQueue) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.followUpQueue = q
|
|
}
|
|
}
|
|
|
|
func WithSessionCompaction(sessionCompaction bool) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.sessionCompaction = sessionCompaction
|
|
}
|
|
}
|
|
|
|
func WithProviderRegistry(registry *provider.Registry) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if registry != nil {
|
|
r.providerRegistry = registry
|
|
}
|
|
}
|
|
}
|
|
|
|
func WithModelStore(store ModelStore) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.modelsStore = store
|
|
}
|
|
}
|
|
|
|
func WithSessionStore(store session.Store) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.sessionStore = store
|
|
}
|
|
}
|
|
|
|
// WithWorkingDir sets the working directory for hooks execution
|
|
func WithWorkingDir(dir string) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.workingDir = dir
|
|
}
|
|
}
|
|
|
|
// WithEnv sets the environment variables for hooks execution
|
|
func WithEnv(env []string) Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.env = env
|
|
}
|
|
}
|
|
|
|
// WithClock replaces the runtime's clock. Defaults to time.Now. Tests that
|
|
// need deterministic timestamps (assistant message CreatedAt, fallback
|
|
// cooldown windows, tool-call latency) can pass a fake clock so assertions
|
|
// don't depend on wall-clock advancement.
|
|
func WithClock(now func() time.Time) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if now != nil {
|
|
r.now = now
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithTelemetry replaces the runtime's Telemetry sink. Defaults to a
|
|
// pass-through to the package-level pkg/telemetry helpers. Tests pass a
|
|
// recorder to assert that the runtime emitted the expected lifecycle
|
|
// events without setting up an OTel client.
|
|
func WithTelemetry(t Telemetry) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if t != nil {
|
|
r.telemetry = t
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithMaxOverflowCompactions overrides how many consecutive context-overflow
|
|
// auto-compactions the run loop is allowed to attempt before surfacing the
|
|
// error. Defaults to defaultMaxOverflowCompactions (1).
|
|
//
|
|
// Tests use this to exercise both branches of the overflow-recovery code
|
|
// path: pass 0 to verify the failure surface immediately; pass a higher
|
|
// number to verify the loop bounds compaction attempts. Negative values
|
|
// are clamped to 0.
|
|
func WithMaxOverflowCompactions(n int) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if n < 0 {
|
|
n = 0
|
|
}
|
|
r.maxOverflowCompactions = n
|
|
}
|
|
}
|
|
|
|
// WithBudget sets the run-wide budget: the cost, token and wall-clock
|
|
// ceilings that stop a run once crossed, regardless of which agent spends.
|
|
// A nil or all-zero config leaves runs unbudgeted, which is the default.
|
|
func WithBudget(cfg *latest.BudgetConfig) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if cfg.IsZero() {
|
|
return
|
|
}
|
|
r.budgetCfg = cfg
|
|
}
|
|
}
|
|
|
|
// WithNamedBudgets registers the manifest's named budget definitions and
|
|
// the budget names each agent declared. A named budget referenced by
|
|
// several agents is one shared pot, not a copy per agent.
|
|
func WithNamedBudgets(defs map[string]latest.BudgetConfig, agentBudgets map[string][]string) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if len(defs) == 0 || len(agentBudgets) == 0 {
|
|
return
|
|
}
|
|
r.budgetsCfg = defs
|
|
r.agentBudgets = agentBudgets
|
|
}
|
|
}
|
|
|
|
// WithToolListTimeout overrides how long EmitStartupInfo waits for a single
|
|
// toolset to enumerate its tools before skipping it. Defaults to
|
|
// defaultToolListTimeout. A non-positive value is ignored so the default
|
|
// stands. Tests pass a short timeout to exercise the skip path (a toolset
|
|
// whose Tools() blocks) without a real-time wait.
|
|
func WithToolListTimeout(d time.Duration) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if d > 0 {
|
|
r.toolListTimeout = d
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithToolStartTimeout overrides how long EmitStartupInfo waits for a single
|
|
// toolset to start before skipping it. Defaults to defaultToolStartTimeout.
|
|
// A non-positive value is ignored so the default stands. Tests pass a short
|
|
// timeout to exercise the skip path (a toolset whose Start() blocks) without
|
|
// a real-time wait.
|
|
func WithToolStartTimeout(d time.Duration) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if d > 0 {
|
|
r.toolStartTimeout = d
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithRetryOnRateLimit enables automatic retry with backoff for HTTP 429 (rate limit)
|
|
// errors when no fallback models are available. When enabled, the runtime will honor
|
|
// the Retry-After header from the provider's response to determine wait time before
|
|
// retrying, falling back to exponential backoff if the header is absent.
|
|
//
|
|
// This is off by default. It is intended for library consumers that run agents
|
|
// programmatically and prefer to wait for rate limits to clear rather than fail
|
|
// immediately.
|
|
//
|
|
// When fallback models are configured, 429 errors always skip to the next model
|
|
// regardless of this setting.
|
|
func WithRetryOnRateLimit() Opt {
|
|
return func(r *LocalRuntime) {
|
|
r.fallback.retryOnRateLimit = true
|
|
}
|
|
}
|
|
|
|
// WithAutoInjector adds an [builtins.AutoInjector] that augments every
|
|
// per-agent hook configuration during executor build. The canonical
|
|
// use case is the snapshot controller returned by
|
|
// [builtins.RegisterSnapshot]: pass the same controller to the App via
|
|
// app.WithSnapshotController so /undo and friends drive the same
|
|
// instance that captures the checkpoints.
|
|
//
|
|
// Multiple calls accumulate; injectors run in registration order.
|
|
func WithAutoInjector(inj builtins.AutoInjector) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if inj != nil {
|
|
r.autoInjectors = append(r.autoInjectors, inj)
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithHooksRegistry plugs a pre-populated [hooks.Registry] into the
|
|
// runtime instead of letting it allocate a fresh one. Embedders use
|
|
// this to pre-register builtins they own (today snapshot, tomorrow
|
|
// any custom builtin) so the auto-injection chain set up by
|
|
// [WithAutoInjector] resolves against the same registry.
|
|
//
|
|
// The runtime continues to register its own stateless and
|
|
// closure-bound builtins (add_date, max_iterations, cache_response,
|
|
// unload, ...) on top of the supplied registry, so the embedder only
|
|
// needs to install entries that the runtime can't construct itself.
|
|
func WithHooksRegistry(reg *hooks.Registry) Opt {
|
|
return func(r *LocalRuntime) {
|
|
if reg != nil {
|
|
r.hooksRegistry = reg
|
|
}
|
|
}
|
|
}
|
|
|
|
// New creates a runtime ready to drive an agent loop. It is a thin
|
|
// alias for [NewLocalRuntime] returning the [Runtime] interface, kept
|
|
// for source compatibility with callers written before persistence
|
|
// became an [EventObserver]. Persistence is auto-registered against
|
|
// the configured (or default in-memory) session store; pass
|
|
// [WithSessionStore] to override and [WithEventObserver] to layer
|
|
// additional observers (telemetry, audit, ...).
|
|
func New(ctx context.Context, agents *team.Team, opts ...Opt) (Runtime, error) {
|
|
return NewLocalRuntime(ctx, agents, opts...)
|
|
}
|
|
|
|
// NewLocalRuntime creates a new LocalRuntime without the persistence wrapper.
|
|
// This is useful for testing or when persistence is handled externally.
|
|
func NewLocalRuntime(ctx context.Context, agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
|
|
defaultAgent, err := agents.DefaultAgent()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
r := &LocalRuntime{
|
|
ctx: func() context.Context { return context.WithoutCancel(ctx) },
|
|
toolMap: make(map[string]ToolHandlerFunc),
|
|
liveSessions: make(map[string]*liveSessionEntry),
|
|
team: agents,
|
|
agents: newAgentRouter(agents, defaultAgent.Name()),
|
|
resumeChan: make(chan ResumeRequest),
|
|
steerQueue: NewInMemoryMessageQueue(defaultSteerQueueCapacity),
|
|
followUpQueue: NewInMemoryMessageQueue(defaultFollowUpQueueCapacity),
|
|
sessionCompaction: true,
|
|
managedOAuth: true,
|
|
sessionStore: session.NewInMemorySessionStore(),
|
|
fallback: newFallbackExecutor(),
|
|
now: time.Now,
|
|
telemetry: defaultTelemetry{},
|
|
providerRegistry: provider.EmptyRegistry(),
|
|
maxOverflowCompactions: defaultMaxOverflowCompactions,
|
|
toolListTimeout: defaultToolListTimeout,
|
|
toolStartTimeout: defaultToolStartTimeout,
|
|
streamStoppedDeliveryTimeout: defaultStreamStoppedDeliveryTimeout,
|
|
dmrModelLister: dmrmodels.ListModelsWithMetadata,
|
|
}
|
|
r.bgAgents = agenttool.NewHandler(r)
|
|
r.fallback.prepareMessages = r.prepareMessagesForModel
|
|
|
|
// stripUnsupportedModalitiesTransform captures the runtime closure to
|
|
// resolve the agent from Input.AgentName, so it lives here rather
|
|
// than as a stateless builtin in pkg/hooks/builtins. It drops image
|
|
// content for text-only models on every model call.
|
|
//
|
|
// redact_secrets used to live here as a sibling [MessageTransform];
|
|
// it now ships entirely as a [hooks.BuiltinFunc] in
|
|
// pkg/hooks/builtins/redact_secrets.go and is wired into all three
|
|
// of tool_input_transform, before_llm_call, and tool_response_transform via
|
|
// [builtins.ApplyAgentDefaults] (or a user's hooks YAML directly),
|
|
// so the rewrite path is the same for every leak vector and there
|
|
// is no flag-only code path to keep in sync.
|
|
//
|
|
// Ordering matters: strip_generated_media MUST run before
|
|
// strip_unsupported_modalities. The latter strips any image/audio/
|
|
// video-kind document part the resolved model can't accept,
|
|
// regardless of whether that part is a runtime-materialized generated
|
|
// artifact or a user attachment — it has no placeholder logic. If it
|
|
// ran first, a capability-less or unknown model would have a
|
|
// media-only generated-media assistant message stripped down to
|
|
// nothing right there, and strip_generated_media would then see no
|
|
// generated-media part left to react to: its placeholder would never
|
|
// fire, and the turn would silently vanish from outgoing history.
|
|
// Running strip_generated_media first guarantees its placeholder text
|
|
// is already in place — as ordinary text, not a media part — by the
|
|
// time strip_unsupported_modalities runs, so there is nothing left for
|
|
// it to strip from that message.
|
|
r.transforms = append(r.transforms,
|
|
// strip_generated_media has no runtime state to capture (the policy
|
|
// is unconditional), so it registers the free function directly
|
|
// rather than a method value like the transform below.
|
|
registeredTransform{
|
|
name: BuiltinStripGeneratedMedia,
|
|
fn: stripGeneratedMediaTransform,
|
|
},
|
|
registeredTransform{
|
|
name: BuiltinStripUnsupportedModalities,
|
|
fn: r.stripUnsupportedModalitiesTransform,
|
|
},
|
|
)
|
|
|
|
for _, opt := range opts {
|
|
opt(r)
|
|
}
|
|
|
|
// Set up the hooks registry. Use the embedder-supplied registry
|
|
// (via [WithHooksRegistry]) when present so any builtins the
|
|
// embedder pre-registered — typically the snapshot builtin from
|
|
// [builtins.RegisterSnapshot] — are visible to the runtime, then
|
|
// register the runtime-owned builtins on top.
|
|
if r.hooksRegistry == nil {
|
|
r.hooksRegistry = hooks.NewRegistry()
|
|
}
|
|
if err := builtins.Register(r.hooksRegistry); err != nil {
|
|
return nil, fmt.Errorf("register builtin hooks: %w", err)
|
|
}
|
|
registerModelHook(r.hooksRegistry, r.providerRegistry)
|
|
r.hooksRegistry.Register(hooks.HookTypeEvaluator, hooks.NewEvaluatorFactory(func(agentName, name string) (evaluator.Evaluator, bool) {
|
|
a, err := r.team.Agent(agentName)
|
|
if err != nil {
|
|
return nil, false
|
|
}
|
|
var client evaluator.Evaluator
|
|
var ok bool
|
|
if a.HasEvaluatorScope() {
|
|
client, ok = a.Evaluator(name)
|
|
} else {
|
|
client, ok = r.team.Evaluator(name)
|
|
}
|
|
if !ok || client == nil {
|
|
return nil, false
|
|
}
|
|
return &accountedEvaluator{client: client, name: name}, true
|
|
}, hooks.WithRoutingDefaults(r.routingDefaultAgent)))
|
|
|
|
// cache_response is registered here (not in pkg/hooks/builtins)
|
|
// because it needs to capture the runtime to resolve the agent
|
|
// referenced by Input.AgentName. The other builtins are stateless
|
|
// and can stay as package-level functions registered via
|
|
// [builtins.Register] above.
|
|
if err := r.hooksRegistry.RegisterBuiltin(BuiltinCacheResponse, r.cacheResponseBuiltin); err != nil {
|
|
return nil, fmt.Errorf("register %q builtin: %w", BuiltinCacheResponse, err)
|
|
}
|
|
|
|
// Build cooldowns after opts so they pick up [WithClock].
|
|
r.fallback.cooldowns = newCooldownManager(r.now)
|
|
|
|
// Default the runtime's working directory to the process CWD when no
|
|
// caller supplied one. This matches the session's default and ensures
|
|
// builtin hooks that look up files (add_prompt_files) can find them
|
|
// without the embedder having to remember to call WithWorkingDir.
|
|
if r.workingDir == "" {
|
|
if cwd, err := os.Getwd(); err == nil {
|
|
r.workingDir = cwd
|
|
}
|
|
}
|
|
|
|
if r.modelsStore == nil {
|
|
// Precedence: an explicit WithModelStore (already set above) wins; then a
|
|
// store carried on the ModelSwitcherConfig (the team loader shares the
|
|
// one it warmed so the first /model open skips the cold catalog parse);
|
|
// otherwise a lazy store constructed on first use.
|
|
if r.modelSwitcherCfg != nil && r.modelSwitcherCfg.ModelsStore != nil {
|
|
r.modelsStore = r.modelSwitcherCfg.ModelsStore
|
|
} else {
|
|
r.modelsStore = &lazyModelStore{}
|
|
}
|
|
}
|
|
|
|
// Validate that the current agent exists and has a model
|
|
// (the router's current name might have been changed by WithCurrentAgent)
|
|
defaultAgent, err = r.team.Agent(r.agents.Name())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if defaultAgent.Model(ctx) == nil && !defaultAgent.HasHarness() {
|
|
return nil, fmt.Errorf("agent %s has no valid model", defaultAgent.Name())
|
|
}
|
|
|
|
// Register runtime-managed tool handlers once during construction.
|
|
// This avoids concurrent map writes when multiple goroutines call
|
|
// RunStream on the same runtime (e.g. background agent sessions).
|
|
r.registerDefaultTools()
|
|
|
|
// Pre-build per-agent hook executors now that workingDir, env and
|
|
// the team are finalized. Read-only afterwards.
|
|
r.buildHooksExecutors()
|
|
|
|
// Auto-register the stock persistence observer against the
|
|
// (possibly user-supplied) session store. It runs first in the
|
|
// observer chain so any user-supplied observers see the same view
|
|
// of the session that future RunStream calls and store reads will.
|
|
if obs := newPersistenceObserver(r.sessionStore); obs != nil {
|
|
r.observers = append([]EventObserver{obs}, r.observers...)
|
|
}
|
|
|
|
slog.DebugContext(ctx, "Creating new runtime", "agent", r.agents.Name(), "available_agents", agents.Size())
|
|
|
|
return r, nil
|
|
}
|
|
|
|
func (r *LocalRuntime) CurrentAgentName(context.Context) string {
|
|
return r.currentAgentName()
|
|
}
|
|
|
|
// currentAgentName is the context-free internal accessor. LocalRuntime
|
|
// resolves the active agent purely from in-memory state, so its internal
|
|
// callers (event callbacks, logging) use this directly; the exported
|
|
// context-taking method exists to satisfy the Runtime interface, whose
|
|
// remote implementation needs a context for its one-time lookup.
|
|
func (r *LocalRuntime) currentAgentName() string {
|
|
return r.agents.Name()
|
|
}
|
|
|
|
func (r *LocalRuntime) setCurrentAgent(name string) {
|
|
r.agents.Set(name)
|
|
}
|
|
|
|
func (r *LocalRuntime) CurrentAgentInfo(context.Context) CurrentAgentInfo {
|
|
currentAgent := r.CurrentAgent()
|
|
|
|
return CurrentAgentInfo{
|
|
Name: currentAgent.Name(),
|
|
Description: currentAgent.Description(),
|
|
Commands: currentAgent.Commands(),
|
|
}
|
|
}
|
|
|
|
func (r *LocalRuntime) SetCurrentAgent(_ context.Context, agentName string) error {
|
|
return r.agents.SetValidated(agentName)
|
|
}
|
|
|
|
func (r *LocalRuntime) CurrentAgentCommands(context.Context) types.Commands {
|
|
return r.CurrentAgent().Commands()
|
|
}
|
|
|
|
// CurrentAgentTools returns the tools available to the current agent.
|
|
// This starts toolsets if needed; returned handlers enforce this agent's skill
|
|
// policy without permitting embedded commands outside a tool call.
|
|
func (r *LocalRuntime) CurrentAgentTools(ctx context.Context) ([]tools.Tool, error) {
|
|
a := r.CurrentAgent()
|
|
agentTools, err := a.Tools(ctx)
|
|
if err != nil {
|
|
return agentTools, err
|
|
}
|
|
return bindSkillRuntime(agentTools, skillRuntime{runtime: r, agent: a}), nil
|
|
}
|
|
|
|
// ToolsetState is the coarse lifecycle bucket the agent inspector renders as a
|
|
// status glyph. It collapses the full lifecycle.State machine into the three
|
|
// distinctions a reader cares about: serving (started), not running (stopped),
|
|
// or broken (error).
|
|
type ToolsetState string
|
|
|
|
const (
|
|
ToolsetStarted ToolsetState = "started"
|
|
ToolsetStopped ToolsetState = "stopped"
|
|
ToolsetError ToolsetState = "error"
|
|
)
|
|
|
|
// ToolsetDetail describes one configured toolset for the agent inspector: its
|
|
// display name, kind, lifecycle bucket and the tools it exposes. Tools holds
|
|
// the live tool names when the toolset is started, otherwise the declared
|
|
// `tools:` allow-list from the retained config, and is empty when neither is
|
|
// available (a not-yet-started toolset with no explicit allow-list).
|
|
type ToolsetDetail struct {
|
|
Name string
|
|
Kind string
|
|
State ToolsetState
|
|
Tools []string
|
|
}
|
|
|
|
// AgentConfigInfo is the static-plus-live dataset behind the read-only agent
|
|
// inspector modal. The static parts (sub-agents, handoffs, fallbacks, skills,
|
|
// limits, option flags, declared toolset allow-lists) are derived from the
|
|
// resolved *agent.Agent and the retained config without starting any toolset;
|
|
// the live parts (per-toolset lifecycle state, started tool names, IsCurrent)
|
|
// reflect the running team. Remote runtimes (which hold no local team) return
|
|
// the zero value, so the modal omits every config-derived section.
|
|
type AgentConfigInfo struct {
|
|
SubAgents []string // sub-agent names, sorted
|
|
Handoffs []string // handoff target names, sorted
|
|
Fallbacks []string // fallback model ids ("provider/model"), in priority order
|
|
Skills []string // configured skill names (inline + included + used), sorted
|
|
|
|
MaxIterations int // 0 when unset
|
|
NumHistoryItems int // 0 when unset
|
|
MaxConsecutiveToolCalls int // 0 when unset
|
|
|
|
Options []string // enabled option flags (add-date, redact-secrets, ...)
|
|
Toolsets []ToolsetDetail // per-toolset live state + tools, in declaration order
|
|
|
|
IsCurrent bool // true when this is the live current agent
|
|
}
|
|
|
|
// AgentConfigInfo returns the named agent's inspector dataset. It inspects the
|
|
// resolved agent and the retained config, reading live tool names only from
|
|
// already-started toolsets (never starting one), so it is safe to call for any
|
|
// agent whether or not it has run. Unknown agents yield the zero value, so the
|
|
// modal omits the corresponding sections.
|
|
func (r *LocalRuntime) AgentConfigInfo(ctx context.Context, agentName string) AgentConfigInfo {
|
|
a, err := r.team.Agent(agentName)
|
|
if err != nil || a == nil {
|
|
return AgentConfigInfo{}
|
|
}
|
|
|
|
cfg, hasCfg := r.team.AgentConfig(agentName)
|
|
|
|
var subAgents []string
|
|
for _, sub := range a.SubAgents() {
|
|
if sub != nil {
|
|
subAgents = append(subAgents, sub.Name())
|
|
}
|
|
}
|
|
|
|
var handoffs []string
|
|
for _, h := range a.Handoffs() {
|
|
if h != nil {
|
|
handoffs = append(handoffs, h.Name())
|
|
}
|
|
}
|
|
|
|
var fallbacks []string
|
|
for _, p := range a.FallbackModels() {
|
|
if p != nil {
|
|
fallbacks = append(fallbacks, p.ID().String())
|
|
}
|
|
}
|
|
|
|
info := AgentConfigInfo{
|
|
SubAgents: uniqueNames(subAgents, true),
|
|
Handoffs: uniqueNames(handoffs, true),
|
|
Fallbacks: uniqueNames(fallbacks, false),
|
|
MaxIterations: a.MaxIterations(),
|
|
NumHistoryItems: a.NumHistoryItems(),
|
|
MaxConsecutiveToolCalls: a.MaxConsecutiveToolCalls(),
|
|
Options: agentOptionFlags(a, cfg, hasCfg),
|
|
Toolsets: toolsetDetails(ctx, a, cfg, hasCfg),
|
|
IsCurrent: r.agents != nil && r.agents.Name() == agentName,
|
|
}
|
|
if hasCfg {
|
|
info.Skills = configSkillNames(cfg)
|
|
}
|
|
return info
|
|
}
|
|
|
|
// agentOptionFlags lists the agent's enabled boolean options as stable,
|
|
// hyphenated display names. AddDate/AddEnvironmentInfo/RedactSecrets read the
|
|
// agent's effective (resolved) values; CodeModeTools is only knowable from the
|
|
// retained config because the runtime folds it into a single wrapper toolset.
|
|
func agentOptionFlags(a *agent.Agent, cfg latest.AgentConfig, hasCfg bool) []string {
|
|
var opts []string
|
|
if a.AddDate() {
|
|
opts = append(opts, "add-date")
|
|
}
|
|
if a.AddEnvironmentInfo() {
|
|
opts = append(opts, "add-environment-info")
|
|
}
|
|
if a.RedactSecrets() {
|
|
opts = append(opts, "redact-secrets")
|
|
}
|
|
if hasCfg && cfg.CodeModeTools {
|
|
opts = append(opts, "code-mode-tools")
|
|
}
|
|
return opts
|
|
}
|
|
|
|
// configSkillNames lists the cleanly-resolvable skill names from the agent
|
|
// config: inline skill names, the Include allow-list, and referenced skill
|
|
// groups (UseSkills). Skills auto-discovered from Sources (e.g. a local
|
|
// directory) are not enumerated here because doing so would require loading
|
|
// them from disk; the inspector notes this rather than starting that work.
|
|
func configSkillNames(cfg latest.AgentConfig) []string {
|
|
var names []string
|
|
for _, s := range cfg.Skills.Inline {
|
|
names = append(names, s.Name)
|
|
}
|
|
names = append(names, cfg.Skills.Include...)
|
|
names = append(names, cfg.UseSkills...)
|
|
return uniqueNames(names, true)
|
|
}
|
|
|
|
// toolsetDetails builds one ToolsetDetail per live toolset of the agent,
|
|
// combining the side-effect-free lifecycle status with tool names: live names
|
|
// for started toolsets, otherwise the declared `tools:` allow-list keyed by the
|
|
// same name the registry assigns (cmp.Or(name, type)).
|
|
func toolsetDetails(ctx context.Context, a *agent.Agent, cfg latest.AgentConfig, hasCfg bool) []ToolsetDetail {
|
|
toolSets := a.ToolSets()
|
|
if len(toolSets) == 0 {
|
|
return nil
|
|
}
|
|
var declared map[string][]string
|
|
if hasCfg {
|
|
declared = declaredToolNames(cfg)
|
|
}
|
|
infos := make([]ToolsetDetail, 0, len(toolSets))
|
|
for _, ts := range toolSets {
|
|
status := toolsetStatusFor(ts)
|
|
info := ToolsetDetail{
|
|
Name: status.Name,
|
|
Kind: status.Kind,
|
|
State: toolsetStateBucket(status.State),
|
|
}
|
|
if live, ok := startedToolNames(ctx, ts); ok {
|
|
info.Tools = live
|
|
} else if names, ok := declared[status.Name]; ok {
|
|
info.Tools = names
|
|
}
|
|
infos = append(infos, info)
|
|
}
|
|
return infos
|
|
}
|
|
|
|
// toolsetStateBucket collapses the lifecycle state machine into the inspector's
|
|
// three buckets. Failed is the only error state; Stopped (including a
|
|
// not-yet-started toolset) is stopped; everything else (Ready, Degraded,
|
|
// Starting, Restarting) reads as started/serving.
|
|
func toolsetStateBucket(s lifecycle.State) ToolsetState {
|
|
switch s {
|
|
case lifecycle.StateFailed:
|
|
return ToolsetError
|
|
case lifecycle.StateStopped:
|
|
return ToolsetStopped
|
|
default:
|
|
return ToolsetStarted
|
|
}
|
|
}
|
|
|
|
// startedToolNames returns the live tool names of ts, but only when it is a
|
|
// started toolset; the boolean is false for a not-yet-started toolset, or one
|
|
// whose Start is still in flight, so the caller can fall back to the declared
|
|
// allow-list without blocking on a long-running Start (e.g. RAG indexing).
|
|
// context.TODO is safe here: the toolset is already started, so listing
|
|
// returns its cached tools without a cancellable round-trip.
|
|
func startedToolNames(ctx context.Context, ts tools.ToolSet) ([]string, bool) {
|
|
s, ok := tools.As[*tools.StartableToolSet](ts)
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
if started, _ := s.TryState(); !started {
|
|
return nil, false
|
|
}
|
|
tl, err := s.Tools(ctx)
|
|
if err != nil {
|
|
return nil, false
|
|
}
|
|
names := make([]string, 0, len(tl))
|
|
for i := range tl {
|
|
if tl[i].Name != "" {
|
|
names = append(names, tl[i].Name)
|
|
}
|
|
}
|
|
return names, true
|
|
}
|
|
|
|
// declaredToolNames maps each configured toolset's display name to its declared
|
|
// `tools:` allow-list. The key matches the registry's naming
|
|
// (cmp.Or(name, type)) so it lines up with the live toolset's Name. Toolsets
|
|
// with no explicit allow-list (serve all tools) are omitted.
|
|
func declaredToolNames(cfg latest.AgentConfig) map[string][]string {
|
|
m := make(map[string][]string, len(cfg.Toolsets))
|
|
for i := range cfg.Toolsets {
|
|
t := cfg.Toolsets[i]
|
|
key := cmp.Or(t.Name, t.Type)
|
|
if key == "" || len(t.Tools) == 0 {
|
|
continue
|
|
}
|
|
m[key] = slices.Clone(t.Tools)
|
|
}
|
|
return m
|
|
}
|
|
|
|
// uniqueNames drops empty entries and de-duplicates names. By default it keeps
|
|
// the first-seen order (meaningful for declaration/priority order); when sorted
|
|
// is true it orders the result case-insensitively instead.
|
|
func uniqueNames(names []string, sorted bool) []string {
|
|
seen := make(map[string]struct{}, len(names))
|
|
out := make([]string, 0, len(names))
|
|
for _, name := range names {
|
|
if name == "" {
|
|
continue
|
|
}
|
|
if _, dup := seen[name]; dup {
|
|
continue
|
|
}
|
|
seen[name] = struct{}{}
|
|
out = append(out, name)
|
|
}
|
|
if sorted {
|
|
slices.SortFunc(out, func(a, b string) int {
|
|
return strings.Compare(strings.ToLower(a), strings.ToLower(b))
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
// CurrentAgentToolsetStatuses returns one ToolsetStatus per toolset of the
|
|
// active agent. The list is in declaration order. Toolsets that wrap
|
|
// another (StartableToolSet, Multiplexer) are unwrapped so the inner
|
|
// supervisor's state is visible.
|
|
func (r *LocalRuntime) CurrentAgentToolsetStatuses() []tools.ToolsetStatus {
|
|
a := r.CurrentAgent()
|
|
if a == nil {
|
|
return nil
|
|
}
|
|
toolSets := a.ToolSets()
|
|
statuses := make([]tools.ToolsetStatus, 0, len(toolSets))
|
|
for _, ts := range toolSets {
|
|
statuses = append(statuses, toolsetStatusFor(ts))
|
|
}
|
|
return statuses
|
|
}
|
|
|
|
// AgentToolsetStatuses returns one ToolsetStatus per toolset of the named
|
|
// agent, in declaration order. It mirrors CurrentAgentToolsetStatuses but for
|
|
// an arbitrary agent and is side-effect-free: it reads each toolset's live
|
|
// lifecycle (state, kind, last error, restart count) without starting it.
|
|
// Unknown agents yield nil, so callers degrade gracefully.
|
|
func (r *LocalRuntime) AgentToolsetStatuses(name string) []tools.ToolsetStatus {
|
|
a, err := r.team.Agent(name)
|
|
if err != nil || a == nil {
|
|
return nil
|
|
}
|
|
toolSets := a.ToolSets()
|
|
statuses := make([]tools.ToolsetStatus, 0, len(toolSets))
|
|
for _, ts := range toolSets {
|
|
statuses = append(statuses, toolsetStatusFor(ts))
|
|
}
|
|
return statuses
|
|
}
|
|
|
|
// RestartToolset locates the named toolset on the active agent and
|
|
// asks its canonical lifecycle wrapper to restart it in place. The wrapper
|
|
// serializes the operation with start/stop and keeps its lifecycle state in
|
|
// sync with the underlying supervisor.
|
|
//
|
|
// Returns an error when:
|
|
// - no toolset matches name (matching uses the same logic as the
|
|
// /tools dialog: the toolset's Name() if any, otherwise its
|
|
// description),
|
|
// - no matching toolset supports restart,
|
|
// - the supervisor itself returned an error (timeout, classified
|
|
// transport failure, etc.).
|
|
//
|
|
// Display names are not guaranteed unique (a YAML name: can shadow a
|
|
// built-in), so non-restartable matches are skipped rather than aborting:
|
|
// the first restartable toolset with the given name wins.
|
|
func (r *LocalRuntime) RestartToolset(ctx context.Context, name string) error {
|
|
a := r.CurrentAgent()
|
|
if a == nil {
|
|
return errors.New("no active agent")
|
|
}
|
|
found := false
|
|
for _, ts := range a.ToolSets() {
|
|
if nameFor(ts, tools.DescribeToolSet(ts)) != name {
|
|
continue
|
|
}
|
|
found = true
|
|
startable, ok := ts.(*tools.StartableToolSet)
|
|
if ok && startable.CanRestart() {
|
|
return startable.RestartIfSupported(ctx)
|
|
}
|
|
}
|
|
if found {
|
|
return fmt.Errorf("toolset %q does not support restart", name)
|
|
}
|
|
return fmt.Errorf("toolset %q not found", name)
|
|
}
|
|
|
|
// toolsetStatusFor builds a ToolsetStatus for ts. tools.As walks the
|
|
// wrapper chain so Statable/Describer can live anywhere in the stack.
|
|
func toolsetStatusFor(ts tools.ToolSet) tools.ToolsetStatus {
|
|
status := tools.ToolsetStatus{
|
|
Description: tools.DescribeToolSet(ts),
|
|
}
|
|
if kinder, ok := tools.As[tools.Kinder](ts); ok {
|
|
status.Kind = kinder.Kind()
|
|
}
|
|
if statable, ok := tools.As[tools.Statable](ts); ok {
|
|
info := statable.State()
|
|
status.State = info.State
|
|
status.LastError = info.LastError
|
|
status.RestartCount = info.RestartCount
|
|
} else {
|
|
// Toolsets without a supervisor are considered ready by default;
|
|
// the StartableToolSet wrapper would have surfaced an error
|
|
// earlier if Start failed.
|
|
status.State = lifecycleStateForUnsupervised(ts)
|
|
}
|
|
if startable, ok := ts.(*tools.StartableToolSet); ok {
|
|
status.Restartable = startable.CanRestart()
|
|
} else {
|
|
_, status.Restartable = tools.As[tools.Restartable](ts)
|
|
}
|
|
status.Name = nameFor(ts, status.Description)
|
|
return status
|
|
}
|
|
|
|
// lifecycleStateForUnsupervised reports the lifecycle state of a toolset with
|
|
// no Statable supervisor, using the wrapper's own non-blocking status probe
|
|
// so a toolset whose Start is legitimately still running (e.g. RAG indexing
|
|
// a large knowledge base) never stalls a status/introspection caller.
|
|
func lifecycleStateForUnsupervised(ts tools.ToolSet) lifecycle.State {
|
|
s, ok := ts.(*tools.StartableToolSet)
|
|
if !ok {
|
|
return lifecycle.StateReady
|
|
}
|
|
started, inFlight := s.TryState()
|
|
switch {
|
|
case inFlight:
|
|
return lifecycle.StateStarting
|
|
case !started:
|
|
return lifecycle.StateStopped
|
|
default:
|
|
return lifecycle.StateReady
|
|
}
|
|
}
|
|
|
|
// nameFor picks a stable, user-visible name for a toolset. We look for
|
|
// any inner toolset that implements tools.Named (walked via tools.As so
|
|
// wrappers like StartableToolSet are transparent). The registry adds a
|
|
// WithName wrapper for every built-in toolset so this is reachable for
|
|
// almost every toolset; fallback uses the description ("mcp(stdio cmd=...)"),
|
|
// which is still better than the Go type name.
|
|
func nameFor(ts tools.ToolSet, fallback string) string {
|
|
if name := tools.GetName(ts); name != "" {
|
|
return name
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
// mcpPromptToolset is the subset of the MCP toolset's API the runtime needs
|
|
// for prompt discovery and execution. It is an interface (satisfied by
|
|
// *mcp.Toolset) so pkg/runtime does not import the MCP toolset package,
|
|
// which would link its whole dependency tree into every embedder.
|
|
type mcpPromptToolset interface {
|
|
ListPrompts(ctx context.Context) ([]tools.PromptInfo, error)
|
|
GetPrompt(ctx context.Context, name string, arguments map[string]string) (*mcp.GetPromptResult, error)
|
|
}
|
|
|
|
// CurrentMCPPrompts returns the available MCP prompts from all active MCP toolsets
|
|
// for the current agent. It discovers prompts by calling ListPrompts on each MCP toolset
|
|
// and aggregates the results into a map keyed by prompt name.
|
|
func (r *LocalRuntime) CurrentMCPPrompts(ctx context.Context) map[string]tools.PromptInfo {
|
|
prompts := make(map[string]tools.PromptInfo)
|
|
|
|
// Get the current agent to access its toolsets
|
|
currentAgent := r.CurrentAgent()
|
|
if currentAgent == nil {
|
|
slog.WarnContext(ctx, "No current agent available for MCP prompt discovery")
|
|
return prompts
|
|
}
|
|
|
|
// Iterate through all toolsets of the current agent
|
|
for _, toolset := range currentAgent.ToolSets() {
|
|
mcpToolsets := tools.FindAll[mcpPromptToolset](toolset)
|
|
if len(mcpToolsets) == 0 {
|
|
slog.DebugContext(ctx, "Toolset contains no MCP prompt provider", "type", fmt.Sprintf("%T", toolset))
|
|
continue
|
|
}
|
|
for _, mcpToolset := range mcpToolsets {
|
|
slog.DebugContext(ctx, "Found MCP toolset", "toolset", mcpToolset)
|
|
maps.Copy(prompts, r.discoverMCPPrompts(ctx, mcpToolset))
|
|
}
|
|
}
|
|
|
|
slog.DebugContext(ctx, "Discovered MCP prompts", "agent", currentAgent.Name(), "prompt_count", len(prompts))
|
|
return prompts
|
|
}
|
|
|
|
// discoverMCPPrompts queries an MCP toolset for available prompts and converts them
|
|
// to PromptInfo structures. This method handles the MCP protocol communication
|
|
// and gracefully handles any errors during prompt discovery.
|
|
func (r *LocalRuntime) discoverMCPPrompts(ctx context.Context, toolset mcpPromptToolset) map[string]tools.PromptInfo {
|
|
mcpPrompts, err := toolset.ListPrompts(ctx)
|
|
if err != nil {
|
|
slog.WarnContext(ctx, "Failed to list MCP prompts from toolset", "error", err)
|
|
return nil
|
|
}
|
|
|
|
prompts := make(map[string]tools.PromptInfo, len(mcpPrompts))
|
|
for _, mcpPrompt := range mcpPrompts {
|
|
promptInfo := tools.PromptInfo{
|
|
Name: mcpPrompt.Name,
|
|
Description: mcpPrompt.Description,
|
|
Arguments: make([]tools.PromptArgument, 0, len(mcpPrompt.Arguments)),
|
|
}
|
|
|
|
for _, arg := range mcpPrompt.Arguments {
|
|
promptInfo.Arguments = append(promptInfo.Arguments, tools.PromptArgument{
|
|
Name: arg.Name,
|
|
Description: arg.Description,
|
|
Required: arg.Required,
|
|
})
|
|
}
|
|
|
|
prompts[mcpPrompt.Name] = promptInfo
|
|
slog.DebugContext(ctx, "Discovered MCP prompt", "name", mcpPrompt.Name, "args_count", len(promptInfo.Arguments))
|
|
}
|
|
|
|
return prompts
|
|
}
|
|
|
|
// CurrentAgent returns the current agent
|
|
func (r *LocalRuntime) CurrentAgent() *agent.Agent {
|
|
return r.agents.Current()
|
|
}
|
|
|
|
// resolveSessionAgent returns the agent for the given session. Delegates to
|
|
// agentRouter.ResolveSession; kept on LocalRuntime for the existing callsites
|
|
// in loop.go and elsewhere.
|
|
func (r *LocalRuntime) resolveSessionAgent(sess *session.Session) *agent.Agent {
|
|
return r.agents.ResolveSession(sess)
|
|
}
|
|
|
|
// CurrentAgentSkillsToolset returns the skills toolset for the current agent, or nil if not enabled.
|
|
func (r *LocalRuntime) CurrentAgentSkillsToolset() *skills.ToolSet {
|
|
return agentSkillsToolset(r.CurrentAgent())
|
|
}
|
|
|
|
// agentSkillsToolset returns the skills toolset configured on a, or nil if
|
|
// a is nil or has none. Session-aware callers (e.g. RunSkillFork) use it
|
|
// directly so a pinned session's agent is honoured instead of the shared
|
|
// current agent.
|
|
func agentSkillsToolset(a *agent.Agent) *skills.ToolSet {
|
|
if a == nil {
|
|
return nil
|
|
}
|
|
for _, ts := range a.ToolSets() {
|
|
if st, ok := tools.Find[*skills.ToolSet](ts); ok {
|
|
return st
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ExecuteMCPPrompt executes an MCP prompt with provided arguments and returns the content.
|
|
func (r *LocalRuntime) ExecuteMCPPrompt(ctx context.Context, promptName string, arguments map[string]string) (string, error) {
|
|
currentAgent := r.CurrentAgent()
|
|
if currentAgent == nil {
|
|
return "", errors.New("no current agent available")
|
|
}
|
|
|
|
for _, toolset := range currentAgent.ToolSets() {
|
|
for _, mcpToolset := range tools.FindAll[mcpPromptToolset](toolset) {
|
|
result, err := mcpToolset.GetPrompt(ctx, promptName, arguments)
|
|
if err != nil {
|
|
// If error is "prompt not found", continue to next toolset
|
|
if err.Error() == "prompt not found" {
|
|
continue
|
|
}
|
|
return "", fmt.Errorf("error executing prompt '%s': %w", promptName, err)
|
|
}
|
|
|
|
// Convert the MCP result to a string format
|
|
if len(result.Messages) == 0 {
|
|
return "No content returned from MCP prompt", nil
|
|
}
|
|
|
|
var content strings.Builder
|
|
for i, message := range result.Messages {
|
|
if i > 0 {
|
|
content.WriteString("\n\n")
|
|
}
|
|
if textContent, ok := message.Content.(*mcp.TextContent); ok {
|
|
content.WriteString(textContent.Text)
|
|
} else {
|
|
fmt.Fprintf(&content, "[Non-text content: %T]", message.Content)
|
|
}
|
|
}
|
|
return content.String(), nil
|
|
}
|
|
}
|
|
|
|
return "", fmt.Errorf("MCP prompt '%s' not found in any active toolset", promptName)
|
|
}
|
|
|
|
// TitleGenerator returns a title generator for automatic session title
|
|
// generation, or nil when no configured model is a usable title candidate
|
|
// (see [sessiontitle.New]).
|
|
func (r *LocalRuntime) TitleGenerator(ctx context.Context) *sessiontitle.Generator {
|
|
a := r.CurrentAgent()
|
|
if a == nil {
|
|
return nil
|
|
}
|
|
return sessiontitle.New(a.TitleModels(ctx)...)
|
|
}
|
|
|
|
// getAgentModelID returns the model ID for an agent. The zero ID is
|
|
// returned when no model is configured.
|
|
func getAgentModelID(ctx context.Context, a *agent.Agent) modelsdev.ID {
|
|
if a == nil {
|
|
return modelsdev.ID{}
|
|
}
|
|
if model := a.Model(ctx); model != nil {
|
|
return model.ID()
|
|
}
|
|
return modelsdev.ID{}
|
|
}
|
|
|
|
// getEffectiveModelID returns the currently active model ID for an agent, accounting
|
|
// for any active fallback cooldown. During a cooldown period, this returns the fallback
|
|
// model ID instead of the configured primary model, so the UI reflects the actual model in use.
|
|
func (r *LocalRuntime) getEffectiveModelID(ctx context.Context, a *agent.Agent) modelsdev.ID {
|
|
cooldownState := r.fallback.cooldowns.Get(a.Name())
|
|
if cooldownState != nil {
|
|
fallbacks := a.FallbackModels()
|
|
if cooldownState.fallbackIndex >= 0 && cooldownState.fallbackIndex < len(fallbacks) {
|
|
return fallbacks[cooldownState.fallbackIndex].ID()
|
|
}
|
|
}
|
|
return getAgentModelID(ctx, a)
|
|
}
|
|
|
|
// agentDetailsFromTeam converts team agent info to AgentDetails for events.
|
|
// It accounts for active fallback cooldowns, returning the effective model
|
|
// instead of the configured model when a fallback is in effect.
|
|
func (r *LocalRuntime) agentDetailsFromTeam(ctx context.Context) []AgentDetails {
|
|
agentsInfo := r.team.AgentsInfo(ctx)
|
|
details := make([]AgentDetails, len(agentsInfo))
|
|
for i, info := range agentsInfo {
|
|
providerName := info.Provider
|
|
modelName := info.Model
|
|
var thinking string
|
|
|
|
// Get the agent to access fallbacks and the effective thinking level.
|
|
if a, err := r.team.Agent(info.Name); err == nil && a != nil {
|
|
// Check if this agent has an active fallback cooldown
|
|
cooldownState := r.fallback.cooldowns.Get(info.Name)
|
|
if cooldownState != nil {
|
|
fallbacks := a.FallbackModels()
|
|
if cooldownState.fallbackIndex >= 0 && cooldownState.fallbackIndex < len(fallbacks) {
|
|
fb := fallbacks[cooldownState.fallbackIndex].ID()
|
|
providerName = fb.Provider
|
|
modelName = fb.Model
|
|
}
|
|
}
|
|
thinking = r.agentThinkingLabel(ctx, a)
|
|
}
|
|
|
|
details[i] = AgentDetails{
|
|
Name: info.Name,
|
|
Description: info.Description,
|
|
Provider: providerName,
|
|
Model: modelName,
|
|
Thinking: thinking,
|
|
Commands: info.Commands,
|
|
}
|
|
}
|
|
return details
|
|
}
|
|
|
|
// agentThinkingLabel returns a short, user-facing label for the effective
|
|
// thinking-effort level of the agent's current model: the effort level (e.g.
|
|
// "high"), "adaptive" for adaptive budgets, the decimal token count for
|
|
// token-based budgets, "off" when thinking is disabled on a reasoning-capable
|
|
// model, or "" when the model has no selectable thinking configuration to
|
|
// display.
|
|
func (r *LocalRuntime) agentThinkingLabel(ctx context.Context, a *agent.Agent) string {
|
|
models := a.EffectiveModels()
|
|
if len(models) == 0 {
|
|
return ""
|
|
}
|
|
cfg := models[0].BaseConfig().ModelConfig
|
|
// Only models that can actually reason get a thinking line.
|
|
if !r.modelSupportsThinking(ctx, &cfg) {
|
|
return ""
|
|
}
|
|
budget := cfg.ThinkingBudget
|
|
if budget == nil || budget.IsDisabled() {
|
|
return "off"
|
|
}
|
|
if l, ok := budget.EffortLevel(); ok {
|
|
return l.String()
|
|
}
|
|
if budget.IsAdaptive() {
|
|
return "adaptive"
|
|
}
|
|
return strconv.Itoa(budget.Tokens) // token-based budget
|
|
}
|
|
|
|
// SessionStore returns the session store for browsing/loading past sessions.
|
|
func (r *LocalRuntime) SessionStore() session.Store {
|
|
return r.sessionStore
|
|
}
|
|
|
|
// Close releases resources held by the runtime. The session store is
|
|
// *not* closed here: it is provided by the embedder via
|
|
// WithSessionStore and may be shared with other runtimes (e.g. the
|
|
// TUI's session spawner reuses the same store across spawned
|
|
// sessions, so closing it here would break later session lookups).
|
|
// Embedders that own the store are responsible for closing it once
|
|
// when their process is shutting down.
|
|
func (r *LocalRuntime) Close() error {
|
|
r.bgAgents.StopAll()
|
|
r.OnToolsChanged(nil)
|
|
return nil
|
|
}
|
|
|
|
// UpdateSessionTitle persists the session title via the session store.
|
|
func (r *LocalRuntime) UpdateSessionTitle(ctx context.Context, sess *session.Session, title string) error {
|
|
sess.SetTitle(title)
|
|
if r.sessionStore != nil {
|
|
return r.sessionStore.UpdateSession(ctx, sess)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// PermissionsInfo returns the team-level permission patterns.
|
|
// Returns nil if no permissions are configured.
|
|
func (r *LocalRuntime) PermissionsInfo() *PermissionsInfo {
|
|
permChecker := r.team.Permissions()
|
|
if permChecker == nil || permChecker.IsEmpty() {
|
|
return nil
|
|
}
|
|
return &PermissionsInfo{
|
|
Allow: permChecker.AllowPatterns(),
|
|
Ask: permChecker.AskPatterns(),
|
|
Deny: permChecker.DenyPatterns(),
|
|
}
|
|
}
|
|
|
|
// ResetStartupInfo resets the startup info emission flag.
|
|
// This should be called when replacing a session to allow re-emission of
|
|
// agent, team, and toolset info to the UI.
|
|
func (r *LocalRuntime) ResetStartupInfo() {
|
|
r.startupInfoEmitted.Store(false)
|
|
}
|
|
|
|
// OnToolsChanged registers a handler that is called when an MCP toolset
|
|
// reports a tool list change outside of a RunStream. This allows the UI
|
|
// to update the tool count immediately. The runtime subscribes to the
|
|
// team's toolsets on behalf of the handler, replacing only its own previous
|
|
// subscriptions so runtimes sharing a toolset do not displace each other;
|
|
// a nil handler unsubscribes.
|
|
func (r *LocalRuntime) OnToolsChanged(handler func(Event)) {
|
|
r.toolsChangedMu.Lock()
|
|
r.onToolsChanged = handler
|
|
previous := r.toolsChangedUnsubs
|
|
r.toolsChangedUnsubs = nil
|
|
if handler != nil {
|
|
r.toolsChangedUnsubs = r.subscribeToolsChanged()
|
|
}
|
|
r.toolsChangedMu.Unlock()
|
|
|
|
for _, unsub := range previous {
|
|
unsub()
|
|
}
|
|
}
|
|
|
|
// emitToolsChanged is the callback registered on MCP toolsets. It re-reads
|
|
// the current agent's full tool list and pushes a ToolsetInfo event.
|
|
func (r *LocalRuntime) emitToolsChanged() {
|
|
r.toolsChangedMu.RLock()
|
|
handler := r.onToolsChanged
|
|
r.toolsChangedMu.RUnlock()
|
|
if handler == nil {
|
|
return
|
|
}
|
|
ctx, cancel := context.WithTimeout(r.ctx(), toolsChangedTimeout)
|
|
defer cancel()
|
|
a := r.CurrentAgent()
|
|
agentTools, err := a.StartedTools(ctx)
|
|
if err != nil {
|
|
return
|
|
}
|
|
handler(ToolsetInfo(len(agentTools), false, r.currentAgentName()))
|
|
}
|
|
|
|
// OnBackgroundEvent registers a handler that receives events surfaced from
|
|
// detached background work — today the TokenUsageEvents of background agent
|
|
// tasks (run_background_agent), whose sub-sessions run on their own goroutine
|
|
// with no live event channel. This lets the UI keep per-agent context
|
|
// accounting for background agents.
|
|
func (r *LocalRuntime) OnBackgroundEvent(handler func(Event)) {
|
|
r.backgroundEventMu.Lock()
|
|
defer r.backgroundEventMu.Unlock()
|
|
r.onBackgroundEvent = handler
|
|
r.onBackgroundContext = nil
|
|
}
|
|
|
|
// OnBackgroundEventWithContext preserves detached producer ownership.
|
|
func (r *LocalRuntime) OnBackgroundEventWithContext(handler func(context.Context, Event)) {
|
|
r.backgroundEventMu.Lock()
|
|
defer r.backgroundEventMu.Unlock()
|
|
r.onBackgroundContext = handler
|
|
r.onBackgroundEvent = nil
|
|
}
|
|
|
|
// emitBackgroundEvent forwards an event from detached background work to the
|
|
// registered handler, if any.
|
|
func (r *LocalRuntime) emitBackgroundEvent(ctx context.Context, event Event) {
|
|
r.backgroundEventMu.RLock()
|
|
handler := r.onBackgroundEvent
|
|
contextual := r.onBackgroundContext
|
|
r.backgroundEventMu.RUnlock()
|
|
if contextual != nil {
|
|
contextual(ctx, event)
|
|
return
|
|
}
|
|
if handler != nil {
|
|
handler(event)
|
|
}
|
|
}
|
|
|
|
// emitAgentAndTeamInfo sends the AgentInfo and TeamInfo events that drive the
|
|
// sidebar's agent/model/thinking display. It returns false when sending was
|
|
// aborted (e.g. the context was cancelled). Shared by EmitStartupInfo and
|
|
// EmitAgentInfo so both render identical model labels.
|
|
func (r *LocalRuntime) emitAgentAndTeamInfo(ctx context.Context, a *agent.Agent, send func(Event) bool) bool {
|
|
modelID := r.getEffectiveModelID(ctx, a)
|
|
modelLabel := modelID.String()
|
|
contextLimit := r.contextLimitForAgentModel(ctx, a, modelID)
|
|
var compactionModel string
|
|
var primaryLimit int64
|
|
if a.CompactionModel() != nil {
|
|
var compactionLimit int64
|
|
compactionModel, primaryLimit, compactionLimit = r.compactionCapAttribution(ctx, a, modelID)
|
|
if !compactionCaps(primaryLimit, compactionLimit) {
|
|
compactionModel, primaryLimit = "", 0
|
|
}
|
|
}
|
|
if a.HasHarness() {
|
|
modelLabel = agentModelLabel(ctx, a)
|
|
contextLimit = 0
|
|
compactionModel, primaryLimit = "", 0
|
|
}
|
|
info := AgentInfo(a.Name(), modelLabel, a.Description(), a.WelcomeMessage(), contextLimit).(*AgentInfoEvent)
|
|
info.CompactionModel = compactionModel
|
|
info.PrimaryContextLimit = primaryLimit
|
|
if !send(info) {
|
|
return false
|
|
}
|
|
return send(TeamInfo(r.agentDetailsFromTeam(ctx), r.currentAgentName()))
|
|
}
|
|
|
|
func (r *LocalRuntime) contextLimitForAgentModel(ctx context.Context, a *agent.Agent, modelID modelsdev.ID) int64 {
|
|
if a == nil || modelID.IsZero() {
|
|
return 0
|
|
}
|
|
return r.effectiveContextLimit(ctx, a, r.resolveContextLimit(ctx, a.Model(ctx), modelID))
|
|
}
|
|
|
|
// EmitAgentInfo implements [Runtime.EmitAgentInfo]: it refreshes the agent and
|
|
// team display without touching toolset discovery.
|
|
func (r *LocalRuntime) EmitAgentInfo(ctx context.Context, events EventSink) {
|
|
a := r.CurrentAgent()
|
|
if a == nil {
|
|
return
|
|
}
|
|
r.emitAgentAndTeamInfo(ctx, a, func(event Event) bool {
|
|
if ctx.Err() != nil {
|
|
return false
|
|
}
|
|
events.Emit(event)
|
|
return true
|
|
})
|
|
}
|
|
|
|
// EmitStartupInfo emits initial agent, team, and toolset information for immediate sidebar display.
|
|
// When sess is non-nil and contains token data, a TokenUsageEvent is also emitted so that the
|
|
// sidebar can display context usage percentage on session restore.
|
|
func (r *LocalRuntime) EmitStartupInfo(ctx context.Context, sess *session.Session, events EventSink) {
|
|
// CompareAndSwap makes the check-and-set atomic: App.Start emits from a
|
|
// spawned goroutine while reEmitStartupInfo (e.g. on /new session) resets
|
|
// and re-emits from another, and a plain bool check-then-set race would
|
|
// let both goroutines through or miss an emission.
|
|
if !r.startupInfoEmitted.CompareAndSwap(false, true) {
|
|
return
|
|
}
|
|
|
|
a := r.CurrentAgent()
|
|
|
|
// Helper to send events with context check
|
|
send := func(event Event) bool {
|
|
if ctx.Err() != nil {
|
|
return false
|
|
}
|
|
events.Emit(event)
|
|
return true
|
|
}
|
|
|
|
// Emit agent and team information immediately for fast sidebar display.
|
|
// Use getEffectiveModelID to account for active fallback cooldowns.
|
|
modelID := r.getEffectiveModelID(ctx, a)
|
|
if !r.emitAgentAndTeamInfo(ctx, a, send) {
|
|
return
|
|
}
|
|
|
|
// When restoring a session that already has token data, emit a
|
|
// TokenUsageEvent so the sidebar can show the context usage percentage.
|
|
// The context limit comes from the model definition (models.dev), which
|
|
// is a model property — not persisted in the session.
|
|
//
|
|
// SessionUsage folds embedded sub-session costs into the parent's cost,
|
|
// which is what a restore/branch context needs: those sub-sessions
|
|
// won't emit their own events.
|
|
var restoredInput, restoredOutput int64
|
|
if sess != nil {
|
|
restoredInput, restoredOutput = sess.Usage()
|
|
}
|
|
if restoredInput > 0 || restoredOutput > 0 {
|
|
contextLimit := r.contextLimitForAgentModel(ctx, a, modelID)
|
|
usage := SessionUsage(sess, contextLimit, a.CompactionThreshold())
|
|
|
|
// Reconstruct LastMessage from the parent session's last assistant
|
|
// message so that FinishReason (and other per-message fields) are
|
|
// available on session restore. We intentionally iterate
|
|
// sess.Messages (not GetAllMessages) so the result reflects the
|
|
// parent agent's state: this event carries the parent session_id,
|
|
// and sub-agents emit their own token_usage events with their own
|
|
// session_id during live streaming.
|
|
//
|
|
// MessagesSnapshot takes sess.mu so this cannot race a concurrent
|
|
// AddMessage/ApplyCompaction (e.g. a live HTTP AddMessage arriving
|
|
// while startup info is being emitted for a restored session).
|
|
items := sess.MessagesSnapshot()
|
|
for i := range slices.Backward(items) {
|
|
item := &items[i]
|
|
if !item.IsMessage() || item.Message.Message.Role != chat.MessageRoleAssistant {
|
|
continue
|
|
}
|
|
msg := &item.Message.Message
|
|
lm := &MessageUsage{
|
|
Model: msg.Model,
|
|
Cost: msg.Cost,
|
|
FinishReason: msg.FinishReason,
|
|
}
|
|
if msg.Usage != nil {
|
|
lm.Usage = *msg.Usage
|
|
}
|
|
usage.LastMessage = lm
|
|
break
|
|
}
|
|
|
|
send(NewTokenUsageEvent(sess.ID, r.currentAgentName(), usage))
|
|
}
|
|
|
|
// Tool loading can be slow (MCP servers need to start). Mark the
|
|
// context as non-interactive so toolsets that require user-driven
|
|
// flows (e.g. an OAuth elicitation for a remote MCP server) fail
|
|
// fast with a recognisable error rather than blocking on a dialog
|
|
// the TUI is not yet ready to render. The actual prompt happens on
|
|
// the first RunStream when the user is interacting with the agent.
|
|
nonInteractiveCtx := tools.WithoutInteractivePrompts(ctx)
|
|
r.emitToolsProgressively(nonInteractiveCtx, a, send)
|
|
|
|
// Flush any agent warnings: load-time warnings recorded at agent
|
|
// construction (WithLoadTimeWarnings) and per-toolset warnings recorded
|
|
// during startup above (e.g. a remote MCP server returning 4xx during
|
|
// initialize). Surfacing them as WarningEvents lets the TUI show a
|
|
// persistent notice with the actual server-side explanation — otherwise
|
|
// the user only sees the toolset disappear from the sidebar with no clue
|
|
// as to why.
|
|
r.emitAgentWarnings(a, events)
|
|
}
|
|
|
|
// emitToolsProgressively loads tools from each toolset and emits progress updates.
|
|
// This allows the UI to show the tool count incrementally as each toolset loads,
|
|
// with a spinner indicating that more tools may be coming.
|
|
func (r *LocalRuntime) emitToolsProgressively(ctx context.Context, a *agent.Agent, send func(Event) bool) {
|
|
toolsets := a.ToolSets()
|
|
totalToolsets := len(toolsets)
|
|
|
|
// If no toolsets, emit final state immediately
|
|
if totalToolsets == 0 {
|
|
send(ToolsetInfo(0, false, r.currentAgentName()))
|
|
return
|
|
}
|
|
|
|
// Emit initial loading state
|
|
if !send(ToolsetInfo(0, true, r.currentAgentName())) {
|
|
return
|
|
}
|
|
|
|
// Start scheduling and outcome classification are shared with the turn path.
|
|
// Channels retain configuration order while allowing early toolsets to be
|
|
// listed without waiting for slower later ones.
|
|
starts := a.StartToolSets(ctx, r.toolStartTimeout)
|
|
|
|
// Load tools from each toolset and emit progress
|
|
var totalTools int
|
|
for i, toolset := range toolsets {
|
|
// Check context before potentially slow operations
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
|
|
isLast := i == totalToolsets-1
|
|
|
|
outcome := <-starts[i]
|
|
desc := tools.DescribeToolSet(outcome.ToolSet.ToolSet)
|
|
switch outcome.Kind {
|
|
case tools.StartCanceled:
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
slog.DebugContext(ctx, "Toolset start canceled; skipping", "agent", a.Name(), "toolset", desc, "cause", outcome.Err)
|
|
continue
|
|
case tools.StartTimedOut:
|
|
slog.WarnContext(ctx, "Toolset start timed out; skipping", "agent", a.Name(), "toolset", desc, "timeout", r.toolStartTimeout)
|
|
a.AddToolWarning(fmt.Sprintf("%s is taking too long to start (>%s) — it keeps starting in the background and its tools appear once it is ready", desc, r.toolStartTimeout))
|
|
continue
|
|
case tools.StartAuthorizationRequired:
|
|
if outcome.ReportRecovery {
|
|
slog.WarnContext(ctx, "Toolset needs re-authentication after background token rejection", "agent", a.Name(), "toolset", desc)
|
|
a.AddToolWarning(desc + " needs re-authentication — it will prompt on your next message, or use /toolset-restart")
|
|
} else {
|
|
slog.DebugContext(ctx, "Toolset deferred until first message", "agent", a.Name(), "toolset", desc, "reason", outcome.Err)
|
|
}
|
|
continue
|
|
case tools.StartFailed:
|
|
if outcome.ReportFailure {
|
|
slog.WarnContext(ctx, "Toolset start failed; skipping", "agent", a.Name(), "toolset", desc, "error", outcome.Err)
|
|
a.AddToolWarning(fmt.Sprintf("%s start failed: %v", desc, outcome.Err))
|
|
} else {
|
|
slog.DebugContext(ctx, "Toolset still unavailable; skipping", "agent", a.Name(), "toolset", desc, "error", outcome.Err)
|
|
}
|
|
continue
|
|
case tools.StartPartial:
|
|
if outcome.ReportFailure {
|
|
slog.WarnContext(ctx, "Toolset start failed; using healthy subset", "agent", a.Name(), "toolset", desc, "error", outcome.Err)
|
|
a.AddToolWarning(fmt.Sprintf("%s start failed: %v", desc, outcome.Err))
|
|
}
|
|
case tools.StartInFlight:
|
|
slog.DebugContext(ctx, "Toolset start already in flight; skipping", "agent", a.Name(), "toolset", desc)
|
|
continue
|
|
case tools.StartReady:
|
|
}
|
|
|
|
// Get tools from this toolset under a bounded deadline. A toolset
|
|
// whose Tools() blocks indefinitely would otherwise stall the whole
|
|
// loop: the terminal ToolsetInfo{Loading:false} below is never sent,
|
|
// so the sidebar stays on "Loading tools…" forever and /quit appears
|
|
// to hang. Time it out, skip it, and move on. The skip is only for
|
|
// this startup sidebar pass — the toolset is not torn down, so a
|
|
// slow-but-responsive one is listed (and counted) again on the next
|
|
// turn or tool-change refresh.
|
|
ts, err := listToolsWithTimeout(ctx, toolset, r.toolListTimeout)
|
|
if err != nil {
|
|
slog.WarnContext(ctx, "Failed to list tools from toolset; skipping",
|
|
"agent", a.Name(), "toolset", tools.DescribeToolSet(toolset), "error", err)
|
|
continue
|
|
}
|
|
|
|
totalTools += len(ts)
|
|
|
|
// Emit progress update - still loading unless this is the last toolset
|
|
if !send(ToolsetInfo(totalTools, !isLast, r.currentAgentName())) {
|
|
return
|
|
}
|
|
}
|
|
|
|
// Emit final state (not loading)
|
|
send(ToolsetInfo(totalTools, false, r.currentAgentName()))
|
|
}
|
|
|
|
// listToolsWithTimeout enumerates a toolset's tools under a bounded deadline.
|
|
// The (potentially blocking) Tools call runs in a goroutine and we select on
|
|
// either its completion or the timeout, so a toolset whose Tools() ignores
|
|
// context cancellation — e.g. a wedged MCP stdio subprocess — cannot block
|
|
// startup. On timeout it returns the context error; the orphaned goroutine
|
|
// sends into a buffered channel and exits if the call ever returns, so it does
|
|
// not leak past the eventual (or never) return of Tools().
|
|
func listToolsWithTimeout(ctx context.Context, toolset tools.ToolSet, timeout time.Duration) ([]tools.Tool, error) {
|
|
// Defend against a zero/negative timeout (e.g. a directly-constructed
|
|
// LocalRuntime that bypassed NewLocalRuntime) so we never collapse to an
|
|
// already-expired context that skips every toolset.
|
|
if timeout <= 0 {
|
|
timeout = defaultToolListTimeout
|
|
}
|
|
toolCtx, cancel := context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
|
|
type listResult struct {
|
|
tools []tools.Tool
|
|
err error
|
|
}
|
|
done := make(chan listResult, 1) // buffered so a late send never blocks
|
|
go func() {
|
|
ts, err := toolset.Tools(toolCtx)
|
|
done <- listResult{tools: ts, err: err}
|
|
}()
|
|
|
|
select {
|
|
case <-toolCtx.Done():
|
|
return nil, toolCtx.Err()
|
|
case res := <-done:
|
|
return res.tools, res.err
|
|
}
|
|
}
|
|
|
|
func (r *LocalRuntime) Resume(_ context.Context, req ResumeRequest) {
|
|
slog.Debug("Resuming runtime", "agent", r.currentAgentName(), "type", req.Type, "reason", req.Reason)
|
|
|
|
// Defensive validation:
|
|
//
|
|
// The runtime may be resumed by multiple entry points (API, CLI, TUI, tests).
|
|
// Even if upstream layers perform validation, the runtime must never assume
|
|
// the ResumeType is valid. Accepting invalid values here leads to confusing
|
|
// downstream behavior where tool execution fails without a clear cause.
|
|
if !IsValidResumeType(req.Type) {
|
|
slog.Warn(
|
|
"Invalid resume type received; ignoring resume request",
|
|
"agent", r.currentAgentName(),
|
|
"confirmation_type", req.Type,
|
|
"valid_types", ValidResumeTypes(),
|
|
)
|
|
return
|
|
}
|
|
// Normalize legacy verbs (approve-session, approve-safe,
|
|
// approve-safer) so every downstream consumer sees the current set.
|
|
req.Type = NormalizeResumeType(req.Type)
|
|
|
|
// Attempt to deliver the resume signal to the execution loop.
|
|
//
|
|
// The channel is non-blocking by design to avoid deadlocks if the runtime
|
|
// is not currently waiting for a confirmation (e.g. already resumed,
|
|
// canceled, or shutting down).
|
|
select {
|
|
case r.resumeChan <- req:
|
|
slog.Debug("Resume signal sent", "agent", r.currentAgentName())
|
|
default:
|
|
slog.Debug(
|
|
"Resume channel not ready; resume signal dropped",
|
|
"agent", r.currentAgentName(),
|
|
"confirmation_type", req.Type,
|
|
)
|
|
}
|
|
}
|
|
|
|
var (
|
|
// ErrSteerQueueFull indicates that no more steering messages can be queued.
|
|
ErrSteerQueueFull = errors.New("steer queue full")
|
|
// ErrFollowUpQueueFull indicates that no more follow-up messages can be queued.
|
|
ErrFollowUpQueueFull = errors.New("follow-up queue full")
|
|
)
|
|
|
|
// Steer enqueues a user message for urgent mid-turn injection into the
|
|
// running agent loop. The message will be picked up after the current batch
|
|
// of tool calls finishes but before the loop checks whether to stop.
|
|
func (r *LocalRuntime) Steer(ctx context.Context, msg QueuedMessage) error {
|
|
if !r.steerQueue.Enqueue(ctx, msg) {
|
|
return ErrSteerQueueFull
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// FollowUp enqueues a message to be processed after the current agent turn
|
|
// finishes. Unlike Steer, follow-ups are popped one at a time and each gets
|
|
// a full undivided agent turn.
|
|
func (r *LocalRuntime) FollowUp(ctx context.Context, msg QueuedMessage) error {
|
|
if !r.followUpQueue.Enqueue(ctx, msg) {
|
|
return ErrFollowUpQueueFull
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *LocalRuntime) CancelSteer(_ context.Context, id string) bool {
|
|
q, ok := r.steerQueue.(cancelableMessageQueue)
|
|
return ok && q.Cancel(id)
|
|
}
|
|
|
|
func (r *LocalRuntime) CancelFollowUp(_ context.Context, id string) bool {
|
|
q, ok := r.followUpQueue.(cancelableMessageQueue)
|
|
return ok && q.Cancel(id)
|
|
}
|
|
|
|
// SetRecallHandler registers an embedder-owned wake-up path for tool recalls.
|
|
// When unset, recalls fall back to the steer queue and are consumed by the next
|
|
// active RunStream.
|
|
func (r *LocalRuntime) SetRecallHandler(handler RecallHandler) {
|
|
r.recallMu.Lock()
|
|
defer r.recallMu.Unlock()
|
|
r.recallHandler = handler
|
|
}
|
|
|
|
func (r *LocalRuntime) recall(ctx context.Context, msg QueuedMessage) error {
|
|
if r.activeRootStreams.Load() > 0 {
|
|
return r.Steer(ctx, msg)
|
|
}
|
|
|
|
r.recallMu.RLock()
|
|
handler := r.recallHandler
|
|
r.recallMu.RUnlock()
|
|
if handler == nil {
|
|
return r.Steer(ctx, msg)
|
|
}
|
|
if !handler(ctx, msg) {
|
|
return errors.New("recall handler rejected message")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *LocalRuntime) QueueStatus() QueueStatus {
|
|
status := QueueStatus{}
|
|
if steerQ, ok := r.steerQueue.(*inMemoryMessageQueue); ok {
|
|
status.SteerDepth, status.SteerCapacity = steerQ.status()
|
|
}
|
|
if followupQ, ok := r.followUpQueue.(*inMemoryMessageQueue); ok {
|
|
status.FollowupDepth, status.FollowupCapacity = followupQ.status()
|
|
}
|
|
return status
|
|
}
|
|
|
|
// Run starts the agent's interaction loop
|
|
|
|
func (r *LocalRuntime) startSpan(ctx context.Context, name string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
|
|
if r.tracer == nil {
|
|
return ctx, trace.SpanFromContext(ctx)
|
|
}
|
|
return r.tracer.Start(ctx, name, opts...)
|
|
}
|
|
|
|
// Summarize generates a summary for the session based on the conversation history.
|
|
// The additionalPrompt parameter allows users to provide additional instructions
|
|
// for the summarization (e.g., "focus on code changes" or "include action items").
|
|
//
|
|
// Summarize is the public entry point used by user-driven /compact actions; it
|
|
// reports compactionReasonManual to BeforeCompaction / AfterCompaction hooks
|
|
// and "manual" to PreCompact hooks.
|
|
// Internal callers (proactive threshold, overflow recovery) use
|
|
// [LocalRuntime.compactWithReason] directly to forward a more specific reason.
|
|
func (r *LocalRuntime) Summarize(ctx context.Context, sess *session.Session, additionalPrompt string, events EventSink) {
|
|
// Manual compaction bypasses RunStream, but native compaction may start tools.
|
|
ctx = genai.WithConversationID(ctx, sess.ID)
|
|
ctx = tools.WithHandlerScope(ctx, tools.HandlerScope{
|
|
Elicitation: r.elicitationHandler,
|
|
Sampling: r.samplingHandler,
|
|
SamplingWithTools: r.samplingWithToolsHandler,
|
|
OAuthSuccess: func() {
|
|
events.Emit(Authorization(tools.ElicitationActionAccept, r.resolveSessionAgent(sess).Name()))
|
|
},
|
|
})
|
|
r.configureToolsetHandlers(r.resolveSessionAgent(sess))
|
|
r.compactWithReason(ctx, sess, additionalPrompt, compactionReasonManual, events)
|
|
}
|
|
|
|
// compactWithReason runs a session compaction with the supplied reason and
|
|
// emits a TokenUsageEvent so the UI immediately reflects the new context
|
|
// pressure.
|
|
//
|
|
// reason is reported to BeforeCompaction / AfterCompaction hooks as
|
|
// CompactionReason. Use [compactionReasonThreshold] for proactive
|
|
// threshold-of-context triggers (90% by default, tunable via
|
|
// compaction_threshold), [compactionReasonOverflow] for post-overflow
|
|
// auto-recovery, or [compactionReasonManual] for user-invoked compactions.
|
|
//
|
|
// PreCompact hooks fire first via the legacy [hooks.Input.Source] field
|
|
// ("auto" / "tool_overflow" / "overflow" / "manual"); they may cancel the
|
|
// compaction or contribute additional steering text. BeforeCompaction
|
|
// hooks then fire inside [LocalRuntime.doCompact] with [Input.CompactionReason]
|
|
// set to the canonical reason; they may veto or supply a custom summary.
|
|
func (r *LocalRuntime) compactWithReason(ctx context.Context, sess *session.Session, additionalPrompt, reason string, events EventSink) {
|
|
// Stamp the session ID on ctx so the compaction LLM call carries
|
|
// `X-Cagent-Session-Id` to the gateway. Manual compaction
|
|
// (via `Summarize` from the App) bypasses `runStreamLoop`'s seed;
|
|
// internal callers (proactive threshold, overflow recovery) already
|
|
// run with a stamped ctx, but re-stamping is idempotent.
|
|
ctx = httpclient.ContextWithSessionID(ctx, sess.ID)
|
|
a := r.resolveSessionAgent(sess)
|
|
|
|
source := preCompactSourceFor(reason)
|
|
skip, msg, extraPrompt := r.executePreCompactHooks(ctx, sess, a, source, events)
|
|
if skip {
|
|
slog.WarnContext(ctx, "pre_compact hook signalled skip",
|
|
"agent", a.Name(), "session_id", sess.ID, "source", source, "reason", msg)
|
|
if msg != "" {
|
|
events.Emit(Warning(msg, a.Name()))
|
|
}
|
|
return
|
|
}
|
|
additionalPrompt = joinPrompts(additionalPrompt, extraPrompt)
|
|
|
|
r.doCompact(ctx, sess, a, additionalPrompt, reason, events)
|
|
|
|
// Emit a TokenUsageEvent so the sidebar immediately reflects the
|
|
// compaction: tokens drop to the summary size, context % drops, and
|
|
// cost increases by the summary generation cost.
|
|
modelID := r.getEffectiveModelID(ctx, a)
|
|
contextLimit := r.effectiveContextLimit(ctx, a, r.resolveContextLimit(ctx, a.Model(ctx), modelID))
|
|
events.Emit(NewTokenUsageEvent(sess.ID, a.Name(), SessionUsage(sess, contextLimit, a.CompactionThreshold())))
|
|
}
|
|
|
|
// preCompactSourceFor maps the canonical compaction reason
|
|
// ([compactionReasonThreshold] / [compactionReasonOverflow] /
|
|
// [compactionReasonManual]) onto the [hooks.Input.Source] string
|
|
// surfaced by the pre_compact hook ("auto" / "overflow" / "manual").
|
|
// Unknown reasons fall through unchanged so future, more specific
|
|
// reasons (e.g. "tool_overflow") can be forwarded verbatim without
|
|
// touching this map.
|
|
func preCompactSourceFor(reason string) string {
|
|
switch reason {
|
|
case compactionReasonThreshold:
|
|
return "auto"
|
|
case compactionReasonOverflow:
|
|
return "overflow"
|
|
case compactionReasonManual:
|
|
return "manual"
|
|
default:
|
|
return reason
|
|
}
|
|
}
|
|
|
|
// joinPrompts concatenates two non-empty prompt fragments with a blank
|
|
// line, returning whichever is non-empty when the other isn't. Used by
|
|
// compactWithReason to splice pre_compact's additional_context into
|
|
// the caller's additionalPrompt without having to special-case empty
|
|
// strings at the callsite.
|
|
func joinPrompts(a, b string) string {
|
|
switch {
|
|
case a == "":
|
|
return b
|
|
case b == "":
|
|
return a
|
|
default:
|
|
return a + "\n\n" + b
|
|
}
|
|
}
|