55 lines
3.6 KiB
Python
55 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""A Kie task survives a failed poll and resumes in a new PHP process without re-submission."""
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from threading import Thread
|
|
import json
|
|
import subprocess
|
|
import tempfile
|
|
from support import copy_core, write_json
|
|
|
|
class Provider(BaseHTTPRequestHandler):
|
|
def log_message(self, *_): pass
|
|
def reply(self, status, body, kind='application/json'):
|
|
if isinstance(body, dict): body = json.dumps(body).encode()
|
|
self.send_response(status); self.send_header('Content-Type', kind); self.send_header('Content-Length', str(len(body))); self.end_headers(); self.wfile.write(body)
|
|
def do_POST(self):
|
|
self.rfile.read(int(self.headers['Content-Length']))
|
|
self.server.submissions += 1
|
|
self.reply(200, {'data': {'taskId': 'persistent-task'}})
|
|
def do_GET(self):
|
|
if self.path == '/image.png': self.reply(200, self.server.png, 'image/png'); return
|
|
self.server.polls += 1
|
|
if self.server.polls == 1: self.reply(503, {'message': 'Temporary poll outage'}); return
|
|
self.reply(200, {'data': {'state': 'success', 'resultJson': json.dumps({'resultUrls': [self.server.url + '/image.png']})}})
|
|
|
|
with tempfile.TemporaryDirectory(prefix='raiju-ai-resume-') as folder:
|
|
root = copy_core(folder)
|
|
server = ThreadingHTTPServer(('127.0.0.1', 0), Provider)
|
|
server.url = f'http://127.0.0.1:{server.server_port}'
|
|
server.submissions = server.polls = 0
|
|
server.png = subprocess.check_output(['php', '-r', '$im=imagecreatetruecolor(16,16);imagepng($im);'])
|
|
thread = Thread(target=server.serve_forever, daemon=True); thread.start()
|
|
try:
|
|
write_json(root / 'content/ai-settings.json', {'enabled': True, 'provider': 'kie', 'baseUrl': server.url, 'allowLocalProvider': True, 'apiKey': 'fixture', 'features': {'imageGeneration': True}})
|
|
prefix = "require 'includes/ai/ai-helper.php'; "
|
|
def php(code): return subprocess.check_output(['php', '-r', prefix + code], cwd=root, text=True)
|
|
options = "['_requestId'=>'image_job_0123456789abcdef','model'=>'gpt-image-2','size'=>'auto','count'=>1]"
|
|
first = php("try { raijuAiGenerateImage('Fixture',"+options+"); } catch(Throwable $e) { echo $e->getMessage(); }")
|
|
assert 'outage' in first and server.submissions == 1
|
|
usage = json.loads((root / 'content/ai-usage.json').read_text())
|
|
entry = usage['reservations']['image_job_0123456789abcdef']
|
|
assert entry['status'] == 'uncertain' and entry['tasks'][0]['id'] == 'persistent-task'
|
|
php("raijuAiSaveImageJob(['id'=>'job_0123456789abcdef','status'=>'running','startedAt'=>gmdate('c',time()-2000)]);")
|
|
state = json.loads(php("echo json_encode(raijuAiRefreshImageJobState(raijuAiLoadImageJob('job_0123456789abcdef')));"))
|
|
assert state['status'] == 'queued', state
|
|
result = json.loads(php("echo json_encode(raijuAiGenerateImage('Fixture',"+options+"));"))
|
|
assert server.submissions == 1 and server.polls == 2 and (root / result['path'].lstrip('/')).is_file()
|
|
# Re-entering after a successful provider response reuses saved local outputs.
|
|
again = json.loads(php("echo json_encode(raijuAiGenerateImage('Fixture',"+options+"));"))
|
|
assert again['path'] == result['path'] and server.submissions == 1
|
|
usage = json.loads((root / 'content/ai-usage.json').read_text())
|
|
assert next(iter(usage['days'].values()))['requests'] == 1
|
|
print('PASS persistent Kie task: failed polling, fresh-process resume, stale worker recovery, one provider submission and one usage settlement')
|
|
finally:
|
|
server.shutdown(); server.server_close(); thread.join()
|