starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-21 00:02:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: address review round 1 on dynamic ai agent toolsets

Browse Source 
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
hugocasa
2026-09-09 17:39:57 +02:00
co-authored by Claude Opus 5
parent 4509c96de7
commit 76e70d6a8f
10 changed files with 245 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-worker/src/ai/utils.rs
+9 -1
View File
@@ -26,6 +26,11 @@ use windmill_queue::{flow_status::get_step_of_flow_status, MiniPulledJob};
use crate::parse_sig_of_lang;
/// What every tool an MCP server exposes is advertised under, ahead of the server and tool names
/// (`convert_mcp_tools_to_windmill_tools`). Outside the `mcp` gate: a run narrowing its tools reads
/// it to tell a name that could belong to a server from one that could not.
pub const MCP_TOOL_NAME_PREFIX: &str = "mcp_";
pub async fn parse_raw_script_schema(
content: &str,
language: &ScriptLang,
@@ -377,7 +382,10 @@ fn convert_mcp_tools_to_windmill_tools(
.iter()
.map(|mcp_tool| {
let sanitized_resource_name = sanitize_tool_name_part(resource_name);
let tool_name = format!("mcp_{}_{}", sanitized_resource_name, mcp_tool.name);
let tool_name = format!(
"{}{}_{}",
MCP_TOOL_NAME_PREFIX, sanitized_resource_name, mcp_tool.name
);
let mut schema_value = serde_json::to_value(&*mcp_tool.input_schema)
.context("Failed to convert MCP schema to JSON value")?;
backend/windmill-worker/src/ai_executor.rs
+197 -66
View File
@@ -4,7 +4,7 @@ use crate::ai::utils::{
filter_schema_by_input_transforms, find_unique_tool_name, get_flow_context,
get_flow_job_runnable_and_raw_flow, get_step_name_from_flow, load_mcp_tools,
parse_raw_script_schema, update_flow_status_module_with_actions,
update_flow_status_module_with_actions_success,
update_flow_status_module_with_actions_success, MCP_TOOL_NAME_PREFIX,
};
use crate::memory_oss::{read_from_memory, write_to_memory};
use crate::worker_flow::{get_previous_job_result, get_transform_context};
@@ -246,13 +246,52 @@ fn overlay_tool_inputs(
}
}
/// Write a linked step's own inputs over the brain its agent supplied.
///
/// Called only after the resource has been interpolated: these values are caller-controlled and
/// `build_args_map` has already resolved them, so passing them through it again would expand
/// contextual values — `$WM_TOKEN` in a user message would reach the model provider.
///
/// `user_message`/`user_attachments` are the step's whatever it holds, blank included. The two
/// below them are only an override when the step actually holds a value: an unfilled field arrives
/// as null, and writing that through would drop the `memory` of an agent saved back when memory was
/// part of the brain, ending the conversations it holds without saying so.
fn overlay_flow_local_args(
brain: &mut serde_json::Map<String, serde_json::Value>,
local_args: &HashMap<String, Box<RawValue>>,
) {
for key in ["user_message", "user_attachments"] {
if let Some(v) = local_args.get(key) {
brain.insert(
key.to_string(),
serde_json::from_str(v.get()).unwrap_or(serde_json::Value::Null),
);
}
}
for key in ["memory", "enabled_tools"] {
let Some(v) = local_args
.get(key)
.and_then(|v| serde_json::from_str::<serde_json::Value>(v.get()).ok())
.filter(|v| !v.is_null())
else {
continue;
};
brain.insert(key.to_string(), v);
}
}
/// The roster a run advertises, given the tool names it enabled, plus the resource paths of the
/// MCP entries it named outright — those enable every tool of that server, which only
/// `load_mcp_tools` can enumerate.
///
/// `None` advertises the whole roster, which is what every agent written before the field existed
/// relies on; an empty list advertises nothing. MCP entries always survive this stage: they
/// advertise `mcp_<server>_<tool>` names that do not exist until the server has answered.
/// relies on; an empty list advertises nothing.
///
/// An MCP entry the run did not name survives only while some enabled name could still turn out to
/// be one of its tools, i.e. carries the `mcp_` prefix every one of them is advertised under. A
/// server nothing can match is dropped here rather than after `load_mcp_tools`: resolving one reads
/// its resource, refreshes its token and opens a client, each of which can fail the whole run — a
/// run that switched that server off must not be brought down by it.
fn narrow_roster(
tools: Vec<AgentTool>,
enabled_tools: Option<&[String]>,
@@ -260,6 +299,7 @@ fn narrow_roster(
let Some(enabled) = enabled_tools else {
return (tools, HashSet::new());
};
let names_a_server_tool = enabled.iter().any(|n| n.starts_with(MCP_TOOL_NAME_PREFIX));
let mut enabled_mcp_paths = HashSet::new();
let tools = tools
.into_iter()
@@ -274,7 +314,7 @@ fn narrow_roster(
enabled_mcp_paths
.insert(mcp.resource_path.trim_start_matches("$res:").to_string());
}
true
named || names_a_server_tool
}
_ => named,
}
@@ -283,15 +323,57 @@ fn narrow_roster(
(tools, enabled_mcp_paths)
}
/// Names in `enabled_tools` that name nothing on the agent. The list is an input transform, so it
/// can be computed per run; a name that has since been renamed away must not fail the step, but it
/// would otherwise silently narrow the agent, so the caller logs what it dropped.
fn unmatched_enabled_tools(enabled_tools: &[String], advertised: &[&str]) -> Vec<String> {
enabled_tools
/// Whether a tool an MCP server exposes is advertised: the run named it directly, or named the
/// server entry it came from.
///
/// The two sides of that second test are different types — a roster entry's `resource_path` and the
/// `McpToolSource` the loader builds — so they are matched on the string both strip to
/// (`ai/utils.rs`). Getting that wrong advertises nothing and says nothing, which is why it is
/// pinned by a test rather than left inline.
fn mcp_tool_enabled(tool: &Tool, enabled: &[String], enabled_mcp_paths: &HashSet<String>) -> bool {
let Some(source) = &tool.mcp_source else {
// A Windmill tool, already settled by `narrow_roster`.
return true;
};
enabled.iter().any(|n| n == &tool.def.function.name)
|| enabled_mcp_paths.contains(&source.resource_path)
}
/// How many unmatched names are worth naming in the log. The list is an input transform, so its
/// length is the flow author's to choose; the log line exists to point at a typo, and the count
/// carries the rest.
const MAX_LOGGED_UNMATCHED_TOOLS: usize = 20;
/// The log line for names in `enabled_tools` that name nothing on the agent, if any. The list can
/// be computed per run, so a name that has since been renamed away must not fail the step — but it
/// would otherwise silently narrow the agent, so the run says what it dropped.
fn unmatched_enabled_tools_message(
enabled_tools: &[String],
advertised: &[&str],
) -> Option<String> {
let unmatched: Vec<&str> = enabled_tools
.iter()
.filter(|name| !advertised.contains(&name.as_str()))
.map(|name| name.as_str())
.filter(|name| !advertised.contains(name))
.collect();
if unmatched.is_empty() {
return None;
}
let shown = unmatched
.iter()
.take(MAX_LOGGED_UNMATCHED_TOOLS)
.cloned()
.collect()
.collect::<Vec<_>>()
.join(", ");
let rest = unmatched.len().saturating_sub(MAX_LOGGED_UNMATCHED_TOOLS);
let suffix = if rest > 0 {
format!(" and {rest} more")
} else {
String::new()
};
Some(format!(
"--- ENABLED TOOLS: {shown}{suffix} named no tool of this agent and had no effect ---\n"
))
}
pub async fn handle_ai_agent_job(
@@ -429,9 +511,10 @@ pub async fn handle_ai_agent_job(
));
};
// A linked step takes its brain and tools from the resource and keeps only the flow-local
// inputs (user_message/user_attachments) of its own; both stay rigid, so the one thing it may
// bind to this flow is the tools' inputs, overlaid from `tool_inputs` below.
// A linked step takes its brain and tools from the resource and keeps only its own flow-local
// inputs. The brain and the roster stay rigid; what the step binds to this flow is the message
// it asks, which of those tools this use may call, the conversation it is part of, and the
// tools' own inputs — the last overlaid from `tool_inputs` below.
let (args, tools): (AIAgentArgs, Vec<AgentTool>) = if let Some(agent_ref) = agent.as_deref() {
let agent_path = agent_ref
.trim_start_matches("$res:")
@@ -482,30 +565,7 @@ pub async fn handle_ai_agent_job(
)))
}
};
// Only after interpolating the resource: these are caller-controlled and already resolved by
// build_args_map, so passing them through it again would expand contextual values —
// `$WM_TOKEN` in a user message would reach the model provider.
for key in ["user_message", "user_attachments"] {
if let Some(v) = local_args.get(key) {
brain.insert(
key.to_string(),
serde_json::from_str(v.get()).unwrap_or(serde_json::Value::Null),
);
}
}
// Flow-local like the two above, but only when the step actually holds a value: an unset
// field arrives as null, and writing that through would drop the `memory` of an agent saved
// back when memory was part of the brain, silently ending the conversations it holds.
for key in ["memory", "enabled_tools"] {
let Some(v) = local_args
.get(key)
.and_then(|v| serde_json::from_str::<serde_json::Value>(v.get()).ok())
.filter(|v| !v.is_null())
else {
continue;
};
brain.insert(key.to_string(), v);
}
overlay_flow_local_args(&mut brain, &local_args);
let args = serde_json::from_value::<AIAgentArgs>(serde_json::Value::Object(brain))
.map_err(|e| {
Error::internal_err(format!(
@@ -758,15 +818,8 @@ pub async fn handle_ai_agent_job(
};
if let Some(enabled) = enabled_tools {
// The other half of the narrowing above: an MCP tool is advertised when the run names it
// directly, or names the server entry it came from.
tools.retain(|t| match &t.mcp_source {
Some(source) => {
enabled.iter().any(|n| n == &t.def.function.name)
|| enabled_mcp_paths.contains(&source.resource_path)
}
None => true,
});
// The other half of the narrowing above, now that the servers kept by it have answered.
tools.retain(|t| mcp_tool_enabled(t, enabled, &enabled_mcp_paths));
let mut matchable: Vec<&str> = roster_names.iter().map(|s| s.as_str()).collect();
matchable.extend(
tools
@@ -783,18 +836,8 @@ pub async fn handle_ai_agent_job(
"tools"
},
);
let unmatched = unmatched_enabled_tools(enabled, &matchable);
if !unmatched.is_empty() {
append_logs(
&job.id,
&job.workspace_id,
format!(
"--- ENABLED TOOLS: {} named no tool of this agent and was ignored ---\n",
unmatched.join(", ")
),
conn,
)
.await;
if let Some(message) = unmatched_enabled_tools_message(enabled, &matchable) {
append_logs(&job.id, &job.workspace_id, message, conn).await;
}
}
@@ -2001,24 +2044,112 @@ mod tests {
assert_eq!(names(&all), ["get_user", "send_email", "github"]);
assert!(paths.is_empty());
// An empty list is a list: it advertises nothing, bar the deferred MCP entry.
// An empty list is a list: nothing is advertised, and no server is resolved to find that
// out — one that is down must not fail a run that switched it off.
let (none, paths) = narrow_roster(roster(), Some(&[]));
assert_eq!(names(&none), ["github"]);
assert!(names(&none).is_empty());
assert!(paths.is_empty());
let enabled = ["get_user".to_string(), "renamed_away".to_string()];
let (kept, paths) = narrow_roster(roster(), Some(&enabled));
assert_eq!(names(&kept), ["get_user", "github"]);
// Named nothing, so the MCP entry only survives to be settled against its own tool names.
assert_eq!(names(&kept), ["get_user"]);
assert!(paths.is_empty());
assert_eq!(
unmatched_enabled_tools(&enabled, &["get_user", "send_email", "github"]),
["renamed_away"]
unmatched_enabled_tools_message(&enabled, &["get_user", "send_email", "github"])
.unwrap(),
"--- ENABLED TOOLS: renamed_away named no tool of this agent and had no effect ---\n"
);
// Naming the entry enables every tool of that server, keyed by the path load_mcp_tools uses.
let (_, paths) = narrow_roster(roster(), Some(&["github".to_string()]));
let (kept, paths) = narrow_roster(roster(), Some(&["github".to_string()]));
assert_eq!(names(&kept), ["github"]);
assert_eq!(paths.into_iter().collect::<Vec<_>>(), ["u/test/gh"]);
// A name that could only be one of a server's own tools keeps every server in, since which
// one it belongs to is not knowable until they answer.
let (kept, paths) = narrow_roster(roster(), Some(&["mcp_github_create_issue".to_string()]));
assert_eq!(names(&kept), ["github"]);
assert!(paths.is_empty());
}
/// The two sides of the server-entry match are different types, and getting it wrong advertises
/// nothing and says nothing.
#[test]
fn mcp_tools_match_their_own_name_or_their_server_entry() {
fn mcp_tool(name: &str, resource_path: &str) -> Tool {
Tool {
def: ToolDef {
r#type: "function".to_string(),
function: ToolDefFunction {
name: name.to_string(),
description: None,
parameters: to_raw_value(&serde_json::json!({})),
},
},
module: None,
mcp_source: Some(McpToolSource {
name: "github".to_string(),
tool_name: name.to_string(),
resource_path: resource_path.to_string(),
}),
}
}
let create = mcp_tool("mcp_github_create_issue", "u/test/gh");
let list = mcp_tool("mcp_github_list_issues", "u/test/gh");
// Named directly: that tool only.
let enabled = ["mcp_github_create_issue".to_string()];
assert!(mcp_tool_enabled(&create, &enabled, &HashSet::new()));
assert!(!mcp_tool_enabled(&list, &enabled, &HashSet::new()));
// The server entry named instead: everything it exposes, whatever the names turned out to be.
let paths = HashSet::from(["u/test/gh".to_string()]);
assert!(mcp_tool_enabled(&create, &[], &paths));
assert!(mcp_tool_enabled(&list, &[], &paths));
assert!(!mcp_tool_enabled(
&mcp_tool("mcp_gitlab_list_issues", "u/test/gl"),
&[],
&paths
));
}
/// The rule the whole back-compat story rests on: an unfilled step field arrives as null, and
/// writing it through would end the conversations a legacy agent's own `memory` holds.
#[test]
fn flow_local_args_override_the_agent_only_when_set() {
fn raw(json: &str) -> Box<RawValue> {
RawValue::from_string(json.to_string()).unwrap()
}
let mut brain = serde_json::Map::new();
brain.insert("system_prompt".to_string(), serde_json::json!("from agent"));
brain.insert("memory".to_string(), serde_json::json!({ "kind": "auto" }));
let local_args = HashMap::from([
(
"user_message".to_string(),
raw("\"ask the flow's question\""),
),
("memory".to_string(), raw("null")),
("enabled_tools".to_string(), raw("null")),
]);
overlay_flow_local_args(&mut brain, &local_args);
assert_eq!(
brain["user_message"],
serde_json::json!("ask the flow's question")
);
assert_eq!(brain["system_prompt"], serde_json::json!("from agent"));
// Unfilled, so the agent's own is what runs.
assert_eq!(brain["memory"], serde_json::json!({ "kind": "auto" }));
assert!(!brain.contains_key("enabled_tools"));
let local_args = HashMap::from([
("memory".to_string(), raw("{\"kind\":\"off\"}")),
("enabled_tools".to_string(), raw("[\"get_user\"]")),
]);
overlay_flow_local_args(&mut brain, &local_args);
assert_eq!(brain["memory"], serde_json::json!({ "kind": "off" }));
assert_eq!(brain["enabled_tools"], serde_json::json!(["get_user"]));
}
#[test]
docs/feature-telemetry.md
+1 -1
View File
@@ -4,7 +4,7 @@
anonymous usage-stats payload. It answers "does anyone use this, and which variant do they pick"
without any identifying data leaving the instance.
It currently carries 48 registered actions across eighteen features (`ai_session`, `ai_chat`,
It currently carries 49 registered actions across eighteen features (`ai_session`, `ai_chat`,
`ai_fix`, `ai_agent`, `ai_agent_eval`, `app_sandbox`, `datatable`, `flow_editor`, `flow_run`,
`flow_step`, `home`, `run_form`, `debugger`, `trigger`, `command_script`, `hub_script`,
`usage_meter`, `sso_groups_claim`). Nearly all of the
frontend/src/lib/components/ModuleTest.svelte
+8 -2
View File
@@ -173,8 +173,14 @@
// `args` is built from the whole AI agent schema whatever the step is, so on a linked step
// it carries every brain key as undefined even though the form renders only the flow-local
// ones (`flowLocalAgentSchema`). Overlaying those would shadow the brain the draft just
// supplied with nothing, so an inlined step takes only the inputs its form actually offers.
const formKeys = draft ? (AGENT_FLOW_LOCAL_KEYS as readonly string[]) : Object.keys(args)
// supplied with nothing, so an inlined step takes only the inputs its form actually offers
// — and of those, only the ones it was given a value for. An unfilled field must inherit
// what the agent carries, the way a deployed run does: an agent saved before `memory`
// became a step input still holds one, and a test that blanked it would answer without the
// history the same step answers with when the flow runs.
const formKeys = draft
? AGENT_FLOW_LOCAL_KEYS.filter((key) => args?.[key] !== undefined)
: Object.keys(args)
// The test form only covers the schema it was given, and for a standalone agent that may be
// the flow-local one (the agent editor shows the brain in its own form, not here). Take the
frontend/src/lib/components/copilot/chat/flow/openFlow.json
+1 -1
View File
File diff suppressed because one or more lines are too long
frontend/src/lib/components/copilot/chat/flow/openFlowZod.gen.ts
Generated
+2 -2
View File
File diff suppressed because one or more lines are too long
frontend/src/lib/components/flows/agentFormFields.ts
+1 -1
View File
@@ -120,7 +120,7 @@ export const AGENT_FIELDS: AgentFieldSpec[] = [
group: 'tools',
label: 'Enabled tools',
tooltip:
'Narrows the tools above to the ones named here, so a run only carries what it needs. Set it to an expression to decide per run. An MCP server named here enables all of its tools.',
'Narrows the tools above to the ones named here, so a run only carries what it needs. Leave it empty for no tools at all, or set it to an expression to decide per run. An MCP server named here enables every tool it exposes; an expression can name one of them on its own, as mcp_<server>_<tool>.',
defaultHint: 'Default: all of them'
},
{
frontend/src/lib/components/flows/agentResourceUtils.ts
+17 -4
View File
@@ -83,6 +83,9 @@ export function flowLocalInputs(
function transformIsSet(transform: InputTransform | undefined): boolean {
if (!transform) return false
const t = transform as any
// Same reading of "set" as `agentFieldIsSet`: an emptied expression is a field being written,
// not one holding a value.
if (t.type === 'javascript') return Boolean(t.expr)
if (t.type !== 'static') return true
return t.value !== undefined && t.value !== null
}
@@ -161,7 +164,7 @@ export function inputTransformsToAgentConfig(
/**
* Reduce the AI agent schema to only the flow-local inputs. Used when a step is linked to a saved
* agent: the brain fields come from the resource, so only user_message/user_attachments stay editable.
* agent: the brain fields come from the resource, so only `AGENT_FLOW_LOCAL_KEYS` stay editable.
*/
export function flowLocalAgentSchema(schema: any): any {
if (!schema?.properties) {
@@ -203,18 +206,28 @@ export function transformValuedBrainKeys(args: Record<string, any> | undefined):
})
}
/** Flatten a saved agent's brain config into human-readable label/value rows for a read-only
* display on a linked step. Only set fields are returned, in the canonical brain-key order. */
/**
* Flatten a saved agent's brain config into human-readable label/value rows for a read-only display
* on a linked step. Only set fields are returned, in the canonical brain-key order.
*
* `memory` is listed after them although it is no longer a brain field, because an agent saved
* while it was one still carries a config the worker honours. Nothing writes one any more, so the
* row only ever appears on such an agent — and where it does, the step's own Memory field would
* otherwise be the only thing on screen saying anything about memory, while reading "off".
*/
export function summarizeAgentBrain(
config: AIAgentConfig | undefined
): { label: string; value: string }[] {
const rows: { label: string; value: string }[] = []
for (const key of AGENT_BRAIN_KEYS) {
for (const key of [...AGENT_BRAIN_KEYS, 'memory']) {
const v = (config as any)?.[key]
if (v === undefined || v === null || v === '') continue
let value: string
if (key === 'provider') {
value = [v.kind, v.model].filter(Boolean).join(' · ') || 'configured'
} else if (key === 'memory') {
// Memory configs are serialized with a `kind` tag (serde tag = "kind").
value = typeof v === 'object' ? (v.kind ?? 'configured') : String(v)
} else if (key === 'output_schema') {
value = 'configured'
} else if (typeof v === 'boolean') {
frontend/src/lib/components/flows/content/AgentResourceBar.svelte
+8 -5
View File
@@ -437,11 +437,14 @@
const local = overridingFlowLocalInputs(inputTransforms)
// `memory` is the step's, but an agent saved back when it was part of the brain still carries
// one that the worker honours while the step holds none. Forking is where that ends, so it
// comes across as the step's own rather than being dropped with the link.
const legacyMemory: Record<string, InputTransform> =
cfg.memory != undefined && !local.memory
? { memory: { type: 'static', value: cfg.memory } as InputTransform }
: {}
// comes across as the step's own rather than being dropped with the link — without the
// `memory_id`, so that saving mints one per step (`cleanInputs`) instead of leaving every
// flow that forked this agent answering from a single shared conversation.
const legacyMemory: Record<string, InputTransform> = {}
if (cfg.memory != undefined && !local.memory) {
const { memory_id: _minted, ...rest } = cfg.memory as Record<string, unknown>
legacyMemory.memory = { type: 'static', value: rest } as InputTransform
}
const forkedInputs = { ...agentConfigToInputTransforms(cfg), ...legacyMemory, ...local }
const forkedTools = cfg.tools ?? []
inputTransforms = forkedInputs
openflow.openapi.yaml
+1 -1
View File
@@ -1115,7 +1115,7 @@ components:
Path of a reusable `ai_agent` resource (hybrid linking). When set, the agent brain
config (provider/model/system prompt/etc.) and tool set are resolved at runtime from
that resource; the module's input_transforms then only carry the flow-local inputs
(user_message/user_attachments).
(user_message/user_attachments/memory/enabled_tools).
tool_inputs:
type: object
description: |