agent-runtime-definition-builder 0.1.1

use std::{path::Component, sync::Arc};

use async_trait::async_trait;
use runtime_extension_api::{
    ContextInvalidation, ExtensionError, SideEffect, Tool, ToolContext, ToolDescriptor, ToolOutcome,
};
use serde::Deserialize;
use serde_json::{Value, json};

/// Standard Workspace Tool composition for an Agent Definition.
///
/// `read_only` is suitable for Skill + Workspace agents. `read_write` adds one
/// bounded write Tool, and `with_command_execution` explicitly opts into the
/// higher-risk process capability. Caller capability filtering remains
/// authoritative at execution time.
pub struct WorkspaceToolSet {
    write: bool,
    execute: bool,
}

impl WorkspaceToolSet {
    pub fn read_only() -> Self {
        Self {
            write: false,
            execute: false,
        }
    }

    pub fn read_write() -> Self {
        Self {
            write: true,
            execute: false,
        }
    }

    pub fn with_command_execution(mut self) -> Self {
        self.execute = true;
        self
    }

    pub(crate) fn into_tools(self) -> Vec<Arc<dyn Tool>> {
        let mut tools: Vec<Arc<dyn Tool>> =
            vec![Arc::new(ReadTool), Arc::new(ListTool), Arc::new(SearchTool)];
        if self.write {
            tools.push(Arc::new(WriteTool));
        }
        if self.execute {
            tools.push(Arc::new(ExecuteTool));
        }
        tools
    }
}

fn descriptor(
    name: &str,
    description: &str,
    input_schema: Value,
    side_effect: SideEffect,
    capability: &str,
) -> ToolDescriptor {
    ToolDescriptor {
        name: name.into(),
        description: description.into(),
        input_schema,
        required_capabilities: vec![capability.into()],
        side_effect,
    }
}

fn parse<T: for<'de> Deserialize<'de>>(value: Value) -> Result<T, ExtensionError> {
    serde_json::from_value(value).map_err(|error| ExtensionError::Invalid(error.to_string()))
}

fn validate_path(path: &str, root_allowed: bool) -> Result<(), ExtensionError> {
    if (!root_allowed && path.is_empty())
        || std::path::Path::new(path).is_absolute()
        || std::path::Path::new(path)
            .components()
            .any(|part| !matches!(part, Component::Normal(_)))
    {
        return Err(ExtensionError::Invalid(
            "workspace path must be relative and must not contain traversal".into(),
        ));
    }
    Ok(())
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PathInput {
    path: String,
}

struct ReadTool;
#[async_trait]
impl Tool for ReadTool {
    fn descriptor(&self) -> ToolDescriptor {
        descriptor(
            "workspace.read",
            "Read one UTF-8 workspace file.",
            json!({"type":"object","properties":{"path":{"type":"string","minLength":1}},"required":["path"],"additionalProperties":false}),
            SideEffect::ReadOnly,
            "files:read",
        )
    }

    async fn invoke(
        &self,
        context: ToolContext,
        value: Value,
    ) -> Result<ToolOutcome, ExtensionError> {
        let input: PathInput = parse(value)?;
        validate_path(&input.path, false)?;
        Ok(ToolOutcome::success(
            context.workspace.read_text(&input.path).await?,
        ))
    }
}

struct ListTool;
#[async_trait]
impl Tool for ListTool {
    fn descriptor(&self) -> ToolDescriptor {
        descriptor(
            "workspace.list",
            "List entries below a workspace-relative path.",
            json!({"type":"object","properties":{"path":{"type":"string"}},"required":["path"],"additionalProperties":false}),
            SideEffect::ReadOnly,
            "files:read",
        )
    }

    async fn invoke(
        &self,
        context: ToolContext,
        value: Value,
    ) -> Result<ToolOutcome, ExtensionError> {
        let input: PathInput = parse(value)?;
        validate_path(&input.path, true)?;
        let entries = context.workspace.list(&input.path).await?;
        Ok(ToolOutcome::success(
            entries
                .into_iter()
                .map(|entry| format!("{}{}", entry.path, if entry.is_dir { "/" } else { "" }))
                .collect::<Vec<_>>()
                .join("\n"),
        ))
    }
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SearchInput {
    path: String,
    query: String,
    #[serde(default = "default_limit")]
    limit: usize,
}

fn default_limit() -> usize {
    100
}

struct SearchTool;
#[async_trait]
impl Tool for SearchTool {
    fn descriptor(&self) -> ToolDescriptor {
        descriptor(
            "workspace.search",
            "Search workspace text with a bounded result count.",
            json!({"type":"object","properties":{"path":{"type":"string"},"query":{"type":"string","minLength":1},"limit":{"type":"integer","minimum":1,"maximum":1000}},"required":["path","query"],"additionalProperties":false}),
            SideEffect::ReadOnly,
            "files:read",
        )
    }

    async fn invoke(
        &self,
        context: ToolContext,
        value: Value,
    ) -> Result<ToolOutcome, ExtensionError> {
        let input: SearchInput = parse(value)?;
        validate_path(&input.path, true)?;
        if input.query.is_empty() || !(1..=1000).contains(&input.limit) {
            return Err(ExtensionError::Invalid("query or limit is invalid".into()));
        }
        Ok(ToolOutcome::success(
            context
                .workspace
                .search(&input.path, &input.query, input.limit)
                .await?
                .join("\n"),
        ))
    }
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct WriteInput {
    path: String,
    content: String,
}

struct WriteTool;
#[async_trait]
impl Tool for WriteTool {
    fn descriptor(&self) -> ToolDescriptor {
        descriptor(
            "workspace.write",
            "Create or replace one UTF-8 workspace file.",
            json!({"type":"object","properties":{"path":{"type":"string","minLength":1},"content":{"type":"string"}},"required":["path","content"],"additionalProperties":false}),
            SideEffect::WorkspaceWrite,
            "files:write",
        )
    }

    async fn invoke(
        &self,
        context: ToolContext,
        value: Value,
    ) -> Result<ToolOutcome, ExtensionError> {
        let input: WriteInput = parse(value)?;
        validate_path(&input.path, false)?;
        context
            .workspace
            .write_text(&input.path, &input.content)
            .await?;
        Ok(ToolOutcome::success(format!(
            "wrote {} bytes to {}",
            input.content.len(),
            input.path
        ))
        .with_context_invalidation(ContextInvalidation::RUNTIME_STATE))
    }
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ExecuteInput {
    command: String,
    cwd: Option<String>,
}

struct ExecuteTool;
#[async_trait]
impl Tool for ExecuteTool {
    fn descriptor(&self) -> ToolDescriptor {
        descriptor(
            "workspace.exec",
            "Execute one command through the managed Workspace service.",
            json!({"type":"object","properties":{"command":{"type":"string","minLength":1},"cwd":{"type":["string","null"]}},"required":["command"],"additionalProperties":false}),
            SideEffect::Process,
            "commands:execute",
        )
    }

    async fn invoke(
        &self,
        context: ToolContext,
        value: Value,
    ) -> Result<ToolOutcome, ExtensionError> {
        let input: ExecuteInput = parse(value)?;
        if input.command.trim().is_empty() {
            return Err(ExtensionError::Invalid("command must not be empty".into()));
        }
        if let Some(cwd) = &input.cwd {
            validate_path(cwd, true)?;
        }
        let result = context
            .workspace
            .execute(&input.command, input.cwd.as_deref())
            .await?;
        Ok(ToolOutcome::success(format!(
            "exit={}\nstdout:\n{}\nstderr:\n{}",
            result.exit_code, result.stdout, result.stderr
        ))
        .with_context_invalidation(ContextInvalidation::RUNTIME_STATE))
    }
}