1	import Base
2	import ./json.bend as Json
3	import ./clock.bend as Clock
4	import ./usage/norm.bend as UNorm
5	import ./usage/req.bend as UReq
6	import ./usage/cred.bend as UCred
7	import ./usage/gate.bend as UGate
8	
9	# Usage reads, as the reference's oauth.py (claude-swap 0.27.2 plus patch
10	# 0003): the usage endpoint for full-login tokens, the one-token inference
11	# probe (rate-limit headers) for setup-tokens and on a usage-endpoint 429.
12	# Pure parts first (header mapping, normalisation, classification, request
13	# text), then the effects. Nothing here logs: tokens live only in a curl
14	# config on stdin and in the credential text the caller owns.
15	
16	# ---- fetch_usage -----------------------------------------------------------
17	
18	# usage: the normalised result (JNull when there is none); error: the kind
19	# (http-429, http-403, timeout, network, bad-response, no-access-token,
20	# invalid_grant, no_refresh_token, invalid_client, refresh-failed, ...);
21	# retry: Retry-After seconds; struck: the credential text whose refresh token
22	# the server rejected for good (the store quarantines on it).
23	type Outcome is Data:
24	  Outcome{usage: Json.Json, error: Maybe<&2, String>, retry: Maybe<&2, U32>, struck: Maybe<&2, String>}
25	
26	# credential: the rotated credential when a refresh succeeded, for the store
27	# to persist.
28	type Fetched is Data:
29	  Fetched{outcome: Outcome, credential: Maybe<&2, String>}
30	
31	def done(o: Outcome, upd: Maybe<&2, String>) -> IO(Fetched):
32	  IO.pure(Fetched, Fetched{o, upd})
33	
34	def out.ok(+data: Json.Json, +now: U32, +off: U32) -> Outcome:
35	  Outcome{UNorm.build_usage_result(data, now, off), None{}, None{}, None{}}
36	
37	def out.err(e: UReq.Err) -> Outcome:
38	  match e:
39	    case UReq.Err{kind, retry, body}:
40	      Outcome{Json.JNull{}, Some{kind}, retry, None{}}
41	
42	def out.kind(k: String, struck: Maybe<&2, String>) -> Outcome:
43	  Outcome{Json.JNull{}, Some{k}, None{}, struck}
44	
45	def dead(+k: String) -> Bool:
46	  String.eq(k, "invalid_grant") || String.eq(k, "no_refresh_token")
47	
48	def final.pick(r: Result<&2, &2, UReq.Err, Json.Json>, +now: U32, +off: U32, upd: Maybe<&2, String>) -> IO(Fetched):
49	  match r:
50	    case Done{d}:
51	      done(out.ok(d, now, off), upd)
52	    case Fail{e}:
53	      done(out.err(e), upd)
54	
55	# Refusals that will not resolve by retrying this pass: the caller must not
56	# fall through to the usage endpoint with the known-expired token.
57	def deterministic(+k: String) -> Bool:
58	  String.eq(k, "invalid_client") || String.eq(k, "consume-busy") || String.eq(k, "store-unmirrored")
59	
60	def retry.k(+k: String) -> String:
61	  Clock.sel.s(dead(k) || deterministic(k), k, "refresh-failed")
62	
63	def struck.pick(d: Bool, cred: String) -> Maybe<&2, String>:
64	  match d:
65	    case True{}:
66	      Some{cred}
67	    case False{}:
68	      None{}
69	
70	def retry.struck(+k: String, +cred: String) -> Maybe<&2, String>:
71	  struck.pick(dead(k), cred)
72	
73	def retry.tok(m: Maybe<&2, String>, new: String, +now: U32, +off: U32) -> IO(Fetched):
74	  match m:
75	    case None{}:
76	      done(out.kind("refresh-failed", None{}), Some{new})
77	    case Some{t}:
78	      IO.bind(Result<&2, &2, UReq.Err, Json.Json>, Fetched, UReq.request_usage_data(t, False{}),
79	        r => final.pick(r, now, off, Some{new}))
80	
81	def retry.fail(e: Maybe<&2, String>, consumed: Maybe<&2, String>, +cred: String, upd: Maybe<&2, String>) -> IO(Fetched):
82	  match e:
83	    case Some{+k}:
84	      done(out.kind(retry.k(k), retry.struck(k, Maybe.default(&2, String, consumed, cred))), upd)
85	    case None{}:
86	      done(out.kind("refresh-failed", None{}), upd)
87	
88	def retry.got(c: Maybe<&2, String>, e: Maybe<&2, String>, consumed: Maybe<&2, String>, +cred: String, +now: U32, +off: U32, upd: Maybe<&2, String>) -> IO(Fetched):
89	  match c:
90	    case Some{+new}:
91	      retry.tok(UCred.access_token(UCred.oauth.of(new)), new, now, off)
92	    case None{}:
93	      retry.fail(e, consumed, cred, upd)
94	
95	# After a 401 and a refresh: failures keep their identity only when the
96	# reference does (dead lineage, rejected client, consume-busy), else
97	# `refresh-failed`.
98	def retry.pick(g: UGate.Gated, +cred: String, +now: U32, +off: U32, upd: Maybe<&2, String>) -> IO(Fetched):
99	  match g:
100	    case UGate.Gated{c, e, consumed, acct}:
101	      retry.got(c, e, consumed, cred, now, off, upd)
102	
103	def fail.refresh(retry: Bool, e: UReq.Err, gate: UGate.Gate, +cred: String, +now: U32, +off: U32, upd: Maybe<&2, String>) -> IO(Fetched):
104	  match retry:
105	    case True{}:
106	      IO.bind(UGate.Gated, Fetched, UGate.refresh_via(gate, cred), g => retry.pick(g, cred, now, off, upd))
107	    case False{}:
108	      done(out.err(e), upd)
109	
110	def is_401(e: UReq.Err) -> Bool:
111	  match e:
112	    case UReq.Err{kind, retry, body}:
113	      String.eq(kind, "http-401")
114	
115	# 401 on an inactive account with a refresh token: refresh once, retry once.
116	def fetch.after(r: Result<&2, &2, UReq.Err, Json.Json>, gate: UGate.Gate, +cred: String, +o: Json.Json, +active: Bool, +now: U32, +off: U32, upd: Maybe<&2, String>) -> IO(Fetched):
117	  match r:
118	    case Done{d}:
119	      done(out.ok(d, now, off), upd)
120	    case Fail{+e}:
121	      fail.refresh(is_401(e) && Bool.not(active) && UCred.has_refresh(o), e, gate, cred, now, off, upd)
122	
123	def fetch.req(+gate: UGate.Gate, +cred: String, +o: Json.Json, +token: String, +active: Bool, +now: U32, +off: U32, upd: Maybe<&2, String>, inf: Bool) -> IO(Fetched):
124	  IO.bind(Result<&2, &2, UReq.Err, Json.Json>, Fetched, UReq.request_usage_data(token, inf),
125	    r => fetch.after(r, gate, cred, o, active, now, off, upd))
126	
127	def pre.stop(stop: Bool, +k: String, gate: UGate.Gate, +cred: String, +o: Json.Json, +token: String, +active: Bool, +now: U32, +off: U32, consumed: Maybe<&2, String>) -> IO(Fetched):
128	  match stop:
129	    case True{}:
130	      done(out.kind(k, retry.struck(k, Maybe.default(&2, String, consumed, cred))), None{})
131	    case False{}:
132	      fetch.req(gate, cred, o, token, active, now, off, None{}, UCred.inference_only(o))
133	
134	def pre.fail(e: Maybe<&2, String>, consumed: Maybe<&2, String>, gate: UGate.Gate, +cred: String, +o: Json.Json, +token: String, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
135	  match e:
136	    case Some{+k}:
137	      pre.stop(dead(k) || deterministic(k), k, gate, cred, o, token, active, now, off, consumed)
138	    case None{}:
139	      pre.stop(False{}, "transient", gate, cred, o, token, active, now, off, consumed)
140	
141	# Before the request: an expired token of an inactive account is refreshed. A
142	# dead lineage or a deterministic refusal stops here; a transient failure goes
143	# on with the old token. A refreshed token is a new token: the usage endpoint
144	# first, even for a setup-token scope.
145	def pre.new(gate: UGate.Gate, +new: String, +token: String, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
146	  fetch.req(gate, new, UCred.oauth.of(new), Maybe.default(&2, String, UCred.access_token(UCred.oauth.of(new)), token), active, now, off, Some{new}, False{})
147	
148	def pre.got(c: Maybe<&2, String>, e: Maybe<&2, String>, consumed: Maybe<&2, String>, gate: UGate.Gate, +cred: String, +o: Json.Json, +token: String, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
149	  match c:
150	    case Some{new}:
151	      pre.new(gate, new, token, active, now, off)
152	    case None{}:
153	      pre.fail(e, consumed, gate, cred, o, token, active, now, off)
154	
155	def fetch.pre(g: UGate.Gated, gate: UGate.Gate, +cred: String, +o: Json.Json, +token: String, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
156	  match g:
157	    case UGate.Gated{c, e, consumed, acct}:
158	      pre.got(c, e, consumed, gate, cred, o, token, active, now, off)
159	
160	def gate.pick(needs: Bool, +gate: UGate.Gate, +cred: String, +o: Json.Json, +token: String, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
161	  match needs:
162	    case True{}:
163	      IO.bind(UGate.Gated, Fetched, UGate.refresh_via(gate, cred), g => fetch.pre(g, gate, cred, o, token, active, now, off))
164	    case False{}:
165	      fetch.req(gate, cred, o, token, active, now, off, None{}, UCred.inference_only(o))
166	
167	def fetch.gate(+gate: UGate.Gate, +cred: String, +o: Json.Json, token: String, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
168	  gate.pick(Bool.not(active) && UCred.has_refresh(o) && UCred.expired(o, now), gate, cred, o, token, active, now, off)
169	
170	def fetch.start(m: Maybe<&2, String>, gate: UGate.Gate, +cred: String, +o: Json.Json, +active: Bool, +now: U32, +off: U32) -> IO(Fetched):
171	  match m:
172	    case None{}:
173	      done(out.kind("no-access-token", None{}), None{})
174	    case Some{t}:
175	      fetch.gate(gate, cred, o, t, active, now, off)
176	
177	# The local UTC offset, biased by 86400 (86400 = UTC when `date` is unusable).
178	def offset.of(oe: String & String) -> U32:
179	  (o, e) = oe
180	  Maybe.default(&2, U32, Clock.tz(String.trim(o)), 86400)
181	
182	def offset.run(r: U32 & (String & String)) -> U32:
183	  (st, oe) = r
184	  offset.of(oe)
185	
186	def offset.pick(r: Result<&1, &1, U32 & String, U32 & (String & String)>) -> U32:
187	  match r:
188	    case Fail{e}:
189	      86400
190	    case Done{v}:
191	      offset.run(v)
192	
193	def local_offset() -> IO(U32):
194	  IO.bind(Result<&1, &1, U32 & String, U32 & (String & String)>, U32, Process.run("date", ["+%z"], "", 64, 2000),
195	    r => IO.pure(U32, offset.pick(r)))
196	
197	# fetch_usage plus the 401 refresh of try_fetch_usage_for_account, over the
198	# credential JSON text. Refreshes only inactive accounts. The rotated
199	# credential comes back in `Fetched.credential` for the store to persist.
200	# `gate` picks the refresh path: NoGate{} posts directly (tests); Gate{slot,
201	# email, root} runs it under the slot's consume lock.
202	def fetch_with(gate: UGate.Gate, +credential: String, is_active: Bool) -> IO(Fetched):
203	  do IO<Fetched>:
204	    now : U32 <- Clock.now_s()
205	    off : U32 <- local_offset()
206	    fetch.start(UCred.access_token(UCred.oauth.of(credential)), gate, credential, UCred.oauth.of(credential), is_active, now, off)
207	
208	def fetch_usage(credential: String, is_active: Bool) -> IO(Fetched):
209	  fetch_with(UGate.NoGate{}, credential, is_active)
210	
211	# The same, refreshing through the consume gate on <root>/credentials.
212	def fetch_usage_gated(credential: String, is_active: Bool, slot: String, email: String, root: String) -> IO(Fetched):
213	  fetch_with(UGate.Gate{slot, email, root}, credential, is_active)