chore: refactor code

rustfmt.toml

@@ -1,10 +1,19 @@
+# Language edition (parsing) + style guide edition (formatting rules).
 edition = "2024"
-comment_width = 100
-format_code_in_doc_comments = true
+style_edition = "2024"
+max_width = 100
+
+# Imports: one merged `use` per crate, grouped std → external → crate.
 imports_granularity = "Crate"
-imports_layout = "Horizontal"
+group_imports = "StdExternalCrate"
+
+# Comments: wrap to width and format code samples inside doc comments.
 wrap_comments = true
-brace_style = "SameLineWhere"
-fn_params_layout = "Vertical"
-fn_call_width = 30
-chain_width = 1
+comment_width = 100
+format_code_in_doc_comments = true
+
+# Tidy-ups.
+use_field_init_shorthand = true
+reorder_impl_items = true
+match_block_trailing_comma = true
+overflow_delimited_expr = true

src/agent/mod.rs

@@ -0,0 +1,7 @@
+//! LLM layer: the routing orchestrator, the optimization/verification agents, and
+//! the rig `Tool` they call to compile candidate rewrites.
+
+pub(crate) mod orchestrator;
+pub(crate) mod rig_agent;
+pub(crate) mod tools;
+pub(crate) mod verify_agent;

src/orchestrator.rs → src/agent/orchestrator.rs

@@ -11,11 +11,13 @@
 
 use std::sync::{Arc, Mutex};
 
-use rig_core::agent::{PromptHook, ToolCallHookAction};
-use rig_core::client::CompletionClient;
-use rig_core::completion::{CompletionModel, Prompt, ToolDefinition};
-use rig_core::providers::deepseek;
-use rig_core::tool::Tool;
+use rig_core::{
+    agent::{PromptHook, ToolCallHookAction},
+    client::CompletionClient,
+    completion::{CompletionModel, Prompt, ToolDefinition},
+    providers::deepseek,
+    tool::Tool,
+};
 use serde::Deserialize;
 use serde_json::json;
 use tracing::info;
@@ -67,17 +69,14 @@ struct PlanArgs {
 
 /// Route a contract from its skeleton. On any failure (no tool call, unparseable
 /// args), returns `Err` so the caller can fall back to full per-function fan-out.
-pub async fn route(
-    client: &deepseek::Client,
-    skeleton: &str,
-) -> Result<Route, String> {
+pub async fn route(client: &deepseek::Client, skeleton: &str) -> Result<Route, String> {
     let slot: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
     let hook = CaptureHook { slot: slot.clone() };
 
     // The loop terminates inside the hook the moment the tool is called, so the
     // result is a cancellation error — the decision lives in `slot`.
     let _ = client
-        .agent(crate::rig_agent::MODEL)
+        .agent(super::rig_agent::MODEL)
         .preamble(ROUTER_PROMPT)
         .context(skeleton)
         .temperature(0.0)
@@ -98,19 +97,8 @@ pub async fn route(
     let parsed: PlanArgs =
         serde_json::from_str(&args).map_err(|e| format!("router args parse failed: {e}"))?;
 
-    if parsed
-        .mode
-        .eq_ignore_ascii_case("decompose")
-        && !parsed
-            .tasks
-            .is_empty()
-    {
-        info!(
-            "  router: decompose ({} task(s))",
-            parsed
-                .tasks
-                .len()
-        );
+    if parsed.mode.eq_ignore_ascii_case("decompose") && !parsed.tasks.is_empty() {
+        info!("  router: decompose ({} task(s))", parsed.tasks.len());
         Ok(Route::Decompose(parsed.tasks))
     } else {
         info!("  router: oneshot");
@@ -132,10 +120,7 @@ impl<M: CompletionModel> PromptHook<M> for CaptureHook {
         _internal_call_id: &str,
         args: &str,
     ) -> ToolCallHookAction {
-        *self
-            .slot
-            .lock()
-            .unwrap() = Some(args.to_string());
+        *self.slot.lock().unwrap() = Some(args.to_string());
         ToolCallHookAction::terminate("plan captured")
     }
 }
@@ -148,16 +133,13 @@ pub struct PlanError(String);
 struct SubmitPlanTool;
 
 impl Tool for SubmitPlanTool {
-    const NAME: &'static str = "submit_plan";
-
-    type Error = PlanError;
     type Args = PlanArgs;
+    type Error = PlanError;
     type Output = String;
 
-    async fn definition(
-        &self,
-        _prompt: String,
-    ) -> ToolDefinition {
+    const NAME: &'static str = "submit_plan";
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
         ToolDefinition {
             name: Self::NAME.to_string(),
             description: "Submit the optimization routing decision for this contract.".to_string(),
@@ -196,10 +178,7 @@ impl Tool for SubmitPlanTool {
         }
     }
 
-    async fn call(
-        &self,
-        _args: Self::Args,
-    ) -> Result<Self::Output, Self::Error> {
+    async fn call(&self, _args: Self::Args) -> Result<Self::Output, Self::Error> {
         // Unreachable in practice: the hook terminates the loop before the tool body
         // executes. Present only so the model is offered the tool.
         Ok("ok".to_string())

src/rig_agent.rs → src/agent/rig_agent.rs

@@ -19,8 +19,8 @@ use rig_core::{
 };
 use tracing::info;
 
-
-use crate::{retrieval::GasliteIndex, tools::FunctionForgeTool, utils::strip_code_fences};
+use super::tools::FunctionForgeTool;
+use crate::{kb::retrieval::GasliteIndex, utils::strip_code_fences};
 
 /// The single model id used by every agent (router, oneshot, per-function, verify).
 /// `deepseek-chat` is the **non-thinking** mode of `deepseek-v4-flash` — it skips the
@@ -151,9 +151,7 @@ pub async fn optimize_function(
     let hook = TimingHook::new(fn_name);
     let result = if use_forge {
         builder
-            .tool(FunctionForgeTool::new(
-                original, fn_start, fn_end,
-            ))
+            .tool(FunctionForgeTool::new(original, fn_start, fn_end))
             .build()
             .prompt(user)
             .with_hook(hook.clone())
@@ -196,10 +194,7 @@ pub async fn optimize_function(
 
 /// Render the file-level declarations block for agent context (empty when none).
 fn decls_block(file_decls: &str) -> String {
-    if file_decls
-        .trim()
-        .is_empty()
-    {
+    if file_decls.trim().is_empty() {
         String::new()
     } else {
         format!(
@@ -282,45 +277,26 @@ impl TimingHook {
     fn new(label: &str) -> Self {
         Self {
             label: Arc::from(label),
-            state: Arc::new(Mutex::new(
-                TurnTimer::default(),
-            )),
+            state: Arc::new(Mutex::new(TurnTimer::default())),
             captured: Arc::new(Mutex::new(None)),
         }
     }
 
     /// `(turns, total LLM time, total in-loop tool time)`.
     fn summary(&self) -> (usize, Duration, Duration) {
-        let s = self
-            .state
-            .lock()
-            .unwrap();
-        (
-            s.turn,
-            s.llm_total,
-            s.tool_total,
-        )
+        let s = self.state.lock().unwrap();
+        (s.turn, s.llm_total, s.tool_total)
     }
 
     /// The verified candidate captured on early exit, if any.
     fn captured(&self) -> Option<String> {
-        self.captured
-            .lock()
-            .unwrap()
-            .clone()
+        self.captured.lock().unwrap().clone()
     }
 }
 
 impl<M: CompletionModel> PromptHook<M> for TimingHook {
-    async fn on_completion_call(
-        &self,
-        _prompt: &Message,
-        _history: &[Message],
-    ) -> HookAction {
-        let mut s = self
-            .state
-            .lock()
-            .unwrap();
+    async fn on_completion_call(&self, _prompt: &Message, _history: &[Message]) -> HookAction {
+        let mut s = self.state.lock().unwrap();
         s.turn += 1;
         s.llm_start = Some(Instant::now());
         HookAction::cont()
@@ -331,22 +307,13 @@ impl<M: CompletionModel> PromptHook<M> for TimingHook {
         _prompt: &Message,
         _response: &CompletionResponse<M::Response>,
     ) -> HookAction {
-        let mut s = self
-            .state
-            .lock()
-            .unwrap();
-        if let Some(start) = s
-            .llm_start
-            .take()
-        {
+        let mut s = self.state.lock().unwrap();
+        if let Some(start) = s.llm_start.take() {
             let d = start.elapsed();
             s.llm_total += d;
             let turn = s.turn;
             drop(s);
-            info!(
-                "  [{}] turn {turn}: LLM {d:.2?}",
-                self.label
-            );
+            info!("  [{}] turn {turn}: LLM {d:.2?}", self.label);
         }
         HookAction::cont()
     }
@@ -358,10 +325,7 @@ impl<M: CompletionModel> PromptHook<M> for TimingHook {
         _internal_call_id: &str,
         _args: &str,
     ) -> ToolCallHookAction {
-        self.state
-            .lock()
-            .unwrap()
-            .tool_start = Some(Instant::now());
+        self.state.lock().unwrap().tool_start = Some(Instant::now());
         ToolCallHookAction::cont()
     }
 
@@ -374,22 +338,13 @@ impl<M: CompletionModel> PromptHook<M> for TimingHook {
         result: &str,
     ) -> HookAction {
         {
-            let mut s = self
-                .state
-                .lock()
-                .unwrap();
-            if let Some(start) = s
-                .tool_start
-                .take()
-            {
+            let mut s = self.state.lock().unwrap();
+            if let Some(start) = s.tool_start.take() {
                 let d = start.elapsed();
                 s.tool_total += d;
                 let turn = s.turn;
                 drop(s);
-                info!(
-                    "  [{}] turn {turn}: tool {tool_name} {d:.2?}",
-                    self.label
-                );
+                info!("  [{}] turn {turn}: tool {tool_name} {d:.2?}", self.label);
             }
         }
 
@@ -412,10 +367,7 @@ impl<M: CompletionModel> PromptHook<M> for TimingHook {
                     .and_then(|v| v.get("optimized_function"))
                     .and_then(|x| x.as_str());
                 if let Some(f) = func {
-                    *self
-                        .captured
-                        .lock()
-                        .unwrap() = Some(strip_code_fences(f));
+                    *self.captured.lock().unwrap() = Some(strip_code_fences(f));
                     return HookAction::terminate("forge_verify compiled — skipping final turn");
                 }
             }
@@ -424,4 +376,3 @@ impl<M: CompletionModel> PromptHook<M> for TimingHook {
         HookAction::cont()
     }
 }
-