From 9fc1ee5d40b71a0ee4ac4f020d07af4714c89c3a Mon Sep 17 00:00:00 2001 From: prajwal-395 Date: Sat, 15 Aug 2026 13:01:46 -0400 Subject: [PATCH] feat: Add quota memory extraction for agy harness Extracts model, quota, and reset time from agy pane footer and stores it persistently in state/ so dispatch rules can read it. Fails open on ambiguous reads and resets on staleness. --- bin/fm-agy-quota-lib.sh | 130 +++++++++++++++++++++++++++++++++ bin/fm-watch.sh | 5 ++ tests/fm-agy-quota-lib.test.sh | 80 ++++++++++++++++++++ 3 files changed, 215 insertions(+) create mode 100644 bin/fm-agy-quota-lib.sh create mode 100755 tests/fm-agy-quota-lib.test.sh diff --git a/bin/fm-agy-quota-lib.sh b/bin/fm-agy-quota-lib.sh new file mode 100644 index 0000000000..e172d2953a --- /dev/null +++ b/bin/fm-agy-quota-lib.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +# fm-agy-quota-lib.sh - Quota memory observation and read helpers. +# Usage: . bin/fm-agy-quota-lib.sh + +# fm_agy_parse_reset_time: converts '4h 24m', '1d 2h', etc to seconds. +fm_agy_parse_reset_time() { + local ts="$1" + local total=0 + local d=0 h=0 m=0 s=0 + + case "$ts" in + *d*) d="${ts%%d*}"; d="${d##* }"; total=$((total + d * 86400)) ;; + esac + case "$ts" in + *h*) h="${ts%%h*}"; h="${h##* }"; total=$((total + h * 3600)) ;; + esac + case "$ts" in + *m*) m="${ts%%m*}"; m="${m##* }"; total=$((total + m * 60)) ;; + esac + case "$ts" in + *s*) s="${ts%%s*}"; s="${s##* }"; total=$((total + s)) ;; + esac + echo "$total" +} + +# fm_agy_quota_observe: extract and record quota from agy pane footer. +# Format: Gemini 3.1 Pro (High) | ctx: 10.5% | quota: 94.7% (4h 24m) +fm_agy_quota_observe() { # + local text="$1" state_dir="$2" + case "$text" in + *" | quota: "*) ;; + *) return 0 ;; + esac + + local pre="${text%% | quota: *}" + local nl=' +' + local line_start="${pre##*"$nl"}" + local model="${line_start%% | ctx: *}" + if [ "$model" = "$line_start" ]; then + return 0 + fi + + local post="${text#* | quota: }" + local line_end="${post%%"$nl"*}" + local quota="${line_end%% (*}" + local reset_time="${line_end#*(}" + reset_time="${reset_time%)}" + + if [ -z "$model" ] || [ -z "$quota" ] || [ -z "$reset_time" ]; then + return 0 + fi + + local safe_model + if command -v md5 >/dev/null 2>&1; then + safe_model=$(printf '%s' "$model" | md5 -q) + else + safe_model=$(printf '%s' "$model" | md5sum | cut -d' ' -f1) + fi + + local now + now=$(date +%s) + + # Only write if it actually changed, but doing this requires reading. + # Since this is already conditionally executed, writing directly is fine. + printf '%s|%s|%s|%s\n' "$model" "$quota" "$reset_time" "$now" > "$state_dir/.agy-quota-$safe_model" +} + +# fm_agy_quota_read: returns the last known value for a model together with its age. +fm_agy_quota_read() { # [] + local model="$1" state_dir="$2" now_override="$3" + local safe_model + if command -v md5 >/dev/null 2>&1; then + safe_model=$(printf '%s' "$model" | md5 -q) + else + safe_model=$(printf '%s' "$model" | md5sum | cut -d' ' -f1) + fi + + local file="$state_dir/.agy-quota-$safe_model" + if [ ! -f "$file" ]; then + echo "unknown" + return 0 + fi + + local line + line=$(cat "$file" 2>/dev/null || true) + if [ -z "$line" ]; then + echo "unknown" + return 0 + fi + + local file_model="${line%%|*}" + local rest="${line#*|}" + local quota="${rest%%|*}" + rest="${rest#*|}" + local reset_time_str="${rest%%|*}" + local obs_ts="${rest#*|}" + + if [ "$file_model" != "$model" ] || [ -z "$quota" ] || [ -z "$reset_time_str" ] || [ -z "$obs_ts" ]; then + echo "unknown" + return 0 + fi + + local now + if [ -n "$now_override" ]; then + now="$now_override" + else + now=$(date +%s) + fi + + local age=$((now - obs_ts)) + if [ "$age" -lt 0 ]; then + age=0 + fi + + local reset_seconds + reset_seconds=$(fm_agy_parse_reset_time "$reset_time_str") + if [ "$reset_seconds" -eq 0 ]; then + echo "unknown" + return 0 + fi + + if [ "$age" -ge "$reset_seconds" ]; then + echo "unknown" + return 0 + fi + + quota="${quota%%%*}" + echo "$quota $age" +} diff --git a/bin/fm-watch.sh b/bin/fm-watch.sh index 3f4a57afd6..cdf93e491c 100755 --- a/bin/fm-watch.sh +++ b/bin/fm-watch.sh @@ -88,6 +88,8 @@ mkdir -p "$STATE" . "$SCRIPT_DIR/fm-pending-reply-lib.sh" # shellcheck source=bin/fm-busy-lib.sh . "$SCRIPT_DIR/fm-busy-lib.sh" +# shellcheck source=bin/fm-agy-quota-lib.sh +. "$SCRIPT_DIR/fm-agy-quota-lib.sh" WATCH_LOCK="$STATE/.watch.lock" WATCH_PATH="$SCRIPT_DIR/fm-watch.sh" @@ -1033,6 +1035,9 @@ EOF # content cannot suppress stale detection. Read once per window per poll and # reused below so a busy verdict is consistent within one cycle. if window_is_busy "$w" "$tail40"; then busy_now=0; else busy_now=1; fi + if [ "$h" != "$prev" ]; then + fm_agy_quota_observe "$tail40" "$STATE" + fi if [ "$h" = "$prev" ]; then n=$(( $(cat "$cf" 2>/dev/null || echo 0) + 1 )) echo "$n" > "$cf" diff --git a/tests/fm-agy-quota-lib.test.sh b/tests/fm-agy-quota-lib.test.sh new file mode 100755 index 0000000000..e262c9ea5a --- /dev/null +++ b/tests/fm-agy-quota-lib.test.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +set -u + +# shellcheck source=tests/lib.sh +. "$(dirname "${BASH_SOURCE[0]}")/lib.sh" +# shellcheck source=bin/fm-agy-quota-lib.sh +. "$ROOT/bin/fm-agy-quota-lib.sh" + +TMP_ROOT=$(fm_test_tmproot fm-agy-quota-lib) + +assert_eq() { + local expected=$1 actual=$2 + if [ "$expected" != "$actual" ]; then + fail "expected '$expected' but got '$actual'" + fi +} + +echo "Testing fm_agy_parse_reset_time" +assert_eq "15840" "$(fm_agy_parse_reset_time '4h 24m')" +assert_eq "2700" "$(fm_agy_parse_reset_time '45m')" +assert_eq "93600" "$(fm_agy_parse_reset_time '1d 2h')" +assert_eq "30" "$(fm_agy_parse_reset_time '30s')" +pass "fm_agy_parse_reset_time logic works" + +echo "Testing observe and read valid" +state="$TMP_ROOT/state1" +mkdir -p "$state" + +footer="Gemini 3.1 Pro (High) | ctx: 10.5% | quota: 94.7% (4h 24m)" +fm_agy_quota_observe "$footer" "$state" + +now=$(date +%s) +res=$(fm_agy_quota_read "Gemini 3.1 Pro (High)" "$state" "$now") +assert_eq "94.7 0" "$res" + +# Advance time by 10 seconds +res=$(fm_agy_quota_read "Gemini 3.1 Pro (High)" "$state" "$((now + 10))") +assert_eq "94.7 10" "$res" +pass "observe and read valid handles normal cases" + +echo "Testing read older than reset window" +state="$TMP_ROOT/state2" +mkdir -p "$state" + +footer="Claude Opus 4.6 (Thinking) | ctx: 5% | quota: 14.7% (1h 0m)" +fm_agy_quota_observe "$footer" "$state" + +now=$(date +%s) + +# 3600 seconds is the window. +res=$(fm_agy_quota_read "Claude Opus 4.6 (Thinking)" "$state" "$now") +assert_eq "14.7 0" "$res" + +# 3600 seconds later -> reset +res=$(fm_agy_quota_read "Claude Opus 4.6 (Thinking)" "$state" "$((now + 3600))") +assert_eq "unknown" "$res" +pass "older than window returns unknown" + +echo "Testing missing, unparseable, ambiguous" +state="$TMP_ROOT/state3" +mkdir -p "$state" + +# Missing +res=$(fm_agy_quota_read "No Such Model" "$state" "$(date +%s)") +assert_eq "unknown" "$res" + +# Unparseable reset time +footer="Bad Model | ctx: 0% | quota: 50% (forever)" +fm_agy_quota_observe "$footer" "$state" +res=$(fm_agy_quota_read "Bad Model" "$state" "$(date +%s)") +assert_eq "unknown" "$res" + +# Ambiguous formatting (missing ctx but has quota, observe will reject it) +footer="Ambiguous Model | quota: 50% (1h)" +fm_agy_quota_observe "$footer" "$state" +res=$(fm_agy_quota_read "Ambiguous Model" "$state" "$(date +%s)") +assert_eq "unknown" "$res" +pass "missing unparseable ambiguous fail open" + +echo "ALL TESTS PASSED"