aoh

commit 60c368b7a1149aa69983081d3e1627a64beb2562
parent 05ba00427778f7ea726a14992ccbb4c3060d72fe
Author: Asher Morgan <asher@ashermorgan.net>
Date:   Sun, 23 Aug 2026 19:01:42 -0700

Improve client error handling

Diffstat:
M
aoh/client.py
 | 
41
++++++++++++++++++++++++++++-------------
M
aoh/server.py
 | 
10
+++++++---

2 files changed, 35 insertions(+), 16 deletions(-)
diff --git a/aoh/client.py b/aoh/client.py
@@ -20,6 +20,10 @@ PASSWORD_PROMPTS = {
 }
 
 
+class ClientError(Exception):
+    """Raised for miscellaneous client errors."""
+
+
 def exec(args):
     """Run a command on the local host."""
 
@@ -88,8 +92,12 @@ def send_json_request(url, req_data=None, method='GET', cookies=''):
         headers = e.headers
         data = e.fp.read()
 
-    # We assume that all responses will include JSON bodies
-    assert headers.get('Content-Type') == 'application/json'
+    if headers.get('Content-Type') != 'application/json':
+        if status_code >= 400:
+            raise ClientError(f'Server responded with {status_code}')
+        else:
+            raise ClientError('Server responded with unexpected content type: '
+                              f"{headers.get('Content-Type')}")
 
     return status_code, headers, json.loads(data)
 
@@ -98,11 +106,10 @@ def runner_loop(runner_url, cookies):
     """Main execution loop."""
 
     req = {}
+
     while True:
         status, _, res = send_json_request(runner_url, req, 'PUT', cookies)
 
-        assert status == 200
-
         if 'logs' in res:
             # Strip duplicate password prompts
             logs = res['logs']
@@ -112,19 +119,22 @@ def runner_loop(runner_url, cookies):
 
             print(logs, end='')
 
-        req = {}
         if 'err' in res:
-            print(res['err'])
-            sys.exit(1)
+            raise ClientError(res['err'])
+        elif status != 200:
+            raise ClientError(f'Server responded with {status}')
+
         if res.get('finished'):
-            break
-        elif 'exec' in res:
+            return
+
+        if 'exec' in res:
             req = exec(res['exec'])
         elif 'put' in res:
             req = put(res['put'])
         elif 'fetch' in res:
             req = fetch(res['fetch'])
         else:
+            req = {}
             time.sleep(0.1)
 
 
@@ -148,9 +158,10 @@ def create_runner(playbook, args):
                     getpass.getpass(PASSWORD_PROMPTS[pw_type])
         status, headers, res = send_json_request(url, req, 'POST')
 
-    if status != 201:
-        print(res['err'])
-        sys.exit(1)
+    if 'err' in res:
+        raise ClientError(res['err'])
+    elif status != 201:
+        raise ClientError(f'Server responded with {status}')
 
     # We assume that only one cookie will be set
     session_token = headers['Set-Cookie'].split(';')[0]
@@ -190,7 +201,11 @@ def cli(args):
         playbook = 'main.yml'
         aoh_args = args[1:]
 
-    create_runner(playbook, aoh_args)
+    try:
+        create_runner(playbook, aoh_args)
+    except ClientError as e:
+        print(f'Error: {e}')
+        sys.exit(1)
 
 
 if __name__ == '__main__':
diff --git a/aoh/server.py b/aoh/server.py
@@ -1,6 +1,6 @@
 import os
 
-from flask import Flask, abort, make_response, render_template, request, session
+from flask import Flask, make_response, render_template, request, session
 from flask_apscheduler import APScheduler
 
 from .playbook import get_playbook
@@ -64,11 +64,15 @@ def new_runner():
 @app.put('/runners/<id>')
 def runner_update(id):
     if id != session.get('runner'):
-        abort(401)
+        return make_response({
+            'err': 'Invalid runner ID',
+        }, 401)
 
     runner = _RUNNERS.get(id)
     if not runner:
-        abort(400)
+        return make_response({
+            'err': 'Runner not found',
+        }, 404)
 
     res = runner.process_client_request(request.json)