Add TKMind platform extensions, H5/MindSpace stack, and deployment tooling.
Deploy Documentation / deploy (push) Has been cancelled
Canary / Prepare Version (push) Has been cancelled
Canary / build-cli (push) Has been cancelled
Canary / Upload Install Script (push) Has been cancelled
Canary / bundle-desktop (push) Has been cancelled
Canary / bundle-desktop-intel (push) Has been cancelled
Canary / bundle-desktop-linux (push) Has been cancelled
Canary / bundle-desktop-windows (push) Has been cancelled
Canary / bundle-desktop-windows-cuda (push) Has been cancelled
Canary / Release (push) Has been cancelled
Unused Dependencies / machete (push) Has been cancelled
CI / changes (push) Has been cancelled
CI / Check Rust Code Format (push) Has been cancelled
CI / Build and Test Rust Project (push) Has been cancelled
CI / Build Rust Project on Windows (push) Has been cancelled
CI / Check MSRV (push) Has been cancelled
CI / Lint Rust Code (push) Has been cancelled
CI / Check Generated Schemas are Up-to-Date (push) Has been cancelled
CI / Test and Lint Electron Desktop App (push) Has been cancelled
CI / H5 Plaza Tests and Build (push) Has been cancelled
Live Provider Tests / check-fork (push) Has been cancelled
Live Provider Tests / changes (push) Has been cancelled
Live Provider Tests / Build Binary (push) Has been cancelled
Live Provider Tests / Smoke Tests (push) Has been cancelled
Live Provider Tests / Smoke Tests (Code Execution) (push) Has been cancelled
Live Provider Tests / Compaction Tests (push) Has been cancelled
Live Provider Tests / goose server HTTP integration tests (push) Has been cancelled
Publish Ask AI Bot Docker Image / docker (push) Has been cancelled
Publish Docker Image / docker (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled

Fork goose with custom MCP widgets, platform extensions (aider, git, web, search),
MindSpace H5 backend/frontend, Plaza/Ops UIs, and deploy scripts for tkmind.cn.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
john
2026-06-14 21:30:20 +08:00
parent e5fd568e01
commit 4e21ca937a
359 changed files with 70658 additions and 56 deletions
@@ -0,0 +1,308 @@
use std::path::{Path, PathBuf};
use std::time::Duration;
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::platform_extensions::developer::workspace_path::resolve_path_in_workspace;
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use indoc::indoc;
use rmcp::model::{
CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult,
ServerCapabilities, Tool, ToolAnnotations,
};
use schemars::{schema_for, JsonSchema};
use serde::Deserialize;
use tokio::io::AsyncReadExt;
use tokio::process::Command;
use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "aider";
const DEFAULT_TIMEOUT_SECS: u64 = 600;
#[derive(Debug, Deserialize, JsonSchema)]
struct AiderCodeParams {
/// Coding task for Aider: what to change, fix, or implement.
task: String,
/// Project directory. Defaults to the session working directory.
project_dir: Option<String>,
}
pub struct AiderClient {
info: InitializeResult,
}
impl AiderClient {
pub fn new(_context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(
Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string())
.with_title("Aider"),
)
.with_instructions(indoc! {"
Delegate multi-file coding work to Aider instead of using developer write/edit.
Prefer this tool for:
- Bug fixes and refactors across multiple files
- New feature implementation
- Test-failure fixes that need code changes
After Aider finishes, verify with developer shell (tests, git status).
For one-line or single-file tweaks, developer edit/write is fine.
"});
Ok(Self { info })
}
fn schema<T: JsonSchema>() -> JsonObject {
serde_json::to_value(schema_for!(T))
.expect("schema serialization should succeed")
.as_object()
.expect("schema should serialize to an object")
.clone()
}
fn get_tools() -> Vec<Tool> {
vec![Tool::new(
"code".to_string(),
"Run Aider to implement a coding task in the project directory. \
Aider edits files directly and uses git. Returns stdout/stderr and exit code."
.to_string(),
Self::schema::<AiderCodeParams>(),
)
.annotate(ToolAnnotations::from_raw(
Some("Aider Code".to_string()),
Some(false),
Some(true),
Some(false),
Some(true),
))]
}
fn resolve_aider_bin() -> Result<PathBuf, String> {
for key in ["GOOSE_AIDER_BIN", "AIDER_BIN"] {
if let Ok(path) = std::env::var(key) {
let candidate = PathBuf::from(&path);
if candidate.is_file() {
return Ok(candidate);
}
return Err(format!("{key} is set but not a file: {path}"));
}
}
for candidate in [
"/root/aider/.venv/bin/aider",
"/Users/john/PycharmProjects/aider/.venv/bin/aider",
] {
let path = PathBuf::from(candidate);
if path.is_file() {
return Ok(path);
}
}
Err(
"Aider binary not found. Set GOOSE_AIDER_BIN or install Aider under /root/aider."
.to_string(),
)
}
fn resolve_coding_router() -> Option<PathBuf> {
std::env::var("GOOSE_CODING_ROUTER")
.ok()
.map(PathBuf::from)
.filter(|path| path.is_file())
}
fn timeout_secs() -> u64 {
std::env::var("GOOSE_AIDER_TIMEOUT_SECS")
.ok()
.and_then(|value| value.parse().ok())
.filter(|secs| *secs > 0)
.unwrap_or(DEFAULT_TIMEOUT_SECS)
}
async fn run_aider(
task: &str,
project_dir: &Path,
cancel_token: CancellationToken,
) -> CallToolResult {
if task.trim().is_empty() {
return CallToolResult::error(vec![Content::text("task cannot be empty")]);
}
if !project_dir.is_dir() {
return CallToolResult::error(vec![Content::text(format!(
"project directory does not exist: {}",
project_dir.display()
))]);
}
let timeout_secs = Self::timeout_secs();
let mut command = if let Some(router) = Self::resolve_coding_router() {
let mut cmd = Command::new(router);
cmd.arg("aider")
.arg(project_dir)
.arg(task)
.current_dir(project_dir);
cmd
} else {
let aider_bin = match Self::resolve_aider_bin() {
Ok(path) => path,
Err(error) => return CallToolResult::error(vec![Content::text(error)]),
};
let mut cmd = Command::new(aider_bin);
cmd.arg("--yes-always")
.arg("--message")
.arg(task)
.current_dir(project_dir);
cmd
};
command.stdout(std::process::Stdio::piped());
command.stderr(std::process::Stdio::piped());
let mut child = match command.spawn() {
Ok(child) => child,
Err(error) => {
return CallToolResult::error(vec![Content::text(format!(
"Failed to start Aider: {error}"
))]);
}
};
let stdout = child.stdout.take();
let stderr = child.stderr.take();
let wait_result = tokio::select! {
() = cancel_token.cancelled() => {
let _ = child.kill().await;
return CallToolResult::error(vec![Content::text("Aider run cancelled")]);
}
result = tokio::time::timeout(Duration::from_secs(timeout_secs), child.wait()) => result,
};
let exit_code = match wait_result {
Ok(Ok(status)) => status.code(),
Ok(Err(error)) => {
return CallToolResult::error(vec![Content::text(format!(
"Failed waiting on Aider: {error}"
))]);
}
Err(_) => {
return CallToolResult::error(vec![Content::text(format!(
"Aider timed out after {timeout_secs}s"
))]);
}
};
let mut stdout_text = String::new();
let mut stderr_text = String::new();
if let Some(mut stdout) = stdout {
let _ = stdout.read_to_string(&mut stdout_text).await;
}
if let Some(mut stderr) = stderr {
let _ = stderr.read_to_string(&mut stderr_text).await;
}
let body = format!(
"project_dir: {}\nexit_code: {}\n\n--- stdout ---\n{}\n\n--- stderr ---\n{}",
project_dir.display(),
exit_code
.map(|code| code.to_string())
.unwrap_or_else(|| "null".to_string()),
stdout_text.trim_end(),
stderr_text.trim_end(),
);
if exit_code == Some(0) {
CallToolResult::success(vec![Content::text(body)])
} else {
CallToolResult::error(vec![Content::text(body)])
}
}
}
#[async_trait]
impl McpClientTrait for AiderClient {
async fn list_tools(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Ok(ListToolsResult {
tools: Self::get_tools(),
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
if name != "code" {
return Ok(CallToolResult::error(vec![Content::text(format!(
"Unknown tool: {name}"
))]));
}
let Some(value) = arguments.map(serde_json::Value::Object) else {
return Ok(CallToolResult::error(vec![Content::text(
"Missing arguments",
)]));
};
let params: AiderCodeParams = match serde_json::from_value(value) {
Ok(params) => params,
Err(error) => {
return Ok(CallToolResult::error(vec![Content::text(format!(
"Failed to parse arguments: {error}"
))]));
}
};
let project_dir = match params.project_dir.as_deref() {
Some(dir) => match resolve_path_in_workspace(dir, ctx.working_dir.as_deref()) {
Ok(path) => path,
Err(error) => {
return Ok(CallToolResult::error(vec![Content::text(error)]));
}
},
None => match ctx.working_dir.clone() {
Some(path) => path,
None => {
return Ok(CallToolResult::error(vec![Content::text(
"working_dir is required",
)]));
}
},
};
Ok(Self::run_aider(&params.task, &project_dir, cancellation_token).await)
}
fn get_info(&self) -> Option<&InitializeResult> {
Some(&self.info)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn aider_exposes_code_tool() {
let names: Vec<String> = AiderClient::get_tools()
.into_iter()
.map(|tool| tool.name.to_string())
.collect();
assert_eq!(names, vec!["code"]);
}
}
@@ -2,6 +2,8 @@ use std::fs;
use std::path::{Path, PathBuf};
use rmcp::model::{CallToolResult, Content};
use super::workspace_path::resolve_path_in_workspace;
use schemars::JsonSchema;
use serde::Deserialize;
@@ -43,7 +45,12 @@ impl EditTools {
params: FileReadParams,
working_dir: Option<&Path>,
) -> CallToolResult {
let path = resolve_path(&params.path, working_dir);
let path = match resolve_path_in_workspace(&params.path, working_dir) {
Ok(path) => path,
Err(error) => {
return CallToolResult::error(vec![Content::text(error).with_priority(0.0)]);
}
};
match fs::read_to_string(&path) {
Ok(content) => {
@@ -67,7 +74,12 @@ impl EditTools {
params: FileWriteParams,
working_dir: Option<&Path>,
) -> CallToolResult {
let path = resolve_path(&params.path, working_dir);
let path = match resolve_path_in_workspace(&params.path, working_dir) {
Ok(path) => path,
Err(error) => {
return CallToolResult::error(vec![Content::text(error).with_priority(0.0)]);
}
};
if let Some(parent) = path.parent() {
if !parent.as_os_str().is_empty() && !parent.exists() {
@@ -111,7 +123,12 @@ impl EditTools {
params: FileEditParams,
working_dir: Option<&Path>,
) -> CallToolResult {
let path = resolve_path(&params.path, working_dir);
let path = match resolve_path_in_workspace(&params.path, working_dir) {
Ok(path) => path,
Err(error) => {
return CallToolResult::error(vec![Content::text(error).with_priority(0.0)]);
}
};
let content = match fs::read_to_string(&path) {
Ok(c) => c,
@@ -9,7 +9,7 @@ use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::json;
use super::edit::resolve_path;
use super::workspace_path::resolve_path_in_workspace;
const MAX_IMAGE_BYTES: u64 = 20 * 1024 * 1024;
@@ -171,10 +171,16 @@ async fn load_image_bytes(source: &str, working_dir: Option<&Path>) -> Result<Ve
.map_err(|_| "invalid file URL".to_string())?;
load_file_bytes(path)
}
_ => load_file_bytes(resolve_path(source, working_dir)),
_ => match resolve_path_in_workspace(source, working_dir) {
Ok(path) => load_file_bytes(path),
Err(error) => Err(error),
},
}
} else {
load_file_bytes(resolve_path(source, working_dir))
match resolve_path_in_workspace(source, working_dir) {
Ok(path) => load_file_bytes(path),
Err(error) => Err(error),
}
}
}
@@ -2,6 +2,7 @@ pub mod edit;
pub mod image;
pub mod shell;
pub mod tree;
pub mod workspace_path;
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
@@ -46,6 +47,12 @@ fn developer_instructions() -> &'static str {
and file sizes. When you need to search, prefer findstr or Select-String (via shell).
Then use type or Get-Content to gather the context you need, always reading before
editing. Use write and edit to efficiently make changes. Test and verify as appropriate.
When a session working directory is set, every tree/shell/write/edit/read_image path must
stay inside that directory. Default file search starts at `.` (the working directory root).
Treat user-named areas (for example `oa/`) as subdirectories under the working directory.
Never search parent directories, MindSpace root, sibling user folders, or paths outside the
session scope. Do not tell the user you will search outside their workspace.
"}
} else {
indoc! {"
@@ -61,6 +68,12 @@ fn developer_instructions() -> &'static str {
content. Then use cat or sed to gather the context you need, always reading before editing.
Use write and edit to efficiently make changes. Test and verify as appropriate.
When a session working directory is set, every tree/shell/write/edit/read_image path must
stay inside that directory. Default file search starts at `.` (the working directory root).
Treat user-named areas (for example `oa/`) as subdirectories under the working directory.
Never search parent directories, MindSpace root, sibling user folders, the project root, or
the host home directory. Do not tell the user you will search outside their workspace.
When running Python scripts or commands, always use `python3` instead of `python`.
"}
}
@@ -19,6 +19,7 @@ use tokio::sync::OnceCell;
use tokio::task::JoinHandle;
use tokio_stream::{wrappers::SplitStream, StreamExt};
use super::workspace_path::ensure_shell_command_within_workspace;
use crate::subprocess::SubprocessExt;
/// Check if the current process is running inside a Flatpak sandbox.
@@ -355,6 +356,12 @@ impl ShellTool {
return Self::error_result("Command cannot be empty.", None);
}
if let Some(base) = working_dir {
if let Err(error) = ensure_shell_command_within_workspace(&params.command, base) {
return Self::error_result(&error, None);
}
}
#[cfg(not(windows))]
let login_path = self.login_path.get().await;
#[cfg(not(windows))]
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Component, Path, PathBuf};
use super::workspace_path::ensure_within_workspace;
use ignore::WalkBuilder;
use rmcp::model::{CallToolResult, Content};
use schemars::JsonSchema;
@@ -41,6 +42,11 @@ impl TreeTool {
.unwrap_or_else(|| PathBuf::from("."))
.join(path)
};
if let Some(base) = working_dir {
if let Err(error) = ensure_within_workspace(&root, base) {
return CallToolResult::error(vec![Content::text(error).with_priority(0.0)]);
}
}
self.tree_at(root, params.depth)
}
@@ -0,0 +1,316 @@
use std::path::{Component, Path, PathBuf};
use super::edit::resolve_path;
pub fn resolve_path_in_workspace(
path: &str,
working_dir: Option<&Path>,
) -> Result<PathBuf, String> {
let Some(base) = working_dir else {
return Ok(resolve_path(path, None));
};
let resolved = resolve_path(path, Some(base));
ensure_within_workspace(&resolved, base)?;
Ok(resolved)
}
pub fn ensure_shell_command_within_workspace(
command: &str,
working_dir: &Path,
) -> Result<(), String> {
if command.contains("../") {
return Err(format!(
"Shell command must not use parent paths (../). All file access must stay within {}.",
working_dir.display()
));
}
for pattern in ["/Users/", "/home/", "$HOME", "${HOME}"] {
if command.contains(pattern) {
return Err(format!(
"Shell command must not reference host home paths ({pattern}). \
All file access must stay within {}.",
working_dir.display()
));
}
}
let lower = command.to_ascii_lowercase();
for token in [
"curl ",
"curl\t",
"wget ",
"wget\t",
"fetch http",
"fetch https",
] {
if lower.contains(token) {
return Err(format!(
"Shell network fetch is disabled in user workspace sessions. \
Use local commands (find, ls, cat) within {}.",
working_dir.display()
));
}
}
let base = canonicalize_existing(working_dir).map_err(|error| {
format!(
"Failed to resolve working directory {}: {error}",
working_dir.display()
)
})?;
let base_norm = lexical_normalize(&base);
for token in shell_path_tokens(command) {
let resolved = resolve_shell_path_token(&token, working_dir);
let Some(resolved) = resolved else {
continue;
};
let resolved_norm =
if resolved.exists() {
lexical_normalize(&canonicalize_existing(&resolved).map_err(|error| {
format!("Failed to resolve {}: {error}", resolved.display())
})?)
} else {
lexical_normalize(&resolved)
};
if !resolved_norm.starts_with(&base_norm) {
return Err(format!(
"Shell command references path outside the session working directory: {token}. \
All file access must stay within {}.",
working_dir.display()
));
}
}
Ok(())
}
fn shell_path_tokens(command: &str) -> Vec<String> {
let mut tokens = Vec::new();
let mut current = String::new();
let mut in_single = false;
let mut in_double = false;
let mut escape = false;
for ch in command.chars() {
if escape {
current.push(ch);
escape = false;
continue;
}
match ch {
'\\' if in_double => escape = true,
'\'' if !in_double => in_single = !in_single,
'"' if !in_single => in_double = !in_double,
c if c.is_whitespace() && !in_single && !in_double => {
if !current.is_empty() {
tokens.push(std::mem::take(&mut current));
}
}
_ => current.push(ch),
}
}
if !current.is_empty() {
tokens.push(current);
}
tokens
}
fn looks_like_shell_path_token(token: &str) -> bool {
let token = token.trim_matches(|c| "'\"".contains(c));
if token.is_empty() || token.starts_with('-') {
return false;
}
if token.starts_with("2>")
|| token.starts_with("1>")
|| token.starts_with(">>")
|| token == ">"
|| token.contains("://")
{
return false;
}
!matches!(token, "|" | "&&" | "||" | ";" | "(" | ")")
&& (token.starts_with('/')
|| token.starts_with("./")
|| token.starts_with("../")
|| token.starts_with('~')
|| token == "."
|| token == ".."
|| token.contains('/')
|| token.contains('\\'))
}
fn resolve_shell_path_token(token: &str, working_dir: &Path) -> Option<PathBuf> {
let token = token.trim_matches(|c| "'\"".contains(c));
if !looks_like_shell_path_token(token) {
return None;
}
if token == "." {
return Some(working_dir.to_path_buf());
}
if token == ".." {
return working_dir.parent().map(Path::to_path_buf);
}
let expanded = if let Some(rest) = token.strip_prefix('~') {
dirs::home_dir()?.join(rest.trim_start_matches(['/', '\\']))
} else {
PathBuf::from(token)
};
Some(if expanded.is_absolute() {
expanded
} else {
working_dir.join(expanded)
})
}
pub fn ensure_within_workspace(resolved: &Path, base: &Path) -> Result<(), String> {
let base_canon = canonicalize_existing(base).map_err(|error| {
format!(
"Failed to resolve working directory {}: {error}",
base.display()
)
})?;
let resolved_canon = workspace_check_path(resolved, &base_canon)?;
let base_check = lexical_normalize(&base_canon);
let resolved_check = lexical_normalize(&resolved_canon);
if !resolved_check.starts_with(&base_check) {
return Err(format!(
"Path {} is outside the session working directory {}",
resolved.display(),
base.display()
));
}
Ok(())
}
fn workspace_check_path(resolved: &Path, base_canon: &Path) -> Result<PathBuf, String> {
if resolved.exists() {
return canonicalize_existing(resolved);
}
let mut anchor = lexical_normalize(resolved);
while !anchor.exists() {
match anchor.parent() {
Some(parent) if parent.as_os_str().is_empty() => break,
Some(parent) => anchor = parent.to_path_buf(),
None => break,
}
}
if anchor.exists() {
let anchor_canon = canonicalize_existing(&anchor)?;
let suffix = resolved
.strip_prefix(&anchor)
.unwrap_or_else(|_| resolved.as_ref());
return Ok(anchor_canon.join(suffix));
}
if resolved.is_absolute() {
Ok(lexical_normalize(resolved))
} else {
Ok(base_canon.join(resolved))
}
}
fn canonicalize_existing(path: &Path) -> Result<PathBuf, String> {
path.canonicalize()
.map_err(|error| format!("Failed to resolve {}: {error}", path.display()))
}
fn lexical_normalize(path: &Path) -> PathBuf {
let mut out = PathBuf::new();
for component in path.components() {
match component {
Component::CurDir => {}
Component::ParentDir => {
out.pop();
}
Component::Prefix(prefix) => out.push(prefix.as_os_str()),
Component::RootDir => out.push(component.as_os_str()),
Component::Normal(part) => out.push(part),
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::tempdir;
#[test]
fn allows_relative_paths_inside_workspace() {
let dir = tempdir().unwrap();
let base = dir.path();
let resolved =
resolve_path_in_workspace("notes/report.html", Some(base)).expect("inside workspace");
assert_eq!(resolved, base.join("notes/report.html"));
}
#[test]
fn blocks_parent_escape() {
let dir = tempdir().unwrap();
let base = dir.path();
let error = resolve_path_in_workspace("../outside.txt", Some(base))
.expect_err("should block escape");
assert!(error.contains("outside the session working directory"));
}
#[test]
fn blocks_absolute_paths_outside_workspace() {
let dir = tempdir().unwrap();
let base = dir.path();
let outside = std::env::temp_dir().join("goose-workspace-outside.txt");
let error = resolve_path_in_workspace(outside.to_str().unwrap(), Some(base))
.expect_err("should block outside absolute path");
assert!(error.contains("outside the session working directory"));
}
#[test]
fn blocks_parent_path_segments_in_shell_commands() {
let dir = tempdir().unwrap();
let base = dir.path();
let error = ensure_shell_command_within_workspace("find .. -name '*.csv'", base)
.expect_err("blocked");
assert!(
error.contains("outside the session working directory")
|| error.contains("parent paths")
);
}
#[test]
fn blocks_shell_commands_with_paths_outside_workspace() {
let dir = tempdir().unwrap();
let base = dir.path();
let outside = std::env::temp_dir().join("goose-shell-outside.txt");
let command = format!("find {} -name '*.csv'", outside.display());
let error = ensure_shell_command_within_workspace(&command, base).expect_err("blocked");
assert!(error.contains("outside the session working directory"));
}
#[test]
fn allows_shell_commands_within_workspace() {
let dir = tempdir().unwrap();
let base = dir.path();
fs::create_dir_all(base.join("oa")).unwrap();
ensure_shell_command_within_workspace("find . -name '*.csv'", base).expect("allowed");
ensure_shell_command_within_workspace("rg export csv oa/", base).expect("allowed");
}
#[test]
fn allows_absolute_paths_inside_workspace() {
let dir = tempdir().unwrap();
let base = dir.path();
let inside = base.join("inside.txt");
fs::write(&inside, "ok").unwrap();
let resolved =
resolve_path_in_workspace(inside.to_str().unwrap(), Some(base)).expect("allowed");
assert_eq!(resolved, inside);
}
}
@@ -0,0 +1,358 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{
CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult,
ServerCapabilities, Tool, ToolAnnotations,
};
use schemars::{schema_for, JsonSchema};
use serde::{Deserialize, Serialize};
use std::process::Command;
use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "git";
const MAX_OUTPUT_CHARS: usize = 50_000;
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct GitDiffParams {
path: Option<String>,
staged: bool,
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct GitLogParams {
#[schemars(default = "default_limit")]
limit: u32,
path: Option<String>,
}
fn default_limit() -> u32 {
20
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct GitBlameParams {
path: String,
start_line: Option<u32>,
end_line: Option<u32>,
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct GitCommitParams {
message: String,
paths: Vec<String>,
}
pub struct GitClient {
info: InitializeResult,
#[allow(dead_code)]
context: PlatformExtensionContext,
}
impl GitClient {
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(
Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string())
.with_title("Git"),
)
.with_instructions(
"Git operations: status, diff, log, blame, and commit changes.".to_string(),
);
Ok(Self { info, context })
}
fn run_command(args: &[&str]) -> Result<String, String> {
let output = Command::new("git")
.args(args)
.output()
.map_err(|e| format!("Failed to run git: {}", e))?;
let stdout = String::from_utf8_lossy(&output.stdout);
let stderr = String::from_utf8_lossy(&output.stderr);
let combined = if output.status.success() {
stdout.into_owned()
} else {
format!("{}{}", stdout, stderr)
};
if combined.len() > MAX_OUTPUT_CHARS {
Ok(format!(
"{}\n[output truncated]",
&combined[..MAX_OUTPUT_CHARS]
))
} else {
Ok(combined)
}
}
fn handle_git_status(_arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let status = Self::run_command(&["status", "--short"])?;
let branch = Self::run_command(&["branch", "--show-current"])?;
Ok(vec![Content::text(format!(
"Branch: {}\n{}",
branch.trim(),
status
))])
}
fn handle_git_diff(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let staged = arguments
.as_ref()
.and_then(|a| a.get("staged"))
.and_then(|v| v.as_bool())
.unwrap_or(false);
let path = arguments
.as_ref()
.and_then(|a| a.get("path"))
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let mut args = vec!["diff".to_string()];
if staged {
args.push("--staged".to_string());
}
if let Some(ref p) = path {
args.push(p.clone());
}
let arg_refs: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
let output = Self::run_command(&arg_refs)?;
Ok(vec![Content::text(output)])
}
fn handle_git_log(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let limit = arguments
.as_ref()
.and_then(|a| a.get("limit"))
.and_then(|v| v.as_u64())
.unwrap_or(20)
.to_string();
let path = arguments
.as_ref()
.and_then(|a| a.get("path"))
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let mut args = vec![
"log".to_string(),
"--oneline".to_string(),
"-n".to_string(),
limit,
];
if let Some(ref p) = path {
args.push(p.clone());
}
let arg_refs: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
let output = Self::run_command(&arg_refs)?;
Ok(vec![Content::text(output)])
}
fn handle_git_blame(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let path = arguments
.as_ref()
.ok_or("Missing arguments")?
.get("path")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: path")?
.to_string();
let start_line = arguments
.as_ref()
.and_then(|a| a.get("start_line"))
.and_then(|v| v.as_u64())
.map(|v| v as u32);
let end_line = arguments
.as_ref()
.and_then(|a| a.get("end_line"))
.and_then(|v| v.as_u64())
.map(|v| v as u32);
let mut args = vec!["blame".to_string()];
if let (Some(start), Some(end)) = (start_line, end_line) {
args.push(format!("-L {},{}", start, end));
} else if let Some(start) = start_line {
args.push(format!("-L {},{}", start, start));
}
args.push(path);
let arg_refs: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
let output = Self::run_command(&arg_refs)?;
Ok(vec![Content::text(output)])
}
fn handle_git_commit(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let message = arguments
.as_ref()
.ok_or("Missing arguments")?
.get("message")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: message")?
.to_string();
let paths: Vec<String> = arguments
.as_ref()
.and_then(|a| a.get("paths"))
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str())
.map(|s| s.to_string())
.collect()
})
.unwrap_or_default();
if paths.is_empty() {
return Err("No paths provided for commit".to_string());
}
let mut add_args = vec!["add".to_string()];
add_args.extend(paths.iter().cloned());
let add_arg_refs: Vec<&str> = add_args.iter().map(|s| s.as_str()).collect();
let add_output = Command::new("git")
.args(&add_arg_refs)
.output()
.map_err(|e| format!("Failed to run git add: {}", e))?;
if !add_output.status.success() {
return Err(format!(
"git add failed: {}",
String::from_utf8_lossy(&add_output.stderr)
));
}
let commit_output = Command::new("git")
.args(["commit", "-s", "-m", &message])
.output()
.map_err(|e| format!("Failed to run git commit: {}", e))?;
let stdout = String::from_utf8_lossy(&commit_output.stdout);
let stderr = String::from_utf8_lossy(&commit_output.stderr);
Ok(vec![Content::text(format!("{}{}", stdout, stderr))])
}
fn get_tools() -> Vec<Tool> {
let diff_schema = serde_json::to_value(schema_for!(GitDiffParams))
.expect("Failed to serialize GitDiffParams schema");
let log_schema = serde_json::to_value(schema_for!(GitLogParams))
.expect("Failed to serialize GitLogParams schema");
let blame_schema = serde_json::to_value(schema_for!(GitBlameParams))
.expect("Failed to serialize GitBlameParams schema");
let commit_schema = serde_json::to_value(schema_for!(GitCommitParams))
.expect("Failed to serialize GitCommitParams schema");
let empty_schema = serde_json::json!({"type": "object", "properties": {}});
vec![
Tool::new(
"git_status".to_string(),
"Show current git branch and working tree status (short format)".to_string(),
empty_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Git Status".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
Tool::new(
"git_diff".to_string(),
"Show git diff. Optionally specify a path and whether to show staged changes."
.to_string(),
diff_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Git Diff".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
Tool::new(
"git_log".to_string(),
"Show git commit log in oneline format. Optionally limit number and filter by path."
.to_string(),
log_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Git Log".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
Tool::new(
"git_blame".to_string(),
"Show git blame for a file, optionally limited to a line range.".to_string(),
blame_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Git Blame".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
Tool::new(
"git_commit".to_string(),
"Stage specified paths and create a signed git commit with the given message."
.to_string(),
commit_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Git Commit".to_string()),
Some(false),
Some(false),
Some(false),
Some(false),
)),
]
}
}
#[async_trait]
impl McpClientTrait for GitClient {
async fn list_tools(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Ok(ListToolsResult {
tools: Self::get_tools(),
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
_ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
"git_status" => Self::handle_git_status(arguments),
"git_diff" => Self::handle_git_diff(arguments),
"git_log" => Self::handle_git_log(arguments),
"git_blame" => Self::handle_git_blame(arguments),
"git_commit" => Self::handle_git_commit(arguments),
_ => Err(format!("Unknown tool: {}", name)),
};
match content {
Ok(content) => Ok(CallToolResult::success(content)),
Err(error) => Ok(CallToolResult::error(vec![Content::text(format!(
"Error: {}",
error
))])),
}
}
fn get_info(&self) -> Option<&InitializeResult> {
Some(&self.info)
}
}
@@ -1,3 +1,4 @@
pub mod aider;
pub mod analyze;
pub mod apps;
pub mod chatrecall;
@@ -5,11 +6,16 @@ pub mod chatrecall;
pub mod code_execution;
pub mod developer;
pub mod ext_manager;
pub mod git;
pub mod orchestrator;
pub mod projectmemory;
pub mod search;
pub mod summarize;
pub mod summon;
pub mod test_runner;
pub mod todo;
pub mod tom;
pub mod web;
use std::collections::HashMap;
@@ -29,6 +35,20 @@ pub static PLATFORM_EXTENSIONS: Lazy<HashMap<&'static str, PlatformExtensionDef>
|| {
let mut map = HashMap::new();
map.insert(
aider::EXTENSION_NAME,
PlatformExtensionDef {
name: aider::EXTENSION_NAME,
display_name: "Aider",
description:
"Delegate multi-file coding tasks to Aider instead of developer write/edit",
default_enabled: true,
unprefixed_tools: false,
hidden: false,
client_factory: |ctx| Box::new(aider::AiderClient::new(ctx).unwrap()),
},
);
map.insert(
analyze::EXTENSION_NAME,
PlatformExtensionDef {
@@ -85,6 +105,21 @@ pub static PLATFORM_EXTENSIONS: Lazy<HashMap<&'static str, PlatformExtensionDef>
},
);
map.insert(
projectmemory::EXTENSION_NAME,
PlatformExtensionDef {
name: projectmemory::EXTENSION_NAME,
display_name: "Project Memory",
description: "Persistent project context injected via harness bootstrap and MOIM",
default_enabled: false,
unprefixed_tools: false,
hidden: false,
client_factory: |ctx| {
Box::new(projectmemory::ProjectMemoryClient::new(ctx).unwrap())
},
},
);
map.insert(
"extensionmanager",
PlatformExtensionDef {
@@ -202,6 +237,60 @@ pub static PLATFORM_EXTENSIONS: Lazy<HashMap<&'static str, PlatformExtensionDef>
},
);
map.insert(
git::EXTENSION_NAME,
PlatformExtensionDef {
name: git::EXTENSION_NAME,
display_name: "Git",
description: "Git operations: status, diff, log, blame, branch management",
default_enabled: true,
unprefixed_tools: false,
hidden: false,
client_factory: |ctx| Box::new(git::GitClient::new(ctx).unwrap()),
},
);
map.insert(
search::EXTENSION_NAME,
PlatformExtensionDef {
name: search::EXTENSION_NAME,
display_name: "Search",
description:
"Search code with ripgrep: find patterns, symbols, and text across files",
default_enabled: true,
unprefixed_tools: false,
hidden: false,
client_factory: |ctx| Box::new(search::SearchClient::new(ctx).unwrap()),
},
);
map.insert(
test_runner::EXTENSION_NAME,
PlatformExtensionDef {
name: test_runner::EXTENSION_NAME,
display_name: "Test Runner",
description:
"Run tests and parse results: cargo test with structured failure reporting",
default_enabled: true,
unprefixed_tools: false,
hidden: false,
client_factory: |ctx| Box::new(test_runner::TestRunnerClient::new(ctx).unwrap()),
},
);
map.insert(
web::EXTENSION_NAME,
PlatformExtensionDef {
name: web::EXTENSION_NAME,
display_name: "Web",
description: "Fetch web pages and search the internet",
default_enabled: true,
unprefixed_tools: false,
hidden: false,
client_factory: |ctx| Box::new(web::WebClient::new(ctx).unwrap()),
},
);
map
},
);
@@ -0,0 +1,127 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use crate::session::extension_data::ExtensionState;
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{
CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult,
ServerCapabilities,
};
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "projectmemory";
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ProjectMemoryState {
pub summary: String,
pub source: Option<String>,
pub updated_at: Option<String>,
}
impl ExtensionState for ProjectMemoryState {
const EXTENSION_NAME: &'static str = EXTENSION_NAME;
const VERSION: &'static str = "v0";
}
pub struct ProjectMemoryClient {
info: InitializeResult,
context: PlatformExtensionContext,
}
impl ProjectMemoryClient {
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().build()).with_server_info(
Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string())
.with_title("Project Memory"),
);
Ok(Self { info, context })
}
}
#[async_trait]
impl McpClientTrait for ProjectMemoryClient {
async fn list_tools(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Ok(ListToolsResult {
tools: vec![],
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
_ctx: &ToolCallContext,
name: &str,
_arguments: Option<JsonObject>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
Ok(CallToolResult::error(vec![Content::text(format!(
"projectmemory has no tools (called: {name})"
))]))
}
fn get_info(&self) -> Option<&InitializeResult> {
Some(&self.info)
}
async fn get_moim(&self, session_id: &str) -> Option<String> {
let session = self
.context
.session_manager
.get_session(session_id, false)
.await
.ok()?;
let state = ProjectMemoryState::from_extension_data(&session.extension_data)?;
let summary = state.summary.trim();
if summary.is_empty() {
return None;
}
let mut lines = vec![
"Project memory bootstrap: treat this as background context for the current project."
.to_string(),
"Prefer newer user instructions if anything conflicts.".to_string(),
];
if let Some(source) = state.source.as_deref() {
lines.push(format!("Source: {source}"));
}
if let Some(updated_at) = state.updated_at.as_deref() {
lines.push(format!("Updated: {updated_at}"));
}
Some(format!("{}\n\n{}", lines.join("\n"), summary))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::session::ExtensionData;
#[test]
fn project_memory_state_round_trips_through_extension_data() {
let state = ProjectMemoryState {
summary: "Use the H5 client for mobile conversations.".to_string(),
source: Some("harness".to_string()),
updated_at: Some("2026-06-10T00:00:00Z".to_string()),
};
let mut extension_data = ExtensionData::default();
state.to_extension_data(&mut extension_data).unwrap();
assert_eq!(
ProjectMemoryState::from_extension_data(&extension_data)
.unwrap()
.summary,
state.summary
);
}
}
@@ -0,0 +1,317 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{
CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult,
ServerCapabilities, Tool, ToolAnnotations,
};
use schemars::{schema_for, JsonSchema};
use serde::{Deserialize, Serialize};
use std::process::Command;
use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "search";
const MAX_OUTPUT_CHARS: usize = 50_000;
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct SearchTextParams {
pattern: String,
path: Option<String>,
case_sensitive: bool,
file_glob: Option<String>,
#[schemars(default = "default_max_results")]
max_results: u32,
}
fn default_max_results() -> u32 {
100
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct SearchFilesParams {
name_pattern: String,
path: Option<String>,
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct SearchSymbolParams {
symbol: String,
path: Option<String>,
}
pub struct SearchClient {
info: InitializeResult,
#[allow(dead_code)]
context: PlatformExtensionContext,
}
impl SearchClient {
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(
Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string())
.with_title("Search"),
)
.with_instructions(
"Search code with ripgrep: find patterns, symbols, and text across files."
.to_string(),
);
Ok(Self { info, context })
}
fn truncate(s: String) -> String {
if s.len() > MAX_OUTPUT_CHARS {
format!("{}\n[output truncated]", &s[..MAX_OUTPUT_CHARS])
} else {
s
}
}
fn handle_search_text(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let args = arguments.as_ref().ok_or("Missing arguments")?;
let pattern = args
.get("pattern")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: pattern")?
.to_string();
let path = args
.get("path")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let case_sensitive = args
.get("case_sensitive")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let file_glob = args
.get("file_glob")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let max_results = args
.get("max_results")
.and_then(|v| v.as_u64())
.unwrap_or(100)
.to_string();
// Try rg first, fall back to grep
let rg_available = Command::new("which")
.arg("rg")
.output()
.map(|o| o.status.success())
.unwrap_or(false);
let output = if rg_available {
let mut cmd_args = vec![
"--line-number".to_string(),
"-m".to_string(),
max_results.clone(),
];
if !case_sensitive {
cmd_args.push("--ignore-case".to_string());
}
if let Some(ref glob) = file_glob {
cmd_args.push("-g".to_string());
cmd_args.push(glob.clone());
}
cmd_args.push(pattern.clone());
if let Some(ref p) = path {
cmd_args.push(p.clone());
}
let arg_refs: Vec<&str> = cmd_args.iter().map(|s| s.as_str()).collect();
let out = Command::new("rg")
.args(&arg_refs)
.output()
.map_err(|e| format!("Failed to run rg: {}", e))?;
String::from_utf8_lossy(&out.stdout).into_owned()
} else {
let mut cmd_args = vec!["-rn".to_string()];
if !case_sensitive {
cmd_args.push("-i".to_string());
}
if let Some(ref glob) = file_glob {
cmd_args.push("--include".to_string());
cmd_args.push(glob.clone());
}
cmd_args.push(pattern.clone());
if let Some(ref p) = path {
cmd_args.push(p.clone());
} else {
cmd_args.push(".".to_string());
}
let arg_refs: Vec<&str> = cmd_args.iter().map(|s| s.as_str()).collect();
let out = Command::new("grep")
.args(&arg_refs)
.output()
.map_err(|e| format!("Failed to run grep: {}", e))?;
String::from_utf8_lossy(&out.stdout).into_owned()
};
Ok(vec![Content::text(Self::truncate(output))])
}
fn handle_search_files(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let args = arguments.as_ref().ok_or("Missing arguments")?;
let name_pattern = args
.get("name_pattern")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: name_pattern")?
.to_string();
let search_path = args
.get("path")
.and_then(|v| v.as_str())
.unwrap_or(".")
.to_string();
let out = Command::new("find")
.args([&search_path, "-name", &name_pattern])
.output()
.map_err(|e| format!("Failed to run find: {}", e))?;
let output = String::from_utf8_lossy(&out.stdout).into_owned();
Ok(vec![Content::text(Self::truncate(output))])
}
fn handle_search_symbol(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let args = arguments.as_ref().ok_or("Missing arguments")?;
let symbol = args
.get("symbol")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: symbol")?
.to_string();
let path = args
.get("path")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let rg_available = Command::new("which")
.arg("rg")
.output()
.map(|o| o.status.success())
.unwrap_or(false);
let output = if rg_available {
let mut cmd_args = vec![
"--line-number".to_string(),
"-w".to_string(),
symbol.clone(),
];
if let Some(ref p) = path {
cmd_args.push(p.clone());
}
let arg_refs: Vec<&str> = cmd_args.iter().map(|s| s.as_str()).collect();
let out = Command::new("rg")
.args(&arg_refs)
.output()
.map_err(|e| format!("Failed to run rg: {}", e))?;
String::from_utf8_lossy(&out.stdout).into_owned()
} else {
let mut cmd_args = vec!["-rn".to_string(), "-w".to_string(), symbol.clone()];
if let Some(ref p) = path {
cmd_args.push(p.clone());
} else {
cmd_args.push(".".to_string());
}
let arg_refs: Vec<&str> = cmd_args.iter().map(|s| s.as_str()).collect();
let out = Command::new("grep")
.args(&arg_refs)
.output()
.map_err(|e| format!("Failed to run grep: {}", e))?;
String::from_utf8_lossy(&out.stdout).into_owned()
};
Ok(vec![Content::text(Self::truncate(output))])
}
fn get_tools() -> Vec<Tool> {
let text_schema = serde_json::to_value(schema_for!(SearchTextParams))
.expect("Failed to serialize SearchTextParams schema");
let files_schema = serde_json::to_value(schema_for!(SearchFilesParams))
.expect("Failed to serialize SearchFilesParams schema");
let symbol_schema = serde_json::to_value(schema_for!(SearchSymbolParams))
.expect("Failed to serialize SearchSymbolParams schema");
vec![
Tool::new(
"search_text".to_string(),
"Search for a text pattern in files using ripgrep (falls back to grep). Supports case sensitivity and file glob filters.".to_string(),
text_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Search Text".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
Tool::new(
"search_files".to_string(),
"Find files by name pattern using the find command.".to_string(),
files_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Search Files".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
Tool::new(
"search_symbol".to_string(),
"Find exact symbol matches (whole word) in code using ripgrep or grep.".to_string(),
symbol_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Search Symbol".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
]
}
}
#[async_trait]
impl McpClientTrait for SearchClient {
async fn list_tools(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Ok(ListToolsResult {
tools: Self::get_tools(),
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
_ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
"search_text" => Self::handle_search_text(arguments),
"search_files" => Self::handle_search_files(arguments),
"search_symbol" => Self::handle_search_symbol(arguments),
_ => Err(format!("Unknown tool: {}", name)),
};
match content {
Ok(content) => Ok(CallToolResult::success(content)),
Err(error) => Ok(CallToolResult::error(vec![Content::text(format!(
"Error: {}",
error
))])),
}
}
fn get_info(&self) -> Option<&InitializeResult> {
Some(&self.info)
}
}
@@ -0,0 +1,313 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{
CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult,
ServerCapabilities, Tool, ToolAnnotations,
};
use schemars::{schema_for, JsonSchema};
use serde::{Deserialize, Serialize};
use std::process::Command;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "test_runner";
const MAX_OUTPUT_CHARS: usize = 50_000;
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct RunTestsParams {
package: Option<String>,
test_filter: Option<String>,
#[schemars(default = "default_timeout")]
timeout_secs: u32,
}
fn default_timeout() -> u32 {
120
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct ListTestsParams {
package: Option<String>,
}
pub struct TestRunnerClient {
info: InitializeResult,
#[allow(dead_code)]
context: PlatformExtensionContext,
}
impl TestRunnerClient {
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(
Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string())
.with_title("Test Runner"),
)
.with_instructions(
"Run tests and parse results: cargo test with structured failure reporting."
.to_string(),
);
Ok(Self { info, context })
}
fn truncate(s: String) -> String {
if s.len() > MAX_OUTPUT_CHARS {
format!("{}\n[output truncated]", &s[..MAX_OUTPUT_CHARS])
} else {
s
}
}
fn parse_test_summary(output: &str) -> String {
let mut passed = 0u32;
let mut failed = 0u32;
let mut ignored = 0u32;
let mut failures = Vec::new();
let mut in_failures = false;
for line in output.lines() {
if line.contains("test result:") {
if let Some(ok_idx) = line.find("ok.") {
let rest = &line[ok_idx + 3..];
if let Some(p) = Self::extract_number(rest, "passed") {
passed += p;
}
if let Some(f) = Self::extract_number(rest, "failed") {
failed += f;
}
if let Some(i) = Self::extract_number(rest, "ignored") {
ignored += i;
}
} else if line.contains("FAILED") {
if let Some(rest) = line.find("FAILED.").map(|i| &line[i + 7..]) {
if let Some(f) = Self::extract_number(rest, "failed") {
failed += f;
}
}
}
}
if line.contains("failures:") && line.trim() == "failures:" {
in_failures = true;
} else if in_failures && line.starts_with(" ") {
let test_name = line.trim();
if !test_name.is_empty() {
failures.push(test_name.to_string());
}
} else if in_failures && line.trim().is_empty() {
// continue
} else if in_failures && !line.starts_with(" ") && !line.trim().is_empty() {
in_failures = false;
}
}
let mut summary = format!(
"Test Summary: {} passed, {} failed, {} ignored\n",
passed, failed, ignored
);
if !failures.is_empty() {
summary.push_str("\nFailed tests:\n");
for f in &failures {
summary.push_str(&format!(" - {}\n", f));
}
}
summary
}
fn extract_number(text: &str, label: &str) -> Option<u32> {
if let Some(label_pos) = text.find(&format!(" {}", label)) {
let before = &text[..label_pos];
let num_str: String = before
.chars()
.rev()
.take_while(|c| c.is_ascii_digit())
.collect();
let reversed: String = num_str.chars().rev().collect();
reversed.parse().ok()
} else {
None
}
}
fn handle_run_tests(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let package = arguments
.as_ref()
.and_then(|a| a.get("package"))
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let test_filter = arguments
.as_ref()
.and_then(|a| a.get("test_filter"))
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let timeout_secs = arguments
.as_ref()
.and_then(|a| a.get("timeout_secs"))
.and_then(|v| v.as_u64())
.unwrap_or(120) as u64;
let mut cmd = Command::new("cargo");
cmd.arg("test");
if let Some(ref pkg) = package {
cmd.args(["-p", pkg]);
}
if let Some(ref filter) = test_filter {
cmd.arg(filter);
}
cmd.args(["--", "--nocapture"]);
cmd.stderr(std::process::Stdio::piped());
cmd.stdout(std::process::Stdio::piped());
let child = cmd
.spawn()
.map_err(|e| format!("Failed to spawn cargo test: {}", e))?;
// Use a thread to enforce the timeout
let timeout = Duration::from_secs(timeout_secs);
let result = std::thread::spawn(move || {
let mut child = child;
let start = std::time::Instant::now();
loop {
match child.try_wait() {
Ok(Some(_)) => break,
Ok(None) => {
if start.elapsed() > timeout {
let _ = child.kill();
return Err(format!(
"cargo test timed out after {} seconds",
timeout_secs
));
}
std::thread::sleep(Duration::from_millis(100));
}
Err(e) => return Err(format!("Error waiting for cargo test: {}", e)),
}
}
child
.wait_with_output()
.map_err(|e| format!("Failed to get output: {}", e))
})
.join()
.map_err(|_| "cargo test thread panicked".to_string())?;
let output = match result {
Ok(out) => out,
Err(msg) => return Ok(vec![Content::text(msg)]),
};
let stdout = String::from_utf8_lossy(&output.stdout).into_owned();
let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
let combined = format!("{}{}", stdout, stderr);
let summary = Self::parse_test_summary(&combined);
let full = format!("{}\n{}", summary, combined);
Ok(vec![Content::text(Self::truncate(full))])
}
fn handle_list_tests(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let package = arguments
.as_ref()
.and_then(|a| a.get("package"))
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let mut cmd_args = vec!["test".to_string()];
if let Some(ref pkg) = package {
cmd_args.push("-p".to_string());
cmd_args.push(pkg.clone());
}
cmd_args.push("--".to_string());
cmd_args.push("--list".to_string());
let arg_refs: Vec<&str> = cmd_args.iter().map(|s| s.as_str()).collect();
let out = Command::new("cargo")
.args(&arg_refs)
.output()
.map_err(|e| format!("Failed to run cargo test --list: {}", e))?;
let stdout = String::from_utf8_lossy(&out.stdout).into_owned();
let stderr = String::from_utf8_lossy(&out.stderr).into_owned();
let combined = format!("{}{}", stdout, stderr);
Ok(vec![Content::text(Self::truncate(combined))])
}
fn get_tools() -> Vec<Tool> {
let run_schema = serde_json::to_value(schema_for!(RunTestsParams))
.expect("Failed to serialize RunTestsParams schema");
let list_schema = serde_json::to_value(schema_for!(ListTestsParams))
.expect("Failed to serialize ListTestsParams schema");
vec![
Tool::new(
"run_tests".to_string(),
"Run cargo tests with optional package filter, test name filter, and timeout. Returns pass/fail counts and failure details.".to_string(),
run_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Run Tests".to_string()),
Some(false),
Some(false),
Some(false),
Some(false),
)),
Tool::new(
"list_tests".to_string(),
"List all available tests in the project or a specific package.".to_string(),
list_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("List Tests".to_string()),
Some(false),
Some(true),
Some(false),
Some(false),
)),
]
}
}
#[async_trait]
impl McpClientTrait for TestRunnerClient {
async fn list_tools(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Ok(ListToolsResult {
tools: Self::get_tools(),
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
_ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
"run_tests" => Self::handle_run_tests(arguments),
"list_tests" => Self::handle_list_tests(arguments),
_ => Err(format!("Unknown tool: {}", name)),
};
match content {
Ok(content) => Ok(CallToolResult::success(content)),
Err(error) => Ok(CallToolResult::error(vec![Content::text(format!(
"Error: {}",
error
))])),
}
}
fn get_info(&self) -> Option<&InitializeResult> {
Some(&self.info)
}
}
@@ -0,0 +1,337 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{
CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult,
ServerCapabilities, Tool, ToolAnnotations,
};
use schemars::{schema_for, JsonSchema};
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "web";
const MAX_OUTPUT_CHARS: usize = 50_000;
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct FetchUrlParams {
url: String,
extract_text: bool,
}
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct WebSearchParams {
query: String,
#[schemars(default = "default_num_results")]
num_results: u32,
}
fn default_num_results() -> u32 {
10
}
pub struct WebClient {
info: InitializeResult,
#[allow(dead_code)]
context: PlatformExtensionContext,
}
impl WebClient {
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(
Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string())
.with_title("Web"),
)
.with_instructions(
"Fetch web pages and search the internet using DuckDuckGo.".to_string(),
);
Ok(Self { info, context })
}
fn truncate(s: String) -> String {
if s.len() > MAX_OUTPUT_CHARS {
format!("{}\n[output truncated]", &s[..MAX_OUTPUT_CHARS])
} else {
s
}
}
fn strip_html(html: &str) -> String {
// Remove script and style blocks first
let mut text = html.to_string();
// Remove <script>...</script>
while let Some(start) = text.to_lowercase().find("<script") {
if let Some(end) = text.to_lowercase()[start..].find("</script>") {
text = format!("{}{}", &text[..start], &text[start + end + 9..]);
} else {
break;
}
}
// Remove <style>...</style>
while let Some(start) = text.to_lowercase().find("<style") {
if let Some(end) = text.to_lowercase()[start..].find("</style>") {
text = format!("{}{}", &text[..start], &text[start + end + 8..]);
} else {
break;
}
}
// Strip remaining HTML tags
let mut result = String::new();
let mut in_tag = false;
for ch in text.chars() {
match ch {
'<' => in_tag = true,
'>' => in_tag = false,
c if !in_tag => result.push(c),
_ => {}
}
}
// Clean up whitespace
let lines: Vec<&str> = result.lines().map(|l| l.trim()).collect();
lines
.iter()
.filter(|l| !l.is_empty())
.cloned()
.collect::<Vec<_>>()
.join("\n")
}
fn handle_fetch_url(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let args = arguments.as_ref().ok_or("Missing arguments")?;
let url = args
.get("url")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: url")?
.to_string();
let extract_text = args
.get("extract_text")
.and_then(|v| v.as_bool())
.unwrap_or(true);
let client = reqwest::blocking::Client::builder()
.user_agent("Mozilla/5.0 (compatible; Goose/1.0)")
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| format!("Failed to build HTTP client: {}", e))?;
let response = client
.get(&url)
.send()
.map_err(|e| format!("Failed to fetch URL: {}", e))?;
let body = response
.text()
.map_err(|e| format!("Failed to read response body: {}", e))?;
let output = if extract_text {
Self::strip_html(&body)
} else {
body
};
Ok(vec![Content::text(Self::truncate(output))])
}
fn handle_web_search(arguments: Option<JsonObject>) -> Result<Vec<Content>, String> {
let args = arguments.as_ref().ok_or("Missing arguments")?;
let query = args
.get("query")
.and_then(|v| v.as_str())
.ok_or("Missing required parameter: query")?
.to_string();
let num_results = args
.get("num_results")
.and_then(|v| v.as_u64())
.unwrap_or(10) as usize;
let encoded_query: String = query
.chars()
.map(|c| match c {
' ' => '+',
c if c.is_ascii_alphanumeric() || "-_.~".contains(c) => c,
_ => '+',
})
.collect();
let url = format!("https://html.duckduckgo.com/html/?q={}", encoded_query);
let client = reqwest::blocking::Client::builder()
.user_agent("Mozilla/5.0 (compatible; Goose/1.0)")
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| format!("Failed to build HTTP client: {}", e))?;
let response = client
.get(&url)
.send()
.map_err(|e| format!("Failed to search: {}", e))?;
let html = response
.text()
.map_err(|e| format!("Failed to read search response: {}", e))?;
// Parse DuckDuckGo HTML results
let results = Self::parse_ddg_results(&html, num_results);
Ok(vec![Content::text(Self::truncate(results))])
}
fn parse_ddg_results(html: &str, max: usize) -> String {
let mut results = Vec::new();
let lower = html.to_lowercase();
// Find result blocks — DuckDuckGo uses class="result"
let mut pos = 0;
while results.len() < max {
// Find the next result link anchor
let search_str = "class=\"result__a\"";
let Some(link_start) = lower[pos..].find(search_str).map(|i| i + pos) else {
break;
};
// Extract href from the preceding <a tag
let before = &html[..link_start + search_str.len()];
let a_start = before.rfind('<').unwrap_or(0);
let a_block = &html[a_start
..link_start
+ search_str.len()
+ 100.min(html.len() - link_start - search_str.len())];
let href = Self::extract_attr(a_block, "href").unwrap_or_default();
// Get link text
let after_a_tag_end = link_start + search_str.len();
let close_offset = html[after_a_tag_end..].find('>').unwrap_or(0);
let content_start = after_a_tag_end + close_offset + 1;
let close_a = html[content_start..].find("</a>").unwrap_or(0);
let title = Self::strip_html(&html[content_start..content_start + close_a]);
// Try to find snippet
let snippet_search = "class=\"result__snippet\"";
let snippet = if let Some(snip_pos) = lower[pos..].find(snippet_search).map(|i| i + pos)
{
let after = snip_pos + snippet_search.len();
if let Some(close) = html[after..].find('>') {
let text_start = after + close + 1;
if let Some(end_div) = html[text_start..].find("</a>") {
Self::strip_html(&html[text_start..text_start + end_div])
} else {
String::new()
}
} else {
String::new()
}
} else {
String::new()
};
if !title.trim().is_empty() {
results.push(format!(
"{}. {}\n URL: {}\n {}",
results.len() + 1,
title.trim(),
href,
snippet.trim()
));
}
pos = link_start + 1;
}
if results.is_empty() {
"No results found.".to_string()
} else {
results.join("\n\n")
}
}
fn extract_attr(tag: &str, attr: &str) -> Option<String> {
let search = format!("{}=\"", attr);
let lower = tag.to_lowercase();
let start = lower.find(&search)? + search.len();
let end = tag[start..].find('"')?;
Some(tag[start..start + end].to_string())
}
fn get_tools() -> Vec<Tool> {
let fetch_schema = serde_json::to_value(schema_for!(FetchUrlParams))
.expect("Failed to serialize FetchUrlParams schema");
let search_schema = serde_json::to_value(schema_for!(WebSearchParams))
.expect("Failed to serialize WebSearchParams schema");
vec![
Tool::new(
"fetch_url".to_string(),
"Fetch the content of a URL. If extract_text is true, strips HTML tags to return readable text.".to_string(),
fetch_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Fetch URL".to_string()),
Some(false),
Some(false),
Some(false),
Some(false),
)),
Tool::new(
"web_search".to_string(),
"Search the web using DuckDuckGo and return a list of result titles, URLs, and snippets.".to_string(),
search_schema.as_object().unwrap().clone(),
)
.annotate(ToolAnnotations::from_raw(
Some("Web Search".to_string()),
Some(false),
Some(false),
Some(false),
Some(false),
)),
]
}
}
#[async_trait]
impl McpClientTrait for WebClient {
async fn list_tools(
&self,
_session_id: &str,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Ok(ListToolsResult {
tools: Self::get_tools(),
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
_ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
"fetch_url" => Self::handle_fetch_url(arguments),
"web_search" => Self::handle_web_search(arguments),
_ => Err(format!("Unknown tool: {}", name)),
};
match content {
Ok(content) => Ok(CallToolResult::success(content)),
Err(error) => Ok(CallToolResult::error(vec![Content::text(format!(
"Error: {}",
error
))])),
}
}
fn get_info(&self) -> Option<&InitializeResult> {
Some(&self.info)
}
}
+1 -3
View File
@@ -160,9 +160,7 @@ impl<'a> SystemPromptBuilder<'a, PromptManager> {
} else {
prompt_template::render_template("system.md", &context)
}
.unwrap_or_else(|_| {
"You are a general-purpose AI agent called goose, created by Block".to_string()
});
.unwrap_or_else(|_| "You are a general-purpose AI agent called TKMind".to_string());
let mut system_prompt_extras = self.manager.system_prompt_extras.clone();
@@ -3,8 +3,8 @@ source: crates/goose/src/agents/prompt_manager.rs
assertion_line: 458
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
goose is being developed as an open-source software project.
You are a general-purpose AI agent called TKMind, created by TKMind.
TKMind is an intelligent assistant for coding, analysis, and task automation.
# Extensions
@@ -46,6 +46,19 @@ Analyze code structure using tree-sitter AST parsing. Three auto-selected modes:
For large codebases, delegate analysis to a subagent and retain only the summary.
## aider
### Instructions
Delegate multi-file coding work to Aider instead of using developer write/edit.
Prefer this tool for:
- Bug fixes and refactors across multiple files
- New feature implementation
- Test-failure fixes that need code changes
After Aider finishes, verify with developer shell (tests, git status).
For one-line or single-file tweaks, developer edit/write is fine.
## apps
apps supports resources.
@@ -2,8 +2,8 @@
source: crates/goose/src/agents/prompt_manager.rs
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
goose is being developed as an open-source software project.
You are a general-purpose AI agent called TKMind, created by TKMind.
TKMind is an intelligent assistant for coding, analysis, and task automation.
# Extensions
@@ -2,8 +2,8 @@
source: crates/goose/src/agents/prompt_manager.rs
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
goose is being developed as an open-source software project.
You are a general-purpose AI agent called TKMind, created by TKMind.
TKMind is an intelligent assistant for coding, analysis, and task automation.
# Extensions
@@ -2,8 +2,8 @@
source: crates/goose/src/agents/prompt_manager.rs
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
goose is being developed as an open-source software project.
You are a general-purpose AI agent called TKMind, created by TKMind.
TKMind is an intelligent assistant for coding, analysis, and task automation.
# Extensions
+2
View File
@@ -7,6 +7,7 @@ use std::{
use crate::config::paths::Paths;
use crate::hints::import_files::read_referenced_files;
pub const TKMIND_HINTS_FILENAME: &str = ".tkmindhints";
pub const GOOSE_HINTS_FILENAME: &str = ".goosehints";
pub const AGENTS_MD_FILENAME: &str = "AGENTS.md";
@@ -17,6 +18,7 @@ pub fn get_context_filenames() -> Vec<String> {
.get_param::<Vec<String>>("CONTEXT_FILE_NAMES")
.unwrap_or_else(|_| {
vec![
TKMIND_HINTS_FILENAME.to_string(),
GOOSE_HINTS_FILENAME.to_string(),
AGENTS_MD_FILENAME.to_string(),
]
+1 -1
View File
@@ -3,5 +3,5 @@ pub mod load_hints;
pub use load_hints::{
build_gitignore, get_context_filenames, load_hint_files, SubdirectoryHintTracker,
AGENTS_MD_FILENAME, GOOSE_HINTS_FILENAME,
AGENTS_MD_FILENAME, GOOSE_HINTS_FILENAME, TKMIND_HINTS_FILENAME,
};
+1 -1
View File
@@ -1,4 +1,4 @@
You are a specialized subagent within the goose AI framework, created by AAIF (Agentic AI Foundation). You were spawned by the main goose agent to handle a specific task efficiently.
You are a specialized subagent within the TKMind AI framework. You were spawned by the main TKMind agent to handle a specific task efficiently.
# Your Role
You are an autonomous subagent with these characteristics:
+2 -2
View File
@@ -1,5 +1,5 @@
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
goose is being developed as an open-source software project.
You are a general-purpose AI agent called TKMind, created by TKMind.
TKMind is an intelligent assistant for coding, analysis, and task automation.
{% if not code_execution_mode %}
# Extensions
@@ -1,4 +1,4 @@
You are goose, an autonomous AI agent created by AAIF (Agentic AI Foundation). You act on the user's
You are TKMind, an autonomous AI agent. You act on the user's
behalf — you do not explain how to do things, you DO them directly.
The OS is {{os}}, the shell is {{shell}}, and the working directory is {{working_directory}}
+26 -1
View File
@@ -15,6 +15,16 @@ use std::fs::read_to_string;
use std::path::PathBuf;
use std::time::Duration;
/// Cap how long establishing the TCP/TLS connection may take, independent of the
/// (much larger) total request timeout used for inference. Without this, pointing
/// a provider at an unreachable host (e.g. a wrong LAN Ollama address) hangs until
/// the full request timeout elapses, which looks like the app freezing.
const CONNECT_TIMEOUT_SECS: u64 = 10;
fn connect_timeout(total_timeout: Duration) -> Duration {
Duration::from_secs(CONNECT_TIMEOUT_SECS).min(total_timeout)
}
pub struct ApiClient {
client: Client,
host: String,
@@ -292,7 +302,9 @@ impl ApiClient {
}
pub fn with_timeout(host: String, auth: AuthMethod, timeout: Duration) -> Result<Self> {
let mut client_builder = Client::builder().timeout(timeout);
let mut client_builder = Client::builder()
.timeout(timeout)
.connect_timeout(connect_timeout(timeout));
// Configure TLS if needed
let tls_config = TlsConfig::from_config()?;
@@ -316,6 +328,7 @@ impl ApiClient {
fn rebuild_client(&mut self) -> Result<()> {
let mut client_builder = Client::builder()
.timeout(self.timeout)
.connect_timeout(connect_timeout(self.timeout))
.default_headers(self.default_headers.clone());
// Configure TLS if needed
@@ -699,4 +712,16 @@ mod tests {
assert_eq!(actual, expected);
});
}
#[test]
fn test_connect_timeout_is_capped_by_total_timeout() {
assert_eq!(
connect_timeout(Duration::from_secs(600)),
Duration::from_secs(CONNECT_TIMEOUT_SECS)
);
assert_eq!(
connect_timeout(Duration::from_secs(3)),
Duration::from_secs(3)
);
}
}