1  import { mkdir, readdir, readFile, unlink, writeFile } from 'fs/promises'
  2  import { join } from 'path'
  3  import { z } from 'zod/v4'
  4  import { getCwd } from '../../utils/cwd.js'
  5  import { logForDebugging } from '../../utils/debug.js'
  6  import { lazySchema } from '../../utils/lazySchema.js'
  7  import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'
  8  import { type AgentMemoryScope, getAgentMemoryDir } from './agentMemory.js'
  9  
 10  const SNAPSHOT_BASE = 'agent-memory-snapshots'
 11  const SNAPSHOT_JSON = 'snapshot.json'
 12  const SYNCED_JSON = '.snapshot-synced.json'
 13  
 14  const snapshotMetaSchema = lazySchema(() =>
 15    z.object({
 16      updatedAt: z.string().min(1),
 17    }),
 18  )
 19  
 20  const syncedMetaSchema = lazySchema(() =>
 21    z.object({
 22      syncedFrom: z.string().min(1),
 23    }),
 24  )
 25  type SyncedMeta = z.infer<ReturnType<typeof syncedMetaSchema>>
 26  
 27  /**
 28   * Returns the path to the snapshot directory for an agent in the current project.
 29   * e.g., <cwd>/.claude/agent-memory-snapshots/<agentType>/
 30   */
 31  export function getSnapshotDirForAgent(agentType: string): string {
 32    return join(getCwd(), '.claude', SNAPSHOT_BASE, agentType)
 33  }
 34  
 35  function getSnapshotJsonPath(agentType: string): string {
 36    return join(getSnapshotDirForAgent(agentType), SNAPSHOT_JSON)
 37  }
 38  
 39  function getSyncedJsonPath(agentType: string, scope: AgentMemoryScope): string {
 40    return join(getAgentMemoryDir(agentType, scope), SYNCED_JSON)
 41  }
 42  
 43  async function readJsonFile<T>(
 44    path: string,
 45    schema: z.ZodType<T>,
 46  ): Promise<T | null> {
 47    try {
 48      const content = await readFile(path, { encoding: 'utf-8' })
 49      const result = schema.safeParse(jsonParse(content))
 50      return result.success ? result.data : null
 51    } catch {
 52      return null
 53    }
 54  }
 55  
 56  async function copySnapshotToLocal(
 57    agentType: string,
 58    scope: AgentMemoryScope,
 59  ): Promise<void> {
 60    const snapshotMemDir = getSnapshotDirForAgent(agentType)
 61    const localMemDir = getAgentMemoryDir(agentType, scope)
 62  
 63    await mkdir(localMemDir, { recursive: true })
 64  
 65    try {
 66      const files = await readdir(snapshotMemDir, { withFileTypes: true })
 67      for (const dirent of files) {
 68        if (!dirent.isFile() || dirent.name === SNAPSHOT_JSON) continue
 69        const content = await readFile(join(snapshotMemDir, dirent.name), {
 70          encoding: 'utf-8',
 71        })
 72        await writeFile(join(localMemDir, dirent.name), content)
 73      }
 74    } catch (e) {
 75      logForDebugging(`Failed to copy snapshot to local agent memory: ${e}`)
 76    }
 77  }
 78  
 79  async function saveSyncedMeta(
 80    agentType: string,
 81    scope: AgentMemoryScope,
 82    snapshotTimestamp: string,
 83  ): Promise<void> {
 84    const syncedPath = getSyncedJsonPath(agentType, scope)
 85    const localMemDir = getAgentMemoryDir(agentType, scope)
 86    await mkdir(localMemDir, { recursive: true })
 87    const meta: SyncedMeta = { syncedFrom: snapshotTimestamp }
 88    try {
 89      await writeFile(syncedPath, jsonStringify(meta))
 90    } catch (e) {
 91      logForDebugging(`Failed to save snapshot sync metadata: ${e}`)
 92    }
 93  }
 94  
 95  /**
 96   * Check if a snapshot exists and whether it's newer than what we last synced.
 97   */
 98  export async function checkAgentMemorySnapshot(
 99    agentType: string,
100    scope: AgentMemoryScope,
101  ): Promise<{
102    action: 'none' | 'initialize' | 'prompt-update'
103    snapshotTimestamp?: string
104  }> {
105    const snapshotMeta = await readJsonFile(
106      getSnapshotJsonPath(agentType),
107      snapshotMetaSchema(),
108    )
109  
110    if (!snapshotMeta) {
111      return { action: 'none' }
112    }
113  
114    const localMemDir = getAgentMemoryDir(agentType, scope)
115  
116    let hasLocalMemory = false
117    try {
118      const dirents = await readdir(localMemDir, { withFileTypes: true })
119      hasLocalMemory = dirents.some(d => d.isFile() && d.name.endsWith('.md'))
120    } catch {
121      // Directory doesn't exist
122    }
123  
124    if (!hasLocalMemory) {
125      return { action: 'initialize', snapshotTimestamp: snapshotMeta.updatedAt }
126    }
127  
128    const syncedMeta = await readJsonFile(
129      getSyncedJsonPath(agentType, scope),
130      syncedMetaSchema(),
131    )
132  
133    if (
134      !syncedMeta ||
135      new Date(snapshotMeta.updatedAt) > new Date(syncedMeta.syncedFrom)
136    ) {
137      return {
138        action: 'prompt-update',
139        snapshotTimestamp: snapshotMeta.updatedAt,
140      }
141    }
142  
143    return { action: 'none' }
144  }
145  
146  /**
147   * Initialize local agent memory from a snapshot (first-time setup).
148   */
149  export async function initializeFromSnapshot(
150    agentType: string,
151    scope: AgentMemoryScope,
152    snapshotTimestamp: string,
153  ): Promise<void> {
154    logForDebugging(
155      `Initializing agent memory for ${agentType} from project snapshot`,
156    )
157    await copySnapshotToLocal(agentType, scope)
158    await saveSyncedMeta(agentType, scope, snapshotTimestamp)
159  }
160  
161  /**
162   * Replace local agent memory with the snapshot.
163   */
164  export async function replaceFromSnapshot(
165    agentType: string,
166    scope: AgentMemoryScope,
167    snapshotTimestamp: string,
168  ): Promise<void> {
169    logForDebugging(
170      `Replacing agent memory for ${agentType} with project snapshot`,
171    )
172    // Remove existing .md files before copying to avoid orphans
173    const localMemDir = getAgentMemoryDir(agentType, scope)
174    try {
175      const existing = await readdir(localMemDir, { withFileTypes: true })
176      for (const dirent of existing) {
177        if (dirent.isFile() && dirent.name.endsWith('.md')) {
178          await unlink(join(localMemDir, dirent.name))
179        }
180      }
181    } catch {
182      // Directory may not exist yet
183    }
184    await copySnapshotToLocal(agentType, scope)
185    await saveSyncedMeta(agentType, scope, snapshotTimestamp)
186  }
187  
188  /**
189   * Mark the current snapshot as synced without changing local memory.
190   */
191  export async function markSnapshotSynced(
192    agentType: string,
193    scope: AgentMemoryScope,
194    snapshotTimestamp: string,
195  ): Promise<void> {
196    await saveSyncedMeta(agentType, scope, snapshotTimestamp)
197  }