Skip to content

Commit fe4ad89

Browse files
Merge pull request #3833 from springfall2008/copilot/fix-octopus-direct-token-refresh
Fix Octopus Direct: refresh JWT on HTTP 401/403 (intelligent go slots)
2 parents 58159d7 + 5928e38 commit fe4ad89

2 files changed

Lines changed: 133 additions & 6 deletions

File tree

‎apps/predbat/octopus.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1558,6 +1558,21 @@ async def async_graphql_query(self, query, request_context, returns_data=True, i
15581558
headers = {"Authorization": f"{auth_prefix}{self.graphql_token}", integration_context_header: request_context}
15591559
self.log("OctopusAPI: Making GraphQL request to {} payload {} headers {}".format(url, payload, headers))
15601560
async with client.post(url, json=payload, headers=headers) as response:
1561+
# Check for HTTP-level 401/403 (transport-level auth failure) and retry once.
1562+
# This handles cases where the JWT has been revoked server-side and the server
1563+
# returns a bare 401/403 status rather than a GraphQL error body — which would
1564+
# otherwise loop forever without ever refreshing the token.
1565+
if response.status in [401, 403] and _retry_count == 0:
1566+
self.log(f"OctopusAPI: HTTP {response.status} for graphql query {request_context}, forcing token refresh and retry")
1567+
record_api_call("octopus", False, "auth_error")
1568+
self.graphql_token = None
1569+
retry_token = await self.async_refresh_token()
1570+
if retry_token is None:
1571+
self.failures_total += 1
1572+
self.log(f"Warn: OctopusAPI: Failed to refresh token for retry of graphql query {request_context}")
1573+
return None
1574+
return await self.async_graphql_query(query, request_context, returns_data=returns_data, ignore_errors=ignore_errors, _retry_count=1, use_backend=use_backend)
1575+
15611576
# Process response (which reads the text)
15621577
response_body = await self.async_read_response_retry(response, url, ignore_errors=ignore_errors)
15631578

‎apps/predbat/tests/test_octopus_rate_limit.py‎

Lines changed: 118 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -559,11 +559,123 @@ async def mock_refresh():
559559
else:
560560
print("PASS: All auth error codes (KT-CT-1139, KT-CT-1111, KT-CT-1143) trigger auto-refresh")
561561

562-
# Summary
563-
if failed:
564-
print("\n**** Octopus API rate limit tests FAILED ****")
565-
raise Exception("Octopus API rate limit tests failed")
566-
else:
567-
print("\n**** All Octopus API rate limit tests PASSED ****")
562+
# Test 10: HTTP 401 response triggers token refresh and retry (Octopus Direct intelligent slots)
563+
# Regression test for: repeated 401/403 HTTP responses (e.g. when fetching intelligent go slots)
564+
# not refreshing the token, requiring an add-on restart to recover.
565+
print("\n*** Test 10: HTTP 401 response triggers token refresh and retry ***")
566+
for http_status in [401, 403]:
567+
api = OctopusAPI(my_predbat, key="test-key", account_id="test-account", automatic=False)
568+
api.graphql_token = "stale-token"
569+
570+
# Mock client
571+
mock_client_10 = AsyncMock()
572+
api.api.async_create_client_session = AsyncMock(return_value=mock_client_10)
573+
574+
# First response: bare HTTP 401/403 with no GraphQL body
575+
# Second response (after token refresh): HTTP 200 with data
576+
call_count_10 = [0]
577+
578+
def make_mock_http_response(status):
579+
"""Create a mock aiohttp response with the given status."""
580+
r = AsyncMock()
581+
r.status = status
582+
r.__aenter__ = AsyncMock(return_value=r)
583+
r.__aexit__ = AsyncMock(return_value=None)
584+
return r
585+
586+
success_data = {"data": {"intelligentDispatches": []}}
587+
588+
# async_read_response is only reached on the retry (status 200)
589+
api.async_read_response = AsyncMock(return_value=success_data)
590+
591+
# POST returns 401 on first call, 200 on second
592+
def post_side_effect_10(*args, **kwargs):
593+
call_count_10[0] += 1
594+
if call_count_10[0] == 1:
595+
return make_mock_http_response(http_status)
596+
return make_mock_http_response(200)
597+
598+
mock_client_10.post = MagicMock(side_effect=post_side_effect_10)
599+
600+
# Token refresh: on retry, return a fresh token
601+
refresh_calls_10 = [0]
602+
603+
async def mock_token_refresh_with_retry():
604+
"""Mock token refresh returning a new token on the second call."""
605+
refresh_calls_10[0] += 1
606+
if refresh_calls_10[0] == 1:
607+
# Initial check — stale token still set, acts as "valid enough to try"
608+
return "stale-token"
609+
# Second call — forced refresh after 401
610+
api.graphql_token = "fresh-token"
611+
return "fresh-token"
612+
613+
api.async_refresh_token = AsyncMock(side_effect=mock_token_refresh_with_retry)
614+
615+
result = await api.async_graphql_query("query { intelligentDispatches }", f"test-http-{http_status}", returns_data=True, ignore_errors=True)
616+
617+
if result != success_data["data"]:
618+
print(f"ERROR: HTTP {http_status} did not trigger token refresh and retry. Got {result}")
619+
failed = True
620+
break
621+
elif call_count_10[0] != 2:
622+
print(f"ERROR: HTTP {http_status} - expected 2 POST calls (original + retry), got {call_count_10[0]}")
623+
failed = True
624+
break
625+
elif refresh_calls_10[0] < 2:
626+
print(f"ERROR: HTTP {http_status} - token refresh not called enough times. Got {refresh_calls_10[0]}")
627+
failed = True
628+
break
629+
elif api.graphql_token != "fresh-token":
630+
print(f"ERROR: HTTP {http_status} - token not updated after refresh. Got {api.graphql_token}")
631+
failed = True
632+
break
633+
else:
634+
print("PASS: HTTP 401 and 403 responses trigger token refresh and retry")
635+
636+
# Test 10b: HTTP 401 with failed token refresh returns None without infinite loop
637+
print("\n*** Test 10b: HTTP 401 with failed token refresh returns None ***")
638+
api = OctopusAPI(my_predbat, key="test-key", account_id="test-account", automatic=False)
639+
api.graphql_token = "stale-token"
640+
641+
mock_client_10b = AsyncMock()
642+
api.api.async_create_client_session = AsyncMock(return_value=mock_client_10b)
643+
644+
mock_resp_10b = AsyncMock()
645+
mock_resp_10b.status = 401
646+
mock_resp_10b.__aenter__ = AsyncMock(return_value=mock_resp_10b)
647+
mock_resp_10b.__aexit__ = AsyncMock(return_value=None)
648+
mock_client_10b.post = MagicMock(return_value=mock_resp_10b)
649+
650+
# Token refresh fails
651+
refresh_calls_10b = [0]
652+
653+
async def mock_token_refresh_with_failure():
654+
"""Mock token refresh that fails on the forced retry."""
655+
refresh_calls_10b[0] += 1
656+
if refresh_calls_10b[0] == 1:
657+
return "stale-token"
658+
return None # Refresh fails
659+
660+
api.async_refresh_token = AsyncMock(side_effect=mock_token_refresh_with_failure)
661+
initial_failures_10b = api.failures_total
662+
663+
result = await api.async_graphql_query("query { intelligentDispatches }", "test-http-401-no-refresh", returns_data=True, ignore_errors=False)
664+
665+
if result is not None:
666+
print(f"ERROR: Expected None when refresh fails after HTTP 401, got {result}")
667+
failed = True
668+
elif api.failures_total != initial_failures_10b + 1:
669+
print(f"ERROR: failures_total not incremented after failed refresh. Expected {initial_failures_10b + 1}, got {api.failures_total}")
670+
failed = True
671+
else:
672+
print("PASS: HTTP 401 with failed token refresh returns None and increments failures_total")
673+
674+
# Summary
675+
if failed:
676+
print("\n**** Octopus API rate limit tests FAILED ****")
677+
raise Exception("Octopus API rate limit tests failed")
678+
else:
679+
print("\n**** All Octopus API rate limit tests PASSED ****")
568680

569681
return failed

0 commit comments

Comments
 (0)