atto.go

// MIT License
//
// Copyright (c) 2026 Arsene Tochemey Gandote
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//

// Package atto is the user-facing entry point of the atto agent
// runtime. Construct a [Runtime] with [New] and drive invocations
// through [Runtime.Run]; the actor system, the per-session actors,
// the model actor and the retry policy are wired automatically.
//
// Typical use:
//
//	rt, err := atto.New(ctx, model, func(m llm.LLM) agent.Agent {
//	    return agent.NewLLM(
//	        agent.WithName("assistant"),
//	        agent.WithModel(m),
//	        agent.WithInstruction("..."),
//	        agent.WithTools(tools...),
//	    )
//	})
//	if err != nil { ... }
//	defer rt.Stop(context.Background())
//
//	for ev, err := range rt.Run(ctx, sessionID, session.UserText("...")) {
//	    ...
//	}
//
// The [AgentBuilder] receives the [llm.LLM] that atto routes through
// the model actor; the agent built against it inherits retry,
// passivation and (in cluster mode) placement transparently. Users
// who need a shared actor system or goakt features beyond the
// [Option] set pair [WithActorSystem] with [Extension].
package atto

import (
	"context"
	"fmt"
	"iter"
	"sync"
	"time"

	gactor "github.com/tochemey/goakt/v4/actor"
	gaktlog "github.com/tochemey/goakt/v4/log"
	"github.com/tochemey/goakt/v4/passivation"
	"github.com/tochemey/goakt/v4/remote"

	"github.com/tochemey/atto/agent"
	attoactor "github.com/tochemey/atto/internal/actor"
	internaldiscovery "github.com/tochemey/atto/internal/discovery"
	internalrunner "github.com/tochemey/atto/internal/runner"
	"github.com/tochemey/atto/llm"
	"github.com/tochemey/atto/session"
	"github.com/tochemey/atto/store/inmemory"
)

// defaultSessionAskTimeout caps how long [Runtime.AskSession] waits for
// the session actor to reply. The session actor's task-projection
// handlers are non-blocking (single state-map write plus a store Save);
// anything longer indicates either a store-level stall or a stuck
// mailbox and should fail fast so the A2A wire can surface backpressure.
const defaultSessionAskTimeout = 10 * time.Second

// AgentBuilder is the caller-supplied closure that constructs the
// root [agent.Agent]. It receives the model-actor-backed [llm.LLM]
// that atto built from the actor system; the agent must use that
// instance (via [agent.WithModel] or equivalent) so completions go
// through the model actor's retry and placement machinery.
//
// Returning nil is treated as [ErrNoAgent].
type AgentBuilder func(model llm.LLM) agent.Agent

// Runtime is the front door of atto. It owns the goakt actor system
// (unless one is supplied via [WithActorSystem]), the per-session
// session actors and the model actor that fronts the [llm.LLM].
// Construct one with [New] and drive invocations with [Run].
//
// A Runtime is safe for concurrent use across distinct session IDs.
type Runtime struct {
	sys              gactor.ActorSystem
	ownsSystem       bool
	root             agent.Agent
	eventBufferSize  int
	passivationAfter time.Duration

	mu       sync.Mutex
	stopped  bool
	sessions map[string]*gactor.PID
}

// New constructs and starts a [Runtime].
//
// By default New builds and starts an internal goakt actor system,
// registers atto's runtime extension with the supplied model, store
// and retry policy, spawns the per-process model actor and returns
// a Runtime ready to serve [Run] calls. The supplied build closure
// receives the model-actor-backed [llm.LLM] and must hand it to the
// agent it returns.
//
// When [WithActorSystem] is supplied the supplied system is adopted
// as-is: it must already be started and must have a runtime
// extension registered (see [Extension]). [Runtime.Stop] then
// shuts down the per-session actors but leaves the system itself
// running for the caller.
//
// Errors:
//   - [ErrNoAgent] — build is nil or returns nil
//   - [ErrNoModel] — model is nil and [WithActorSystem] was not used
//   - [ErrRuntimeExtensionMissing] — [WithActorSystem] supplied a
//     system with no atto runtime extension registered
//   - [ErrClusterWithActorSystem] — both [WithCluster] and
//     [WithActorSystem] were supplied
//   - [ErrClusterNeedsSharedStore] — [WithCluster] was supplied
//     without an explicit [WithStore]; the in-memory default is
//     private per node and would silently break session affinity
//   - any goakt error from constructing, starting or spawning on
//     the actor system
func New(ctx context.Context, model llm.LLM, build AgentBuilder, opts ...Option) (*Runtime, error) {
	if build == nil {
		return nil, ErrNoAgent
	}

	cfg := defaultConfig()

	for _, o := range opts {
		o(&cfg)
	}

	if err := validateConfig(&cfg); err != nil {
		return nil, err
	}

	sys, ownsSystem, err := resolveActorSystem(ctx, model, cfg)
	if err != nil {
		return nil, err
	}

	var modelLLMOpts []internalrunner.ModelLLMOption
	if cfg.modelAskTimeout > 0 {
		modelLLMOpts = append(modelLLMOpts, internalrunner.WithAskTimeout(cfg.modelAskTimeout))
	}

	wrapped, err := internalrunner.NewModelLLM(ctx, sys, modelLLMOpts...)
	if err != nil {
		teardown(ctx, sys, ownsSystem)
		return nil, err
	}

	if _, err := attoactor.EnsureTaskRegistry(ctx, sys); err != nil {
		teardown(ctx, sys, ownsSystem)
		return nil, err
	}

	root := build(wrapped)
	if root == nil {
		teardown(ctx, sys, ownsSystem)
		return nil, ErrNoAgent
	}

	return &Runtime{
		sys:              sys,
		ownsSystem:       ownsSystem,
		root:             root,
		eventBufferSize:  cfg.eventBufferSize,
		passivationAfter: cfg.passivationAfter,
		sessions:         make(map[string]*gactor.PID),
	}, nil
}

// validateConfig enforces invariants that depend on the *combination*
// of supplied options, before any actor system construction begins.
// Catching these here keeps the error path cheap and saves the
// caller from goakt-flavoured failure messages later.
func validateConfig(cfg *config) error {
	if cfg.sys != nil && cfg.cluster != nil {
		return ErrClusterWithActorSystem
	}

	if cfg.cluster != nil && cfg.store == nil {
		return ErrClusterNeedsSharedStore
	}

	// Cluster mode does not use the user-facing default store
	// (each node would have its own private map). Every other path
	// gets the in-memory default when the caller did not supply one.
	if cfg.store == nil && cfg.sys == nil {
		cfg.store = inmemory.New()
	}

	return nil
}

// resolveActorSystem returns the goakt actor system the [Runtime]
// will use, plus a flag indicating whether the runtime owns the
// lifecycle. Three branches:
//
//   - [WithActorSystem]: validate the caller's system has the atto
//     extension and adopt it.
//   - [WithCluster]: build a fresh cluster-enabled system with atto's
//     kinds, serialisables and runtime extension wired automatically.
//   - default: build a fresh local system with atto's runtime
//     extension wired automatically.
func resolveActorSystem(ctx context.Context, model llm.LLM, cfg config) (gactor.ActorSystem, bool, error) {
	switch {
	case cfg.sys != nil:
		ext := cfg.sys.Extension(attoactor.RuntimeExtensionID)
		if ext == nil {
			return nil, false, ErrRuntimeExtensionMissing
		}

		// Reject extensions with no session store up-front. The
		// default in-memory store cannot be auto-installed on a
		// caller-supplied system without overriding configuration
		// the caller intended; failing here surfaces the
		// misconfiguration at the API boundary rather than later
		// during session-actor PreStart.
		rt, ok := ext.(attoactor.Runtime)
		if !ok || rt.Store() == nil {
			return nil, false, ErrRuntimeStoreMissing
		}

		return cfg.sys, false, nil

	case cfg.cluster != nil:
		return buildClusterSystem(ctx, model, cfg)

	default:
		return buildLocalSystem(ctx, model, cfg)
	}
}

// buildLocalSystem constructs a single-process actor system seeded
// with atto's runtime extension and no remote/cluster wiring. PubSub
// is enabled so the bridge in [github.com/tochemey/atto/a2a] can use
// per-task topics for streaming without requiring cluster mode;
// cluster mode spawns the topic actor on its own.
func buildLocalSystem(ctx context.Context, model llm.LLM, cfg config) (gactor.ActorSystem, bool, error) {
	if model == nil {
		return nil, false, ErrNoModel
	}

	sys, err := gactor.NewActorSystem(cfg.systemName,
		gactor.WithLogger(goaktLoggerFromConfig(cfg)),
		gactor.WithPubSub(),
		gactor.WithExtensions(attoactor.NewRuntime(runtimeOptionsFromConfig(model, cfg)...)),
	)
	if err != nil {
		return nil, false, err
	}

	if err := sys.Start(ctx); err != nil {
		return nil, false, err
	}

	return sys, true, nil
}

// buildClusterSystem constructs a cluster-enabled actor system.
// Atto's cluster kinds, remote serialisables and runtime extension
// are wired automatically; the caller never imports goakt.
func buildClusterSystem(ctx context.Context, model llm.LLM, cfg config) (gactor.ActorSystem, bool, error) {
	if model == nil {
		return nil, false, ErrNoModel
	}

	cc := cfg.cluster

	clusterCfg := gactor.NewClusterConfig().
		WithDiscovery(internaldiscovery.Wrap(cc.provider)).
		WithKinds(ClusterKinds()...).
		WithDiscoveryPort(cc.discoveryPort).
		WithPeersPort(cc.peersPort).
		WithMinimumPeersQuorum(cc.quorum).
		WithReplicaCount(cc.replicaCount).
		WithBootstrapTimeout(cc.bootstrapTimeout)

	if cc.partitions > 0 {
		clusterCfg = clusterCfg.WithPartitionCount(cc.partitions)
	}

	remoteCfg := remote.NewConfig(cc.bindHost, cc.remotingPort,
		remote.WithSerializables(RemoteSerializables()...),
	)

	sys, err := gactor.NewActorSystem(cfg.systemName,
		gactor.WithLogger(goaktLoggerFromConfig(cfg)),
		gactor.WithRemote(remoteCfg),
		gactor.WithCluster(clusterCfg),
		gactor.WithExtensions(attoactor.NewRuntime(runtimeOptionsFromConfig(model, cfg)...)),
	)
	if err != nil {
		return nil, false, err
	}

	if err := sys.Start(ctx); err != nil {
		return nil, false, err
	}

	return sys, true, nil
}

// runtimeOptionsFromConfig maps the user-facing [config] to the
// [attoactor.RuntimeOption] slice consumed by the internal runtime
// extension.
func runtimeOptionsFromConfig(model llm.LLM, cfg config) []attoactor.RuntimeOption {
	opts := []attoactor.RuntimeOption{
		attoactor.WithLLM(model),
	}

	if cfg.store != nil {
		opts = append(opts, attoactor.WithStore(cfg.store))
	}

	if cfg.stashBound > 0 {
		opts = append(opts, attoactor.WithStashBound(cfg.stashBound))
	}

	if cfg.modelMaxRetries > 0 {
		opts = append(opts, attoactor.WithModelMaxRetries(cfg.modelMaxRetries))
	}

	if cfg.modelBaseBackoff > 0 {
		opts = append(opts, attoactor.WithModelBaseBackoff(cfg.modelBaseBackoff))
	}

	if cfg.modelMaxBackoff > 0 {
		opts = append(opts, attoactor.WithModelMaxBackoff(cfg.modelMaxBackoff))
	}

	return opts
}

// goaktLoggerFromConfig converts the user-supplied [slog.Logger]
// into the [gaktlog.Logger] goakt expects. A nil logger maps to
// [gaktlog.DiscardLogger] so atto stays silent unless [WithLogger]
// was used. The slog handler's level decides which records are
// emitted; we pass [gaktlog.DebugLevel] so goakt's own gating does
// not filter records the user's handler considered relevant.
func goaktLoggerFromConfig(cfg config) gaktlog.Logger {
	if cfg.logger == nil {
		return gaktlog.DiscardLogger
	}

	return gaktlog.NewSlogFrom(cfg.logger, gaktlog.DebugLevel)
}

// teardown stops an actor system that [New] owns when a later
// construction step fails. Adoption-mode systems are left alone.
func teardown(ctx context.Context, sys gactor.ActorSystem, ownsSystem bool) {
	if !ownsSystem || sys == nil {
		return
	}

	_ = sys.Stop(ctx)
}

// Run drives one invocation of the configured root agent against
// sessionID, supplying input as the user message. It returns an
// iterator over the resulting events. The iterator terminates when
// the invocation completes, errors or the supplied context is
// cancelled.
func (r *Runtime) Run(ctx context.Context, sessionID string, input session.Message) iter.Seq2[*session.Event, error] {
	return func(yield func(*session.Event, error) bool) {
		pid, err := r.sessionPID(ctx, sessionID)
		if err != nil {
			yield(nil, err)
			return
		}

		events := make(chan *session.Event, r.eventBufferSize)
		done := make(chan error, 1)

		go attoactor.RunWorker(ctx, attoactor.WorkerConfig{
			Agent:      r.root,
			SessionID:  sessionID,
			SessionPID: pid,
			Input:      input,
			Events:     events,
			Done:       done,
		})

		drainAndYield(events, done, yield)
	}
}

// Stop releases the runtime's resources: every per-session actor is
// shut down and, when [Runtime] owns the actor system, the system
// itself is stopped. Subsequent calls to [Run] return
// [ErrRuntimeStopped]. Stop is idempotent.
func (r *Runtime) Stop(ctx context.Context) error {
	r.mu.Lock()

	if r.stopped {
		r.mu.Unlock()
		return nil
	}

	r.stopped = true
	pids := make([]*gactor.PID, 0, len(r.sessions))

	for _, p := range r.sessions {
		pids = append(pids, p)
	}

	r.sessions = nil
	r.mu.Unlock()

	for _, p := range pids {
		_ = p.Shutdown(ctx)
	}

	if r.ownsSystem {
		return r.sys.Stop(ctx)
	}

	return nil
}

// ActorSystem returns the goakt actor system that backs the runtime.
// The A2A bridge uses it to spawn per-task topic subscribers and to
// reach the system topic actor. Callers should treat the handle as
// read-only.
//
// Always non-nil for a [Runtime] returned by [New]: the runtime owns
// the system in default mode and adopts the caller's system in
// [WithActorSystem] mode.
func (r *Runtime) ActorSystem() gactor.ActorSystem {
	return r.sys
}

// AskSession sends msg to the session actor for sessionID and returns
// its reply. The runtime spawns the actor on first use (mirroring
// [Run]'s behaviour) so callers do not pre-warm sessions.
//
// AskSession is the integration point the A2A bridge uses to drive
// task-projection writes through the actor mailbox instead of writing
// to the session store directly: the actor is already the
// serialisation point for the session, so layering task projection
// reads and writes onto the same mailbox preserves the
// "single-source-of-truth" guarantee for [store.Snapshot.State]
// without a parallel store. Bridge code outside this repository should
// treat AskSession as the supported channel for any state interaction
// outside the agent loop.
//
// Returns [ErrRuntimeStopped] when the runtime has been stopped, and
// surfaces goakt's Ask errors otherwise (timeout, remote failure,
// undelivered reply).
func (r *Runtime) AskSession(ctx context.Context, sessionID string, msg any) (any, error) {
	pid, err := r.sessionPID(ctx, sessionID)
	if err != nil {
		return nil, err
	}

	return gactor.Ask(ctx, pid, msg, defaultSessionAskTimeout)
}

// sessionPID returns the PID of the session actor for sessionID,
// spawning one on first use.
//
// The runtime mutex is held across [gactor.ActorSystem.Spawn] so
// that concurrent invocations for the same sessionID serialise on
// the spawn path. goakt's Spawn is itself idempotent — if a live
// actor with the requested name already exists (e.g. spawned
// earlier by another runtime attached to the same system) it
// returns the existing PID rather than an error — but holding the
// lock keeps the internal session cache consistent without relying
// on that.
func (r *Runtime) sessionPID(ctx context.Context, sessionID string) (*gactor.PID, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.stopped {
		return nil, ErrRuntimeStopped
	}

	if pid, ok := r.sessions[sessionID]; ok {
		if pid.IsRunning() {
			return pid, nil
		}

		// The cached actor has been passivated or otherwise stopped.
		// Drop it and re-spawn below so the next turn re-hydrates
		// from the configured store.
		delete(r.sessions, sessionID)
	}

	actorName := attoactor.SessionActorName(sessionID)

	// In cluster mode the session actor may already live on a peer pod
	// (consistent-hash placement or a prior request that won the race).
	// Look it up first so we hand back the existing PID instead of
	// failing the local Spawn with ErrActorAlreadyExists.
	if pid, err := r.sys.ActorOf(ctx, actorName); err == nil && pid != nil {
		r.sessions[sessionID] = pid
		return pid, nil
	}

	pid, err := r.sys.Spawn(ctx, actorName, attoactor.NewSession(), r.spawnOptions()...)
	if err != nil {
		// Lost the race to a concurrent spawn on this or a peer pod —
		// resolve the live PID and return it.
		if lookup, lookupErr := r.sys.ActorOf(ctx, actorName); lookupErr == nil && lookup != nil {
			r.sessions[sessionID] = lookup
			return lookup, nil
		}

		return nil, fmt.Errorf("atto: spawn session actor %q: %w", actorName, err)
	}

	r.sessions[sessionID] = pid

	return pid, nil
}

// spawnOptions returns the [gactor.SpawnOption] list applied to
// every session-actor spawn: stashing for in-flight serialisation
// and (when configured) a time-based passivation strategy that
// releases idle sessions back to the configured store.
func (r *Runtime) spawnOptions() []gactor.SpawnOption {
	opts := []gactor.SpawnOption{gactor.WithStashing()}

	if r.passivationAfter > 0 {
		opts = append(opts, gactor.WithPassivationStrategy(passivation.NewTimeBasedStrategy(r.passivationAfter)))
	}

	return opts
}

// drainAndYield reads events until the channel is closed, then
// reads the terminal error from done and yields it as the
// iterator's final value (when non-nil).
func drainAndYield(
	events <-chan *session.Event,
	done <-chan error,
	yield func(*session.Event, error) bool,
) {
	for ev := range events {
		if !yield(ev, nil) {
			drain(events, done)
			return
		}
	}

	if err, ok := <-done; ok && err != nil {
		yield(nil, err)
	}
}

// drain consumes any remaining events and the terminal error so
// the worker goroutine can exit when the consumer breaks early.
func drain(events <-chan *session.Event, done <-chan error) {
	//revive:disable-next-line:empty-block // discarding events during drain
	for range events {
	}

	<-done
}