feat(adk): dream improvements

Author: N3kox

adk/middlewares/automemory/dream/api.go

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2026 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dream
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/cloudwego/eino/adk"
+)
+
+// This file holds the operations a caller invokes directly on a dream: starting a
+// run, querying its status, and canceling it. The scheduled middleware is created
+// with New in middleware.go.
+
+// RunRequest parameterizes a manual dream invocation.
+type RunRequest struct {
+	// SessionID identifies the current session.
+	// Optional. When empty, Config.SessionID is used.
+	SessionID string
+}
+
+// Run executes a dream immediately, without schedule gating, and returns the job id.
+//
+// The job's lifecycle is observable through GetDreamStatus and cancelable through
+// CancelDream using the returned id (the in-process default store only supports
+// same-process queries; inject a shared store for cross-process visibility).
+//
+// Run acquires the per-memory-directory run lock so a manual dream does not write
+// concurrently with a scheduled run sharing the same store; if the lock is held, Run
+// returns an empty id and a nil error. The input memory directory is never modified:
+// the model edits a staged working copy that is promoted to the output directory on
+// success.
+func Run[M adk.MessageType](ctx context.Context, cfg *Config[M], req *RunRequest) (string, error) {
+	cfg = cloneConfig(cfg)
+	if err := applyCoreDefaults(ctx, cfg); err != nil {
+		return "", err
+	}
+	m, err := newMiddleware(cfg, nil)
+	if err != nil {
+		return "", err
+	}
+	if req == nil {
+		req = &RunRequest{}
+	}
+	sessionID := strings.TrimSpace(req.SessionID)
+	if sessionID == "" {
+		sessionID = strings.TrimSpace(cfg.SessionID)
+	}
+
+	if store := m.store(); store != nil {
+		unlock, ok, lockErr := store.AcquireLock(ctx, runLockKey(m.resolvedMemoryDir), m.lockTTL())
+		if lockErr != nil || !ok {
+			return "", lockErr
+		}
+		defer func() { _ = unlock(ctx) }()
+	}
+
+	job := m.newJob(sessionID, nil)
+	m.persistJob(ctx, job)
+	if err := m.executeJob(ctx, job, sessionID, nil); err != nil {
+		return job.ID, err
+	}
+	return job.ID, nil
+}
+
+// GetDreamStatus returns the current DreamJob record for jobID. It returns (nil, nil)
+// when no such job exists (for example after the retention TTL elapsed).
+func GetDreamStatus(ctx context.Context, store KVStore, jobID string) (*DreamJob, error) {
+	if store == nil {
+		return nil, fmt.Errorf("dream: nil store")
+	}
+	return getJob(ctx, store, jobID)
+}
+
+// CancelDream requests cancellation of a pending or running dream job. It marks the
+// job canceled in the store and signals any in-process run to abort. Canceling a job
+// that has already reached a terminal state is a no-op. Cross-process runs observe
+// the canceled status on their next iteration check and stop best-effort.
+func CancelDream(ctx context.Context, store KVStore, jobID string) error {
+	if store == nil {
+		return fmt.Errorf("dream: nil store")
+	}
+	job, err := getJob(ctx, store, jobID)
+	if err != nil {
+		return err
+	}
+	if job == nil {
+		return fmt.Errorf("dream: job not found: %s", jobID)
+	}
+	if job.Status.IsTerminal() {
+		return nil
+	}
+	job.Status = DreamStatusCanceled
+	job.EndedAt = time.Now()
+	if err := setJob(ctx, store, job, jobTTL); err != nil {
+		return err
+	}
+	signalCancel(jobID)
+	return nil
+}

adk/middlewares/automemory/dream/config.go

@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-// Package dream provides scheduled consolidation middleware built on top of
-// automemory-managed session files.
 package dream
 
 import (
@@ -24,41 +22,108 @@ import (
 	"time"
 
 	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/adk/filesystem"
 	"github.com/cloudwego/eino/adk/middlewares/automemory"
 	"github.com/cloudwego/eino/components/model"
 )
 
 const (
-	defaultMinInterval       = 24 * time.Hour
-	defaultMinTouchedSession = 5
-	defaultScanInterval      = 10 * time.Minute
-	defaultLockTTL           = time.Hour
+	defaultMinInterval            = 24 * time.Hour
+	defaultMinTouchedSession      = 5
+	defaultScanInterval           = 10 * time.Minute
+	defaultLockTTL                = time.Hour
+	defaultMaxConsecutiveFailures = 3
+	defaultMaxIterations          = 24
 )
 
-// OnError handles non-fatal dream errors.
+// errLocalKVStoreSingleProcess is surfaced through OnError when Store falls back to
+// the in-process default. It is a warning, not a fatal error: dreams still run, but
+// coordination (run lock, job status, scheduler touch counting) is not shared across
+// instances.
+var errLocalKVStoreSingleProcess = fmt.Errorf(
+	"dream: Store not set, using in-process KVStore; this is single-process only, so run " +
+		"locks and job status are not shared across instances and scheduled dreams will not " +
+		"trigger when the middleware is constructed per session across instances " +
+		"(inject a shared, durable KVStore in production)")
+
+// errCountGateNeedsSessionID is surfaced through OnError when the scheduled count
+// gate is active (MinTouchedSession > 1) but no SessionID is configured. The touched
+// set is keyed by SessionID, so without one every run records the same empty member,
+// the distinct-session count is stuck at 1, and the gate never opens. It is a
+// warning, not a fatal error.
+var errCountGateNeedsSessionID = fmt.Errorf(
+	"dream: MinTouchedSession > 1 but Config.SessionID is empty; touched sessions are " +
+		"counted by SessionID, so the count stays at 1 and scheduled dreams will never " +
+		"trigger (set a per-session SessionID, or set MinTouchedSession to 1 to gate on " +
+		"MinInterval alone)")
+
+// OnError handles non-fatal dream errors. The stage argument is one of the
+// OnErrorStage* constants identifying where the failure occurred.
 // Optional. Nil means ignore the error.
-type OnError func(ctx context.Context, stage string, err error)
+type OnError func(ctx context.Context, stage ErrorStage, err error)
 
 // HandleIterator handles the dream sub-agent event stream.
 // Optional. Nil means dream drains the iterator itself.
 type HandleIterator[M adk.MessageType] func(ctx context.Context, iter *adk.AsyncIterator[*adk.TypedAgentEvent[M]]) error
 
-// Config configures auto dream for both `New(...)` and `Run(...)`.
+// Config holds the fields a single dream run needs. It is the input to Run(...) and
+// is embedded by MiddlewareConfig for the scheduled path, so a one-shot caller never
+// sees the scheduler knobs.
 type Config[M adk.MessageType] struct {
 	// MemoryDirectory is the memory root directory.
 	// Required. Relative paths are resolved during init.
+	// The dream run reads from this directory but never modifies it: the model
+	// edits a working copy under StagingDirectory, which is then promoted to
+	// OutputDirectory.
 	MemoryDirectory string
 
+	// OutputDirectory is where consolidated memory is written when a run succeeds.
+	// Optional. Default: MemoryDirectory (in-place consolidation).
+	//
+	// When equal to MemoryDirectory, the staged result is promoted over the source
+	// after the run. The source is left untouched until promotion, so a failed or
+	// canceled run never leaves the source half-processed. Promotion is a
+	// best-effort copy (not atomic across files); a warning is reported via OnError.
+	OutputDirectory string
+
+	// StagingDirectory is the root under which per-run working copies are created
+	// (one subdirectory per job id). The model's writes are bounded to the staging
+	// subdirectory; the source MemoryDirectory is copied in before the run.
+	// Optional. Default: a directory under os.TempDir().
+	StagingDirectory string
+
 	// MemoryBackend reads and updates memory files.
 	// Required.
 	MemoryBackend automemory.Backend
 
+	// Shell, when set, is used only to clean up the staging subdirectory after a
+	// successful promotion (rm -rf).
+	// Optional. When nil, the Shell is auto-derived from MemoryBackend if that value
+	// also implements filesystem.Shell (the common case for sandbox/filesystem
+	// backends that satisfy both interfaces in one struct), so it need not be
+	// configured twice. Set this field only to override that, or to supply a Shell
+	// when MemoryBackend is not one. When neither yields a Shell, staging directories
+	// are left in place under StagingDirectory.
+	Shell filesystem.Shell
+
 	// Model is the model used by the internal dream agent.
 	// Required.
 	Model model.BaseModel[M]
 
+	// MaxIterations caps the dream agent's tool-call loop.
+	// Consolidation reads the index, skims and reads multiple topic files, then
+	// writes/edits several files plus the index, so this needs headroom on large
+	// memory directories.
+	// Optional. Default: 24.
+	MaxIterations int
+
 	// SessionID is the current logical session ID.
 	// Optional. When empty, dream runs without cross-turn session grouping.
+	//
+	// For the scheduled middleware, the trigger counts distinct SessionIDs that have
+	// touched the memory directory. A per-session SessionID is therefore required for
+	// MinTouchedSession > 1 to be meaningful; without one the count stays at 1. New
+	// warns through OnError when this is misconfigured.
 	SessionID string
 
 	// OnError handles non-fatal runtime errors.
@@ -77,86 +142,115 @@ type Config[M adk.MessageType] struct {
 	//   - manual `Run(...)` searches the provided/current session only
 	SessionStore adk.SessionEventStore[M]
 
-	// Schedule controls middleware-triggered runs only.
-	// Optional. `Run(...)` ignores it.
-	Schedule *ScheduleConfig
+	// Store persists job records and the per-memory-directory run lock, enabling
+	// GetDreamStatus/CancelDream and preventing concurrent writes to the same memory
+	// directory. The scheduled path additionally uses it for touched sessions and
+	// schedule state.
+	// Optional. Default: in-process store from NewLocalKVStore (single-process only;
+	// a warning is emitted through OnError when this default is used).
+	//
+	// See KVStore for why production deployments MUST inject a shared, durable store.
+	Store KVStore
+
+	// LockTTL is the lease for the per-memory-directory run lock.
+	// It must comfortably exceed the longest expected dream runtime: if a run
+	// outlives the lease, another process may acquire the lock and write the same
+	// memory directory concurrently.
+	// Optional. Default: 1h.
+	LockTTL time.Duration
 
 	// HandleIterator overrides iterator consumption.
 	// Optional. Default: nil.
 	HandleIterator HandleIterator[M]
 }
 
-// ScheduleConfig controls middleware-triggered runs.
-type ScheduleConfig struct {
+// MiddlewareConfig configures the scheduled dream middleware created by New(...).
+// It embeds Config and adds the trigger knobs that only apply to automatic,
+// middleware-driven runs.
+type MiddlewareConfig[M adk.MessageType] struct {
+	Config[M]
+
 	// MinInterval is the minimum interval between successful runs.
 	// Optional. Default: 24h.
 	MinInterval time.Duration
 
-	// MinTouchedSession is the minimum touched-session count before a run.
+	// MinTouchedSession is the minimum number of distinct sessions (by Config.SessionID)
+	// that must touch the memory directory before a run. Set it to 1 to gate on
+	// MinInterval alone. Values > 1 require a per-session Config.SessionID.
 	// Optional. Default: 5.
 	MinTouchedSession int
 
 	// ScanInterval is the retry delay when the session threshold is not met.
 	// Optional. Default: 10m.
 	ScanInterval time.Duration
 
-	// LockTTL is the lease for the per-memory-directory run lock.
-	// Optional. Default: 1h.
-	LockTTL time.Duration
-
-	// Store persists touched sessions, schedule state, and run locks.
-	// Optional. Default: in-process `LocalStore`.
-	Store Store
+	// MaxConsecutiveFailures caps how many times a failing run is retried against
+	// the same unconsolidated window before the window is advanced to avoid
+	// replaying the same sessions forever.
+	// Optional. Default: 3.
+	MaxConsecutiveFailures int
 
 	// RunInline runs triggered dreams in the `AfterAgent` call path.
 	// Optional. Default: false.
 	RunInline bool
 }
 
-func applyCoreDefaults[M adk.MessageType](cfg *Config[M]) error {
+// scheduleParams carries the resolved scheduler knobs onto the middleware. It is nil
+// for a Run(...)-only construction, which never triggers on a schedule.
+type scheduleParams struct {
+	minInterval            time.Duration
+	minTouchedSession      int
+	scanInterval           time.Duration
+	maxConsecutiveFailures int
+	runInline              bool
+}
+
+func applyCoreDefaults[M adk.MessageType](ctx context.Context, cfg *Config[M]) error {
 	if cfg == nil {
 		return fmt.Errorf("auto dream config: nil")
 	}
 	if cfg.MemoryDirectory == "" || cfg.MemoryBackend == nil || cfg.Model == nil {
 		return fmt.Errorf("auto dream config: invalid")
 	}
+	if cfg.LockTTL <= 0 {
+		cfg.LockTTL = defaultLockTTL
+	}
+	if cfg.Store == nil {
+		cfg.Store = NewLocalKVStore()
+		if cfg.OnError != nil {
+			cfg.OnError(ctx, OnErrorStageInit, errLocalKVStoreSingleProcess)
+		}
+	}
 	return nil
 }
 
 func cloneConfig[M adk.MessageType](cfg *Config[M]) *Config[M] {
 	if cfg == nil {
 		return nil
 	}
-
 	cp := *cfg
-	if cfg.Schedule != nil {
-		scheduleCopy := *cfg.Schedule
-		cp.Schedule = &scheduleCopy
-	}
 	return &cp
 }
 
-func applyScheduleDefaults[M adk.MessageType](cfg *Config[M]) error {
-	if err := applyCoreDefaults(cfg); err != nil {
-		return err
-	}
-	if cfg.Schedule == nil {
-		cfg.Schedule = &ScheduleConfig{}
-	}
-	if cfg.Schedule.MinInterval <= 0 {
-		cfg.Schedule.MinInterval = defaultMinInterval
+func cloneMiddlewareConfig[M adk.MessageType](cfg *MiddlewareConfig[M]) *MiddlewareConfig[M] {
+	if cfg == nil {
+		return nil
 	}
-	if cfg.Schedule.MinTouchedSession <= 0 {
-		cfg.Schedule.MinTouchedSession = defaultMinTouchedSession
+	cp := *cfg
+	return &cp
+}
+
+func applyScheduleDefaults[M adk.MessageType](cfg *MiddlewareConfig[M]) {
+	if cfg.MinInterval <= 0 {
+		cfg.MinInterval = defaultMinInterval
 	}
-	if cfg.Schedule.ScanInterval <= 0 {
-		cfg.Schedule.ScanInterval = defaultScanInterval
+	if cfg.MinTouchedSession <= 0 {
+		cfg.MinTouchedSession = defaultMinTouchedSession
 	}
-	if cfg.Schedule.LockTTL <= 0 {
-		cfg.Schedule.LockTTL = defaultLockTTL
+	if cfg.ScanInterval <= 0 {
+		cfg.ScanInterval = defaultScanInterval
 	}
-	if cfg.Schedule.Store == nil {
-		cfg.Schedule.Store = NewLocalStore()
+	if cfg.MaxConsecutiveFailures <= 0 {
+		cfg.MaxConsecutiveFailures = defaultMaxConsecutiveFailures
 	}
-	return nil
 }