model.go

package agentgo

import (
	"context"
	"encoding/json"
	"sync"
)

// ---------------------------------------------------------------------------
// Agent Context & Loop Config
// ---------------------------------------------------------------------------

// SystemBlock is one segment of a multi-part system prompt.
// Use with AgentContext.SystemBlocks for per-block cache control.
type SystemBlock struct {
	Text         string `json:"text"`
	CacheControl string `json:"cache_control,omitempty"` // e.g. "ephemeral"
}

// AgentContext holds the conversation and request settings used by the loop.
// A stateful Agent refreshes request settings at each turn boundary.
type AgentContext struct {
	artifacts    *memoryArtifactManager // Bound by the runtime, never inherited from a caller.
	SystemPrompt string                 // single-string system prompt (legacy)
	SystemBlocks []SystemBlock          // multi-block system prompt with cache control (takes precedence)
	Messages     []AgentMessage
	Tools        []Tool
}

// LLMRequest carries the inputs to a single LLM call.
type LLMRequest struct {
	Messages []Message  `json:"messages"`
	Tools    []ToolSpec `json:"tools,omitempty"`
}

// LLMResponse carries the result of a single LLM call.
type LLMResponse struct {
	Message Message
}

// ToolSpec describes a tool for the LLM (name + description + JSON schema).
type ToolSpec struct {
	Name         string `json:"name"`
	Description  string `json:"description"`
	Parameters   any    `json:"parameters"`
	DeferLoading bool   `json:"defer_loading,omitempty"`
	// Strict enables provider-side strict schema enforcement (OpenAI strict
	// tool calling / Structured Outputs for arguments). nil leaves the
	// provider default. Set via the optional StrictSchemaTool interface.
	Strict *bool `json:"strict,omitempty"`
}

// LoopConfig configures the agent loop.
type LoopConfig struct {
	artifacts *memoryArtifactManager // Stateful Agent reuses its manager; bare loops create one.
	// Only the stateful Agent refreshes settings; direct loops use their supplied config.
	refreshTurnSettings func(*AgentContext, *LoopConfig)
	Model               ChatModel
	MaxTurns            int           // safety limit, default 100
	MaxRetries          int           // LLM call retry limit for retryable errors, default 3
	MaxToolErrors       int           // consecutive tool failure threshold per tool, 0 = unlimited
	ThinkingLevel       ThinkingLevel // reasoning depth

	// Turn hooks execute once around each logical model/tool turn.
	BeforeTurn BeforeTurnHook
	AfterTurn  AfterTurnHook

	// InitialState is the run baseline supplied by the stateful Agent wrapper.
	// Direct AgentLoop callers may leave it zero-valued.
	InitialState AgentState

	// Context lifecycle. ContextManager drives prompt projection, overflow
	// recovery, and usage reporting. AgentMessage values lower themselves to
	// Message only at the model-call boundary.
	ContextManager ContextManager

	// ToolResultMessageFactory optionally lifts a raw tool result into an
	// application-specific message before it enters the transcript. Nil keeps
	// the built-in model Message representation. The returned message's model
	// projection must preserve the tool role and tool_call_id pairing.
	ToolResultMessageFactory func(ToolCall, ToolResult) AgentMessage

	// CommitContext durably accepts a compaction candidate before publication.
	// The result contains both Messages and the complete staged Artifacts collection.
	// Returning an error rejects both and stops the run. The callback must atomically
	// persist them if durable atomicity is required; external files remain host-owned.
	CommitContext func(ContextCommitResult) error

	// CommitMessage durably records a message before it enters runtime context.
	// Returning an error stops the run; tools requested by that message are not
	// executed. OnMessage remains the post-commit observer hook.
	CommitMessage func(msg AgentMessage) error

	// ToolGate, when non-nil, is called once per tool call after argument
	// validation and the optional Previewer pass. Allowed=false rejects the
	// call (Reason becomes the tool result). The kernel does no permission
	// reasoning of its own.
	ToolGate ToolGate

	// Steering is drained at turn/tool boundaries and after context preparation,
	// before each model request (including retries). Input queued during a slow
	// compaction is committed and included in the immediately following call.
	GetSteeringMessages func() []AgentMessage

	// FollowUp: called when the agent would otherwise stop.
	GetFollowUpMessages func() []AgentMessage

	// ModelMiddlewares wrap each AgentGo-owned physical model attempt, including
	// context summarization. ToolMiddlewares wrap each tool execution. First
	// middleware is outermost.
	ModelMiddlewares []ModelMiddleware
	ToolMiddlewares  []ToolMiddleware

	// MaxToolConcurrency limits parallel tool execution.
	// 0 or 1 = sequential (default, backward compatible).
	// >1 = up to N tools execute concurrently within a single turn.
	MaxToolConcurrency int

	// ShouldEmitAbortMarker reports whether an abort marker message should be
	// emitted when the context is cancelled. When nil or returns false, the
	// cancellation is silent (legacy behavior). Set by Agent.Abort().
	ShouldEmitAbortMarker func() bool

	// StopAfterTool, if non-nil, is called after each successful (non-error)
	// tool execution. If it returns true, the loop exits with EndReasonStop.
	// Use this to let a terminal tool (e.g. commit_chapter) end the loop
	// without wasting turns. The exit passes through StopGuard (with
	// Trigger=StopTriggerAfterTool) like any other normal stop, so a guard
	// can veto a premature terminal-tool exit.
	StopAfterTool func(toolName string) bool

	// StopAfterToolResult is the result-aware variant of StopAfterTool. It is
	// useful when the same tool can be an intermediate step or a terminal step
	// depending on its structured result.
	StopAfterToolResult func(toolName string, result json.RawMessage) bool

	// OnMessage, if non-nil, is called after each committed message is appended
	// to context (assistant, tool result, steering). It is observational; use
	// CommitMessage when persistence failure must stop execution.
	OnMessage func(msg AgentMessage)

	// StopGuard is consulted on every normal stop (end_turn and
	// StopAfterTool exits). Nil (default) means every stop is allowed.
	StopGuard StopGuard

	// LengthRecoveryPrompt overrides the user message injected when the
	// model's output is truncated (max_tokens) with no completed tool calls.
	// Empty uses a built-in default. Each recovery is an extra LLM call on
	// top of normal turn accounting, bounded by an internal cap.
	LengthRecoveryPrompt string

	// AbortMarkerText overrides the marker message recorded when a run is
	// cancelled mid-inference (only when ShouldEmitAbortMarker returns true).
	// Empty uses a built-in default. Lets non-English harnesses localize it.
	AbortMarkerText string
	// AbortMarkerToolText overrides the marker message recorded when
	// cancellation lands during tool execution. Empty uses a built-in default.
	AbortMarkerToolText string

	// CacheLastMessage, when non-empty, instructs the loop to tag the last
	// non-system message in every LLM request with this cache_control value
	// (e.g. "ephemeral"). Providers that support prompt caching place a write
	// breakpoint at that position covering the entire preceding prefix. Empty
	// string (default) leaves messages untouched — keep cache placement under
	// application control.
	//
	// The breakpoint follows the freshest turn (user input, tool_result, or
	// assistant) and skips trailing per-turn system reminders.
	CacheLastMessage string

	// PromptCacheKey, when non-empty, is attached to every LLM request as the
	// provider's prompt-cache routing identity (e.g. OpenAI prompt_cache_key).
	// Use one key per conversation so all requests of a session land on the
	// same cache shard. The adapter drops the hint for providers without
	// key-routed caching. Empty (default) sends nothing.
	PromptCacheKey string
}

// ---------------------------------------------------------------------------
// Context Usage Estimation
// ---------------------------------------------------------------------------

// ContextEstimateFn estimates the current context token consumption from messages.
// Returns total tokens, tokens from LLM Usage, and estimated trailing tokens.
type ContextEstimateFn func(msgs []AgentMessage) (tokens, usageTokens, trailingTokens int)

// ContextUsage represents the current context window occupancy estimate.
type ContextUsage struct {
	Tokens         int     `json:"tokens"`          // estimated total tokens in context
	ContextWindow  int     `json:"context_window"`  // model's context window size
	Percent        float64 `json:"percent"`         // tokens / contextWindow * 100
	UsageTokens    int     `json:"usage_tokens"`    // from last LLM-reported Usage
	TrailingTokens int     `json:"trailing_tokens"` // chars/4 estimate for trailing messages
}

// ---------------------------------------------------------------------------
// Call Options
// ---------------------------------------------------------------------------

// CallOption configures per-call LLM parameters.
type CallOption func(*CallConfig)

// CallConfig holds per-call configuration resolved from CallOptions.
type CallConfig struct {
	ThinkingLevel  ThinkingLevel
	ThinkingBudget int    // max thinking tokens, 0 = use provider default
	APIKey         string // per-call API key override, empty = use model default
	SessionID      string // provider session caching identifier
	PromptCacheKey string // prompt-cache routing identity, empty = no hint
	MaxTokens      int    // per-call max tokens override, 0 = use model default
	ToolChoice     any    // "auto" / "required" / "none" / {"type":"tool","name":"xxx"}, nil = provider default
	ResponseFormat *ResponseFormat
}

const (
	ResponseFormatText       = "text"
	ResponseFormatJSONObject = "json_object"
	ResponseFormatJSONSchema = "json_schema"
)

// ResponseFormat controls provider-native structured output.
//
// JSON object mode asks the model to return valid JSON. JSON schema mode asks
// compatible providers to constrain the final response to the supplied schema.
// Callers should still unmarshal and validate model output like any external
// input.
type ResponseFormat struct {
	Type       string      `json:"type"`
	JSONSchema *JSONSchema `json:"json_schema,omitempty"`
}

// JSONSchema describes the provider-native JSON schema response format.
type JSONSchema struct {
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	Schema      any    `json:"schema"`
	Strict      *bool  `json:"strict,omitempty"`
}

// ResolveCallConfig applies options and returns the resolved config.
func ResolveCallConfig(opts []CallOption) CallConfig {
	var cfg CallConfig
	for _, opt := range opts {
		opt(&cfg)
	}
	return cfg
}

// WithThinking sets the thinking level for a single LLM call.
func WithThinking(level ThinkingLevel) CallOption {
	return func(c *CallConfig) { c.ThinkingLevel = NormalizeThinkingLevel(level) }
}

// WithThinkingBudget sets the max thinking tokens for a single LLM call.
func WithThinkingBudget(tokens int) CallOption {
	return func(c *CallConfig) { c.ThinkingBudget = tokens }
}

// WithAPIKey overrides the API key for a single LLM call.
// Enables key rotation, OAuth short-lived tokens, and multi-tenant scenarios.
func WithAPIKey(key string) CallOption {
	return func(c *CallConfig) { c.APIKey = key }
}

// WithCallPromptCacheKey sets the prompt-cache routing identity for a single
// LLM call. Providers with key-routed prefix caching (OpenAI prompt_cache_key)
// route requests sharing a key to the same cache shard; the adapter drops the
// hint for providers without support.
func WithCallPromptCacheKey(key string) CallOption {
	return func(c *CallConfig) { c.PromptCacheKey = key }
}

// WithCallSessionID sets a session identifier for a single LLM call.
func WithCallSessionID(id string) CallOption {
	return func(c *CallConfig) { c.SessionID = id }
}

// WithMaxTokens overrides the max output tokens for a single LLM call.
func WithMaxTokens(tokens int) CallOption {
	return func(c *CallConfig) { c.MaxTokens = tokens }
}

// WithToolChoice controls whether the model must call a tool.
// Accepted values: "auto" (default), "required" (must call a tool), "none" (no tools).
func WithToolChoice(choice any) CallOption {
	return func(c *CallConfig) { c.ToolChoice = choice }
}

// WithJSONMode requests valid JSON output without enforcing a specific schema.
func WithJSONMode() CallOption {
	return func(c *CallConfig) {
		c.ResponseFormat = &ResponseFormat{Type: ResponseFormatJSONObject}
	}
}

// WithJSONSchema requests structured JSON output constrained by a JSON Schema.
func WithJSONSchema(name, description string, schema any, strict bool) CallOption {
	return func(c *CallConfig) {
		c.ResponseFormat = &ResponseFormat{
			Type: ResponseFormatJSONSchema,
			JSONSchema: &JSONSchema{
				Name:        name,
				Description: description,
				Schema:      schema,
				Strict:      &strict,
			},
		}
	}
}

// WithResponseFormat sets a provider-native response format explicitly.
func WithResponseFormat(format *ResponseFormat) CallOption {
	return func(c *CallConfig) { c.ResponseFormat = format }
}

// ---------------------------------------------------------------------------
// ChatModel Interface
// ---------------------------------------------------------------------------

// ChatModel is the LLM provider interface.
type ChatModel interface {
	Generate(ctx context.Context, messages []Message, tools []ToolSpec, opts ...CallOption) (*LLMResponse, error)
	GenerateStream(ctx context.Context, messages []Message, tools []ToolSpec, opts ...CallOption) (<-chan StreamEvent, error)
	SupportsTools() bool
}

// ProviderNamer is an optional interface for ChatModel implementations
// to expose their provider name (e.g. "openai", "anthropic", "gemini").
// Used by the agent loop to pass provider context to GetApiKey callbacks.
type ProviderNamer interface {
	ProviderName() string
}

// ModelNamer is an optional interface for ChatModel implementations to
// expose their model identifier (e.g. "claude-sonnet-4-6"). Harnesses need
// this for display and telemetry; without it they fall back to concrete-type
// assertions.
type ModelNamer interface {
	ModelName() string
}

// SwappableModel wraps a ChatModel and allows replacing the underlying model
// at runtime. Swaps take effect on the next call.
type SwappableModel struct {
	mu    sync.RWMutex
	model ChatModel
}

func NewSwappableModel(initial ChatModel) *SwappableModel {
	return &SwappableModel{model: initial}
}

func (m *SwappableModel) Swap(next ChatModel) {
	m.mu.Lock()
	m.model = next
	m.mu.Unlock()
}

func (m *SwappableModel) Current() ChatModel {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.model
}

func (m *SwappableModel) Generate(ctx context.Context, messages []Message, tools []ToolSpec, opts ...CallOption) (*LLMResponse, error) {
	model := m.Current()
	if model == nil {
		return nil, ErrNoModel
	}
	return model.Generate(ctx, messages, tools, opts...)
}

func (m *SwappableModel) GenerateStream(ctx context.Context, messages []Message, tools []ToolSpec, opts ...CallOption) (<-chan StreamEvent, error) {
	model := m.Current()
	if model == nil {
		return nil, ErrNoModel
	}
	return model.GenerateStream(ctx, messages, tools, opts...)
}

func (m *SwappableModel) SupportsTools() bool {
	model := m.Current()
	return model != nil && model.SupportsTools()
}

func (m *SwappableModel) ProviderName() string {
	model := m.Current()
	pn, ok := model.(ProviderNamer)
	if !ok {
		return ""
	}
	return pn.ProviderName()
}

func (m *SwappableModel) ModelName() string {
	model := m.Current()
	mn, ok := model.(ModelNamer)
	if !ok {
		return ""
	}
	return mn.ModelName()
}

// ---------------------------------------------------------------------------
// Stream Events (fine-grained)
// ---------------------------------------------------------------------------

// StreamEventType identifies LLM streaming event types.
type StreamEventType string

const (
	// Text content streaming
	StreamEventTextStart StreamEventType = "text_start"
	StreamEventTextDelta StreamEventType = "text_delta"
	StreamEventTextEnd   StreamEventType = "text_end"

	// Thinking/reasoning streaming
	StreamEventThinkingStart StreamEventType = "thinking_start"
	StreamEventThinkingDelta StreamEventType = "thinking_delta"
	StreamEventThinkingEnd   StreamEventType = "thinking_end"

	// Tool call streaming
	StreamEventToolCallStart StreamEventType = "toolcall_start"
	StreamEventToolCallDelta StreamEventType = "toolcall_delta"
	StreamEventToolCallEnd   StreamEventType = "toolcall_end"

	// Terminal events
	StreamEventDone  StreamEventType = "done"
	StreamEventError StreamEventType = "error"
)

// StreamEvent is a streaming event from the LLM. Published payloads are read-only.
// Producers must not mutate reachable message content after sending an event.
//
// +spec=`Each stream event retains its published message while generation continues`
// +case:id=stream_snapshot,desc=`Retain early events and read while the stream advances`,expect=`Stable payloads without data races`
// +link=component://docs/kernel.md
type StreamEvent struct {
	Type         StreamEventType
	ContentIndex int     // which content block is being updated
	Delta        string  // text/thinking/toolcall argument delta
	Message      Message // partial (during streaming) or final (done)
	// ToolID identifies the call for toolcall start/delta/end events. It is
	// empty until the provider supplies an ID.
	ToolID string
	// CompletedToolCall is populated on StreamEventToolCallEnd with the fully
	// reconstructed tool call. It lets the loop start execution immediately
	// without re-parsing the partial assistant message.
	CompletedToolCall *ToolCall
	StopReason        StopReason // finish reason (for done events)
	Err               error      // for error events
}