feat: agent memory system (#1659)

Author: yottahmd

api/v1/api.gen.go

@@ -71,6 +71,7 @@ type API struct {
 	dagStore    exec.DAGStore // For resolving DAG file paths
 	environment EnvironmentInfo
 	hooks       *Hooks
+	memoryStore MemoryStore
 }
 
 // APIConfig contains configuration for the API.
@@ -83,6 +84,7 @@ type APIConfig struct {
 	DAGStore     exec.DAGStore // For resolving DAG file paths
 	Environment  EnvironmentInfo
 	Hooks        *Hooks
+	MemoryStore  MemoryStore
 }
 
 // SessionWithState is a session with its current state.
@@ -110,6 +112,7 @@ func NewAPI(cfg APIConfig) *API {
 		dagStore:    cfg.DAGStore,
 		environment: cfg.Environment,
 		hooks:       cfg.Hooks,
+		memoryStore: cfg.MemoryStore,
 	}
 }
 
@@ -277,13 +280,14 @@ func (a *API) createMessageCallback(id string) func(ctx context.Context, msg Mes
 }
 
 // persistNewSession saves a new session to the store if configured.
-func (a *API) persistNewSession(ctx context.Context, id, userID string, now time.Time) {
+func (a *API) persistNewSession(ctx context.Context, id, userID, dagName string, now time.Time) {
 	if a.store == nil {
 		return
 	}
 	sess := &Session{
 		ID:        id,
 		UserID:    userID,
+		DAGName:   dagName,
 		CreatedAt: now,
 		UpdatedAt: now,
 	}
@@ -340,6 +344,13 @@ func (a *API) handleNewSession(w http.ResponseWriter, r *http.Request) {
 	id := uuid.New().String()
 	now := time.Now()
 
+	// Extract primary DAG name from resolved contexts for per-DAG memory
+	resolved := a.resolveContexts(r.Context(), req.DAGContexts)
+	var dagName string
+	if len(resolved) > 0 {
+		dagName = resolved[0].DAGName
+	}
+
 	mgr := NewSessionManager(SessionManagerConfig{
 		ID:              id,
 		UserID:          userID,
@@ -353,12 +364,14 @@ func (a *API) handleNewSession(w http.ResponseWriter, r *http.Request) {
 		IPAddress:       ipAddress,
 		InputCostPer1M:  modelCfg.InputCostPer1M,
 		OutputCostPer1M: modelCfg.OutputCostPer1M,
+		MemoryStore:     a.memoryStore,
+		DAGName:         dagName,
 	})
 
-	a.persistNewSession(r.Context(), id, userID, now)
+	a.persistNewSession(r.Context(), id, userID, dagName, now)
 	a.sessions.Store(id, mgr)
 
-	messageWithContext := a.formatMessage(r.Context(), req.Message, req.DAGContexts)
+	messageWithContext := formatMessageWithContexts(req.Message, resolved)
 	if err := mgr.AcceptUserMessage(r.Context(), provider, model, modelCfg.Model, messageWithContext); err != nil {
 		a.logger.Error("Failed to accept user message", "error", err)
 		a.respondError(w, http.StatusInternalServerError, api.ErrorCodeInternalError, "Failed to process message")
@@ -609,6 +622,8 @@ func (a *API) reactivateSession(ctx context.Context, id, userID, username, ipAdd
 		Hooks:       a.hooks,
 		Username:    username,
 		IPAddress:   ipAddress,
+		MemoryStore: a.memoryStore,
+		DAGName:     sess.DAGName,
 	})
 	a.sessions.Store(id, mgr)
 

api/v1/api.yaml

@@ -0,0 +1,39 @@
+package agent
+
+import "context"
+
+// MemoryStore provides access to agent memory files.
+// All implementations must be safe for concurrent use.
+type MemoryStore interface {
+	// LoadGlobalMemory reads the global MEMORY.md, truncated to maxLines.
+	LoadGlobalMemory(ctx context.Context) (string, error)
+
+	// LoadDAGMemory reads the MEMORY.md for a specific DAG, truncated to maxLines.
+	LoadDAGMemory(ctx context.Context, dagName string) (string, error)
+
+	// SaveGlobalMemory writes content to the global MEMORY.md.
+	SaveGlobalMemory(ctx context.Context, content string) error
+
+	// SaveDAGMemory writes content to a DAG-specific MEMORY.md.
+	SaveDAGMemory(ctx context.Context, dagName string, content string) error
+
+	// MemoryDir returns the root memory directory path.
+	MemoryDir() string
+
+	// ListDAGMemories returns the names of all DAGs that have memory files.
+	ListDAGMemories(ctx context.Context) ([]string, error)
+
+	// DeleteGlobalMemory removes the global MEMORY.md file.
+	DeleteGlobalMemory(ctx context.Context) error
+
+	// DeleteDAGMemory removes a DAG-specific MEMORY.md file.
+	DeleteDAGMemory(ctx context.Context, dagName string) error
+}
+
+// MemoryContent holds loaded memory for system prompt injection.
+type MemoryContent struct {
+	GlobalMemory string // Contents of global MEMORY.md (truncated)
+	DAGMemory    string // Contents of per-DAG MEMORY.md (truncated)
+	DAGName      string // Name of the DAG (empty if no DAG context)
+	MemoryDir    string // Root memory directory path
+}

internal/agent/api.go

@@ -42,6 +42,8 @@ type SessionManager struct {
 	inputCostPer1M  float64
 	outputCostPer1M float64
 	totalCost       float64
+	memoryStore     MemoryStore
+	dagName         string
 }
 
 // SessionManagerConfig contains configuration for creating a SessionManager.
@@ -61,6 +63,8 @@ type SessionManagerConfig struct {
 	IPAddress       string
 	InputCostPer1M  float64
 	OutputCostPer1M float64
+	MemoryStore     MemoryStore
+	DAGName         string
 }
 
 // NewSessionManager creates a new SessionManager.
@@ -97,6 +101,8 @@ func NewSessionManager(cfg SessionManagerConfig) *SessionManager {
 		ipAddress:       cfg.IPAddress,
 		inputCostPer1M:  cfg.InputCostPer1M,
 		outputCostPer1M: cfg.OutputCostPer1M,
+		memoryStore:     cfg.MemoryStore,
+		dagName:         cfg.DAGName,
 	}
 }
 
@@ -200,6 +206,7 @@ func (sm *SessionManager) GetSession() Session {
 	return Session{
 		ID:        sm.id,
 		UserID:    sm.userID,
+		DAGName:   sm.dagName,
 		CreatedAt: sm.createdAt,
 		UpdatedAt: sm.lastActivity,
 	}
@@ -353,14 +360,15 @@ func (sm *SessionManager) ensureLoop(provider llm.Provider, modelID string, reso
 
 // createLoop creates a new Loop instance with the current configuration.
 func (sm *SessionManager) createLoop(provider llm.Provider, model string, history []llm.Message, safeMode bool) *Loop {
+	memory := sm.loadMemory()
 	return NewLoop(LoopConfig{
 		Provider:         provider,
 		Model:            model,
 		History:          history,
 		Tools:            CreateTools(sm.environment.DAGsDir),
 		RecordMessage:    sm.createRecordMessageFunc(),
 		Logger:           sm.logger,
-		SystemPrompt:     GenerateSystemPrompt(sm.environment, nil),
+		SystemPrompt:     GenerateSystemPrompt(sm.environment, nil, memory),
 		WorkingDir:       sm.workingDir,
 		SessionID:        sm.id,
 		OnWorking:        sm.SetWorking,
@@ -375,6 +383,31 @@ func (sm *SessionManager) createLoop(provider llm.Provider, model string, histor
 	})
 }
 
+// loadMemory loads memory content from the memory store.
+func (sm *SessionManager) loadMemory() MemoryContent {
+	if sm.memoryStore == nil {
+		return MemoryContent{}
+	}
+	ctx := context.Background()
+	global, err := sm.memoryStore.LoadGlobalMemory(ctx)
+	if err != nil {
+		sm.logger.Debug("failed to load global memory", "error", err)
+	}
+	var dagMem string
+	if sm.dagName != "" {
+		dagMem, err = sm.memoryStore.LoadDAGMemory(ctx, sm.dagName)
+		if err != nil {
+			sm.logger.Debug("failed to load DAG memory", "error", err, "dag_name", sm.dagName)
+		}
+	}
+	return MemoryContent{
+		GlobalMemory: global,
+		DAGMemory:    dagMem,
+		DAGName:      sm.dagName,
+		MemoryDir:    sm.memoryStore.MemoryDir(),
+	}
+}
+
 // runLoop executes the session loop and handles cleanup.
 func (sm *SessionManager) runLoop(ctx context.Context, loop *Loop) {
 	defer func() {

internal/agent/memory.go

@@ -34,15 +34,17 @@ type CurrentDAG struct {
 type systemPromptData struct {
 	EnvironmentInfo
 	CurrentDAG *CurrentDAG
+	Memory     MemoryContent
 }
 
-// GenerateSystemPrompt renders the system prompt template with the given environment
-// and optional DAG context.
-func GenerateSystemPrompt(env EnvironmentInfo, currentDAG *CurrentDAG) string {
+// GenerateSystemPrompt renders the system prompt template with the given environment,
+// optional DAG context, and memory content.
+func GenerateSystemPrompt(env EnvironmentInfo, currentDAG *CurrentDAG, memory MemoryContent) string {
 	var buf bytes.Buffer
 	data := systemPromptData{
 		EnvironmentInfo: env,
 		CurrentDAG:      currentDAG,
+		Memory:          memory,
 	}
 	if err := systemPromptTemplate.Execute(&buf, data); err != nil {
 		return fallbackPrompt(env)
@@ -52,5 +54,5 @@ func GenerateSystemPrompt(env EnvironmentInfo, currentDAG *CurrentDAG) string {
 
 // fallbackPrompt returns a basic prompt when template execution fails.
 func fallbackPrompt(env EnvironmentInfo) string {
-	return "You are Hermio, an AI assistant for DAG workflows. DAGs Directory: " + env.DAGsDir
+	return "You are Tsumugi, an AI assistant for DAG workflows. DAGs Directory: " + env.DAGsDir
 }

internal/agent/session.go

@@ -31,6 +31,29 @@ Status: {{.CurrentDAG.Status}}{{end}}
 {{end}}
 </environment>
 
+{{if or .Memory.GlobalMemory .Memory.DAGMemory}}
+<memory>
+{{if .Memory.GlobalMemory}}
+<global_memory>
+{{.Memory.GlobalMemory}}
+</global_memory>
+{{end}}
+{{if .Memory.DAGMemory}}
+<dag_memory dag="{{.Memory.DAGName}}">
+{{.Memory.DAGMemory}}
+</dag_memory>
+{{end}}
+</memory>
+{{end}}
+
+{{if .Memory.MemoryDir}}
+<memory_paths>
+- Memory directory: {{.Memory.MemoryDir}}
+- Global memory: {{.Memory.MemoryDir}}/MEMORY.md
+{{if .Memory.DAGName}}- DAG memory: {{.Memory.MemoryDir}}/dags/{{.Memory.DAGName}}/MEMORY.md
+{{end}}</memory_paths>
+{{end}}
+
 <rules>
 <safety>
 Do not start DAGs unless the user explicitly requests execution—starting a DAG triggers real processes that may have unintended side effects.
@@ -131,6 +154,32 @@ Do not use placeholders unless user explicitly requests dummy data for testing.
 3. Proceed only after user confirms configuration is complete.
 </configuration>
 
+{{if .Memory.MemoryDir}}
+<memory_management>
+Consult your memory files to build on previous experience.
+When you learn something reusable, record it with DAG-first routing.
+
+Paths:
+- Global: {{.Memory.MemoryDir}}/MEMORY.md
+- Per-DAG: {{.Memory.MemoryDir}}/dags/<dag-name>/MEMORY.md
+
+Rules:
+- MEMORY.md is loaded into context automatically; keep under 200 lines
+- Use `read` and `patch` to manage memory files
+- Organize by topic, not chronologically
+- If DAG context is available, save memory to Per-DAG by default (not Global)
+- After creating or updating a DAG, if anything should be remembered, create/update that DAG's memory file
+- DAG memory can include DAG-specific config assumptions, pitfalls, fixes, and debugging playbooks
+- Global memory is only for cross-DAG or user-wide stable preferences/policies
+- If unsure whether knowledge is global or DAG-specific, store it in Per-DAG memory
+- If no DAG context is available, ask the user before writing to Global memory
+- Save stable patterns, environment details, user preferences, and debugging insights
+- Do not save session-specific state, secrets, or unverified info
+- When the user says "remember this", save it immediately
+- When the user says "forget this", remove it
+</memory_management>
+{{end}}
+
 <schema_reference>
 Use `read_schema` to look up DAG YAML structure:
 - Root fields: `path: ""`