feat(llm-plan-executor): add graph engine and ROS submit/status/cancel node

2026-11-11 00:26:37 +10:00
parent 2519514b02
commit 1e8856387b
10 changed files with 2017 additions and 9 deletions
@@ -0,0 +1,384 @@
+"""Graph engine: admission, one-edge transitions, terminals and cancellation."""
+from __future__ import annotations
+
+import threading
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from iiwa_experement.scenario import load_scenario
+from iiwa_llm_plan_executor.engine import GraphEngine
+from iiwa_llm_plan_executor.schema import parse_response
+from iiwa_scene_context.catalog import DEFAULT_CATALOG
+from iiwa_scene_context.snapshot import build_envelope, build_scene
+
+
+PACKAGE = Path(__file__).resolve().parents[1]
+SCENES = PACKAGE.parent / 'iiwa_experement' / 'config' / 'scenes.yaml'
+SAFE = {
+    'collision': 'operator_required',
+    'grasp_failed': 'operator_required',
+    'human_in_scene': 'operator_required',
+    'not_found': 'replan_required',
+    'failed': 'replan_required',
+    'not_ok': 'operator_required',
+    'timeout': 'operator_required',
+    'unknown': 'operator_required',
+}
+PROGRESS = {'success', 'ok', 'found', 'clear', 'empty', 'holding', 'placed'}
+
+
+class FakeContext:
+    def __init__(self, scene_json):
+        self.scene_json = scene_json
+        self.mounted_tool = None
+        self.held_object = None
+        self.scene_revision = 0
+        self.ownership_unknown = False
+
+
+class FakeSkills:
+    def __init__(self):
+        self.calls = []
+        self.events = {}
+        self.on_execute = None
+        self.block = False
+        self.started = threading.Event()
+        self.release = threading.Event()
+
+    def execute(self, node, context):
+        if self.on_execute is not None:
+            self.on_execute(node, context)
+        if self.block:
+            self.started.set()
+            self.release.wait(timeout=5.0)
+        self.calls.append(node.id)
+        return SimpleNamespace(event=self.events.get(node.id, 'success'), details={})
+
+
+def envelope(revision=1, stamp=5.0, experiment_id=1):
+    spec = load_scenario(experiment_id, SCENES)
+    tools = {tool.name: {'state': 'docked'} for tool in spec.tools}
+    scene = build_scene(
+        spec,
+        {
+            'version': revision,
+            'mounted_tool': None,
+            'held_object': None,
+            'ready': True,
+            'cap': {'state': spec.cap_state},
+            'tools': tools,
+            'screws': {},
+            'holes': {},
+        },
+        DEFAULT_CATALOG,
+    )
+    return build_envelope(
+        scene,
+        scene_revision=revision,
+        state_version=revision,
+        stamp=stamp,
+        experiment_id=experiment_id,
+    )
+
+
+def node(node_id, skill, obj='', target='', next_on=None):
+    events = DEFAULT_CATALOG.outcomes(skill)
+    next_on = next_on or {}
+    transitions = []
+    for event in events:
+        if event in next_on:
+            destination = next_on[event]
+        elif event in PROGRESS:
+            destination = 'completed'
+        else:
+            destination = SAFE.get(event, 'operator_required')
+        transitions.append({'event': event, 'next': destination})
+    return {
+        'id': node_id,
+        'skill': skill,
+        'object': obj,
+        'target': target,
+        'transitions': transitions,
+    }
+
+
+def graph(entry, *nodes):
+    return parse_response({
+        'status': 'ready',
+        'entry': entry,
+        'plan': list(nodes),
+    })
+
+
+STOP_GRAPH = lambda: graph(
+    'a',
+    node('a', 'stop', 'robot', next_on={'success': 'b'}),
+    node('b', 'stop', 'robot', next_on={'success': 'completed'}),
+)
+
+GRAB_GRAPH = lambda: graph(
+    'g',
+    node('g', 'grab', 'vacuum', next_on={'success': 't'}),
+    node('t', 'transfer', 'vacuum', 'work',
+         next_on={'success': 'completed'}),
+)
+
+
+def make_engine(skills=None, *, clock=5.0, provider=None, threaded=False,
+                max_steps=512, **kwargs):
+    skills = skills or FakeSkills()
+    clock_box = {'now': clock}
+    state = {'envelope': envelope(1, clock_box['now'])}
+    if provider is not None:
+        state['provider'] = provider
+    def default_provider():
+        return state['envelope']
+    events = {
+        'plans': [],
+        'statuses': [],
+        'replans': [],
+        'stops': [],
+    }
+    engine = GraphEngine(
+        state.get('provider', default_provider),
+        lambda env, scene: FakeContext(scene),
+        skills,
+        clock=lambda: clock_box['now'],
+        on_plan=lambda payload: events['plans'].append(payload),
+        on_status=lambda payload: events['statuses'].append(payload),
+        on_replan=lambda payload: events['replans'].append(payload),
+        on_stop=lambda reason: events['stops'].append(reason),
+        robot_fresh=lambda: (True, 'ok'),
+        threaded=threaded,
+        max_steps=max_steps,
+        **kwargs,
+    )
+    engine.update_readiness(scene_ready=True, ready_at=0.0, motion_ready=True)
+    return engine, skills, events, state, clock_box
+
+
+def test_submission_waits_for_scene_and_the_three_second_delay():
+    engine, _, _, _, clock = make_engine(clock=2.0)
+    rejected = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert rejected.accepted is False
+    assert rejected.code == 'not_ready'
+    assert 'start delay' in rejected.message
+
+    engine, _, _, _, clock = make_engine(clock=3.01)
+    accepted = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert accepted.accepted is True
+    engine.run_accepted('r1')
+
+    engine.update_readiness(scene_ready=False, ready_at=0.0, motion_ready=False)
+    rejected = engine.submit('r2', STOP_GRAPH().as_dict(), 1)
+    assert rejected.code == 'not_ready'
+
+
+def test_rejected_graphs_are_never_published_or_executed():
+    engine, skills, events, _, _ = make_engine()
+    blocked = engine.submit('r1', {'status': 'blocked', 'entry': '', 'plan': []}, 1)
+    assert blocked.code == 'blocked'
+    cap_without_vacuum = graph(
+        'g', node('g', 'grab', 'cap', next_on={'success': 'completed'}),
+    )
+    rejected = engine.submit('r2', cap_without_vacuum.as_dict(), 1)
+    assert rejected.accepted is False
+    assert rejected.code == 'replan_required'
+    assert events['plans'] == []
+    assert skills.calls == []
+    stale = engine.submit('r3', STOP_GRAPH().as_dict(), 2)
+    assert stale.code == 'stale_scene'
+    assert engine.get_status('r2').state == 'rejected'
+
+
+def test_accepted_graph_is_published_once_before_the_first_motion():
+    engine, skills, events, _, _ = make_engine()
+    order = []
+
+    def on_plan(payload):
+        events['plans'].append(payload)
+        order.append(('plan', payload['run_id']))
+
+    engine.on_plan = on_plan
+    skills.on_execute = lambda node, context: order.append(('move', node.id))
+    accepted = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert accepted.accepted
+    engine.run_accepted('r1')
+    assert order[0] == ('plan', 'r1')
+    assert order[1] == ('move', 'a')
+    assert len(events['plans']) == 1
+    status = engine.get_status('r1')
+    assert status.state == 'completed'
+    assert status.found is True
+
+
+def test_same_run_id_is_idempotent_and_a_second_run_is_busy():
+    engine, _, events, _, _ = make_engine()
+    first = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert first.accepted
+    duplicate = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert duplicate.accepted is True
+    assert duplicate.code == 'duplicate'
+    assert len(events['plans']) == 1
+    busy = engine.submit('r2', STOP_GRAPH().as_dict(), 1)
+    assert busy.accepted is False
+    assert busy.code == 'busy'
+    assert len(events['plans']) == 1
+
+
+def test_each_event_selects_exactly_one_edge_and_is_visible_in_status():
+    engine, skills, _, _, _ = make_engine()
+    engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    skills.events = {'a': 'success', 'b': 'success'}
+    engine.run_accepted('r1')
+    status = engine.get_status('r1')
+    assert status.state == 'completed'
+    assert [event['node_id'] for event in status.events] == ['a', 'b']
+    assert [event['next'] for event in status.events] == ['b', 'completed']
+    assert all(event['run_id'] == 'r1' for event in status.events)
+
+
+def test_replan_returns_a_fresh_scene_and_operator_required_stops():
+    engine, skills, events, _, _ = make_engine()
+    replan_graph = graph(
+        'd',
+        node('d', 'detect', 'cap',
+             next_on={'found': 'replan_required'}),
+    )
+    engine.submit('r1', replan_graph.as_dict(), 1)
+    skills.events = {'d': 'found'}
+    engine.run_accepted('r1')
+    assert engine.get_status('r1').state == 'replan_required'
+    assert len(events['replans']) == 1
+    assert events['replans'][0]['scene_revision'] == 1
+
+    engine, skills, events, _, _ = make_engine()
+    operator_graph = graph(
+        'd',
+        node('d', 'detect', 'cap',
+             next_on={'found': 'operator_required'}),
+    )
+    engine.submit('r2', operator_graph.as_dict(), 1)
+    skills.events = {'d': 'found'}
+    engine.run_accepted('r2')
+    assert engine.get_status('r2').state == 'operator_required'
+
+
+def test_unknown_skill_outcome_stops_execution():
+    engine, skills, events, _, _ = make_engine()
+    engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    skills.events = {'a': 'bogus'}
+    engine.run_accepted('r1')
+    status = engine.get_status('r1')
+    assert status.state == 'failed'
+    assert 'unknown skill outcome' in status.reason
+    assert events['stops']
+
+
+def test_lost_freshness_and_missing_scene_stop_motion():
+    engine, _, events, _, _ = make_engine()
+    engine.robot_fresh = lambda: (False, 'joint_states is stale')
+    engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    engine.run_accepted('r1')
+    assert engine.get_status('r1').state == 'failed'
+    assert 'joint_states' in engine.get_status('r1').reason
+    assert events['stops']
+
+    engine, skills, events, state, _ = make_engine()
+    engine.submit('r2', STOP_GRAPH().as_dict(), 1)
+    skills.on_execute = lambda node, context: state.update(envelope=None)
+    engine.run_accepted('r2')
+    assert engine.get_status('r2').state == 'failed'
+    assert 'scene snapshot' in engine.get_status('r2').reason
+
+
+def test_scene_revision_change_stops_non_ownership_nodes_but_not_own_graph_changes():
+    engine, skills, events, state, _ = make_engine()
+    engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    skills.on_execute = lambda node, context: state.update(envelope=envelope(2))
+    engine.run_accepted('r1')
+    assert engine.get_status('r1').state == 'failed'
+    assert 'changed externally' in engine.get_status('r1').reason
+
+    engine, skills, events, state, _ = make_engine()
+    engine.submit('r2', GRAB_GRAPH().as_dict(), 1)
+
+    def own_change(node, context):
+        if node.id == 'g':
+            context.mounted_tool = 'vacuum'
+            state.update(envelope=envelope(2))
+
+    skills.on_execute = own_change
+    engine.run_accepted('r2')
+    status = engine.get_status('r2')
+    assert status.state == 'completed'
+    assert status.scene_revision == 2
+    assert status.mounted_tool == 'vacuum'
+
+
+def test_step_limit_stops_a_recovering_cycle():
+    engine, skills, _, _, _ = make_engine(max_steps=3)
+    cycle = graph(
+        'a',
+        node('a', 'wait_scene_clear', 'workspace',
+             next_on={'clear': 'b'}),
+        node('b', 'wait_scene_clear', 'workspace',
+             next_on={'clear': 'a'}),
+    )
+    engine.submit('r1', cycle.as_dict(), 1)
+    skills.events = {'a': 'clear', 'b': 'clear'}
+    engine.run_accepted('r1')
+    status = engine.get_status('r1')
+    assert status.state == 'failed'
+    assert 'step limit' in status.reason
+
+
+def test_cancellation_stops_the_active_graph_and_preserves_ownership():
+    engine, skills, events, _, _ = make_engine(threaded=True)
+    skills.block = True
+    skills.on_execute = (
+        lambda node, context: setattr(context, 'mounted_tool', node.object)
+        if node.skill == 'grab' else None
+    )
+    engine.submit('r1', GRAB_GRAPH().as_dict(), 1)
+    assert skills.started.wait(timeout=2.0)
+    # A tool grab confirms ownership before the next node.
+    result = engine.cancel('r1')
+    assert result.accepted
+    assert events['stops']
+    skills.release.set()
+    assert engine.wait('r1', timeout=2.0)
+    status = engine.get_status('r1')
+    assert status.state == 'canceled'
+    assert status.mounted_tool == 'vacuum'
+
+
+def test_cancel_rejects_unknown_or_finished_runs():
+    engine, _, _, _, _ = make_engine()
+    assert engine.cancel('missing').accepted is False
+    engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    engine.run_accepted('r1')
+    assert engine.cancel('r1').accepted is False
+
+
+def test_status_query_marks_unknown_ids():
+    engine, _, _, _, _ = make_engine()
+    status = engine.get_status('nope')
+    assert status.found is False
+    assert status.state == 'unknown'
+
+
+def test_not_ready_or_unavailable_scene_is_rejected():
+    engine, _, _, state, _ = make_engine()
+    state['envelope'] = None
+    rejected = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert rejected.code == 'not_ready'
+
+
+def test_stale_scene_timestamp_is_rejected_before_execution():
+    engine, _, _, state, clock = make_engine(clock=10.0)
+    state['envelope'] = envelope(1, stamp=1.0)
+    rejected = engine.submit('r1', STOP_GRAPH().as_dict(), 1)
+    assert rejected.code == 'stale_scene'
@@ -0,0 +1,213 @@
+"""ROS wrapper helpers: request parsing, plan/status publication and logs."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+from iiwa_experement.geometry import load_geometry
+from iiwa_experement.scenario import load_scenario
+from iiwa_llm_plan_executor.engine import AdmissionResult, RunStatus
+from iiwa_llm_plan_executor.node import (
+    PlannerExecutorNode,
+    RunLog,
+    _safe_name,
+    plan_qos,
+)
+from iiwa_msgs.srv import GetRunStatus
+
+
+PACKAGE = Path(__file__).resolve().parents[1]
+ASSETS = PACKAGE.parent / 'iiwa_experement' / 'config'
+
+
+class FakePublisher:
+    def __init__(self):
+        self.messages = []
+
+    def publish(self, message):
+        self.messages.append(message.data)
+
+
+class FakeEngine:
+    def __init__(self, admission=None, status=None, cancel=None):
+        self.admission = admission or AdmissionResult(True, 'accepted', 'ok', 'r1')
+        self.status = status
+        self.cancel_result = cancel
+        self.submitted = []
+        self.cancelled = []
+
+    def submit(self, run_id, graph, revision):
+        self.submitted.append((run_id, graph, revision))
+        return self.admission
+
+    def get_status(self, run_id):
+        return self.status or RunStatus(
+            True, run_id, 'completed', 'done', 'a', 'stop', 'success',
+            'completed', 1, 'completed', 'vacuum', None, (),
+        )
+
+    def cancel(self, run_id):
+        self.cancelled.append(run_id)
+        return self.cancel_result or SimpleNamespace(
+            accepted=True, message='cancel requested',
+        )
+
+
+def fake_node(*, engine=None):
+    node = object.__new__(PlannerExecutorNode)
+    node.engine = engine or FakeEngine()
+    node._plan_publisher = FakePublisher()
+    node._status_publisher = FakePublisher()
+    node._last_plan = None
+    node._last_status = None
+    node.log = RunLog('/tmp/iiwa-plan-log-test')
+    node.get_logger = lambda: SimpleNamespace(
+        info=lambda *a, **k: None,
+        warning=lambda *a, **k: None,
+        error=lambda *a, **k: None,
+    )
+    node._now = lambda: 10.0
+    node._lock = __import__('threading').RLock()
+    return node
+
+
+def test_plan_qos_is_transient_local_depth_one():
+    qos = plan_qos()
+    assert qos.durability.name == 'TRANSIENT_LOCAL'
+    assert qos.depth == 1
+
+
+def test_safe_name_and_run_log_are_machine_readable(tmp_path):
+    assert _safe_name('run/one 1') == 'run_one_1'
+    log = RunLog(str(tmp_path))
+    log.write({'run_id': 'r1', 'event': 'plan_accepted', 'scene_revision': 1})
+    log.write({'run_id': 'r1', 'event': 'completed', 'terminal': 'completed'})
+    lines = (tmp_path / 'r1.jsonl').read_text().splitlines()
+    assert [json.loads(line)['event'] for line in lines] == [
+        'plan_accepted', 'completed',
+    ]
+
+
+def test_plan_is_published_once_and_kept_for_late_subscribers():
+    node = fake_node()
+    payload = {'run_id': 'r1', 'scene_revision': 2, 'response': '{}'}
+    PlannerExecutorNode._publish_plan(node, payload)
+    PlannerExecutorNode._publish_plan(node, payload)
+    assert len(node._plan_publisher.messages) == 2
+    assert node._last_plan == payload
+    assert json.loads(node._plan_publisher.messages[-1])['run_id'] == 'r1'
+
+
+def test_submit_parses_json_and_forwards_the_untouched_response():
+    node = fake_node()
+    published = []
+    node._publish_status = lambda payload: published.append(payload)
+    request = SimpleNamespace(value=json.dumps({
+        'run_id': 'r1',
+        'scene_revision': 3,
+        'response': {'status': 'ready', 'entry': 'a', 'plan': []},
+    }))
+    response = SimpleNamespace(success=False)
+    PlannerExecutorNode._handle_submit(node, request, response)
+    assert response.success is True
+    assert node.engine.submitted[0][0] == 'r1'
+    assert node.engine.submitted[0][1] == {
+        'status': 'ready', 'entry': 'a', 'plan': [],
+    }
+    assert node.engine.submitted[0][2] == 3
+
+
+def test_rejected_submit_returns_false_and_publishes_the_reason():
+    node = fake_node(engine=FakeEngine(
+        admission=AdmissionResult(False, 'replan_required', 'need vacuum', 'r1'),
+    ))
+    published = []
+    node._publish_status = lambda payload: published.append(payload)
+    request = SimpleNamespace(value=json.dumps({
+        'run_id': 'r1', 'scene_revision': 1,
+        'response': {'status': 'blocked', 'entry': '', 'plan': []},
+    }))
+    response = SimpleNamespace(success=True)
+    PlannerExecutorNode._handle_submit(node, request, response)
+    assert response.success is False
+    assert published[-1]['state'] == 'rejected'
+    assert 'replan_required' in published[-1]['reason']
+
+
+def test_invalid_submit_json_is_rejected_without_calling_the_engine():
+    node = fake_node()
+    request = SimpleNamespace(value='{broken')
+    response = SimpleNamespace(success=True)
+    PlannerExecutorNode._handle_submit(node, request, response)
+    assert response.success is False
+    assert node.engine.submitted == []
+
+
+def test_status_service_returns_found_and_json():
+    node = fake_node()
+    request = GetRunStatus.Request(run_id='r1')
+    response = GetRunStatus.Response()
+    PlannerExecutorNode._handle_get_status(node, request, response)
+    assert response.found is True
+    payload = json.loads(response.status_json)
+    assert payload['run_id'] == 'r1'
+    assert payload['state'] == 'completed'
+    assert payload['mounted_tool'] == 'vacuum'
+
+
+def test_cancel_accepts_json_and_plain_run_ids():
+    node = fake_node()
+    response = SimpleNamespace(success=False)
+    PlannerExecutorNode._handle_cancel(
+        node, SimpleNamespace(value=json.dumps({'run_id': 'r1'})), response,
+    )
+    assert response.success is True
+    response = SimpleNamespace(success=False)
+    PlannerExecutorNode._handle_cancel(
+        node, SimpleNamespace(value='r2'), response,
+    )
+    assert node.engine.cancelled == ['r1', 'r2']
+
+
+def test_scene_provider_enforces_freshness():
+    node = fake_node()
+    node._now = lambda: 10.0
+    node.scene_max_age = 2.0
+    node._latest_scene = {
+        'timestamp': 9.0, 'scene_revision': 1, 'scene_json': {},
+    }
+    assert PlannerExecutorNode._scene_provider(node) is not None
+    node._latest_scene['timestamp'] = 7.0
+    assert PlannerExecutorNode._scene_provider(node) is None
+
+
+def test_context_factory_sets_sources_and_inspection_configuration():
+    node = fake_node()
+    node.spec = load_scenario(20, ASSETS / 'scenes.yaml')
+    node.geometry = load_geometry(ASSETS / 'geometry.yaml')
+    node.adapter = 'adapter'
+    node.close_inspection = frozenset({'cap'})
+    node._frame_provider = lambda name: 1.0
+    node._notify_operator = lambda reason, details: True
+    context = PlannerExecutorNode._context_factory(
+        node, {'scene_revision': 4}, {},
+    )
+    assert context.scene_revision == 4
+    assert context.scene_adapter == 'adapter'
+    assert context.close_inspection_objects == frozenset({'cap'})
+    assert context.capture_frame('top') == 1.0
+    assert context.notify('x', {}) is True
+
+
+def test_robot_freshness_uses_the_joint_limit_age():
+    node = fake_node()
+    node.limits = SimpleNamespace(joint_state_max_age_sec=1.0)
+    node._last_joint = None
+    node._last_joint_at = None
+    assert PlannerExecutorNode._robot_fresh(node)[0] is False
+    node._last_joint = (0.0,) * 7
+    node._last_joint_at = 9.5
+    assert PlannerExecutorNode._robot_fresh(node)[0] is True
+    node._last_joint_at = 8.0
+    assert PlannerExecutorNode._robot_fresh(node)[0] is False