Complete Rig tool lifecycle migration

2026-08-04 16:00:20 -05:00
parent 91d8bd0381
commit a3c68e9c30
30 changed files with 1494 additions and 176 deletions
@@ -0,0 +1,319 @@
+use std::time::Duration;
+
+use ai::diff_validation::ParsedDiff;
+use ai::skills::{SkillPathOrigin, SkillReference};
+use galaxy_agent_core::ToolCall;
+use uuid::Uuid;
+
+use crate::ai::agent::task::TaskId;
+use crate::ai::agent::{
+    AIAgentAction, AIAgentActionType, AIAgentPtyWriteMode, AskUserQuestionItem,
+    AskUserQuestionOption, AskUserQuestionType, CreateDocumentsRequest, DocumentDiff,
+    DocumentToCreate, EditDocumentsRequest, FileEdit, FileLocations, ReadDocumentsRequest,
+    ReadFilesRequest, ReadSkillRequest, SearchCodebaseRequest, ShellCommandDelay,
+    StartAgentExecutionMode, StartAgentVersion,
+};
+use crate::ai::document::ai_document_model::AIDocumentId;
+
+pub(super) fn action_from_tool_call(
+    task_id: &str,
+    call: &ToolCall,
+    skill_path_origin: &SkillPathOrigin,
+) -> Result<AIAgentAction, String> {
+    let input = &call.arguments;
+    let action = match call.name.as_str() {
+        "run_shell_command" => AIAgentActionType::RequestCommandOutput {
+            command: string(input, "command"),
+            is_read_only: Some(boolean(input, "is_read_only")),
+            is_risky: Some(boolean(input, "is_risky")),
+            wait_until_completion: boolean(input, "wait_until_complete"),
+            uses_pager: Some(boolean(input, "uses_pager")),
+            rationale: None,
+            citations: Vec::new(),
+        },
+        "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest {
+            locations: input
+                .get("files")
+                .and_then(serde_json::Value::as_array)
+                .into_iter()
+                .flatten()
+                .filter_map(file_location)
+                .collect(),
+        }),
+        "apply_file_diffs" => AIAgentActionType::RequestFileEdits {
+            file_edits: file_edits(input),
+            title: nonempty_string(input, "summary"),
+        },
+        "grep" => AIAgentActionType::Grep {
+            queries: strings(input, "queries"),
+            path: string(input, "path"),
+        },
+        "file_glob" => AIAgentActionType::FileGlob {
+            patterns: strings(input, "patterns"),
+            path: nonempty_string(input, "path"),
+        },
+        "search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest {
+            query: string(input, "query"),
+            partial_paths: nonempty_strings(input, "path_filters"),
+            codebase_path: nonempty_string(input, "path"),
+        }),
+        "write_to_long_running_shell_command" => {
+            AIAgentActionType::WriteToLongRunningShellCommand {
+                block_id: string(input, "command_id").into(),
+                input: string(input, "input").into_bytes().into(),
+                mode: match input.get("mode").and_then(serde_json::Value::as_str) {
+                    Some("line") => AIAgentPtyWriteMode::Line,
+                    Some("block") => AIAgentPtyWriteMode::Block,
+                    Some("raw") | Some(_) | None => AIAgentPtyWriteMode::Raw,
+                },
+            }
+        }
+        "interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand {
+            block_id: string(input, "command_id").into(),
+            input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(),
+            mode: AIAgentPtyWriteMode::Raw,
+        },
+        "read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput {
+            block_id: string(input, "command_id").into(),
+            delay: Some(ShellCommandDelay::Duration(Duration::from_secs(
+                input
+                    .get("wait_seconds")
+                    .and_then(serde_json::Value::as_u64)
+                    .unwrap_or(2)
+                    .min(crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS),
+            ))),
+        },
+        "read_mcp_resource" => AIAgentActionType::ReadMCPResource {
+            server_id: uuid(input, "server_id"),
+            name: String::new(),
+            uri: nonempty_string(input, "uri"),
+        },
+        "read_plan" | "read_documents" | "read_notebook" => {
+            AIAgentActionType::ReadDocuments(ReadDocumentsRequest {
+                document_ids: strings(input, "document_ids")
+                    .into_iter()
+                    .filter_map(|id| AIDocumentId::try_from(id).ok())
+                    .collect(),
+            })
+        }
+        "create_plan" | "create_documents" | "create_notebook" => {
+            AIAgentActionType::CreateDocuments(CreateDocumentsRequest {
+                documents: input
+                    .get("documents")
+                    .and_then(serde_json::Value::as_array)
+                    .into_iter()
+                    .flatten()
+                    .filter_map(|document| {
+                        Some(DocumentToCreate {
+                            title: document.get("title")?.as_str()?.to_string(),
+                            content: document.get("content")?.as_str()?.to_string(),
+                        })
+                    })
+                    .collect(),
+            })
+        }
+        "edit_plan" | "edit_documents" | "edit_notebook" => {
+            AIAgentActionType::EditDocuments(EditDocumentsRequest {
+                diffs: input
+                    .get("diffs")
+                    .and_then(serde_json::Value::as_array)
+                    .into_iter()
+                    .flatten()
+                    .filter_map(|diff| {
+                        Some(DocumentDiff {
+                            document_id: AIDocumentId::try_from(diff.get("document_id")?.as_str()?)
+                                .ok()?,
+                            search: string(diff, "search"),
+                            replace: string(diff, "replace"),
+                        })
+                    })
+                    .collect(),
+            })
+        }
+        "start_agent" => AIAgentActionType::StartAgent {
+            version: StartAgentVersion::V1,
+            name: string(input, "name"),
+            prompt: string(input, "prompt"),
+            execution_mode: StartAgentExecutionMode::local_with_defaults(),
+            lifecycle_subscription: None,
+        },
+        "send_message_to_agent" => AIAgentActionType::SendMessageToAgent {
+            addresses: vec![string(input, "agent_id")],
+            subject: String::new(),
+            message: string(input, "message"),
+        },
+        "ask_user_question" => AIAgentActionType::AskUserQuestion {
+            questions: vec![AskUserQuestionItem {
+                question_id: Uuid::new_v4().to_string(),
+                question: string(input, "question"),
+                question_type: AskUserQuestionType::MultipleChoice {
+                    is_multiselect: false,
+                    options: strings(input, "options")
+                        .into_iter()
+                        .enumerate()
+                        .map(|(index, label)| AskUserQuestionOption {
+                            label,
+                            recommended: index == 0,
+                        })
+                        .collect(),
+                    supports_other: true,
+                },
+            }],
+        },
+        "read_skill" => {
+            let skill = string(input, "skill");
+            let skill = match input
+                .get("reference_type")
+                .and_then(serde_json::Value::as_str)
+            {
+                Some("bundled") => SkillReference::BundledSkillId(skill),
+                Some("path") | Some(_) | None => SkillReference::Path(
+                    skill_path_origin
+                        .location_for_path(skill)
+                        .map_err(|error| error.to_string())?,
+                ),
+            };
+            AIAgentActionType::ReadSkill(ReadSkillRequest { skill })
+        }
+        "fetch_conversation" => AIAgentActionType::FetchConversation {
+            conversation_id: string(input, "conversation_id"),
+        },
+        name if name.starts_with("mcp__") => {
+            let mut parts = name.splitn(3, "__");
+            let _prefix = parts.next();
+            let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok());
+            let name = parts
+                .next()
+                .unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name))
+                .to_string();
+            AIAgentActionType::CallMCPTool {
+                server_id,
+                name,
+                input: input.clone(),
+            }
+        }
+        name => return Err(format!("unsupported Rig tool proposal: {name}")),
+    };
+
+    let tool_name = matches!(
+        call.name.as_str(),
+        "read_notebook" | "create_notebook" | "edit_notebook"
+    )
+    .then(|| "notebook".to_string());
+    Ok(AIAgentAction {
+        id: call.id.clone().into(),
+        task_id: TaskId::new(task_id.to_string()),
+        action,
+        requires_result: true,
+        tool_name,
+    })
+}
+
+fn string(input: &serde_json::Value, key: &str) -> String {
+    input
+        .get(key)
+        .and_then(serde_json::Value::as_str)
+        .unwrap_or_default()
+        .to_string()
+}
+
+fn nonempty_string(input: &serde_json::Value, key: &str) -> Option<String> {
+    let value = string(input, key);
+    (!value.is_empty()).then_some(value)
+}
+
+fn boolean(input: &serde_json::Value, key: &str) -> bool {
+    input
+        .get(key)
+        .and_then(serde_json::Value::as_bool)
+        .unwrap_or(false)
+}
+
+fn strings(input: &serde_json::Value, key: &str) -> Vec<String> {
+    input
+        .get(key)
+        .and_then(serde_json::Value::as_array)
+        .into_iter()
+        .flatten()
+        .filter_map(serde_json::Value::as_str)
+        .map(ToOwned::to_owned)
+        .collect()
+}
+
+fn nonempty_strings(input: &serde_json::Value, key: &str) -> Option<Vec<String>> {
+    let values = strings(input, key);
+    (!values.is_empty()).then_some(values)
+}
+
+fn uuid(input: &serde_json::Value, key: &str) -> Option<Uuid> {
+    input
+        .get(key)
+        .and_then(serde_json::Value::as_str)
+        .and_then(|value| Uuid::parse_str(value).ok())
+}
+
+fn file_location(file: &serde_json::Value) -> Option<FileLocations> {
+    if let Some(name) = file.as_str() {
+        return Some(FileLocations {
+            name: name.to_string(),
+            lines: Vec::new(),
+        });
+    }
+    let name = file
+        .get("path")
+        .or_else(|| file.get("name"))?
+        .as_str()?
+        .to_string();
+    let lines = file
+        .get("line_ranges")
+        .and_then(serde_json::Value::as_array)
+        .into_iter()
+        .flatten()
+        .filter_map(|range| {
+            let start = usize::try_from(range.get("start")?.as_u64()?).ok()?;
+            let end = usize::try_from(range.get("end")?.as_u64()?).ok()?;
+            (start > 0 && end >= start).then_some(start..end)
+        })
+        .collect();
+    Some(FileLocations { name, lines })
+}
+
+fn file_edits(input: &serde_json::Value) -> Vec<FileEdit> {
+    let diffs = input
+        .get("diffs")
+        .and_then(serde_json::Value::as_array)
+        .into_iter()
+        .flatten()
+        .map(|diff| {
+            FileEdit::Edit(ParsedDiff::StrReplaceEdit {
+                file: nonempty_string(diff, "file_path"),
+                search: nonempty_string(diff, "search"),
+                replace: nonempty_string(diff, "replace"),
+            })
+        });
+    let creates = input
+        .get("new_files")
+        .and_then(serde_json::Value::as_array)
+        .into_iter()
+        .flatten()
+        .map(|file| FileEdit::Create {
+            file: nonempty_string(file, "file_path"),
+            content: nonempty_string(file, "content"),
+        });
+    let deletes = input
+        .get("deleted_files")
+        .and_then(serde_json::Value::as_array)
+        .into_iter()
+        .flatten()
+        .map(|file| FileEdit::Delete {
+            file: file
+                .as_str()
+                .map(ToOwned::to_owned)
+                .or_else(|| nonempty_string(file, "file_path")),
+        });
+    diffs.chain(creates).chain(deletes).collect()
+}
+
+#[cfg(test)]
+#[path = "rig_tool_tests.rs"]
+mod tests;