232 lines
14 KiB
Python
232 lines
14 KiB
Python
"""Consume explicit desktop download acknowledgements; never scan or sweep media.
|
|
|
|
Runs as the existing operator user. It uses Docker exec only on configured adapters
|
|
and workers, adds no listener, and never restarts or cancels model jobs.
|
|
"""
|
|
import argparse,contextlib,json,pathlib,sqlite3,subprocess,time
|
|
|
|
# Runs inside the selected adapter so worker credentials never leave its process.
|
|
NORMALIZER_CODE = r'''
|
|
import asyncio,hashlib,json,os,re,sys
|
|
import httpx,app
|
|
request=json.load(sys.stdin)
|
|
def fingerprint(entry):
|
|
return hashlib.sha256(json.dumps(entry,sort_keys=True,separators=(',',':'),ensure_ascii=False).encode()).hexdigest()
|
|
def locators(entry):
|
|
for node in (entry.get('outputs') or {}).values():
|
|
if not isinstance(node,dict):continue
|
|
for field in ('images','gifs','audio'):
|
|
for item in node.get(field) or []:
|
|
if isinstance(item,dict) and item.get('filename'):
|
|
yield {k:item.get(k,'output' if k=='type' else '') for k in ('filename','subfolder','type')}
|
|
async def main():
|
|
jid=request['id'];provider,prompt=app.decode_job(jid)
|
|
if not re.fullmatch(r'[0-9a-fA-F-]{36}',prompt):raise ValueError('Invalid prompt ID')
|
|
if request['mode']=='verify':
|
|
service_headers={'Authorization':'Bearer '+os.getenv('COMFY_ADAPTER_TOKEN','')}
|
|
async with httpx.AsyncClient(timeout=120) as client:
|
|
result=await client.get('http://127.0.0.1:8105/v1/jobs/'+jid,headers=service_headers);result.raise_for_status();state=result.json()
|
|
if state.get('status')!='succeeded':raise ValueError('Video is not successfully completed')
|
|
_,_,entry,outputs=await app.job_state(jid)
|
|
index=request['index']
|
|
if not 0<=index<len(outputs) or outputs[index].get('kind')!='video':raise ValueError('Output is not a video')
|
|
wanted=outputs[index]
|
|
matches={json.dumps(loc,sort_keys=True) for loc in locators(entry) if loc['filename']==wanted['filename']}
|
|
if len(matches)!=1:raise ValueError('Ambiguous output locator')
|
|
locator=json.loads(matches.pop())
|
|
if locator['type']!='output':raise ValueError('Only permanent output videos may be removed')
|
|
digest=hashlib.sha256();size=0
|
|
async with httpx.AsyncClient(timeout=120) as client:
|
|
async with client.stream('GET','http://127.0.0.1:8105/v1/jobs/'+jid+'/outputs/'+str(index),headers=service_headers) as response:
|
|
response.raise_for_status()
|
|
async for chunk in response.aiter_bytes():
|
|
digest.update(chunk);size+=len(chunk)
|
|
if size>request['bytes']:raise ValueError('Video checksum mismatch')
|
|
if size!=request['bytes'] or digest.hexdigest()!=request['sha256']:raise ValueError('Video checksum mismatch')
|
|
_,_,after,_=await app.job_state(jid)
|
|
if fingerprint(after)!=fingerprint(entry):raise ValueError('History changed during verification')
|
|
return {'provider':provider.id,'prompt':prompt,'locator':locator,'historyHash':fingerprint(entry)}
|
|
manifest=request['manifest']
|
|
if manifest['provider']!=provider.id or manifest['prompt']!=prompt:raise ValueError('Worker mismatch')
|
|
headers=app.provider_headers(provider)
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
q=await client.get(provider.base_url+'/queue',headers=headers);q.raise_for_status();queue=q.json()
|
|
for row in queue.get('queue_running',[])+queue.get('queue_pending',[]):
|
|
if row[1]==prompt or manifest['locator']['filename'] in json.dumps(row,ensure_ascii=False):
|
|
raise ValueError('Video is referenced by an active or queued task')
|
|
h=await client.get(provider.base_url+'/history',headers=headers);h.raise_for_status();history=h.json()
|
|
entry=history.get(prompt)
|
|
if entry and fingerprint(entry)!=manifest['historyHash']:raise ValueError('Task history changed')
|
|
if any(key!=prompt and manifest['locator'] in list(locators(value)) for key,value in history.items()):
|
|
raise ValueError('Video is shared by another task')
|
|
# Check permission before touching the file; an empty list removes nothing.
|
|
check=await client.post(provider.base_url+'/history',headers=headers,json={'delete':[]});check.raise_for_status()
|
|
if request['mode']=='remove_history':
|
|
if entry:
|
|
removed=await client.post(provider.base_url+'/history',headers=headers,json={'delete':[prompt]});removed.raise_for_status()
|
|
confirm=await client.get(provider.base_url+'/history/'+prompt,headers=headers);confirm.raise_for_status()
|
|
if prompt in confirm.json():raise ValueError('History deletion not confirmed')
|
|
return {'historyMissing':not bool(entry)}
|
|
try:
|
|
result=asyncio.run(main());print('HB_CLEANUP_RESULT:'+json.dumps({'ok':True,'result':result}))
|
|
except Exception as error:
|
|
# Only our fixed validation messages leave the adapter; not URLs or tokens.
|
|
message=str(error) if type(error) is ValueError else type(error).__name__
|
|
print('HB_CLEANUP_RESULT:'+json.dumps({'ok':False,'error':message[:160]}));sys.exit(1)
|
|
'''
|
|
|
|
# Uses directory descriptors and O_NOFOLLOW; client receipts never supply a path.
|
|
WORKER_CODE = r'''
|
|
import hashlib,json,os,stat,sys
|
|
request=json.load(sys.stdin)
|
|
def execute():
|
|
loc=request['locator'];name=loc['filename'];sub=loc.get('subfolder','')
|
|
if loc.get('type')!='output' or not isinstance(name,str) or not name.lower().endswith(('.mp4','.webm','.mov','.mkv')):
|
|
raise ValueError('Not an output video')
|
|
if any(c in name for c in ('/','\\','\x00')) or name in ('.','..'):raise ValueError('Invalid filename')
|
|
if not isinstance(sub,str) or sub.startswith('/') or '\\' in sub or '\x00' in sub:raise ValueError('Invalid subfolder')
|
|
parts=[part for part in sub.split('/') if part]
|
|
if any(part in ('.','..') for part in parts):raise ValueError('Invalid subfolder')
|
|
directory=os.open('/app/output',os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW)
|
|
try:
|
|
for part in parts:
|
|
child=os.open(part,os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW,dir_fd=directory)
|
|
os.close(directory);directory=child
|
|
try:fd=os.open(name,os.O_RDONLY|os.O_NOFOLLOW,dir_fd=directory)
|
|
except FileNotFoundError:
|
|
if request['mode']=='delete' and request.get('expected'):return {'missing':True}
|
|
raise
|
|
with os.fdopen(fd,'rb') as file:
|
|
before=os.fstat(file.fileno())
|
|
if not stat.S_ISREG(before.st_mode) or before.st_nlink!=1:raise ValueError('Not a unique regular video file')
|
|
digest=hashlib.sha256()
|
|
for block in iter(lambda:file.read(1024*1024),b''):digest.update(block)
|
|
digest=digest.hexdigest();after=os.fstat(file.fileno())
|
|
signature=lambda s:(s.st_dev,s.st_ino,s.st_size,s.st_mtime_ns)
|
|
if signature(before)!=signature(after):raise ValueError('Video changed while hashing')
|
|
result={'sha256':digest,'bytes':after.st_size,'device':after.st_dev,'inode':after.st_ino,'mtime':after.st_mtime_ns}
|
|
if request['mode']=='delete':
|
|
if result!=request['expected']:raise ValueError('Video changed since preparation')
|
|
if signature(os.stat(name,dir_fd=directory,follow_symlinks=False))!=signature(after):raise ValueError('Video was replaced')
|
|
os.unlink(name,dir_fd=directory);os.fsync(directory)
|
|
return result
|
|
finally:os.close(directory)
|
|
try:
|
|
print('HB_CLEANUP_RESULT:'+json.dumps({'ok':True,'result':execute()}))
|
|
except Exception as error:
|
|
message=str(error) if type(error) is ValueError else type(error).__name__
|
|
print('HB_CLEANUP_RESULT:'+json.dumps({'ok':False,'error':message[:160]}));sys.exit(1)
|
|
'''
|
|
|
|
# Reap only byte-for-byte known history records of already cleaned videos. This
|
|
# prevents an operator's pre-upgrade history snapshot from resurrecting them.
|
|
# It never touches media, queues, or unrelated/replaced history records.
|
|
CLEANED_HISTORY_CODE = r'''
|
|
import asyncio,hashlib,json,sys
|
|
import app,httpx
|
|
request=json.load(sys.stdin)
|
|
def fingerprint(value):
|
|
return hashlib.sha256(json.dumps(value,sort_keys=True,separators=(',',':'),ensure_ascii=False).encode()).hexdigest()
|
|
async def main():
|
|
providers=app.providers();removed=[];deferred=[]
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
for provider_id,records in request['providers'].items():
|
|
provider=providers[provider_id];headers=app.provider_headers(provider)
|
|
response=await client.get(provider.base_url+'/history',headers=headers);response.raise_for_status();history=response.json()
|
|
present=[r for r in records if r['prompt'] in history]
|
|
if not present:continue
|
|
response=await client.get(provider.base_url+'/queue',headers=headers);response.raise_for_status();queue=response.json()
|
|
active={r[1] for r in queue.get('queue_running',[])+queue.get('queue_pending',[]) if len(r)>1}
|
|
for record in present:
|
|
prompt=record['prompt']
|
|
if prompt in active or fingerprint(history[prompt]) not in record['hashes']:
|
|
deferred.append(prompt);continue
|
|
# Re-read the exact record immediately before deletion; never clear history.
|
|
response=await client.get(provider.base_url+'/history/'+prompt,headers=headers);response.raise_for_status();entry=response.json().get(prompt)
|
|
if not entry:continue
|
|
if fingerprint(entry) not in record['hashes']:deferred.append(prompt);continue
|
|
response=await client.post(provider.base_url+'/history',headers=headers,json={'delete':[prompt]});response.raise_for_status()
|
|
removed.append(prompt)
|
|
return {'removed':removed,'deferred':deferred}
|
|
try:print('HB_CLEANUP_RESULT:'+json.dumps({'ok':True,'result':asyncio.run(main())}))
|
|
except Exception as error:
|
|
print('HB_CLEANUP_RESULT:'+json.dumps({'ok':False,'error':type(error).__name__}));sys.exit(1)
|
|
'''
|
|
|
|
def docker_python(container,code,payload):
|
|
result=subprocess.run(['docker','exec','-i',container,'python','-c',code],input=json.dumps(payload),capture_output=True,text=True,timeout=240)
|
|
lines=[line[len('HB_CLEANUP_RESULT:'):] for line in result.stdout.splitlines() if line.startswith('HB_CLEANUP_RESULT:')]
|
|
if not lines:raise RuntimeError('Cleanup executor unavailable')
|
|
value=json.loads(lines[-1])
|
|
if result.returncode or not value.get('ok'):raise RuntimeError(value.get('error','Cleanup executor failed'))
|
|
return value['result']
|
|
|
|
def process_ack(db,entry,row,run=docker_python):
|
|
"""Durable prepare -> exact file unlink -> exact history removal -> acknowledgement."""
|
|
jid=row['id'];manifest=json.loads(row['manifest']) if row['manifest'] else None
|
|
if not db.execute('SELECT 1 FROM jobs WHERE id=?',(jid,)).fetchone():raise RuntimeError('Unknown bridge video')
|
|
if not manifest:
|
|
manifest=run(entry['normalizer'],NORMALIZER_CODE,{'mode':'verify','id':jid,'index':row['output_index'],'sha256':row['sha256'],'bytes':row['bytes']})
|
|
worker=entry['workers'].get(manifest['provider'])
|
|
if not worker:raise RuntimeError('Worker not configured for cleanup')
|
|
run(entry['normalizer'],NORMALIZER_CODE,{'mode':'guard','id':jid,'manifest':manifest})
|
|
manifest['worker']=worker
|
|
manifest['fileReceipt']=run(worker,WORKER_CODE,{'mode':'inspect','locator':manifest['locator']})
|
|
# Persist the authority and exact inode/hash before the first destructive action.
|
|
with db:db.execute("UPDATE video_cleanup_acks SET manifest=?,status='prepared',updated=? WHERE id=?",(json.dumps(manifest),time.time(),jid))
|
|
worker=entry['workers'].get(manifest['provider'])
|
|
if not worker or worker!=manifest['worker']:raise RuntimeError('Worker mapping changed; cleanup postponed')
|
|
run(entry['normalizer'],NORMALIZER_CODE,{'mode':'guard','id':jid,'manifest':manifest})
|
|
run(worker,WORKER_CODE,{'mode':'delete','locator':manifest['locator'],'expected':manifest['fileReceipt']})
|
|
run(entry['normalizer'],NORMALIZER_CODE,{'mode':'remove_history','id':jid,'manifest':manifest})
|
|
with db:db.execute("UPDATE video_cleanup_acks SET status='cleaned',error=NULL,updated=? WHERE id=?",(time.time(),jid))
|
|
|
|
def scan(config,run=docker_python):
|
|
for entry in config['bridges']:
|
|
with contextlib.closing(sqlite3.connect(entry['database'],timeout=15)) as db:
|
|
db.row_factory=sqlite3.Row
|
|
if not db.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name='video_cleanup_acks'").fetchone():continue
|
|
rows=db.execute("SELECT * FROM video_cleanup_acks WHERE status!='cleaned' AND next_attempt<=? ORDER BY created LIMIT 5",(time.time(),)).fetchall()
|
|
for row in rows:
|
|
try:
|
|
process_ack(db,entry,row,run)
|
|
print(json.dumps({'event':'video-cleaned','bridge':entry['name'],'id':row['id']}),flush=True)
|
|
except Exception as error:
|
|
# No tokens, signed links or prompts in persistent diagnostics.
|
|
message=str(error) if type(error) is RuntimeError else type(error).__name__
|
|
attempts=row['attempts']+1
|
|
with db:db.execute('UPDATE video_cleanup_acks SET attempts=?,next_attempt=?,error=?,updated=? WHERE id=?',(attempts,time.time()+min(900,30*2**min(attempts,5)),message[:160],time.time(),row['id']))
|
|
print(json.dumps({'event':'cleanup-retry','bridge':entry['name'],'id':row['id'],'reason':message[:160]}),flush=True)
|
|
|
|
def reconcile_cleaned_history(config,run=docker_python):
|
|
for entry in config['bridges']:
|
|
providers={}
|
|
with contextlib.closing(sqlite3.connect(entry['database'],timeout=15)) as db:
|
|
for raw, in db.execute("SELECT manifest FROM video_cleanup_acks WHERE status='cleaned' AND manifest IS NOT NULL"):
|
|
manifest=json.loads(raw);provider=manifest['provider']
|
|
if entry['workers'].get(provider)!=manifest.get('worker'):continue
|
|
hashes=[manifest['historyHash'],*manifest.get('restoredHistoryHashes',[])]
|
|
providers.setdefault(provider,[]).append({'prompt':manifest['prompt'],'hashes':hashes})
|
|
if providers:
|
|
result=run(entry['normalizer'],CLEANED_HISTORY_CODE,{'providers':providers})
|
|
if result['removed'] or result['deferred']:
|
|
print(json.dumps({'event':'cleaned-history-reconciled','bridge':entry['name'],**result}),flush=True)
|
|
|
|
def main():
|
|
import fcntl
|
|
parser=argparse.ArgumentParser();parser.add_argument('config');parser.add_argument('--once',action='store_true');args=parser.parse_args()
|
|
config_path=pathlib.Path(args.config).resolve();config=json.loads(config_path.read_text())
|
|
with (config_path.parent/'cleanup.lock').open('w') as lock:
|
|
fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
|
|
next_history_check=0
|
|
while True:
|
|
try:scan(config)
|
|
except Exception as error:print(json.dumps({'event':'scan-retry','reason':type(error).__name__}),flush=True)
|
|
if time.monotonic()>=next_history_check:
|
|
try:reconcile_cleaned_history(config)
|
|
except Exception as error:print(json.dumps({'event':'history-reconcile-retry','reason':type(error).__name__}),flush=True)
|
|
next_history_check=time.monotonic()+60
|
|
if args.once:return
|
|
time.sleep(10)
|
|
if __name__=='__main__':main()
|