Skip to content

Read codex app-server replies unbuffered in the usage collector #13102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
surim0n wants to merge 1 commit into from
Closed

Read codex app-server replies unbuffered in the usage collector #13102

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 40 additions & 10 deletions bin/omarchy-agent-usage-codex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -474,17 +474,46 @@ def _cached_local_stats(max_age):
return stats


def rpc_request(proc, request_id, method, params=None, timeout=8):
class RpcLineReader:
"""Unbuffered line reader over a pipe fd.

select() watches the fd, but proc.stdout is a buffered text stream: a
reply that shares a chunk with a notification is already inside Python's
buffer when readline() returns the first line, so the fd never becomes
readable again and the request times out with its reply in hand. Reading
the fd raw and splitting lines here keeps leftover bytes where the next
read can see them — and the buffer must outlive one request, since a
reply can arrive together with the previous one.
"""

def __init__(self, stream):
self.fd = stream.fileno()
self.buf = b""

def readline(self, deadline):
while b"\n" not in self.buf:
remaining = deadline - time.time()
if remaining <= 0:
return None
ready, _, _ = select.select([self.fd], [], [], remaining)
if not ready:
return None
chunk = os.read(self.fd, 65536)
if not chunk:
return None
self.buf += chunk
line, self.buf = self.buf.split(b"\n", 1)
return line


def rpc_request(proc, reader, request_id, method, params=None, timeout=8):
payload = {"id": request_id, "method": method, "params": params or {}}
proc.stdin.write(json.dumps(payload) + "\n")
proc.stdin.flush()
deadline = time.time() + timeout
while time.time() < deadline:
ready, _, _ = select.select([proc.stdout], [], [], 0.25)
if not ready:
continue
line = proc.stdout.readline()
if not line:
while True:
line = reader.readline(deadline)
if line is None:
break
try:
message = json.loads(line)
Expand Down Expand Up @@ -541,11 +570,12 @@ def fetch_codex_rpc():
return result

try:
rpc_request(proc, 1, "initialize", {"clientInfo": {"name": "omarchy-agent-usage", "version": "1"}}, timeout=8)
reader = RpcLineReader(proc.stdout)
rpc_request(proc, reader, 1, "initialize", {"clientInfo": {"name": "omarchy-agent-usage", "version": "1"}}, timeout=8)
proc.stdin.write(json.dumps({"method": "initialized", "params": {}}) + "\n")
proc.stdin.flush()
account_msg = rpc_request(proc, 2, "account/read", timeout=4)
limits_msg = rpc_request(proc, 3, "account/rateLimits/read", timeout=4)
account_msg = rpc_request(proc, reader, 2, "account/read", timeout=4)
limits_msg = rpc_request(proc, reader, 3, "account/rateLimits/read", timeout=4)

account = (account_msg.get("result") or {}).get("account") or {}
limits = (limits_msg.get("result") or {}).get("rateLimits") or {}
Expand Down
67 changes: 61 additions & 6 deletions test/shell.d/agent-usage-codex-scanner-test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,17 @@ require_command python3
TEST_HOME=$(mktemp -d)
trap 'rm -rf "$TEST_HOME"' EXIT

# touch -d is GNU-only; set a file's mtime to seconds from now, portably.
set_mtime_offset() {
python3 - "$1" "$2" <<'PY'
import os
import sys
import time

os.utime(sys.argv[1], (time.time() + int(sys.argv[2]),) * 2)
PY
}

mkdir -p "$TEST_HOME/.codex/sessions/$(date +%Y/%m/%d)" "$TEST_HOME/bin"

cat >"$TEST_HOME/bin/codex" <<'EOF'
Expand Down Expand Up @@ -50,7 +61,10 @@ result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" CODEX_ARGS_FILE="$TEST
# NUL-separated, so the assertion sees argument boundaries: a single "-a on-request"
# would flatten to the same text as two arguments but is not a policy codex accepts.
expected_args=(-s read-only -a on-request app-server)
mapfile -d '' -t codex_args <"$TEST_HOME/codex-args"
codex_args=()
while IFS= read -r -d '' arg; do
codex_args+=("$arg")
done <"$TEST_HOME/codex-args"

[[ ${codex_args[*]@Q} == "${expected_args[*]@Q}" ]] ||
fail "Codex collector uses the supported approval policy" "${codex_args[*]@Q}"
Expand Down Expand Up @@ -186,7 +200,8 @@ result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CAC
cache_file=$(ls "$CACHE_HOME/.cache/omarchy/agent-usage/"/codex-scan-*.json 2>/dev/null | head -n 1)
[[ -n $cache_file && -s $cache_file ]] ||
fail "Codex collector leaves a cache file behind" "$result"
[[ $(stat -c %a "$cache_file") == "644" ]] ||
cache_mode=$(stat -c %a "$cache_file" 2>/dev/null || stat -f %Lp "$cache_file")
[[ $cache_mode == "644" ]] ||
fail "Codex collector keeps cache files readable" "$result"
[[ $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "5" ]] ||
fail "Codex collector writes a versioned cache envelope" "$result"
Expand Down Expand Up @@ -278,7 +293,7 @@ conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", (
conn.commit()
conn.close()
PY
touch -d "2 hours ago" "$cache_file"
set_mtime_offset "$cache_file" -7200
result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \
PATH="$CACHE_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only)

Expand Down Expand Up @@ -321,7 +336,7 @@ result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CAC

# 30 seconds is the lowest refreshIntervalSec the widget supports, so a
# cache that old must already be past the no-flag reuse window.
touch -d "30 seconds ago" "$cache_file"
set_mtime_offset "$cache_file" -30
result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \
PATH="$CACHE_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex")

Expand Down Expand Up @@ -353,7 +368,7 @@ conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", (
conn.commit()
conn.close()
PY
touch -d "10 minutes ago" "$cache_file"
set_mtime_offset "$cache_file" -600
result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \
PATH="$CACHE_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only)

Expand Down Expand Up @@ -399,7 +414,7 @@ conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", (
conn.commit()
conn.close()
PY
touch -d "@$(( $(date +%s) + 3600 ))" "$cache_file"
set_mtime_offset "$cache_file" 3600
result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \
PATH="$CACHE_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only)

Expand Down Expand Up @@ -603,3 +618,43 @@ result=$(HOME="$INTERRUPTED_HOME" CODEX_HOME="$INTERRUPTED_HOME/.codex" XDG_CACH
[[ $(jq -r '.todayTotalTokens' <<<"$result") == "9" ]] ||
fail "Codex collector does not reuse a snapshot from an interrupted scan" "$result"
pass "Codex collector does not cache an interrupted opencode scan"

# The app-server emits notifications next to replies, and the pipe chunks them
# however it likes: a notification can share one read with the reply a request
# is waiting on. The collector must not lose the reply to stdio buffering.
CHUNKED_HOME=$(mktemp -d)
trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME" "$MALFORMED_HOME" "$UNWRITABLE_HOME" "$INTERRUPTED_HOME" "$CHUNKED_HOME"' EXIT
mkdir -p "$CHUNKED_HOME/bin"

cat >"$CHUNKED_HOME/bin/codex" <<'EOF'
#!/bin/bash

while read -r request; do
id=$(jq -r '.id // empty' <<<"$request")
method=$(jq -r '.method // empty' <<<"$request")

case "$method" in
initialize)
# One write carrying a notification AND the reply: a buffered reader
# that returns the first line leaves the reply stranded off-fd.
printf '{"method":"remoteControl/status/changed","params":{}}\n{"id":%s,"result":{}}\n' "$id"
;;
account/read)
printf '{"method":"account/updated","params":{}}\n{"id":%s,"result":{"account":{"type":"plus"}}}\n' "$id"
;;
account/rateLimits/read)
jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {planType: "plus", primary: {usedPercent: 40, windowDurationMins: 300, resetsAt: 2000000000}}}}'
;;
esac
done
EOF
chmod +x "$CHUNKED_HOME/bin/codex"

result=$(HOME="$CHUNKED_HOME" CODEX_HOME="$CHUNKED_HOME/.codex" XDG_CACHE_HOME="$CHUNKED_HOME/.cache" XDG_DATA_HOME="$CHUNKED_HOME/.local/share" \
PATH="$CHUNKED_HOME/bin:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --limits-only)

[[ $(jq -r '.usageStatusText' <<<"$result") != "Codex limits unavailable" ]] ||
fail "Codex collector reads a reply chunked together with a notification" "$result"
[[ $(jq -r '.tierLabel' <<<"$result") == "plus" && $(jq -r '.limits[0].percent' <<<"$result") == "0.4" ]] ||
fail "Codex collector reads the limits reply after a coalesced notification" "$result"
pass "Codex collector reads replies chunked together with notifications"