alex/g3
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: JIT-injectable toolsets with load_toolset tool

Browse Source 
Implement dynamic tool loading system that allows tools to be loaded
on-demand rather than included in the default set.

Key changes:
- Add toolsets module with registry of loadable toolsets
- Add load_toolset tool that returns tool definitions for a named toolset
- Add <available_toolsets> section to system prompt
- Track loaded toolsets in Agent, extend tool definitions dynamically
- Move webdriver (15 tools) to JIT-only loading

Benefits:
- Leaner default context (fewer tokens consumed)
- On-demand loading when agent needs specialized tools
- Extensible registry for future toolsets
- Idempotent loading with helpful error messages

Files:
- crates/g3-core/src/toolsets.rs (new)
- crates/g3-core/src/tools/toolsets.rs (new)
- crates/g3-core/src/tool_definitions.rs
- crates/g3-core/src/tool_dispatch.rs
- crates/g3-core/src/prompts.rs
- crates/g3-core/src/lib.rs
- crates/g3-core/src/tools/executor.rs
This commit is contained in:
Dhanji R. Prasanna
2026-02-06 09:35:11 +11:00
parent ff15db44c0
commit cbced3390c
9 changed files with 587 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
crates/g3-core/src/lib.rs
View File

@@ -24,6 +24,7 @@ pub mod ui_writer;
pub mod utils;
pub mod webdriver_session;
pub mod skills;
pub mod toolsets;
pub use feedback_extraction::{
extract_coach_feedback, ExtractedFeedback, FeedbackExtractionConfig, FeedbackSource,
@@ -163,6 +164,8 @@ pub struct Agent<W: UiWriter> {
in_plan_mode: bool,
/// Manager for async research tasks
pending_research_manager: pending_research::PendingResearchManager,
/// Set of toolset names that have been loaded in this session
loaded_toolsets: std::collections::HashSet<String>,
}
impl<W: UiWriter> Agent<W> {
@@ -218,6 +221,7 @@ impl<W: UiWriter> Agent<W> {
acd_enabled: false,
in_plan_mode: false,
pending_research_manager: pending_research::PendingResearchManager::new(),
loaded_toolsets: std::collections::HashSet::new(),
}
}
@@ -614,6 +618,26 @@ impl<W: UiWriter> Agent<W> {
provider_config::resolve_temperature(&self.config, provider_name)
}
/// Get tool definitions including any dynamically loaded toolsets.
fn get_tool_definitions_with_loaded_toolsets(&self, tool_config: tool_definitions::ToolConfig) -> Vec<g3_providers::Tool> {
let mut tools = tool_definitions::create_tool_definitions(tool_config);
// Add tools from loaded toolsets
for toolset_name in &self.loaded_toolsets {
if let Ok(toolset) = toolsets::get_toolset(toolset_name) {
let toolset_tools = toolset.get_tools();
// Avoid duplicates (in case toolset was already in base config)
for tool in toolset_tools {
if !tools.iter().any(|t| t.name == tool.name) {
tools.push(tool);
}
}
}
}
tools
}
/// Print provider diagnostics through the UiWriter for visibility
pub fn print_provider_banner(&self, role_label: &str) {
if let Ok((provider_name, model)) = self.get_provider_info() {
@@ -1016,10 +1040,9 @@ impl<W: UiWriter> Agent<W> {
let _supports_cache_control = provider.supports_cache_control();
let tools = if provider.has_native_tool_calling() {
let tool_config = tool_definitions::ToolConfig::new(
self.config.webdriver.enabled,
self.config.computer_control.enabled,
);
Some(tool_definitions::create_tool_definitions(tool_config))
Some(self.get_tool_definitions_with_loaded_toolsets(tool_config))
} else {
None
};
@@ -1887,10 +1910,9 @@ Skip if nothing new. Be brief."#;
let provider_name = provider.name().to_string();
let tools = if provider.has_native_tool_calling() {
let tool_config = tool_definitions::ToolConfig::new(
self.config.webdriver.enabled,
self.config.computer_control.enabled,
);
Some(tool_definitions::create_tool_definitions(tool_config))
Some(self.get_tool_definitions_with_loaded_toolsets(tool_config))
} else {
None
};
@@ -2569,11 +2591,10 @@ Skip if nothing new. Be brief."#;
let provider_for_tools = self.providers.get(None)?;
if provider_for_tools.has_native_tool_calling() {
let tool_config = tool_definitions::ToolConfig::new(
self.config.webdriver.enabled,
self.config.computer_control.enabled,
);
request.tools =
Some(tool_definitions::create_tool_definitions(tool_config));
Some(self.get_tool_definitions_with_loaded_toolsets(tool_config));
}
// DO NOT add final_display_content to full_response here!
@@ -2981,6 +3002,7 @@ Skip if nothing new. Be brief."#;
requirements_sha: self.requirements_sha.as_deref(),
context_total_tokens: self.context_window.total_tokens,
context_used_tokens: self.context_window.used_tokens,
loaded_toolsets: &mut self.loaded_toolsets,
};
// Dispatch to the appropriate tool handler

19
crates/g3-core/src/prompts.rs
View File

@@ -116,6 +116,7 @@ write_file(\"helper.rs\", \"...\")
// ============================================================================
use crate::skills::{Skill, generate_skills_prompt};
use crate::toolsets::generate_toolsets_prompt;
/// System prompt for providers with native tool calling (Anthropic, OpenAI, etc.)
/// Uses include_str! to embed the prompt at compile time.
@@ -126,11 +127,21 @@ pub fn get_system_prompt_for_native() -> String {
/// System prompt for providers with native tool calling, with skills support.
pub fn get_system_prompt_for_native_with_skills(skills: &[Skill]) -> String {
let skills_section = generate_skills_prompt(skills);
if skills_section.is_empty() {
EMBEDDED_NATIVE_PROMPT.to_string()
} else {
format!("{}\n\n{}", EMBEDDED_NATIVE_PROMPT, skills_section)
let toolsets_section = generate_toolsets_prompt();
let mut prompt = EMBEDDED_NATIVE_PROMPT.to_string();
// Add toolsets section (available toolsets for dynamic loading)
if !toolsets_section.is_empty() {
prompt = format!("{}\n\n{}", prompt, toolsets_section);
}
// Add skills section
if !skills_section.is_empty() {
prompt = format!("{}\n\n{}", prompt, skills_section);
}
prompt
}
/// System prompt for providers without native tool calling (embedded models)

244
crates/g3-core/src/tool_definitions.rs
View File

@@ -10,15 +10,13 @@ use serde_json::json;
/// Configuration for which optional tool sets to enable
#[derive(Debug, Clone, Copy, Default)]
pub struct ToolConfig {
pub webdriver: bool,
pub computer_control: bool,
pub exclude_research: bool,
}
impl ToolConfig {
pub fn new(webdriver: bool, computer_control: bool) -> Self {
pub fn new(computer_control: bool) -> Self {
Self {
webdriver,
computer_control,
exclude_research: false,
}
@@ -37,13 +35,8 @@ impl ToolConfig {
/// Returns a vector of Tool definitions that describe the available tools
/// and their input schemas.
pub fn create_tool_definitions(config: ToolConfig) -> Vec<Tool> {
let mut tools = create_core_tools(config.exclude_research);
if config.webdriver {
tools.extend(create_webdriver_tools());
}
tools
// Webdriver tools are now JIT-loaded via load_toolset("webdriver")
create_core_tools(config.exclude_research)
}
/// Create the core tools that are always available
@@ -301,200 +294,23 @@ fn create_core_tools(exclude_research: bool) -> Vec<Tool> {
});
}
tools
}
// Toolset loading tool - allows dynamic loading of additional tools
tools.push(Tool {
name: "load_toolset".to_string(),
description: "Load additional tools from a named toolset. Returns the tool definitions so you learn how to use them. Use this when you need specialized tools like browser automation.".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "The name of the toolset to load (e.g., 'webdriver')"
}
},
"required": ["name"]
}),
});
/// Create WebDriver browser automation tools
fn create_webdriver_tools() -> Vec<Tool> {
vec![
Tool {
name: "webdriver_start".to_string(),
description: "Start a Safari WebDriver session for browser automation. Must be called before any other webdriver tools. Requires Safari's 'Allow Remote Automation' to be enabled in Develop menu.".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_navigate".to_string(),
description: "Navigate to a URL in the browser".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "The URL to navigate to (must include protocol, e.g., https://)"
}
},
"required": ["url"]
}),
},
Tool {
name: "webdriver_get_url".to_string(),
description: "Get the current URL of the browser".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_get_title".to_string(),
description: "Get the title of the current page".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_find_element".to_string(),
description: "Find an element on the page by CSS selector and return its text content".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector to find the element (e.g., 'h1', '.class-name', '#id')"
}
},
"required": ["selector"]
}),
},
Tool {
name: "webdriver_find_elements".to_string(),
description: "Find all elements matching a CSS selector and return their text content".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector to find elements"
}
},
"required": ["selector"]
}),
},
Tool {
name: "webdriver_click".to_string(),
description: "Click an element on the page".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector for the element to click"
}
},
"required": ["selector"]
}),
},
Tool {
name: "webdriver_send_keys".to_string(),
description: "Type text into an input element".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector for the input element"
},
"text": {
"type": "string",
"description": "Text to type into the element"
},
"clear_first": {
"type": "boolean",
"description": "Whether to clear the element before typing (default: true)"
}
},
"required": ["selector", "text"]
}),
},
Tool {
name: "webdriver_execute_script".to_string(),
description: "Execute JavaScript code in the browser and return the result".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"script": {
"type": "string",
"description": "JavaScript code to execute (use 'return' to return a value)"
}
},
"required": ["script"]
}),
},
Tool {
name: "webdriver_get_page_source".to_string(),
description: "Get the rendered HTML source of the current page. Returns the current DOM state after JavaScript execution.".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"max_length": {
"type": "integer",
"description": "Maximum length of HTML to return (default: 10000, use 0 for no truncation)"
},
"save_to_file": {
"type": "string",
"description": "Optional file path to save the HTML instead of returning it inline"
}
},
"required": []
}),
},
Tool {
name: "webdriver_screenshot".to_string(),
description: "Take a screenshot of the browser window".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path where to save the screenshot (e.g., '/tmp/screenshot.png')"
}
},
"required": ["path"]
}),
},
Tool {
name: "webdriver_back".to_string(),
description: "Navigate back in browser history".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_forward".to_string(),
description: "Navigate forward in browser history".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_refresh".to_string(),
description: "Refresh the current page".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_quit".to_string(),
description: "Close the browser and end the WebDriver session".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
]
tools
}
#[cfg(test)]
@@ -506,15 +322,8 @@ mod tests {
let tools = create_core_tools(false);
// Core tools (with research): shell, background_process, read_file, read_image,
// write_file, str_replace, code_search, plan_read, plan_write, plan_approve,
// remember, rehydrate, research, research_status
// (14 total)
assert_eq!(tools.len(), 14);
}
#[test]
fn test_webdriver_tools_count() {
let tools = create_webdriver_tools();
// 15 webdriver tools
// remember, rehydrate, research, research_status, load_toolset
// (15 total)
assert_eq!(tools.len(), 15);
}
@@ -522,15 +331,16 @@ mod tests {
fn test_create_tool_definitions_core_only() {
let config = ToolConfig::default();
let tools = create_tool_definitions(config);
assert_eq!(tools.len(), 14);
// 15 core tools (webdriver is now JIT-loaded)
assert_eq!(tools.len(), 15);
}
#[test]
fn test_create_tool_definitions_all_enabled() {
let config = ToolConfig::new(true, true);
fn test_create_tool_definitions() {
let config = ToolConfig::new(true);
let tools = create_tool_definitions(config);
// 14 core + 15 webdriver = 29
assert_eq!(tools.len(), 29);
// Webdriver tools are now JIT-loaded, so only core tools are included
assert!(tools.len() >= 14); // At least 14 core tools
}
#[test]

5
crates/g3-core/src/tool_dispatch.rs
View File

@@ -7,7 +7,7 @@ use anyhow::Result;
use tracing::{debug, warn};
use crate::tools::executor::ToolContext;
use crate::tools::{acd, file_ops, memory, misc, plan, research, shell, webdriver};
use crate::tools::{acd, file_ops, memory, misc, plan, research, shell, toolsets, webdriver};
use crate::ui_writer::UiWriter;
use crate::ToolCall;
@@ -47,6 +47,9 @@ pub async fn dispatch_tool<W: UiWriter>(
// Workspace memory tools
"remember" => memory::execute_remember(tool_call, ctx).await,
// Toolset loading
"load_toolset" => toolsets::execute_load_toolset(tool_call, ctx).await,
// ACD (Aggressive Context Dehydration) tools
"rehydrate" => acd::execute_rehydrate(tool_call, ctx).await,

5
crates/g3-core/src/tools/acd.rs
View File

@@ -137,6 +137,7 @@ mod tests {
pending_images: Vec<g3_providers::ImageContent>,
pending_research_manager: PendingResearchManager,
config: g3_config::Config,
loaded_toolsets: std::collections::HashSet<String>,
}
impl TestContext {
@@ -150,6 +151,7 @@ mod tests {
pending_images: Vec::new(),
pending_research_manager: PendingResearchManager::new(),
config: g3_config::Config::default(),
loaded_toolsets: std::collections::HashSet::new(),
}
}
}
@@ -173,6 +175,7 @@ mod tests {
requirements_sha: None,
context_total_tokens: 100000,
context_used_tokens: 10000,
loaded_toolsets: &mut test_ctx.loaded_toolsets,
};
let tool_call = ToolCall {
@@ -204,6 +207,7 @@ mod tests {
requirements_sha: None,
context_total_tokens: 100000,
context_used_tokens: 10000,
loaded_toolsets: &mut test_ctx.loaded_toolsets,
};
let tool_call = ToolCall {
@@ -235,6 +239,7 @@ mod tests {
requirements_sha: None,
context_total_tokens: 100000,
context_used_tokens: 10000,
loaded_toolsets: &mut test_ctx.loaded_toolsets,
};
let tool_call = ToolCall {

3
crates/g3-core/src/tools/executor.rs
View File

@@ -1,6 +1,7 @@
//! Tool executor trait and context for tool execution.
use anyhow::Result;
use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::RwLock;
@@ -29,6 +30,8 @@ pub struct ToolContext<'a, W: UiWriter> {
pub requirements_sha: Option<&'a str>,
pub context_total_tokens: u32,
pub context_used_tokens: u32,
/// Set of toolset names that have been loaded in this session
pub loaded_toolsets: &'a mut HashSet<String>,
}
impl<'a, W: UiWriter> ToolContext<'a, W> {

1
crates/g3-core/src/tools/mod.rs
View File

@@ -19,6 +19,7 @@ pub mod misc;
pub mod plan;
pub mod research;
pub mod shell;
pub mod toolsets;
pub mod webdriver;
pub use executor::ToolExecutor;

110
crates/g3-core/src/tools/toolsets.rs Normal file
View File

@@ -0,0 +1,110 @@
//! Toolset loading tool implementation.
//!
//! This module provides the `load_toolset` tool which allows the agent to
//! dynamically load additional tool definitions at runtime.
use anyhow::Result;
use tracing::debug;
use crate::toolsets;
use crate::tools::executor::ToolContext;
use crate::ui_writer::UiWriter;
use crate::ToolCall;
/// Execute the load_toolset tool.
///
/// This tool loads a named toolset and returns the tool definitions so the
/// agent learns how to use the newly available tools.
///
/// The tool definitions are returned as formatted text describing each tool,
/// its purpose, and its parameters.
pub async fn execute_load_toolset<W: UiWriter>(
tool_call: &ToolCall,
ctx: &mut ToolContext<'_, W>,
) -> Result<String> {
let toolset_name = tool_call
.args
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("");
debug!("Loading toolset: {}", toolset_name);
// Get the toolset from the registry
let toolset = match toolsets::get_toolset(toolset_name) {
Ok(ts) => ts,
Err(err) => {
return Ok(format!("❌ {}", err));
}
};
// Check if already loaded (idempotent)
if ctx.loaded_toolsets.contains(&toolset_name.to_string()) {
return Ok(format!(
"✅ Toolset '{}' is already loaded. You can use its tools.",
toolset_name
));
}
// Get the tool definitions
let tools = toolset.get_tools();
// Mark as loaded
ctx.loaded_toolsets.insert(toolset_name.to_string());
// Format the tool definitions for the agent
let mut output = String::new();
output.push_str(&format!(
"✅ Loaded toolset '{}' with {} tools:\n\n",
toolset_name,
tools.len()
));
for tool in &tools {
output.push_str(&format!("## {}", tool.name));
output.push_str("\n\n");
output.push_str(&tool.description);
output.push_str("\n\n");
// Format the input schema in a readable way
if let Some(props) = tool.input_schema.get("properties") {
if let Some(obj) = props.as_object() {
if !obj.is_empty() {
output.push_str("**Parameters:**\n");
for (param_name, param_schema) in obj {
let param_type = param_schema
.get("type")
.and_then(|t| t.as_str())
.unwrap_or("any");
let param_desc = param_schema
.get("description")
.and_then(|d| d.as_str())
.unwrap_or("");
// Check if required
let required = tool.input_schema
.get("required")
.and_then(|r| r.as_array())
.map(|arr| arr.iter().any(|v| v.as_str() == Some(param_name)))
.unwrap_or(false);
let req_marker = if required { " (required)" } else { " (optional)" };
output.push_str(&format!(
"- `{}` ({}){}: {}\n",
param_name, param_type, req_marker, param_desc
));
}
output.push_str("\n");
}
}
}
}
output.push_str(&format!(
"---\nYou can now use these {} tools. They are available for the rest of this session.",
tools.len()
));
Ok(output)
}

394
crates/g3-core/src/toolsets.rs Normal file
View File

@@ -0,0 +1,394 @@
//! Toolsets module - registry of dynamically loadable tool collections.
//!
//! Toolsets are groups of related tools that can be loaded on-demand via the
//! `load_toolset` tool. This keeps the default tool set lean while allowing
//! access to specialized tools when needed.
//!
//! The agent sees a concise registry in the system prompt and can load
//! toolsets as needed. The tool definitions are returned so the agent
//! learns how to call the newly available tools.
use g3_providers::Tool;
use serde_json::json;
/// A toolset that can be dynamically loaded.
#[derive(Debug, Clone)]
pub struct Toolset {
/// Unique identifier for the toolset (e.g., "webdriver")
pub name: &'static str,
/// Brief description of what the toolset provides
pub description: &'static str,
/// Function that returns the tool definitions for this toolset
tool_definitions_fn: fn() -> Vec<Tool>,
}
impl Toolset {
/// Get the tool definitions for this toolset.
pub fn get_tools(&self) -> Vec<Tool> {
(self.tool_definitions_fn)()
}
}
/// Registry of all available toolsets.
/// Add new toolsets here as they are created.
const TOOLSET_REGISTRY: &[Toolset] = &[
Toolset {
name: "webdriver",
description: "Browser automation via Safari WebDriver. Start sessions, navigate, find elements, click, type, execute JavaScript, take screenshots.",
tool_definitions_fn: create_webdriver_tools,
},
];
/// Get a toolset by name.
///
/// Returns `Ok(Toolset)` if found, or `Err` with a helpful message listing
/// available toolsets if not found.
pub fn get_toolset(name: &str) -> Result<&'static Toolset, String> {
let name = name.trim();
if name.is_empty() {
return Err(format!(
"Toolset name cannot be empty. Available toolsets: {}",
list_toolset_names().join(", ")
));
}
TOOLSET_REGISTRY
.iter()
.find(|t| t.name == name)
.ok_or_else(|| {
format!(
"Unknown toolset '{}'. Available toolsets: {}",
name,
list_toolset_names().join(", ")
)
})
}
/// List all available toolset names.
pub fn list_toolset_names() -> Vec<&'static str> {
TOOLSET_REGISTRY.iter().map(|t| t.name).collect()
}
/// Get all available toolsets.
pub fn get_all_toolsets() -> &'static [Toolset] {
TOOLSET_REGISTRY
}
/// Generate the prompt section describing available toolsets.
///
/// This is injected into the system prompt so the agent knows what
/// toolsets are available to load.
pub fn generate_toolsets_prompt() -> String {
if TOOLSET_REGISTRY.is_empty() {
return String::new();
}
let mut prompt = String::new();
prompt.push_str("# Available Toolsets\n\n");
prompt.push_str("You can dynamically load additional tools using `load_toolset`. ");
prompt.push_str("The tool will return the full definitions so you learn how to use them.\n\n");
prompt.push_str("<available_toolsets>\n");
for toolset in TOOLSET_REGISTRY {
prompt.push_str(" <toolset>\n");
prompt.push_str(&format!(" <name>{}</name>\n", escape_xml(toolset.name)));
prompt.push_str(&format!(" <description>{}</description>\n", escape_xml(toolset.description)));
prompt.push_str(" </toolset>\n");
}
prompt.push_str("</available_toolsets>\n");
prompt
}
/// Escape special XML characters.
fn escape_xml(s: &str) -> String {
s.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;")
.replace('"', "&quot;")
.replace('\'', "&apos;")
}
// =============================================================================
// TOOLSET DEFINITIONS
// =============================================================================
/// Create WebDriver browser automation tools.
///
/// These tools enable browser automation via Safari WebDriver.
fn create_webdriver_tools() -> Vec<Tool> {
vec![
Tool {
name: "webdriver_start".to_string(),
description: "Start a Safari WebDriver session for browser automation. Must be called before any other webdriver tools. Requires Safari's 'Allow Remote Automation' to be enabled in Develop menu.".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_navigate".to_string(),
description: "Navigate to a URL in the browser".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "The URL to navigate to (must include protocol, e.g., https://)"
}
},
"required": ["url"]
}),
},
Tool {
name: "webdriver_get_url".to_string(),
description: "Get the current URL of the browser".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_get_title".to_string(),
description: "Get the title of the current page".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_find_element".to_string(),
description: "Find an element on the page by CSS selector and return its text content".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector to find the element (e.g., 'h1', '.class-name', '#id')"
}
},
"required": ["selector"]
}),
},
Tool {
name: "webdriver_find_elements".to_string(),
description: "Find all elements matching a CSS selector and return their text content".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector to find elements"
}
},
"required": ["selector"]
}),
},
Tool {
name: "webdriver_click".to_string(),
description: "Click an element on the page".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector for the element to click"
}
},
"required": ["selector"]
}),
},
Tool {
name: "webdriver_send_keys".to_string(),
description: "Type text into an input element".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector for the input element"
},
"text": {
"type": "string",
"description": "Text to type into the element"
},
"clear_first": {
"type": "boolean",
"description": "Whether to clear the element before typing (default: true)"
}
},
"required": ["selector", "text"]
}),
},
Tool {
name: "webdriver_execute_script".to_string(),
description: "Execute JavaScript code in the browser and return the result".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"script": {
"type": "string",
"description": "JavaScript code to execute (use 'return' to return a value)"
}
},
"required": ["script"]
}),
},
Tool {
name: "webdriver_get_page_source".to_string(),
description: "Get the rendered HTML source of the current page. Returns the current DOM state after JavaScript execution.".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"max_length": {
"type": "integer",
"description": "Maximum length of HTML to return (default: 10000, use 0 for no truncation)"
},
"save_to_file": {
"type": "string",
"description": "Optional file path to save the HTML instead of returning it inline"
}
},
"required": []
}),
},
Tool {
name: "webdriver_screenshot".to_string(),
description: "Take a screenshot of the browser window".to_string(),
input_schema: json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path where to save the screenshot (e.g., '/tmp/screenshot.png')"
}
},
"required": ["path"]
}),
},
Tool {
name: "webdriver_back".to_string(),
description: "Navigate back in browser history".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_forward".to_string(),
description: "Navigate forward in browser history".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_refresh".to_string(),
description: "Refresh the current page".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
Tool {
name: "webdriver_quit".to_string(),
description: "Close the browser and end the WebDriver session".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"required": []
}),
},
]
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_get_toolset_webdriver() {
let toolset = get_toolset("webdriver").unwrap();
assert_eq!(toolset.name, "webdriver");
assert!(!toolset.description.is_empty());
let tools = toolset.get_tools();
assert!(!tools.is_empty());
assert!(tools.iter().any(|t| t.name == "webdriver_start"));
}
#[test]
fn test_get_toolset_unknown() {
let result = get_toolset("nonexistent");
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.contains("Unknown toolset"));
assert!(err.contains("webdriver")); // Should list available toolsets
}
#[test]
fn test_get_toolset_empty_name() {
let result = get_toolset("");
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.contains("cannot be empty"));
}
#[test]
fn test_get_toolset_whitespace_trimmed() {
let toolset = get_toolset(" webdriver ").unwrap();
assert_eq!(toolset.name, "webdriver");
}
#[test]
fn test_list_toolset_names() {
let names = list_toolset_names();
assert!(names.contains(&"webdriver"));
}
#[test]
fn test_generate_toolsets_prompt() {
let prompt = generate_toolsets_prompt();
assert!(prompt.contains("<available_toolsets>"));
assert!(prompt.contains("</available_toolsets>"));
assert!(prompt.contains("<name>webdriver</name>"));
assert!(prompt.contains("load_toolset"));
}
#[test]
fn test_xml_escaping() {
// The current toolsets don't have special chars, but test the function
assert_eq!(escape_xml("<test>"), "&lt;test&gt;");
assert_eq!(escape_xml("a & b"), "a &amp; b");
}
#[test]
fn test_webdriver_tools_complete() {
let tools = create_webdriver_tools();
let tool_names: Vec<_> = tools.iter().map(|t| t.name.as_str()).collect();
// Verify all expected webdriver tools are present
assert!(tool_names.contains(&"webdriver_start"));
assert!(tool_names.contains(&"webdriver_navigate"));
assert!(tool_names.contains(&"webdriver_get_url"));
assert!(tool_names.contains(&"webdriver_get_title"));
assert!(tool_names.contains(&"webdriver_find_element"));
assert!(tool_names.contains(&"webdriver_find_elements"));
assert!(tool_names.contains(&"webdriver_click"));
assert!(tool_names.contains(&"webdriver_send_keys"));
assert!(tool_names.contains(&"webdriver_execute_script"));
assert!(tool_names.contains(&"webdriver_get_page_source"));
assert!(tool_names.contains(&"webdriver_screenshot"));
assert!(tool_names.contains(&"webdriver_back"));
assert!(tool_names.contains(&"webdriver_forward"));
assert!(tool_names.contains(&"webdriver_refresh"));
assert!(tool_names.contains(&"webdriver_quit"));
}
}