leviath_runtime/pipeline/
tool_stages.rs

1//! Keeping the tool service's view of each agent in step with the stage it is
2//! actually in, including refreshing dynamically advertised tools.
3
4use super::*;
5
6/// Notify the [`ToolService`] of every agent that just entered a stage (tagged
7/// with [`StageJustEntered`] by the transition systems), so it can re-sync that
8/// agent's per-stage tool permissions, then clear the tag. Runs after the
9/// transition systems each tick.
10pub fn sync_tool_stages(
11    service: Res<ToolServiceRes>,
12    entered: Query<(Entity, &StageJustEntered)>,
13    mut commands: Commands,
14) {
15    crate::tick_scope::clear();
16    for (entity, stage) in entered.iter() {
17        crate::tick_scope::enter(entity);
18        service.0.sync_stage(entity, stage.index, &stage.name);
19        commands.entity(entity).remove::<StageJustEntered>();
20    }
21}
22
23/// Re-advertise an agent's tools mid-run: when tagged [`ToolsNeedRefresh`], ask
24/// the tool service for this stage's freshly-resolved tool defs and, if it
25/// returns a set, write it into the live [`StageInference`] (what the next
26/// inference request advertises, read fresh by `build_request`) and the matching
27/// [`StageInferences`] catalog entry (so a later revisit of this stage keeps the
28/// updated set). Always consumes the marker. This is the mechanism behind
29/// mid-run dynamic tool discovery and lazily-listed MCP tools.
30pub fn refresh_advertised_tools(
31    service: Res<ToolServiceRes>,
32    mut agents: Query<
33        (
34            Entity,
35            &StageCursor,
36            &mut StageInference,
37            &mut StageInferences,
38        ),
39        With<ToolsNeedRefresh>,
40    >,
41    mut commands: Commands,
42) {
43    crate::tick_scope::clear();
44    for (entity, cursor, mut si, mut sis) in agents.iter_mut() {
45        crate::tick_scope::enter(entity);
46        if let Some(tools) = service.0.refresh_tools(entity, cursor.index) {
47            si.tools = tools.clone();
48            // Keep the catalog entry in sync so re-entering this stage advertises
49            // the same refreshed set.
50            if let Some(slot) = sis.0.get_mut(cursor.index) {
51                slot.tools = tools;
52            }
53        }
54        commands.entity(entity).remove::<ToolsNeedRefresh>();
55    }
56}
57
58/// Poll each `dynamic_tools` agent for a pending tool re-scan and, when the tool
59/// service reports one, tag it [`ToolsNeedRefresh`] so [`refresh_advertised_tools`]
60/// re-advertises before its next turn. Only agents carrying [`DynamicTools`] are
61/// queried, so static agents (the default) cost nothing.
62pub fn poll_dynamic_tool_refresh(
63    service: Res<ToolServiceRes>,
64    agents: Query<Entity, With<DynamicTools>>,
65    mut commands: Commands,
66) {
67    crate::tick_scope::clear();
68    for entity in agents.iter() {
69        crate::tick_scope::enter(entity);
70        if service.0.wants_refresh(entity) {
71            commands.entity(entity).insert(ToolsNeedRefresh);
72        }
73    }
74}
leviath_runtime/pipeline/tool_stages.rs

leviath_runtime/pipeline/
tool_stages.rs
