feat(policy): /wait awaits local policy reload; demo auto-approves redrafts

zredlined · zredlined · commit 50a6da4871a4 · 2026-05-13T12:22:13.000-07:00
Three things in one commit, all surfaced by running the demo end-to-end against a real gateway and finding the agent had to draft a broader second proposal. 1. /wait race fix. Previously /wait returned `approved` the moment it observed the gateway's chunk status flip, but the local supervisor reloads policy on its own poll cycle (~10s in practice). The agent's retry would race the reload and hit the still-old policy, getting denied. Codex then drafted a broader rule and re-submitted — sound agent behavior, but not what /wait should provoke. Now /wait captures the local policy version at start, and after observed-approved waits for the supervisor to load a strictly-newer version before returning. Bounded by the caller's deadline; best-effort return if the deadline elapses without the version bumping. Two new unit tests pin the happy path and the deadline-clamped fallback. 2. demo.sh auto-approve loop. Replaces approve_when_pending + wait_for_agent with one approve_pending_until_agent_exits function that keeps watching for pending chunks and approving them until the agent process exits (or the configured timeout). Defense in depth against future redraft scenarios for any reason; today (post-fix #1) the agent should only submit one proposal per task, but we don't want to hang silently if it does submit more. 3. UX. Step headers now carry "[t+1.2s]" relative timestamps so reading the run output makes latency visible (the demo's whole point is the wait is cheap — surface that). A spin_wait helper renders an ASCII spinner during the watch loop so the demo never looks frozen on a TTY. Falls back to plain sleep on non-TTY contexts. Closes the race condition diagnosed from the trace timing where the gateway approved at t+0, sandbox observed at t+0.3s, but the supervisor didn't load v2 until t+9.4s — well after the agent had already retried and been denied. Signed-off-by: Alexander Watson <zredlined@gmail.com>
diff --git a/crates/openshell-sandbox/src/policy_local.rs b/crates/openshell-sandbox/src/policy_local.rs
@@ -680,9 +680,28 @@ async fn proposal_wait_response(
     };
     let timeout_secs = parse_timeout_query(query);
     let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(timeout_secs);
+    // Baseline local policy version at /wait start. When a chunk is
+    // approved upstream the gateway merges it immediately, but the local
+    // supervisor reloads policy on its own poll cycle — without this
+    // baseline we can return "approved" before the local rule is active,
+    // and the agent's single retry races the reload.
+    let baseline_policy_version: u32 = ctx
+        .current_policy
+        .read()
+        .await
+        .as_ref()
+        .map_or(0, |p| p.version);
     loop {
         match fetch_chunk(&session, chunk_id).await {
             Ok(Some(chunk)) if is_terminal_status(&chunk.status) => {
+                if chunk.status == "approved" {
+                    // Don't promise the agent its retry will succeed until
+                    // the local supervisor has actually loaded the new
+                    // policy. Bounded by the same deadline so a stuck
+                    // supervisor cannot extend the wait beyond what the
+                    // caller asked for.
+                    wait_for_local_policy_bump(ctx, baseline_policy_version, deadline).await;
+                }
                 // Audit beat: emit at the moment this sandbox observes the
                 // decision so the trace correlates with the proxy events
                 // bracketing the loop. Multiple waiters on the same chunk
@@ -753,6 +772,42 @@ fn is_terminal_status(status: &str) -> bool {
     matches!(status, "approved" | "rejected")
 }
 
+/// After a chunk is approved upstream, wait until the local supervisor has
+/// loaded a policy version strictly newer than the baseline captured at the
+/// start of `/wait`. Bounded by the caller-supplied deadline; returns early
+/// (best-effort) if the deadline passes without the version bumping.
+///
+/// The polling cadence here is faster than `PROPOSAL_WAIT_POLL_INTERVAL`
+/// (which paces upstream gateway calls). This loop only reads in-memory
+/// state, so 200ms gives a responsive handoff to the agent's retry once
+/// the supervisor's own policy poll catches up.
+async fn wait_for_local_policy_bump(
+    ctx: &PolicyLocalContext,
+    baseline_version: u32,
+    deadline: tokio::time::Instant,
+) {
+    const TICK: std::time::Duration = std::time::Duration::from_millis(200);
+    loop {
+        let current: u32 = ctx
+            .current_policy
+            .read()
+            .await
+            .as_ref()
+            .map_or(0, |p| p.version);
+        if current > baseline_version {
+            return;
+        }
+        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
+        if remaining.is_zero() {
+            // Best effort: return approved anyway. The agent's retry may
+            // race the reload, but the alternative (extending past the
+            // caller's deadline) violates the wait contract.
+            return;
+        }
+        tokio::time::sleep(std::cmp::min(remaining, TICK)).await;
+    }
+}
+
 /// Parse `?timeout=<s>` from the query string. Default applies for missing
 /// or unparseable values; bounds clamp to keep the agent's hold ceiling
 /// sane. Re-issue is the right pattern for longer waits.
@@ -1665,6 +1720,64 @@ mod tests {
         assert!(summary.contains("/usr/bin/curl"));
     }
 
+    #[tokio::test]
+    async fn wait_for_local_policy_bump_returns_when_version_advances() {
+        let ctx = PolicyLocalContext::new(
+            Some(ProtoSandboxPolicy {
+                version: 5,
+                ..Default::default()
+            }),
+            None,
+            None,
+        );
+        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
+        let bumped = {
+            let policy = ctx.current_policy.clone();
+            tokio::spawn(async move {
+                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                *policy.write().await = Some(ProtoSandboxPolicy {
+                    version: 6,
+                    ..Default::default()
+                });
+            })
+        };
+        let start = tokio::time::Instant::now();
+        wait_for_local_policy_bump(&ctx, 5, deadline).await;
+        bumped.await.unwrap();
+        // Returned promptly after the bump, well before the 2s deadline.
+        let elapsed = start.elapsed();
+        assert!(
+            elapsed < std::time::Duration::from_millis(800),
+            "should return shortly after version bumps; took {elapsed:?}"
+        );
+    }
+
+    #[tokio::test]
+    async fn wait_for_local_policy_bump_returns_at_deadline_if_no_advance() {
+        let ctx = PolicyLocalContext::new(
+            Some(ProtoSandboxPolicy {
+                version: 5,
+                ..Default::default()
+            }),
+            None,
+            None,
+        );
+        let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(300);
+        let start = tokio::time::Instant::now();
+        wait_for_local_policy_bump(&ctx, 5, deadline).await;
+        let elapsed = start.elapsed();
+        // Best effort: returns at or shortly after the deadline rather
+        // than extending past it.
+        assert!(
+            elapsed >= std::time::Duration::from_millis(250),
+            "should wait until ~deadline; only waited {elapsed:?}"
+        );
+        assert!(
+            elapsed < std::time::Duration::from_millis(800),
+            "should not extend past deadline by much; took {elapsed:?}"
+        );
+    }
+
     #[test]
     fn sanitize_reason_for_audit_strips_control_chars_and_caps_length() {
         // Tabs and newlines are stripped; ordinary printable chars survive;
diff --git a/examples/agent-driven-policy-management/demo.sh b/examples/agent-driven-policy-management/demo.sh
@@ -73,9 +73,50 @@ RESET=$'\033[0m'
 
 AGENT_PID=""
 
-step() { printf "\n${BOLD}${CYAN}==> %s${RESET}\n\n" "$1"; }
+# Wall-clock anchor so each step header can carry a "[t+1.2s]" tag and the
+# reader sees where time is going. `date +%s.%N` works on macOS bash where
+# `${EPOCHREALTIME}` may be unavailable in older bashes.
+DEMO_START_EPOCH="$(date +%s.%N)"
+
+elapsed() {
+    awk -v s="$DEMO_START_EPOCH" -v now="$(date +%s.%N)" \
+        'BEGIN { printf "%.1fs", now - s }'
+}
+
+step() {
+    printf "\n${BOLD}${CYAN}==> [t+%s] %s${RESET}\n\n" "$(elapsed)" "$1"
+}
 info() { printf "  %b\n" "$*"; }
 
+# ASCII spinner for the watch-for-pending loop. Renders only on a TTY so
+# piped runs (CI, tee, etc.) stay clean. spin_wait pairs a message with a
+# bounded sleep so the spinner animates smoothly without polling faster
+# than necessary.
+SPINNER_CHARS=( '⠋' '⠙' '⠹' '⠸' '⠼' '⠴' '⠦' '⠧' '⠇' '⠏' )
+SPINNER_IDX=0
+
+spin_wait() {
+    local message="$1"
+    local duration_secs="${2:-2}"
+    if [[ ! -t 1 ]]; then
+        sleep "$duration_secs"
+        return
+    fi
+    local end=$(( SECONDS + duration_secs ))
+    while (( SECONDS < end )); do
+        printf "\r  ${DIM}%s${RESET} %s  " \
+            "${SPINNER_CHARS[SPINNER_IDX]}" "$message"
+        SPINNER_IDX=$(( (SPINNER_IDX + 1) % ${#SPINNER_CHARS[@]} ))
+        sleep 0.1
+    done
+}
+
+spin_clear() {
+    if [[ -t 1 ]]; then
+        printf "\r%*s\r" "${COLUMNS:-100}" ''
+    fi
+}
+
 # Redact host-side credentials from the agent log tail before printing on
 # failure. Codex shouldn't echo the token, but a misbehaving tool call (e.g.,
 # `curl -v`) could leak it; sanitize before showing the log.
@@ -204,9 +245,19 @@ check_gateway() {
     local raw version
     # `openshell status` colorizes labels with ANSI even when piped, so strip
     # escapes before parsing. Use NO_COLOR as a belt-and-suspenders hint for
-    # libraries that respect it.
-    raw="$(NO_COLOR=1 "$OPENSHELL_BIN" status 2>/dev/null \
-        | sed 's/\x1b\[[0-9;]*m//g')"
+    # libraries that respect it. Capture stderr explicitly so a connection
+    # failure (gateway down, port-forward died after a redeploy) surfaces a
+    # real error message instead of `set -euo pipefail` silently exiting.
+    if ! raw="$(NO_COLOR=1 "$OPENSHELL_BIN" status 2>&1)"; then
+        fail "openshell could not reach the gateway. CLI output:
+${raw}
+
+If you just redeployed, the kubectl port-forward you backgrounded earlier
+probably died with the old pod. Restart it (silenced so its noise doesn't
+bleed into the demo):
+  KUBECONFIG=kubeconfig kubectl -n openshell port-forward svc/openshell 8090:8080 >/dev/null 2>&1 &"
+    fi
+    raw="$(sed 's/\x1b\[[0-9;]*m//g' <<<"$raw")"
     version="$(awk -F': *' '/Version:/ { print $2; exit }' <<<"$raw")"
     [[ -n "$version" ]] \
         || fail "active OpenShell gateway is not reachable; start one with: openshell gateway start"
@@ -357,48 +408,57 @@ EOF
     info "${DIM}Watching for the pending draft on the gateway...${RESET}"
 }
 
-approve_when_pending() {
+approve_pending_until_agent_exits() {
     step "Waiting for the agent to draft a policy proposal"
     narrate_sandbox_workflow
 
-    local start now pending
+    local start now pending approval_count
     start="$(date +%s)"
     pending="${TMP_DIR}/pending.txt"
+    approval_count=0
 
     while true; do
+        # Agent finished? Drain its exit status and we're done.
         if ! kill -0 "$AGENT_PID" >/dev/null 2>&1; then
-            wait "$AGENT_PID" || true
+            spin_clear
+            if ! wait "$AGENT_PID"; then
+                AGENT_PID=""
+                fail "agent run failed"
+            fi
             AGENT_PID=""
-            fail "agent exited before a pending proposal appeared"
+            if (( approval_count == 0 )); then
+                fail "agent exited before any pending proposal appeared"
+            fi
+            info "agent exited after ${approval_count} approval(s)"
+            return
         fi
 
+        # Anything pending? Approve and keep watching — the agent may
+        # redraft if a previous proposal didn't yield the access it needed.
         if "$OPENSHELL_BIN" rule get "$DEMO_SANDBOX_NAME" --status pending >"$pending" 2>/dev/null \
             && grep -q "Chunk:" "$pending" && grep -q "pending" "$pending"; then
+            spin_clear
             info ""
             info "${GREEN}proposal received:${RESET}"
             summarize_pending "$pending"
 
             step "Approving — the agent's /wait will return within ~1s"
             "$OPENSHELL_BIN" rule approve-all "$DEMO_SANDBOX_NAME" \
                 | awk '/approved/ { print "  " $0 }'
-            return
+            approval_count=$((approval_count + 1))
         fi
 
         now="$(date +%s)"
         if (( now - start >= DEMO_APPROVAL_TIMEOUT_SECS )); then
-            fail "timed out waiting for the agent to submit a policy proposal"
+            spin_clear
+            if (( approval_count == 0 )); then
+                fail "timed out waiting for the agent to submit a policy proposal"
+            fi
+            fail "agent did not exit within ${DEMO_APPROVAL_TIMEOUT_SECS}s after ${approval_count} approval(s)"
         fi
-        sleep 2
-    done
-}
 
-wait_for_agent() {
-    if ! wait "$AGENT_PID"; then
-        AGENT_PID=""
-        fail "agent run failed"
-    fi
-    AGENT_PID=""
-    info "agent's /wait returned approved — single PUT retry succeeded"
+        spin_wait "watching for pending proposals (approved ${approval_count} so far)" 2
+    done
 }
 
 verify_github_write() {
@@ -457,8 +517,7 @@ main() {
     show_run_summary
 
     start_agent_sandbox
-    approve_when_pending
-    wait_for_agent
+    approve_pending_until_agent_exits
     verify_github_write
     show_logs
 
