multica/server/internal/daemon/execenv/codex_memory.go

package execenv

import (
	"fmt"
	"log/slog"
	"os"
	"regexp"
	"strings"
)

// Disable native Codex memory in task-local configuration by default to keep
// host or previous-task memories from crossing workspace boundaries. The user's
// global configuration is untouched; MULTICA_CODEX_MEMORY opts into native memory.
//
// Write separate feature and memory switches. Insert overrides inside existing
// TOML tables when present; dotted keys at the root would otherwise redefine
// tables, including tables created implicitly by earlier dotted keys.

// MulticaCodexMemoryEnv is the env var users can set to keep Codex
// native memory enabled inside daemon-managed tasks. Anything truthy
// (1, true, yes, on; case-insensitive) keeps the feature on; everything
// else (including unset) disables it.
const MulticaCodexMemoryEnv = "MULTICA_CODEX_MEMORY"

const (
	multicaMemoryFeatureBeginMarker = "# BEGIN multica-managed memory-feature (do not edit; regenerated by daemon)"
	multicaMemoryFeatureEndMarker   = "# END multica-managed memory-feature"
	multicaMemoryConfigBeginMarker  = "# BEGIN multica-managed memory-config (do not edit; regenerated by daemon)"
	multicaMemoryConfigEndMarker    = "# END multica-managed memory-config"
)

// `\n*` rather than `\n?` so reruns don't accumulate blank lines when
// these blocks coexist with the sandbox / multi-agent blocks.
var memoryFeatureBlockRe = regexp.MustCompile(
	`(?ms)^` + regexp.QuoteMeta(multicaMemoryFeatureBeginMarker) +
		`.*?^` + regexp.QuoteMeta(multicaMemoryFeatureEndMarker) + `\n*`)

var memoryConfigBlockRe = regexp.MustCompile(
	`(?ms)^` + regexp.QuoteMeta(multicaMemoryConfigBeginMarker) +
		`.*?^` + regexp.QuoteMeta(multicaMemoryConfigEndMarker) + `\n*`)

var (
	// matches a top-level `[memories]` table header.
	rootMemoriesTableHeaderRe = regexp.MustCompile(`^\s*\[\s*memories\s*\]\s*(?:#.*)?$`)
	// matches `memories = ...` inside a `[features]` table.
	featuresTableMemoriesRe = regexp.MustCompile(`^\s*memories\s*=`)
	// matches `features.memories = ...` at the TOML root.
	rootDottedFeaturesMemoriesRe = regexp.MustCompile(`^\s*features\s*\.\s*memories\s*=`)
	// matches `generate_memories = ...` / `use_memories = ...` inside a
	// `[memories]` table.
	memoriesTableGenerateRe = regexp.MustCompile(`^\s*generate_memories\s*=`)
	memoriesTableUseRe      = regexp.MustCompile(`^\s*use_memories\s*=`)
	// matches `memories.generate_memories = ...` / `memories.use_memories = ...`
	// at the TOML root.
	rootDottedMemoriesGenerateRe = regexp.MustCompile(`^\s*memories\s*\.\s*generate_memories\s*=`)
	rootDottedMemoriesUseRe      = regexp.MustCompile(`^\s*memories\s*\.\s*use_memories\s*=`)
)

// codexMemoryEnabled reports whether the user opted into keeping Codex
// native memory on for daemon-managed tasks.
func codexMemoryEnabled() bool {
	raw := strings.TrimSpace(os.Getenv(MulticaCodexMemoryEnv))
	switch strings.ToLower(raw) {
	case "1", "true", "yes", "on":
		return true
	}
	return false
}

// renderMulticaMemoryFeatureBlock returns the daemon-managed memory-feature
// block. The body uses `memories = false` when injected inside a
// `[features]` table, and `features.memories = false` otherwise.
func renderMulticaMemoryFeatureBlock(inFeaturesTable bool) string {
	var b strings.Builder
	b.WriteString(multicaMemoryFeatureBeginMarker)
	b.WriteString("\n")
	if inFeaturesTable {
		b.WriteString("memories = false\n")
	} else {
		b.WriteString("features.memories = false\n")
	}
	b.WriteString(multicaMemoryFeatureEndMarker)
	b.WriteString("\n")
	return b.String()
}

// renderMulticaMemoryConfigBlock returns the daemon-managed memory-config
// block. The body uses bare keys when injected inside a `[memories]` table
// and dotted-key form otherwise.
func renderMulticaMemoryConfigBlock(inMemoriesTable bool) string {
	var b strings.Builder
	b.WriteString(multicaMemoryConfigBeginMarker)
	b.WriteString("\n")
	if inMemoriesTable {
		b.WriteString("generate_memories = false\n")
		b.WriteString("use_memories = false\n")
	} else {
		b.WriteString("memories.generate_memories = false\n")
		b.WriteString("memories.use_memories = false\n")
	}
	b.WriteString(multicaMemoryConfigEndMarker)
	b.WriteString("\n")
	return b.String()
}

// stripUserMemoryDirectives removes user-written directives that would
// conflict with the managed blocks:
//
//   - `features.memories = ...` and `memories.{generate,use}_memories = ...`
//     at the TOML root (dotted-key form).
//   - `memories = ...` inside a top-level `[features]` table.
//   - `generate_memories = ...` / `use_memories = ...` inside a top-level
//     `[memories]` table.
//
// Other tables (`[features.experimental]`, `[memories.advanced]`, ...) are
// preserved untouched: they live under their own scope and don't redefine
// the keys the managed blocks set at the root / in the canonical tables.
func stripUserMemoryDirectives(content string) string {
	lines := strings.Split(content, "\n")
	out := make([]string, 0, len(lines))
	currentTable := "" // empty = TOML root
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if rootFeaturesTableHeaderRe.MatchString(line) {
			currentTable = "[features]"
			out = append(out, line)
			continue
		}
		if rootMemoriesTableHeaderRe.MatchString(line) {
			currentTable = "[memories]"
			out = append(out, line)
			continue
		}
		if strings.HasPrefix(trimmed, "[") {
			currentTable = trimmed
			out = append(out, line)
			continue
		}
		switch currentTable {
		case "":
			if rootDottedFeaturesMemoriesRe.MatchString(trimmed) ||
				rootDottedMemoriesGenerateRe.MatchString(trimmed) ||
				rootDottedMemoriesUseRe.MatchString(trimmed) {
				continue
			}
		case "[features]":
			if featuresTableMemoriesRe.MatchString(trimmed) {
				continue
			}
		case "[memories]":
			if memoriesTableGenerateRe.MatchString(trimmed) ||
				memoriesTableUseRe.MatchString(trimmed) {
				continue
			}
		}
		out = append(out, line)
	}
	return strings.Join(out, "\n")
}

// hasRootMemoriesTable reports whether the file contains a top-level
// `[memories]` table header. Sub-tables like `[memories.advanced]` do NOT
// count: they implicitly create `memories` but don't conflict with a
// root-level `memories.generate_memories` dotted key.
func hasRootMemoriesTable(content string) bool {
	for _, line := range strings.Split(content, "\n") {
		if rootMemoriesTableHeaderRe.MatchString(line) {
			return true
		}
	}
	return false
}

// injectMemoryFeatureBlockIntoFeaturesTable inserts the in-table
// memory-feature block immediately after the first `[features]` header
// line. Caller must have already stripped any prior managed block and any
// user-set `memories` directive from inside the table.
func injectMemoryFeatureBlockIntoFeaturesTable(content string) string {
	return injectAfterHeader(content, rootFeaturesTableHeaderRe, renderMulticaMemoryFeatureBlock(true))
}

// injectMemoryConfigBlockIntoMemoriesTable inserts the in-table
// memory-config block immediately after the first `[memories]` header
// line. Caller must have already stripped any prior managed block and any
// user-set `generate_memories` / `use_memories` directive from inside the
// table.
func injectMemoryConfigBlockIntoMemoriesTable(content string) string {
	return injectAfterHeader(content, rootMemoriesTableHeaderRe, renderMulticaMemoryConfigBlock(true))
}

// injectAfterHeader inserts block right after the first line matching
// headerRe. Blank lines that would otherwise sit between the inserted
// block and the next non-blank line are dropped so the END marker is
// butted against the next table / key. Without this normalization, the
// `\n*` quantifier in memoryFeatureBlockRe / memoryConfigBlockRe would
// greedily consume those blank lines on the next strip, and the rewrite
// would no longer be byte-exact idempotent.
//
// The trailing "" element produced by Split when the file ends with `\n`
// is preserved so the rewritten file keeps its EOF newline.
func injectAfterHeader(content string, headerRe *regexp.Regexp, block string) string {
	blockLines := strings.Split(strings.TrimRight(block, "\n"), "\n")
	lines := strings.Split(content, "\n")
	for i, line := range lines {
		if !headerRe.MatchString(line) {
			continue
		}
		tail := lines[i+1:]
		// Drop leading blanks that follow the inserted block, but never
		// strip the lone trailing-newline marker.
		for len(tail) > 1 && tail[0] == "" {
			tail = tail[1:]
		}
		out := make([]string, 0, len(lines[:i+1])+len(blockLines)+len(tail))
		out = append(out, lines[:i+1]...)
		out = append(out, blockLines...)
		out = append(out, tail...)
		return strings.Join(out, "\n")
	}
	return content
}

// ensureCodexMemoryConfig writes the daemon-managed memory blocks into the
// per-task config.toml so Codex native memory stays disabled. Idempotent:
// running it twice produces the same file.
//
// When MULTICA_CODEX_MEMORY is set to a truthy value, the function is a
// no-op — the user has explicitly opted into Codex native memory and
// accepts the leak risk. Toggling the env var across prepare runs is not
// supported: the per-task config is short-lived (recreated per task), so
// users should set the var once at daemon start.
func ensureCodexMemoryConfig(configPath string, logger *slog.Logger) error {
	if codexMemoryEnabled() {
		if logger != nil {
			logger.Info("codex memory: leaving Codex native memory untouched per MULTICA_CODEX_MEMORY",
				"config_path", configPath,
			)
		}
		return nil
	}

	data, err := os.ReadFile(configPath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("read config.toml: %w", err)
	}
	existing := string(data)

	// Always strip any previously written managed blocks (root or in-table
	// form) so reruns and layout transitions stay clean.
	existing = memoryFeatureBlockRe.ReplaceAllString(existing, "")
	existing = memoryConfigBlockRe.ReplaceAllString(existing, "")
	// Strip user-set directives in every encoding; the managed blocks re-add
	// the canonical forms below.
	existing = stripUserMemoryDirectives(existing)

	hasFeatures := hasRootFeaturesTable(existing)
	hasMemories := hasRootMemoriesTable(existing)

	// In-table injections happen first so the root-prepended blocks land
	// at the very top of the file (above any user tables).
	if hasFeatures {
		existing = injectMemoryFeatureBlockIntoFeaturesTable(existing)
	}
	if hasMemories {
		existing = injectMemoryConfigBlockIntoMemoriesTable(existing)
	}

	// Build the root-form blocks (if any) and prepend them.
	var prepend strings.Builder
	if !hasFeatures {
		prepend.WriteString(renderMulticaMemoryFeatureBlock(false))
	}
	if !hasMemories {
		prepend.WriteString(renderMulticaMemoryConfigBlock(false))
	}
	prependStr := prepend.String()

	existing = strings.TrimLeft(existing, "\n")
	if prependStr != "" {
		if existing == "" {
			existing = prependStr
		} else {
			existing = prependStr + "\n" + existing
		}
	}

	if existing == string(data) {
		return nil
	}

	if err := os.WriteFile(configPath, []byte(existing), 0o644); err != nil {
		return fmt.Errorf("write config.toml: %w", err)
	}
	return nil
}