179 lines
7.7 KiB
Python
179 lines
7.7 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import subprocess
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
from datetime import datetime, timezone
|
|
|
|
|
|
def state_path(workspace):
|
|
if not isinstance(workspace, str) or not workspace or any(ord(c) < 32 for c in workspace):
|
|
raise ValueError('Invalid workspace')
|
|
root = Path(workspace)
|
|
if not root.is_absolute() or not root.is_dir():
|
|
raise ValueError('Invalid workspace')
|
|
root = root.resolve()
|
|
directory = root / '.ai-session'
|
|
path = directory / 'state.json'
|
|
if directory.is_symlink() or path.is_symlink() or directory.resolve().parent != root:
|
|
raise ValueError('Invalid state path')
|
|
return path
|
|
|
|
|
|
def show_pointer(path):
|
|
if path.is_file():
|
|
print('Task state is available at ' + str(path) + '; validate its workspace and goal before reuse.')
|
|
|
|
|
|
def git_identity(workspace):
|
|
try:
|
|
def git(*args):
|
|
return subprocess.check_output(['git', '-C', str(workspace), *args], stderr=subprocess.DEVNULL, text=True).strip()
|
|
root = Path(git('rev-parse', '--show-toplevel')).resolve()
|
|
try:
|
|
branch = git('symbolic-ref', '--quiet', '--short', 'HEAD') or 'HEAD'
|
|
except subprocess.CalledProcessError:
|
|
branch = 'HEAD'
|
|
head = git('rev-parse', 'HEAD')
|
|
try:
|
|
repository = git('config', '--get', 'remote.origin.url') or str(root)
|
|
if '://' in repository:
|
|
parsed = urlsplit(repository)
|
|
repository = urlunsplit((parsed.scheme, parsed.hostname or '', parsed.path, '', ''))
|
|
except subprocess.CalledProcessError:
|
|
repository = str(root)
|
|
return {'repository': repository, 'branch': branch, 'head': head}
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
|
|
|
|
def hook_pointer():
|
|
try:
|
|
raw = '' if sys.stdin.isatty() else sys.stdin.buffer.read(1048577).decode('utf-8-sig')
|
|
if len(raw) > 1048576:
|
|
return
|
|
event = json.loads(raw) if raw.strip() else {}
|
|
if not isinstance(event, dict):
|
|
return
|
|
show_pointer(state_path(event.get('cwd', os.getcwd())))
|
|
except (ValueError, OSError, RecursionError):
|
|
return
|
|
|
|
|
|
def update_state(args):
|
|
path = state_path(args.workspace)
|
|
if args.pointer:
|
|
show_pointer(path)
|
|
return
|
|
state = json.loads(path.read_text(encoding='utf-8-sig')) if path.exists() else {}
|
|
if not isinstance(state, dict):
|
|
raise ValueError('Invalid task state')
|
|
for metadata in ('repository', 'branch', 'head'):
|
|
if metadata in state and not isinstance(state[metadata], str):
|
|
raise ValueError('Invalid task state metadata')
|
|
if 'verificationStale' in state and not isinstance(state['verificationStale'], bool):
|
|
raise ValueError('Invalid task state metadata')
|
|
workspace = str(path.parent.parent)
|
|
if state.get('workspace', workspace) != workspace:
|
|
raise ValueError('Task state belongs to another workspace')
|
|
for old, new in (('changedFiles', 'changed'), ('pending', 'open'), ('importantFindings', 'decisions')):
|
|
if old in state:
|
|
state.setdefault(new, state.pop(old))
|
|
for field in ('goal', 'decisions', 'changed', 'verified', 'open', 'risks', 'verification'):
|
|
if field not in state:
|
|
continue
|
|
value = state[field]
|
|
if field == 'goal':
|
|
valid = isinstance(value, str)
|
|
elif field == 'verification':
|
|
valid = isinstance(value, list) and all(isinstance(item, dict) and isinstance(item.get('check'), str) and isinstance(item.get('result'), str) for item in value)
|
|
else:
|
|
valid = isinstance(value, list) and all(isinstance(item, str) for item in value)
|
|
if not valid:
|
|
raise ValueError('Invalid task state field')
|
|
if args.reset:
|
|
state = {}
|
|
elif args.goal is not None and state.get('goal') and state['goal'] != args.goal:
|
|
raise ValueError('A different goal requires reset')
|
|
state['workspace'] = workspace
|
|
previous_identity = {key: state.get(key) for key in ('repository', 'branch', 'head')}
|
|
identity = git_identity(workspace)
|
|
if identity:
|
|
state.update(identity)
|
|
if path.exists() and any(previous_identity.get(key) not in (None, identity[key]) for key in identity):
|
|
state['verificationStale'] = True
|
|
for item in state.get('verification', []):
|
|
item['stale'] = True
|
|
else:
|
|
for key in ('repository', 'branch', 'head'):
|
|
state.pop(key, None)
|
|
for field in ('goal', 'decisions', 'changed', 'verified', 'open', 'risks'):
|
|
value = getattr(args, field)
|
|
if value is not None:
|
|
state[field] = value
|
|
else:
|
|
state.setdefault(field, '' if field == 'goal' else [])
|
|
state.setdefault('verification', [])
|
|
if getattr(args, 'verified', None):
|
|
state['verification'] = [dict(check=check, head=identity['head'], result='pass') for check in args.verified] if identity else state.get('verification', [])
|
|
if identity:
|
|
state['verificationStale'] = False
|
|
state['updatedAt'] = datetime.now(timezone.utc).isoformat()
|
|
path.parent.mkdir(exist_ok=True)
|
|
temporary = None
|
|
try:
|
|
with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', dir=path.parent,
|
|
prefix='.state-', suffix='.tmp', delete=False) as output:
|
|
temporary = output.name
|
|
json.dump(state, output, ensure_ascii=False, indent=2)
|
|
output.write('\n')
|
|
os.replace(temporary, path)
|
|
temporary = None
|
|
finally:
|
|
if temporary is not None:
|
|
os.unlink(temporary)
|
|
|
|
|
|
def main():
|
|
sys.stdout.reconfigure(encoding='utf-8')
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--hook', action='store_true')
|
|
parser.add_argument('--update-json', action='store_true')
|
|
parser.add_argument('--workspace', default=os.getcwd())
|
|
parser.add_argument('--pointer', action='store_true')
|
|
parser.add_argument('--reset', action='store_true')
|
|
parser.add_argument('--goal')
|
|
for field in ('decisions', 'changed', 'verified', 'open', 'risks', 'verification'):
|
|
parser.add_argument('--' + field, nargs='*')
|
|
args = parser.parse_args()
|
|
if args.hook:
|
|
hook_pointer()
|
|
else:
|
|
try:
|
|
if args.update_json:
|
|
values = json.loads(sys.stdin.buffer.read().decode('utf-8-sig'))
|
|
if not isinstance(values, dict):
|
|
raise ValueError('Invalid task state update')
|
|
for field, value in values.items():
|
|
if field not in ('workspace', 'goal', 'decisions', 'changed', 'verified', 'open', 'risks', 'verification'):
|
|
raise ValueError('Invalid task state field')
|
|
if field in ('workspace', 'goal'):
|
|
if not isinstance(value, str):
|
|
raise ValueError('Invalid task state value')
|
|
elif field == 'verification':
|
|
if not isinstance(value, list) or any(not isinstance(item, dict) or not isinstance(item.get('check'), str) or not isinstance(item.get('result'), str) for item in value):
|
|
raise ValueError('Invalid verification list')
|
|
elif not isinstance(value, list) or any(not isinstance(item, str) for item in value):
|
|
raise ValueError('Invalid task state list')
|
|
setattr(args, field, value)
|
|
update_state(args)
|
|
except (ValueError, OSError):
|
|
parser.exit(1, 'Cannot read or update task state for this workspace.\n')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|