File size: 11,809 Bytes
1dfeefd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce3c3bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1dfeefd
 
 
 
 
 
 
 
 
 
ce3c3bd
1dfeefd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce3c3bd
1dfeefd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
#!/bin/sh
# Durable state bridge for the agent harnesses.
#
# Active state always lives on the container's POSIX filesystem. The mounted
# bucket only receives short-lived, closed writes made by rsync (ordinary file
# trees) or SQLite's online backup API (live databases). This is intentional:
# hf-mount's streaming writer holds an open file in memory until close, while
# copying a database and its WAL independently can produce a torn backup.

set -u

MODE="${1:-}"
case "$MODE" in
  restore|checkpoint|checkpoint-final) ;;
  *) echo "usage: $0 restore|checkpoint|checkpoint-final" >&2; exit 2 ;;
esac

: "${DATA_DIR:?DATA_DIR is required}"
: "${AM_LOCAL:?AM_LOCAL is required}"

CODEX_HOME="${CODEX_HOME:-$AM_LOCAL/codex-home}"
CODEX_DURABLE="${CODEX_DURABLE:-$DATA_DIR/state/codex}"
CLAUDE_CONFIG_DIR="${CLAUDE_CONFIG_DIR:-$AM_LOCAL/agent-state/claude}"
CLAUDE_DURABLE="${CLAUDE_DURABLE:-$DATA_DIR/state/claude}"
GEMINI_CLI_HOME="${GEMINI_CLI_HOME:-$AM_LOCAL/agent-state/gemini-home}"
GEMINI_LIVE="${GEMINI_LIVE:-$GEMINI_CLI_HOME/.gemini}"
GEMINI_DURABLE="${GEMINI_DURABLE:-$DATA_DIR/state/gemini}"
OPENCLAW_STATE_DIR="${OPENCLAW_STATE_DIR:-$AM_LOCAL/oc-home/.openclaw}"
OPENCLAW_DURABLE="${OPENCLAW_DURABLE:-$DATA_DIR/state/openclaw-backup}"
OPENCODE_LIVE="${OPENCODE_LIVE:-$AM_LOCAL/opencode-share}"
OPENCODE_DURABLE="${OPENCODE_DURABLE:-$DATA_DIR/state/opencode}"
HERMES_LIVE="${HERMES_LIVE:-$AM_LOCAL/hermes}"
HERMES_DURABLE="${HERMES_DURABLE:-$DATA_DIR/state/hermes}"

LOCK="$AM_LOCAL/agent-state-checkpoint.lock"
mkdir -p "$AM_LOCAL"
exec 9>"$LOCK"
if [ "$MODE" = restore ] || [ "$MODE" = checkpoint-final ]; then
  # A dev-mode restart may overlap the previous app's final checkpoint. Never
  # restore an older durable view while that checkpoint is still being made.
  flock 9
else
  # Timer and shutdown checkpoints can race; the one already holding the lock
  # is sufficient, and a later timer/shutdown will catch subsequent writes.
  flock -n 9 || exit 0
fi

# `find -newer` is a strict comparison. A file written in the same filesystem
# timestamp tick as a checkpoint marker would otherwise be skipped forever.
# Keep a small overlap at every successful boundary; an occasional repeat copy
# is harmless, while a missed transcript/config is not.
mark_checkpoint_floor() {
  touch -d '2 seconds ago' "$1" 2>/dev/null || touch "$1"
}

# rsync exit 24 is "partial transfer due to vanished source files": between
# building its file list and reading them, something deleted files the listing
# still named. Both tree copies are exposed to it — a restore walks the bucket
# for minutes while a harness prunes transcript scratch, and a checkpoint copies
# a live tree the harnesses are still writing. What vanishes is the scratch
# itself, so nothing durable is lost and the copy is otherwise complete.
# Treating it as a failure once refused to boot the Space at all.
# Single-file copies below keep the strict test: there, a vanished source is
# exactly the file we were asked for.
rsync_tree() {
  rsync "$@"; _rs_rc=$?
  [ "$_rs_rc" -eq 24 ] && _rs_rc=0
  return "$_rs_rc"
}

restore_tree() {
  key="$1" durable="$2" live="$3"; shift 3
  mkdir -p "$durable" "$live"
  stamp_dir="$AM_LOCAL/agent-state-stamps"
  stamp="$stamp_dir/$key"
  had_local=false
  [ -n "$(find "$live" -type f -print -quit 2>/dev/null)" ] && had_local=true
  # --update matters for hot/dev restarts: local disk survives those and can be
  # newer than the last completed bucket checkpoint. A fresh container starts
  # with an empty destination, so the same command performs a full restore.
  if rsync_tree -a --update "$@" "$durable/" "$live/"; then
    mkdir -p "$stamp_dir"
    if [ ! -e "$stamp" ]; then
      if [ "$had_local" = true ]; then
        # First deployment over an already-populated local tree: force one
        # checkpoint so newer local files skipped by --update become durable.
        touch -t 197001010000 "$stamp"
      else
        # Fresh container: every local byte came from this durable restore.
        mark_checkpoint_floor "$stamp"
      fi
    fi
  else
    return 1
  fi
}

checkpoint_tree() {
  key="$1" live="$2" durable="$3"; shift 3
  [ -d "$live" ] || return 0
  stamp_dir="$AM_LOCAL/agent-state-stamps"
  mkdir -p "$durable" "$stamp_dir"
  stamp="$stamp_dir/$key"
  if [ ! -e "$stamp" ]; then
    touch -t 197001010000 "$stamp"
  fi

  # Mark the START of this checkpoint. Any file changed during/after its copy
  # is newer than `next` and will therefore be selected again next time.
  next="$stamp.next.$$"
  list="$stamp.files.$$"
  mark_checkpoint_floor "$next"
  (cd "$live" && find . -type f -newer "$stamp" -print0) > "$list"

  if [ -s "$list" ]; then
    # --files-from means rsync walks only changed LOCAL paths. It does not scan
    # the remote tree every 15 seconds — a critical property on the bucket
    # mount. Destination temporaries close before rename, retaining the prior
    # object if this process dies during transfer.
    if ! rsync_tree -a -r --from0 --files-from="$list" --delay-updates \
      "$@" "$live/" "$durable/"; then
      rm -f "$next" "$list"
      return 1
    fi
  fi
  mv "$next" "$stamp"
  rm -f "$list"
}

restore_sqlite_tree() {
  durable="$1" live="$2" db_name="$3" checkpoint_name="$4"
  mkdir -p "$durable" "$live"

  if ! restore_tree "$checkpoint_name-files" "$durable" "$live" \
    --exclude 'checkpoints' --exclude '*.db' --exclude '*.db-*' \
    --exclude '*.sqlite*' --exclude '*-wal' --exclude '*-shm'; then
    return 1
  fi

  checkpoint="$durable/checkpoints/$checkpoint_name"
  target="$live/$db_name"
  if [ -s "$target" ] && [ "$(sqlite3 "$target" 'PRAGMA quick_check;' 2>/dev/null)" = ok ]; then
    # Local disk survives in-container/dev restarts. Its database can be newer
    # than the last checkpoint (including committed rows still in its WAL), so
    # a valid live database is always the restore authority on a hot restart.
    return 0
  fi

  # Retain an invalid local set for diagnosis, then recover from the last
  # known-good checkpoint. These are explicit ephemeral paths, never bucket
  # objects.
  had_invalid=false
  if [ -e "$target" ] || [ -e "$target-wal" ] || [ -e "$target-shm" ]; then
    had_invalid=true
    invalid_dir="$AM_LOCAL/agent-state-invalid/$checkpoint_name.$(date -u +%Y%m%dT%H%M%SZ).$$"
    mkdir -p "$invalid_dir"
    [ -e "$target" ] && mv "$target" "$invalid_dir/$db_name"
    [ -e "$target-wal" ] && mv "$target-wal" "$invalid_dir/$db_name-wal"
    [ -e "$target-shm" ] && mv "$target-shm" "$invalid_dir/$db_name-shm"
  fi

  if [ -s "$checkpoint" ]; then
    # WAL/SHM are ephemeral coordination files, never part of a restored
    # checkpoint. Removing these explicit local paths cannot touch bucket data.
    rm -f "$target-wal" "$target-shm"
    if ! rsync -a "$checkpoint" "$target"; then return 1; fi
  elif [ -s "$durable/$db_name" ]; then
    # One-release compatibility path for state written by the old raw-rsync
    # mechanism. Copy the legacy DB and any WAL so SQLite can recover it; the
    # first successful checkpoint replaces this path as the restore authority.
    if ! rsync -a "$durable/$db_name" "$target"; then return 1; fi
    if [ -f "$durable/$db_name-wal" ] && ! rsync -a "$durable/$db_name-wal" "$target-wal"; then return 1; fi
    if [ -f "$durable/$db_name-shm" ] && ! rsync -a "$durable/$db_name-shm" "$target-shm"; then return 1; fi
  fi

  if [ "$had_invalid" = true ] && [ ! -s "$target" ]; then
    echo "agent-state: invalid local SQLite state has no durable recovery for $target" >&2
    return 1
  fi
  if [ -s "$target" ] && [ "$(sqlite3 "$target" 'PRAGMA quick_check;' 2>/dev/null)" != ok ]; then
    echo "agent-state: restored SQLite state is invalid for $target" >&2
    return 1
  fi
}

checkpoint_sqlite_tree() {
  live="$1" durable="$2" db_name="$3" checkpoint_name="$4"
  source="$live/$db_name"

  # A harness can write ordinary state before it creates its database (Hermes
  # setup files are a real example). Always publish that file tree first. The
  # database backup is optional until the database itself exists.
  if ! checkpoint_tree "$checkpoint_name-files" "$live" "$durable" \
    --exclude 'checkpoints' --exclude '*.db' --exclude '*.db-*' \
    --exclude '*.sqlite*' --exclude '*-wal' --exclude '*-shm'; then
    return 1
  fi
  [ -s "$source" ] || return 0

  sqlite_stamp_dir="$AM_LOCAL/agent-state-stamps"
  sqlite_stamp="$sqlite_stamp_dir/$checkpoint_name-sqlite"
  mkdir -p "$sqlite_stamp_dir"
  # Timestamp-only detection has an equal-tick hole, while deliberately
  # overlapping the marker would rewrite an idle database on rapid successive
  # checkpoints. Record the exact local DB/WAL metadata observed BEFORE the
  # backup instead. A concurrent commit changes the next signature and is
  # therefore picked up by the following checkpoint.
  sqlite_next="$sqlite_stamp.next.$$"
  {
    stat -c 'db|%s|%y|%z' "$source"
    if [ -e "$source-wal" ]; then
      stat -c 'wal|%s|%y|%z' "$source-wal"
    else
      echo 'wal|absent'
    fi
  } > "$sqlite_next"
  if cmp -s "$sqlite_next" "$sqlite_stamp"; then
    rm -f "$sqlite_next"
    return 0
  fi

  stage_dir="$AM_LOCAL/agent-state-snapshots"
  mkdir -p "$stage_dir" "$durable/checkpoints"
  staged="$stage_dir/$checkpoint_name.$$.tmp"
  rm -f "$staged"
  escaped=$(printf '%s' "$staged" | sed "s/'/''/g")

  # .backup observes the main DB and WAL through one consistent SQLite read
  # transaction. The staged file is ordinary local storage and is closed before
  # rsync hands it to the bucket.
  if ! sqlite3 "$source" ".timeout 5000" ".backup '$escaped'"; then
    rm -f "$staged" "$sqlite_next"
    return 1
  fi
  if [ "$(sqlite3 "$staged" 'PRAGMA quick_check;' 2>/dev/null)" != ok ]; then
    echo "agent-state: refusing invalid SQLite checkpoint for $source" >&2
    rm -f "$staged" "$sqlite_next"
    return 1
  fi
  if ! rsync -a --delay-updates "$staged" "$durable/checkpoints/$checkpoint_name"; then
    rm -f "$staged" "$sqlite_next"
    return 1
  fi
  rm -f "$staged"
  mv "$sqlite_next" "$sqlite_stamp"
}

failures=0

if [ "$MODE" = restore ]; then
  restore_tree codex "$CODEX_DURABLE" "$CODEX_HOME" \
    --exclude '*.sqlite*' --exclude '*.db' --exclude '*.db-*' \
    --exclude '*-wal' --exclude '*-shm' --exclude 'db-backups' \
    --exclude 'cache' --exclude '.tmp' --exclude 'mcp-oauth-locks' || failures=$((failures + 1))
  restore_tree claude "$CLAUDE_DURABLE" "$CLAUDE_CONFIG_DIR" || failures=$((failures + 1))
  restore_tree gemini "$GEMINI_DURABLE" "$GEMINI_LIVE" || failures=$((failures + 1))
  restore_tree openclaw "$OPENCLAW_DURABLE" "$OPENCLAW_STATE_DIR" || failures=$((failures + 1))
  restore_sqlite_tree "$OPENCODE_DURABLE" "$OPENCODE_LIVE" opencode.db opencode.db || failures=$((failures + 1))
  restore_sqlite_tree "$HERMES_DURABLE" "$HERMES_LIVE" state.db state.db || failures=$((failures + 1))
else
  checkpoint_tree codex "$CODEX_HOME" "$CODEX_DURABLE" \
    --exclude '*.sqlite*' --exclude '*.db' --exclude '*.db-*' \
    --exclude '*-wal' --exclude '*-shm' --exclude 'db-backups' \
    --exclude 'cache' --exclude '.tmp' --exclude 'mcp-oauth-locks' || failures=$((failures + 1))
  checkpoint_tree claude "$CLAUDE_CONFIG_DIR" "$CLAUDE_DURABLE" || failures=$((failures + 1))
  checkpoint_tree gemini "$GEMINI_LIVE" "$GEMINI_DURABLE" || failures=$((failures + 1))
  checkpoint_tree openclaw "$OPENCLAW_STATE_DIR" "$OPENCLAW_DURABLE" || failures=$((failures + 1))
  checkpoint_sqlite_tree "$OPENCODE_LIVE" "$OPENCODE_DURABLE" opencode.db opencode.db || failures=$((failures + 1))
  checkpoint_sqlite_tree "$HERMES_LIVE" "$HERMES_DURABLE" state.db state.db || failures=$((failures + 1))
fi

[ "$failures" -eq 0 ] || {
  echo "agent-state: $MODE completed with $failures failed adapter(s)" >&2
  exit 1
}