refactor: consolidate first-class tools — colocate compile-time and runtime code per tool

AGENTS.md

@@ -29,7 +29,10 @@ Alongside the correctly generated pipeline yaml, an agent file is generated from
 │   │   ├── common.rs     # Shared helpers across targets
 │   │   ├── standalone.rs # Standalone pipeline compiler
 │   │   ├── onees.rs      # 1ES Pipeline Template compiler
-│   │   ├── extensions.rs # CompilerExtension trait for runtimes/tools
+│   │   ├── extensions/   # CompilerExtension trait and infrastructure extensions
+│   │   │   ├── mod.rs    # Trait, Extension enum, collect_extensions(), re-exports
+│   │   │   ├── github.rs # Always-on GitHub MCP extension
+│   │   │   └── safe_outputs.rs # Always-on SafeOutputs MCP extension
 │   │   └── types.rs      # Front matter grammar and types
 │   ├── init.rs           # Repository initialization for AI-first authoring
 │   ├── execute.rs        # Stage 3 safe output execution
@@ -65,18 +68,26 @@ Alongside the correctly generated pipeline yaml, an agent file is generated from
 │   │   ├── update_wiki_page.rs
 │   │   ├── update_work_item.rs
 │   │   └── upload_attachment.rs
-│   ├── runtimes/         # Runtime environment implementations
+│   ├── runtimes/         # Runtime environment implementations (one dir per runtime)
 │   │   ├── mod.rs        # Module entry point
-│   │   └── lean.rs       # Lean 4 theorem prover runtime
+│   │   └── lean/         # Lean 4 theorem prover runtime
+│   │       ├── mod.rs    # Config types, install helpers
+│   │       └── extension.rs # CompilerExtension impl
 │   ├── data/
 │   │   ├── base.yml          # Base pipeline template for standalone
 │   │   ├── 1es-base.yml      # Base pipeline template for 1ES target
 │   │   ├── ecosystem_domains.json # Network allowlists per ecosystem
 │   │   ├── init-agent.md     # Dispatcher agent template for `init` command
 │   │   └── threat-analysis.md # Threat detection analysis prompt template
-│   └── tools/            # First-class tool implementations (compiler auto-configures)
+│   └── tools/            # First-class tool implementations (one dir per tool)
 │       ├── mod.rs
-│       └── cache_memory.rs
+│       ├── azure_devops/  # Azure DevOps MCP tool
+│       │   ├── mod.rs
+│       │   └── extension.rs # CompilerExtension impl
+│       └── cache_memory/  # Persistent agent memory tool
+│           ├── mod.rs
+│           ├── extension.rs # CompilerExtension impl (compile-time)
+│           └── execute.rs   # Stage 3 runtime (validate/copy)
 ├── examples/             # Example agent definitions
 ├── tests/                # Integration tests and fixtures
 ├── Cargo.toml            # Rust dependencies
@@ -1489,13 +1500,24 @@ When extending the compiler:
 3. **New front matter fields**: Add fields to `FrontMatter` in `src/compile/types.rs`
 4. **New template markers**: Handle replacements in the target-specific compiler (e.g., `standalone.rs` or `onees.rs`)
 5. **New safe-output tools**: Add to `src/safeoutputs/` — implement `ToolResult`, `Executor`, register in `mod.rs`, `mcp.rs`, `execute.rs`
-6. **New first-class tools**: Add to `src/tools/` — extend `ToolsConfig` in `types.rs`, implement `CompilerExtension` trait in `src/compile/extensions.rs`, add collection in `collect_extensions()`
-7. **New runtimes**: Add to `src/runtimes/` — extend `RuntimesConfig` in `types.rs`, implement `CompilerExtension` trait in `src/compile/extensions.rs`, add collection in `collect_extensions()`
+6. **New first-class tools**: Create `src/tools/<name>/` with `mod.rs` and `extension.rs` (CompilerExtension impl). Add `execute.rs` if the tool has Stage 3 runtime logic. Extend `ToolsConfig` in `types.rs`, add collection in `collect_extensions()`
+7. **New runtimes**: Create `src/runtimes/<name>/` with `mod.rs` (config types) and `extension.rs` (CompilerExtension impl). Extend `RuntimesConfig` in `types.rs`, add collection in `collect_extensions()`
 8. **Validation**: Add compile-time validation for safe outputs and permissions
 
+#### Code Organization Principles
+
+The codebase follows a **colocation** principle for tools and runtimes:
+
+- **Tools** (`tools:` front matter) live in `src/tools/<name>/` — one directory per tool, containing both compile-time (`extension.rs`) and runtime (`execute.rs`) code. This means you can look at a single directory to understand everything a tool does.
+- **Runtimes** (`runtimes:` front matter) live in `src/runtimes/<name>/` — one directory per runtime, with config types in `mod.rs` and the `CompilerExtension` impl in `extension.rs`.
+- **Infrastructure extensions** (GitHub MCP, SafeOutputs MCP) that are always-on and not user-configured stay in `src/compile/extensions/`. These are internal plumbing, not user-facing tools.
+- **Safe outputs** (`safe-outputs:` front matter) stay in `src/safeoutputs/` — they follow a different lifecycle (Stage 1 NDJSON → Stage 3 execution) and are not `CompilerExtension` implementations.
+
+The `src/compile/extensions/mod.rs` file owns the `CompilerExtension` trait, the `Extension` enum, and `collect_extensions()`. It re-exports tool/runtime extension types from their colocated homes so the rest of the compiler can import them from a single path.
+
 #### `CompilerExtension` Trait
 
-Runtimes and first-party tools declare their compilation requirements via the `CompilerExtension` trait (`src/compile/extensions.rs`). Instead of scattering special-case `if` blocks across the compiler, each runtime/tool implements this trait and the compiler collects requirements generically:
+Runtimes and first-party tools declare their compilation requirements via the `CompilerExtension` trait (`src/compile/extensions/mod.rs`). Instead of scattering special-case `if` blocks across the compiler, each runtime/tool implements this trait and the compiler collects requirements generically:
 
 ```rust
 pub trait CompilerExtension: Send {
@@ -1509,7 +1531,7 @@ pub trait CompilerExtension: Send {
 }
 ```
 
-To add a new runtime or tool: (1) create a struct implementing `CompilerExtension`, (2) add a collection check in `collect_extensions()`. No other files need modification.
+To add a new runtime or tool: (1) create a directory under `src/tools/` or `src/runtimes/`, (2) implement `CompilerExtension` in `extension.rs`, (3) add a variant to the `Extension` enum and a collection check in `collect_extensions()` in `src/compile/extensions/mod.rs`.
 
 ### Security Considerations
 

src/compile/extensions/mod.rs

@@ -351,16 +351,14 @@ macro_rules! extension_enum {
     };
 }
 
-mod azure_devops;
-mod cache_memory;
 mod github;
-mod lean;
 mod safe_outputs;
 
-pub use azure_devops::{AdoAuthMode, AzureDevOpsExtension};
-pub use cache_memory::CacheMemoryExtension;
+// Re-export tool/runtime extensions from their colocated homes
+pub use crate::tools::azure_devops::{AdoAuthMode, AzureDevOpsExtension};
+pub use crate::tools::cache_memory::CacheMemoryExtension;
 pub use github::GitHubExtension;
-pub use lean::LeanExtension;
+pub use crate::runtimes::lean::LeanExtension;
 pub use safe_outputs::SafeOutputsExtension;
 
 extension_enum! {

src/compile/extensions/tests.rs

@@ -1,5 +1,5 @@
 use super::*;
-use crate::compile::common::{ADO_MCP_SERVER_NAME, parse_markdown};
+use crate::compile::{ADO_MCP_SERVER_NAME, parse_markdown};
 use crate::compile::types::{AzureDevOpsToolConfig, CacheMemoryToolConfig};
 use crate::runtimes::lean::LeanRuntimeConfig;
 

src/compile/mod.rs

@@ -23,7 +23,10 @@ pub use common::generate_copilot_params;
 pub use common::generate_mcpg_config;
 pub use common::MCPG_IMAGE;
 pub use common::MCPG_VERSION;
-pub use common::MCPG_PORT;
+pub use common::ADO_MCP_ENTRYPOINT;
+pub use common::ADO_MCP_IMAGE;
+pub use common::ADO_MCP_PACKAGE;
+pub use common::ADO_MCP_SERVER_NAME;
 pub use types::{CompileTarget, FrontMatter};
 
 /// Trait for pipeline compilers.

src/compile/extensions/lean.rs → src/runtimes/lean/extension.rs

@@ -1,7 +1,7 @@
 // ─── Lean 4 ──────────────────────────────────────────────────────────
 
-use super::{CompileContext, CompilerExtension, ExtensionPhase};
-use crate::runtimes::lean::{self, LEAN_BASH_COMMANDS, LeanRuntimeConfig};
+use crate::compile::extensions::{CompileContext, CompilerExtension, ExtensionPhase};
+use super::{LEAN_BASH_COMMANDS, LeanRuntimeConfig, generate_lean_install};
 use anyhow::Result;
 
 /// Lean 4 runtime extension.
@@ -53,7 +53,7 @@ the toolchain. Lean files use the `.lean` extension.\n"
     }
 
     fn prepare_steps(&self) -> Vec<String> {
-        vec![lean::generate_lean_install(&self.config)]
+        vec![generate_lean_install(&self.config)]
     }
 
     fn validate(&self, ctx: &CompileContext) -> Result<Vec<String>> {

src/runtimes/lean.rs → src/runtimes/lean/mod.rs

@@ -8,6 +8,10 @@
 //! Lean is installed via elan (the Lean toolchain manager) into `$HOME/.elan/bin`,
 //! then symlinked into `/tmp/awf-tools/` for AWF chroot compatibility.
 
+pub mod extension;
+
+pub use extension::LeanExtension;
+
 use ado_aw_derive::SanitizeConfig;
 use serde::Deserialize;
 

src/compile/extensions/azure_devops.rs → src/tools/azure_devops/extension.rs

@@ -1,10 +1,10 @@
 // ─── Azure DevOps MCP ────────────────────────────────────────────────
 
-use super::{
+use crate::compile::extensions::{
     CompileContext, CompilerExtension, ExtensionPhase, McpgServerConfig, PipelineEnvMapping,
 };
 use crate::allowed_hosts::mcp_required_hosts;
-use crate::compile::common::{
+use crate::compile::{
     ADO_MCP_ENTRYPOINT, ADO_MCP_IMAGE, ADO_MCP_PACKAGE, ADO_MCP_SERVER_NAME,
 };
 use crate::compile::types::AzureDevOpsToolConfig;

src/tools/azure_devops/mod.rs

@@ -0,0 +1,9 @@
+//! Azure DevOps first-class tool.
+//!
+//! Compile-time: injects network hosts (ADO domains), MCPG server entry
+//! (containerized ADO MCP), and compile-time validation (org inference,
+//! duplicate MCP).
+
+pub mod extension;
+
+pub use extension::{AdoAuthMode, AzureDevOpsExtension};

src/compile/extensions/cache_memory.rs → src/tools/cache_memory/extension.rs

@@ -1,4 +1,4 @@
-use super::{CompilerExtension, ExtensionPhase};
+use crate::compile::extensions::{CompilerExtension, ExtensionPhase};
 use crate::compile::types::CacheMemoryToolConfig;
 
 /// Cache memory tool extension.

src/tools/cache_memory/mod.rs

@@ -0,0 +1,13 @@
+//! Cache memory first-class tool.
+//!
+//! Compile-time: injects pipeline steps (download/restore previous memory)
+//! and a prompt supplement informing the agent about its memory directory.
+//!
+//! Stage 3 runtime: validates and copies sanitized memory files to the
+//! final safe_outputs artifact for pickup by the next run.
+
+pub mod execute;
+pub mod extension;
+
+pub use execute::{MemoryConfig, process_agent_memory};
+pub use extension::CacheMemoryExtension;

src/tools/mod.rs

@@ -1,10 +1,15 @@
 //! First-class tool implementations for the ado-aw compiler.
 //!
-//! These tools are configured via the `tools:` front-matter section and provide
+//! Each tool is colocated in its own subdirectory containing both
+//! compile-time (`extension.rs` — [`CompilerExtension`] impl) and
+//! runtime (`execute.rs` — Stage 3 logic) code where applicable.
+//!
+//! Tools are configured via the `tools:` front-matter section and provide
 //! built-in functionality that the compiler knows how to auto-configure
 //! (pipeline steps, MCPG entries, network allowlists, etc.).
 //!
 //! This is distinct from `safeoutputs/` which contains safe-output MCP tools
 //! that serialize to NDJSON in Stage 1 and execute in Stage 3.
 
+pub mod azure_devops;
 pub mod cache_memory;